diff --git a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs index dfd9efffd..74a7d3693 100644 --- a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs +++ b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs @@ -877,6 +877,13 @@ impl AuxTradeOutcomeForwardOps { /// /// Block: 256 threads. Grid: `ceil(n_envs / 256)` blocks (per-env /// parallelism, no reduction, no atomicAdd). + /// `out_labels_per_sample_ptr` (Phase B4b 2026-05-14) is the per- + /// (env, t) i32 column the replay buffer scatter consumes. Pass 0 + /// (NULL) to skip the per-(i, t) write — the per-env `out_labels` + /// output is still produced. The kernel-side NULL guard is on the + /// pointer-read site, mirroring the standard NULL-tolerant arg + /// pattern used by `experience_env_step`'s + /// `pnl_vs_target_at_close_per_env` etc. (Phase A3). #[allow(clippy::too_many_arguments)] pub(crate) fn compute_label( &self, @@ -888,6 +895,7 @@ impl AuxTradeOutcomeForwardOps { n_envs: usize, l_timesteps: usize, out_labels_ptr: u64, + out_labels_per_sample_ptr: u64, ) -> Result<(), MLError> { let n_envs_i32 = n_envs as i32; let l_i32 = l_timesteps as i32; @@ -903,6 +911,7 @@ impl AuxTradeOutcomeForwardOps { .arg(&n_envs_i32) .arg(&l_i32) .arg(&out_labels_ptr) + .arg(&out_labels_per_sample_ptr) .launch(LaunchConfig { grid_dim: (grid, 1, 1), block_dim: (block, 1, 1), diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index f8b9cd0ee..946e1b5c2 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1265,6 +1265,26 @@ pub struct GpuExperienceCollector { /// scatter yet. exp_aux_to_label_buf: cudarc::driver::CudaSlice, + /// SP22 H6 vNext Phase B4b (2026-05-14) — per-(env, t) trade-outcome + /// label column. Sized `[alloc_episodes × alloc_timesteps]` i32, + /// row-major: env-major then t. Written by `trade_outcome_label_kernel` + /// at offset `env*L + t` on every rollout step alongside the per-env + /// `exp_aux_to_label_buf` write. Sparse (~95-99% mask=-1). + /// + /// This column threads through the replay-buffer's `aux_sign_labels`- + /// equivalent scatter on `insert_batch` and gather on + /// `sample_proportional` into the trainer's `aux_to_label_buf` + /// (B4b-2 wireup — replay buffer + trainer setter, deferred to + /// next commit). For B4b-1 (this commit) the buffer is allocated + /// and populated by the per-step launch; the consumer wiring is + /// pending. + /// + /// The trade-outcome label kernel's NEW per-sample output arg + /// (added in this commit) writes here. Per-env output stays in + /// `exp_aux_to_label_buf` for Phase C's per-step state-assembly + /// consumer. + exp_aux_to_label_per_sample: cudarc::driver::CudaSlice, + /// SP22 H6 Phase 2 (2026-05-12) — per-env aux directional probability /// cache, RECENTERED encoding. Sized `[alloc_episodes]` f32 (one slot /// per env). Written end-of-step by `exp_aux_softmax_to_per_env_kernel` @@ -2720,6 +2740,18 @@ impl GpuExperienceCollector { "sp22-vnext B3: alloc exp_aux_to_label_buf ({} i32): {e}", alloc_episodes )))?; + // SP22 H6 vNext Phase B4b (2026-05-14): per-(env, t) trade-outcome + // label column for the replay-buffer scatter chain. Sized + // `[alloc_episodes × alloc_timesteps]` — same shape as the K=2 + // head's `aux_sign_labels` ring. Sentinel -1 (mask) populated by + // the alloc_zeros + the per-step kernel writes (mask survives + // until a trade-close event overwrites the env's slot at that t). + let exp_aux_to_label_per_sample = stream + .alloc_zeros::(alloc_episodes * alloc_timesteps) + .map_err(|e| MLError::ModelError(format!( + "sp22-vnext B4b: alloc exp_aux_to_label_per_sample ({} i32): {e}", + alloc_episodes * alloc_timesteps + )))?; // SP22 H6 (2026-05-12): aux→policy state bridge. Load the per-env // softmax extractor (`aux_softmax_to_per_env_kernel`) + the shared // `fill_f32` from the epsilon_greedy cubin (for 0.5 sentinel init at @@ -3309,6 +3341,7 @@ impl GpuExperienceCollector { exp_aux_to_logits_buf, exp_aux_to_softmax_buf, exp_aux_to_label_buf, + exp_aux_to_label_per_sample, // SP22 H6 (2026-05-12): aux→policy state bridge — per-env p_up // cache + extractor kernel + fill_f32 handle for 0.5 sentinel. prev_aux_dir_prob, @@ -6899,6 +6932,11 @@ impl GpuExperienceCollector { { let t_i32 = t as i32; let l_i32 = timesteps as i32; + // Phase B4b (2026-05-14): pass per-(env, t) buffer so the + // kernel writes BOTH the per-env tile (consumed by Phase + // C's per-step state assembly producer) AND the per-(env, + // t) ring column (consumed by the replay-buffer scatter + // chain in B4b-2). self.exp_aux_to_fwd.compute_label( &self.stream, self.trade_close_per_sample.raw_ptr(), @@ -6908,6 +6946,7 @@ impl GpuExperienceCollector { n, // n_envs in this rollout slice timesteps, self.exp_aux_to_label_buf.raw_ptr(), + self.exp_aux_to_label_per_sample.raw_ptr(), )?; let _ = l_i32; // documents the stride contract — consumer reads `env*L + t` } diff --git a/crates/ml/src/cuda_pipeline/trade_outcome_label_kernel.cu b/crates/ml/src/cuda_pipeline/trade_outcome_label_kernel.cu index 64550b042..7a3d7c1a9 100644 --- a/crates/ml/src/cuda_pipeline/trade_outcome_label_kernel.cu +++ b/crates/ml/src/cuda_pipeline/trade_outcome_label_kernel.cu @@ -99,7 +99,16 @@ extern "C" __global__ void trade_outcome_label_kernel( /* Buffer stride: timesteps per env in the per-sample buffer layout. */ int L, /* Per-env i32 output. Written -1 when no trade close, else 0/1/2. */ - int* __restrict__ out_labels + int* __restrict__ out_labels, + /* SP22 H6 vNext Phase B4b (2026-05-14): per-(env, t) i32 output for + * the replay-buffer label ring scatter. Layout `[N, L]` row-major, + * written at the same `env*L + t` offset as `trade_close_per_sample` + * is read. NULL-tolerant — when passed 0 the kernel only emits the + * per-env `out_labels` (Phase A2/A3 contract). The per-(env, t) + * column threads through the replay buffer's `insert_batch` → + * `gather_i32_scalar` → trainer's `aux_to_label_buf` per the + * `aux_sign_labels` pattern. */ + int* __restrict__ out_labels_per_sample ) { int env = blockIdx.x * blockDim.x + threadIdx.x; if (env >= n_episodes) return; @@ -126,4 +135,15 @@ extern "C" __global__ void trade_outcome_label_kernel( } out_labels[env] = label; + + /* Phase B4b (2026-05-14): write to per-(env, t) column too. Same + * `env*L + t` offset as trade_close_per_sample. Persists across + * rollout steps as the per-(i, t) replay column (one write per env + * per step; subsequent steps for the same env write at different + * t-offsets). The bar that produced the trade-close event keeps + * its `0/1/2` label permanently; -1 fills the rest of the env's + * row. */ + if (out_labels_per_sample != NULL) { + out_labels_per_sample[out_off] = label; + } } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index cf5db06e5..ac4f1369d 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17913,3 +17913,22 @@ Fifth Rust-side commit of Phase B. Wires the trade-outcome head's forward + loss **Phase B4b next**: Replay-buffer label scatter so trainer's `aux_to_label_buf` gets populated from the rollout's `exp_aux_to_label_buf` per (i, t) — wires the cold-start degraded "predict Profit everywhere" path to the actual sparse trade-outcome label distribution. **Phase B5 (spec's actual "Phase B")**: Input concat 256→262 with `plan_params`. Requires a kernel-side change to the forward kernel's `h_s2_aux` input + W1 shape bump to `[H, 262]`. Small touch-up vs B0-B4's atomic wireup. + +#### Phase B4b-1 — Per-(env, t) label kernel output + collector buffer (2026-05-14) + +First half of Phase B4b (the replay-buffer label scatter chain). Amends the Phase A2 `trade_outcome_label_kernel` to emit a per-(env, t) output column alongside the existing per-env tile, and adds the collector-side buffer + per-step launch arg. + +**Kernel amendment** (`trade_outcome_label_kernel.cu`): +- Added NULL-tolerant arg `out_labels_per_sample` after the existing `out_labels` arg +- When non-NULL, writes `out_labels_per_sample[env*L + t] = label` at the same offset as the `trade_close_per_sample[env*L + t]` read. NULL = no-op (preserves Phase A2/A3 contract for any caller still passing the old signature). +- Pattern mirrors the existing `aux_sign_labels` per-(i, t) column the K=2 head threads through the replay buffer. + +**Collector field + alloc + launch arg**: +- New struct field `exp_aux_to_label_per_sample: CudaSlice` sized `[alloc_episodes × alloc_timesteps]`. Sentinel -1 (mask) populated by `alloc_zeros` + per-step kernel writes — survives until a trade-close event overwrites the env's slot at that t. +- Updated Phase B3 launcher in `collect_experiences_gpu` to pass `self.exp_aux_to_label_per_sample.raw_ptr()` as the new arg. + +**Why split B4b into B4b-1 + B4b-2**: the full replay-buffer wireup mirrors the K=2 head's `aux_sign_labels` pattern across ~8 distinct code sites (replay-buffer struct field, sample destination buffer, direct-to-trainer pointer, setter method, scatter on insert, gather direct, gather fallback, GpuBatchPtrs field). Splitting lets us validate the per-(i, t) producer in isolation before touching the consumer pipeline. + +**B4b-1 (this commit)** = producer chain complete. The per-(env, t) column is populated correctly every rollout step. Consumer wiring (replay-buffer scatter + trainer setter) is B4b-2's scope. + +**Verification**: `cargo check -p ml` clean. `cargo test -p ml --lib` → 1016/0 on clean runs; the pre-existing `test_dqn_checkpoint_round_trip` NoisyLinear flake still surfaces ~50-70% of full-suite runs (varies `pred1=-1 pred2=1` ↔ `pred1=1 pred2=-1` non-deterministically, even with the `bind_to_thread` fix from B4). Flake predates Phase B4b and is unrelated to the trade-outcome head — `disable_noise()` zeros NoisyLinear ε but apparently leaves some other randomness source intact. Worth dedicated investigation separate from vNext work.