gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(sp22-vnext): Phase B4b-1 — per-(env, t) label kernel output + collector buffer

Browse Source 
First half of Phase B4b (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 `out_labels_per_sample` arg after existing
  `out_labels`. 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 callers passing old signature).
- Pattern mirrors the K=2 head's `aux_sign_labels` per-(i, t) ring
  column that threads through the replay buffer.

Collector field + alloc + launch:
- New struct field `exp_aux_to_label_per_sample: CudaSlice<i32>`
  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. Per-(env, t) column
populated correctly every rollout step. Consumer wiring (replay-
buffer scatter + trainer setter) is B4b-2's scope.

Verification:
- cargo check -p ml clean (21 warnings, none new).
- cargo test -p ml --lib → 1016/0 on clean runs; pre-existing
  test_dqn_checkpoint_round_trip NoisyLinear flake still surfaces
  ~50-70% of full-suite runs (flake predates Phase B4b, unrelated
  to trade-outcome head — disable_noise() zeros ε but leaves some
  other randomness source intact).

Audit: docs/dqn-wire-up-audit.md Phase B4b-1 section.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-14 01:57:23 +02:00
parent 20e1aea27c
commit 491bf7d3e6
4 changed files with 88 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
crates/ml/src/cuda_pipeline/gpu_aux_heads.rs
View File

@@ -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),

39
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -1265,6 +1265,26 @@ pub struct GpuExperienceCollector {
/// scatter yet.
exp_aux_to_label_buf: cudarc::driver::CudaSlice<i32>,
/// 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<i32>,
/// 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::<i32>(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`
}

22
crates/ml/src/cuda_pipeline/trade_outcome_label_kernel.cu
View File

@@ -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;
}
}

19
docs/dqn-wire-up-audit.md
View File

@@ -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<i32>` 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.