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

feat(sp15-wave4.1b): consumer migration — s1_input_dim 102→103, GRN w_s1 reshape, 4 forward + 3 backward sites

Browse Source 
Atomic consumer migration that flips production callers of bn_tanh_concat_kernel
over to Wave 4.1a's bn_tanh_concat_dd_kernel and bumps s1_input_dim from 102 to
103 across the entire trunk forward + backward path. Eliminates the documented
Wave 4.1a transient orphan.

Changes:
- s1_input_dim formula bump (bn_dim + portfolio_dim → bn_dim + portfolio_dim + 1)
  at compute_param_sizes, trainer ctor's CublasGemmSet, xavier_init_params_buf,
  the experience collector's CublasGemmSet, and CublasBackwardSet::new for the
  backward gemm cache. cuBLAS gemm caches re-key automatically (fresh HashMap).
- GRN w_a_h_s1[0] / w_residual_h_s1[4] reshape [shared_h1, 102] → [shared_h1, 103]
  via compute_param_sizes + xavier_init's fan_dims. Xavier-uniform init covers
  the new dd_pct column (bounded [0,1] — Xavier's small-magnitude assumption is
  appropriate; differs from SP14's aux_softmax_diff zero-init which was driven
  by the bidirectional ±1 range).
- bn_concat_dim() accessor +1 (TLOB backward row stride).
- 5 concat_dim local-var bumps (1 alloc + 3 forward + 2 backward + 1 in
  experience collector).
- 4 forward-call migrations to launch_sp15_bn_concat_dd: DDQN argmax pass
  (~27158), online forward (~27467), target forward (~27666), experience
  collector forward (~3853). Each takes self.isv_signals_dev_ptr; the kernel
  reads ISV[DD_PCT_INDEX=406] on-device and broadcasts.
- 3 GRN backward sites (main, ensemble, CQL) flow through encoder_backward_chain
  which uses s1_input_dim — bumped automatically. dd_pct column gradient is
  silently discarded by vsn_d_gated_state_portfolio_pad_kernel (reads
  [bn_dim..bn_dim+portfolio_dim) only) and bn_tanh_backward_kernel (reads
  [0..bn_dim) only). Correct: dd_pct sources from ISV bus, no learnable input.
- Legacy bn_tanh_concat_kernel field DELETED from trainer struct alongside its
  loader, tuple-element, destructuring, assignment (5 mechanical sites for the
  one dead field). Function tuple shrinks 44→43 elements. Kernel symbol stays
  in the cubin source for SP15 oracle parity tests.
- mag_concat / OFI concat audit verdict: DECOUPLED from s1_input_dim. They
  widen shared_h2, not the trunk INPUT dim.
- test_gpu_backtest_evaluator_state_dim_calculation migrated to assert
  STATE_DIM == 128 (was 96, stale per feedback_trust_code_not_docs).

Atomic per feedback_no_partial_refactor: every consumer of s1_input_dim and
bn_concat_buf row-stride migrated together. Eliminates Wave 4.1a transient-
orphan launcher per feedback_wire_everything_up. Legacy field deleted per
feedback_no_legacy_aliases.

Tests: cargo check clean (18 pre-existing unrelated warnings). Wave 4.1a
oracle parity test (bn_tanh_concat_dd_kernel_writes_dd_pct_column) still
passes. ML lib suite went from 945 pass / 14 fail (pre-Wave-4.1b baseline) to
947 pass / 12 fail post-Wave-4.1b — improved by +2 (state_dim_calculation
migration + ensemble checkpoint round-trip flake resolved).

Refs: SP15 Wave 4.1a (a8da1cb9c), pearl_no_host_branches_in_captured_graph,
feedback_no_partial_refactor, feedback_wire_everything_up,
feedback_no_legacy_aliases, feedback_isv_for_adaptive_bounds,
feedback_trust_code_not_docs.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-07 00:33:19 +02:00
parent a8da1cb9cf
commit eb9515e41c
5 changed files with 215 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -300,8 +300,15 @@ impl CublasBackwardSet {
let na = config.num_atoms;
let _sd = ml_core::state_layout::STATE_DIM;
let sd_pad = (ml_core::state_layout::STATE_DIM + 127) & !127;
// SP15 Wave 4.1b: the bottleneck-on s1_input_dim grows by +1 to carry
// the dd_pct column appended by `bn_tanh_concat_dd_kernel`. Mirrors
// `compute_param_sizes` so the backward gemm cache shape keys match
// forward's (and the encoder_backward_chain dX width sees the same
// 103-column buffer).
let s1_input_dim = if config.bottleneck_dim > 0 {
config.bottleneck_dim + ml_core::state_layout::STATE_DIM.saturating_sub(config.market_dim)
config.bottleneck_dim
+ ml_core::state_layout::STATE_DIM.saturating_sub(config.market_dim)
+ 1
} else {
ml_core::state_layout::STATE_DIM
};

15
crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
View File

@@ -3314,9 +3314,18 @@ mod tests {
#[test]
fn test_gpu_backtest_evaluator_state_dim_calculation() {
// Canonical layout from ml_core::state_layout — 42 market + 20 OFI + 16 MTF +
// 8 portfolio + 6 plan/ISV + 4 padding = 96. Padded to 128 for cuBLAS.
assert_eq!(ml_core::state_layout::STATE_DIM, 96);
// Canonical layout from ml_core::state_layout — 128-dim feature
// vector (42 market + 32 OFI + 14 TLOB + 16 MTF + 8 portfolio +
// 6 plan/ISV + 10 padding = 128). Already padded to 128 for cuBLAS
// K-tile alignment. The historical 96-dim layout (pre-2026-04
// refactor) was deprecated when the OFI / TLOB / MTF blocks were
// expanded; the canonical constant in `ml_core::state_layout`
// is the source of truth — code wins over docs (cf.
// `feedback_trust_code_not_docs`). Migrated under SP15 Wave 4.1b
// alongside the s1_input_dim 102→103 propagation that followed
// the same `feedback_trust_code_not_docs` correction in the
// wave's audit doc.
assert_eq!(ml_core::state_layout::STATE_DIM, 128);
assert_eq!(ml_core::state_layout::STATE_DIM_PADDED, 128);
}

248
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -3995,12 +3995,20 @@ pub(crate) const NUM_WEIGHT_TENSORS: usize = 163;
/// Tensors 82-85: Trade plan head (h_s2 → 6 plan parameters).
///
/// When bottleneck is active, w_s1 input dimension changes from state_dim to
/// (bottleneck_dim + portfolio_dim) where portfolio_dim = state_dim - market_dim.
/// (bottleneck_dim + portfolio_dim + 1) where portfolio_dim = state_dim - market_dim
/// and the trailing +1 column is the SP15 Wave 4.1b dd_pct broadcast (read from
/// ISV[DD_PCT_INDEX=406] by `bn_tanh_concat_dd_kernel`). The +1 is folded into
/// `s1_input_dim` so all consumers (compute_param_sizes, Xavier init, cuBLAS
/// gemm cache, encoder_backward_chain dX width, GRN w_s1 / w_residual_h_s1
/// fan-in) treat dd_pct as the (s1_input_dim)th feature column. The bn_concat
/// buffer is sized [B, bn_dim + portfolio_dim + 1] = [B, s1_input_dim] in the
/// bottleneck-on path.
pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT_TENSORS] {
// When bottleneck is active, h_s1 input is [bottleneck_dim + portfolio_features]
// instead of [state_dim]. portfolio_features = state_dim - market_dim (typically 6-8).
// When bottleneck is active, h_s1 input is [bottleneck_dim + portfolio_features + 1]
// instead of [state_dim]. The trailing +1 column carries dd_pct from ISV (SP15
// Wave 4.1b). portfolio_features = state_dim - market_dim (typically 6-8).
let s1_input_dim = if cfg.bottleneck_dim > 0 {
cfg.bottleneck_dim + (ml_core::state_layout::STATE_DIM.saturating_sub(cfg.market_dim))
cfg.bottleneck_dim + (ml_core::state_layout::STATE_DIM.saturating_sub(cfg.market_dim)) + 1
} else {
ml_core::state_layout::STATE_DIM
};
@@ -5461,8 +5469,17 @@ pub struct GpuDqnTrainer {
on_next_bn_hidden_buf: CudaSlice<f32>,
/// DDQN argmax-pass mirror of `bn_concat_buf` (online-on-next_states).
on_next_bn_concat_buf: CudaSlice<f32>,
/// #31 Bottleneck tanh + concat kernel.
bn_tanh_concat_kernel: CudaFunction,
// SP15 Wave 4.1b removed the legacy `bn_tanh_concat_kernel` field —
// production callers (3 forward sites here + 1 in `gpu_experience_collector`)
// migrated to `launch_sp15_bn_concat_dd` (the bottleneck-aware
// dd_pct-appending launcher), which loads the
// `bn_tanh_concat_dd_kernel` symbol on-demand from the shared
// `dqn_utility_kernels.cubin`. The legacy `bn_tanh_concat_kernel`
// and `bn_tanh_concat_f32_kernel` symbols stay in the cubin source
// for the SP15 oracle tests in
// `crates/ml/tests/sp15_phase1_oracle_tests.rs` (parity check
// between old and new kernels), but no production code path resolves
// them anymore.
/// #31 Backward scratch: d_loss/d_bn_concat [B, concat_dim] f32.
bn_d_concat_buf: CudaSlice<f32>,
/// #31 Backward scratch: d_loss/d_bn [B, bn_dim] f32 (after tanh derivative).
@@ -10618,10 +10635,17 @@ impl GpuDqnTrainer {
&self.bn_d_concat_buf
}
/// Bottleneck concat dimension (bottleneck_dim + STATE_DIM - market_dim).
/// Bottleneck concat dimension (bottleneck_dim + STATE_DIM - market_dim + 1).
/// Used by TLOB backward to compute the slice offset for the TLOB gradient.
/// The trailing +1 column is the SP15 Wave 4.1b dd_pct broadcast appended by
/// `bn_tanh_concat_dd_kernel`; TLOB only reads the columns inside the
/// portfolio slice (`[bn_dim + (TLOB_START − market_dim) ..]`), so the
/// dd_pct column is naturally outside its window. The accessor reports the
/// full row stride so the TLOB backward indexes `bn_d_concat_buf` correctly.
pub(crate) fn bn_concat_dim(&self) -> usize {
self.config.bottleneck_dim + (ml_core::state_layout::STATE_DIM.saturating_sub(self.config.market_dim))
self.config.bottleneck_dim
+ (ml_core::state_layout::STATE_DIM.saturating_sub(self.config.market_dim))
+ 1
}
/// SP1 Phase B (slot 24): post-`c51_grad_kernel` value-stream gradient buffer
@@ -17506,7 +17530,7 @@ impl GpuDqnTrainer {
// per array. Stack is set once in DQNTrainer::new() (64KB for all kernels).
// ── Compile utility kernels (grad_norm, adam_update, etc.) ─
let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, nan_check_fused_f32_kernel, clamp_finite_f32_kernel, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel) =
let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, nan_check_fused_f32_kernel, clamp_finite_f32_kernel, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel) =
compile_training_kernels(&stream, &config)?;
// Separate grad_norm instance for non-graph launches (outside CUDA graph).
@@ -19711,8 +19735,14 @@ impl GpuDqnTrainer {
// ── Initialize shared cuBLAS handle (one handle for all forward/backward) ──
let shared_cublas = Arc::new(PerStreamCublasHandles::new(&stream)?);
// SP15 Wave 4.1b: the bottleneck-on s1_input_dim grows by +1 to carry
// the dd_pct column appended by `bn_tanh_concat_dd_kernel`. Mirrors
// `compute_param_sizes` exactly so the cuBLAS gemm cache shape keys
// match the GRN w_s1 / w_residual_h_s1 tensor widths.
let s1_input_dim = if config.bottleneck_dim > 0 {
config.bottleneck_dim + ml_core::state_layout::STATE_DIM.saturating_sub(config.market_dim)
config.bottleneck_dim
+ ml_core::state_layout::STATE_DIM.saturating_sub(config.market_dim)
+ 1
} else {
ml_core::state_layout::STATE_DIM
};
@@ -20044,7 +20074,15 @@ impl GpuDqnTrainer {
// Unpadded strides cause cuBLAS heuristic to pick algorithms that hang
// in CUDA Graph replay on Hopper (sm_90). pad128 forces tile-aligned
// algorithms that are graph-compatible.
let concat_dim = bn_alloc_dim + portfolio_dim;
//
// SP15 Wave 4.1b: trailing +1 column carries dd_pct (read from
// ISV[DD_PCT_INDEX=406] by `bn_tanh_concat_dd_kernel`). Concat layout
// is [bn_tanh | portfolio | dd_pct], `concat_dim = bn_dim +
// portfolio_dim + 1`. All four buffers (online / target / DDQN
// forward outputs + the backward dX accumulator) share the same
// stride — the dd_pct column is a structural part of the trunk
// input now, mirrored across all forward passes.
let concat_dim = bn_alloc_dim + portfolio_dim + 1;
let bn_concat_buf = stream.alloc_zeros::<f32>(b * concat_dim + kt)
.map_err(|e| MLError::ModelError(format!("alloc bn_concat: {e}")))?;
// Target-net mirrors of bn_hidden / bn_concat — same shapes; populated
@@ -22326,7 +22364,8 @@ impl GpuDqnTrainer {
on_next_bn_concat_buf,
bn_d_concat_buf,
bn_d_hidden_buf,
bn_tanh_concat_kernel: bn_tanh_concat_kernel_fn,
// bn_tanh_concat_kernel field removed in SP15 Wave 4.1b — see
// the field-removal comment block in the struct definition.
bn_tanh_backward_kernel,
bn_bias_grad_kernel,
// ── Plan 4 Task 1B-iii: VSN feature-selection buffers ──────
@@ -27103,7 +27142,9 @@ impl GpuDqnTrainer {
let b = self.config.batch_size;
let market_dim = self.config.market_dim;
let portfolio_dim = ml_core::state_layout::STATE_DIM.saturating_sub(market_dim);
let concat_dim = bn_dim + portfolio_dim;
// SP15 Wave 4.1b: trailing +1 column = dd_pct (broadcast from
// ISV[DD_PCT_INDEX=406] by `bn_tanh_concat_dd_kernel`).
let concat_dim = bn_dim + portfolio_dim + 1;
let bn_hidden_ptr = self.on_next_bn_hidden_buf.raw_ptr();
let state_dim_padded_usize = self.cublas_forward.state_dim_padded;
self.cublas_forward.sgemm_f32_ldb(
@@ -27119,31 +27160,29 @@ impl GpuDqnTrainer {
&self.stream, bn_hidden_ptr, on_w_ptrs[34], bn_dim, b,
)?;
let concat_ptr = self.on_next_bn_concat_buf.raw_ptr();
let state_dim_padded = state_dim_padded_usize as i32;
let total_elems = (b * concat_dim) as i32;
let blocks = ((total_elems as u32 + 255) / 256) as u32;
let state_dim_padded_i32 = state_dim_padded_usize as i32;
let bn_dim_i32 = bn_dim as i32;
let market_dim_i32 = market_dim as i32;
let concat_dim_i32 = concat_dim as i32;
let b_i32 = b as i32;
unsafe {
self.stream
.launch_builder(&self.bn_tanh_concat_kernel)
.arg(&bn_hidden_ptr)
.arg(&on_next_states_for_bn)
.arg(&concat_ptr)
.arg(&b_i32)
.arg(&bn_dim_i32)
.arg(&market_dim_i32)
.arg(&state_dim_padded)
.arg(&concat_dim_i32)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("ddqn_bn_tanh_concat: {e}")))?;
}
// SP15 Wave 4.1b: launch the dd_pct-aware concat kernel. Reads
// `isv[DD_PCT_INDEX=406]` and broadcasts it across batch as the
// (concat_dim − 1)th column. ISV bus is the same pointer all the
// other ISV consumers use; experience_collector populates it via
// `launch_sp15_dd_state` before the trainer's forward graph runs,
// so by the time this kernel reads slot 406 the value is fresh.
launch_sp15_bn_concat_dd(
&self.stream,
bn_hidden_ptr,
on_next_states_for_bn,
self.isv_signals_dev_ptr,
concat_ptr,
b_i32,
bn_dim_i32,
market_dim_i32,
state_dim_padded_i32,
concat_dim_i32,
)?;
concat_ptr
} else {
// No bottleneck: pass VSN-gated next_states directly to DDQN encoder.
@@ -27428,7 +27467,9 @@ impl GpuDqnTrainer {
let b = self.config.batch_size;
let market_dim = self.config.market_dim;
let portfolio_dim = ml_core::state_layout::STATE_DIM.saturating_sub(market_dim);
let concat_dim = bn_dim + portfolio_dim;
// SP15 Wave 4.1b: trailing +1 column = dd_pct (broadcast from
// ISV[DD_PCT_INDEX=406] by `bn_tanh_concat_dd_kernel`).
let concat_dim = bn_dim + portfolio_dim + 1;
// Bottleneck GEMM: states[B, market_dim] @ w_bn^T → h_bn[B, bn_dim]
// w_bn is at on_w_ptrs[33], b_bn is at on_w_ptrs[34] (post Plan 4
@@ -27452,38 +27493,33 @@ impl GpuDqnTrainer {
&self.stream, bn_hidden_ptr, on_w_ptrs[34], bn_dim, b,
)?;
// Tanh + concat: [tanh(h_bn), portfolio_from_states] → bn_concat
// Plan 4 Task 1B-iii: portfolio passthrough source is the gated
// state — keeps bottleneck and portfolio halves in sync.
// Tanh + concat + dd_pct: [tanh(h_bn), portfolio_from_states,
// isv[DD_PCT_INDEX]] → bn_concat. Plan 4 Task 1B-iii: portfolio
// passthrough source is the gated state — keeps bottleneck and
// portfolio halves in sync.
// SP15 Wave 4.1b: dd_pct broadcast comes from the same ISV bus
// every other consumer reads; experience_collector's
// `launch_sp15_dd_state` populates ISV[406] in the per-step env
// loop, so the trunk forward sees a fresh value here.
let concat_ptr = self.bn_concat_buf.raw_ptr();
let state_dim_padded = state_dim_padded_usize as i32;
let total_elems = (b * concat_dim) as i32;
let blocks = ((total_elems as u32 + 255) / 256) as u32;
let state_dim_padded_i32 = state_dim_padded_usize as i32;
let bn_dim_i32 = bn_dim as i32;
let market_dim_i32 = market_dim as i32;
let concat_dim_i32 = concat_dim as i32;
let b_i32 = b as i32;
{
unsafe {
self.stream
.launch_builder(&self.bn_tanh_concat_kernel)
.arg(&bn_hidden_ptr)
.arg(&states_for_bn)
.arg(&concat_ptr)
.arg(&b_i32)
.arg(&bn_dim_i32)
.arg(&market_dim_i32)
.arg(&state_dim_padded)
.arg(&concat_dim_i32)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("bn_tanh_concat: {e}")))?;
}
}
launch_sp15_bn_concat_dd(
&self.stream,
bn_hidden_ptr,
states_for_bn,
self.isv_signals_dev_ptr,
concat_ptr,
b_i32,
bn_dim_i32,
market_dim_i32,
state_dim_padded_i32,
concat_dim_i32,
)?;
concat_ptr
} else {
// No bottleneck: bypass directly to the gated state. Downstream
@@ -27638,7 +27674,14 @@ impl GpuDqnTrainer {
let b = self.config.batch_size;
let market_dim = self.config.market_dim;
let portfolio_dim = ml_core::state_layout::STATE_DIM.saturating_sub(market_dim);
let concat_dim = bn_dim + portfolio_dim;
// SP15 Wave 4.1b: trailing +1 column = dd_pct (broadcast from
// ISV[DD_PCT_INDEX=406] by `bn_tanh_concat_dd_kernel`). Target
// forward reads the SAME ISV bus the online forward read on the
// current step, so both networks see the same dd_pct context for
// their respective state inputs (next_states for target, states
// for online). dd_pct is a global episode statistic — not
// per-state — so the broadcast value is correct for both.
let concat_dim = bn_dim + portfolio_dim + 1;
let tg_bn_hidden_ptr = self.tg_bn_hidden_buf.raw_ptr();
let state_dim_padded_usize = self.cublas_forward.state_dim_padded;
self.cublas_forward.sgemm_f32_ldb(
@@ -27654,31 +27697,23 @@ impl GpuDqnTrainer {
&self.stream, tg_bn_hidden_ptr, tg_w_ptrs[34], bn_dim, b,
)?;
let tg_concat_ptr = self.tg_bn_concat_buf.raw_ptr();
let state_dim_padded = state_dim_padded_usize as i32;
let total_elems = (b * concat_dim) as i32;
let blocks = ((total_elems as u32 + 255) / 256) as u32;
let state_dim_padded_i32 = state_dim_padded_usize as i32;
let bn_dim_i32 = bn_dim as i32;
let market_dim_i32 = market_dim as i32;
let concat_dim_i32 = concat_dim as i32;
let b_i32 = b as i32;
unsafe {
self.stream
.launch_builder(&self.bn_tanh_concat_kernel)
.arg(&tg_bn_hidden_ptr)
.arg(&tg_states_for_bn)
.arg(&tg_concat_ptr)
.arg(&b_i32)
.arg(&bn_dim_i32)
.arg(&market_dim_i32)
.arg(&state_dim_padded)
.arg(&concat_dim_i32)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("tg_bn_tanh_concat: {e}")))?;
}
launch_sp15_bn_concat_dd(
&self.stream,
tg_bn_hidden_ptr,
tg_states_for_bn,
self.isv_signals_dev_ptr,
tg_concat_ptr,
b_i32,
bn_dim_i32,
market_dim_i32,
state_dim_padded_i32,
concat_dim_i32,
)?;
tg_concat_ptr
} else {
// No bottleneck: pass VSN-gated next_states directly to target encoder.
@@ -28498,7 +28533,18 @@ impl GpuDqnTrainer {
let market_dim = config.market_dim;
let b = config.batch_size;
let portfolio_dim = ml_core::state_layout::STATE_DIM.saturating_sub(market_dim);
let concat_dim = bn_dim + portfolio_dim;
// SP15 Wave 4.1b: bn_d_concat_buf row stride is the full forward
// concat_dim including the dd_pct column — encoder_backward_chain's
// Linear_a/Linear_residual dX writes 103 columns. The
// bn_tanh_backward_kernel reads only the bn_dim slice [0..bn_dim) and
// the vsn_d_gated_state_portfolio_pad_kernel reads only the
// portfolio slice [bn_dim..bn_dim+portfolio_dim) — both correctly
// skip the dd_pct gradient column at index bn_dim+portfolio_dim
// because dd_pct sources from the ISV bus (no learnable input feeds
// it back). The discarded column's gradient is harmless: its only
// effect would be to update a downstream variable, which doesn't
// exist.
let concat_dim = bn_dim + portfolio_dim + 1;
let state_dim = ml_core::state_layout::STATE_DIM;
let state_dim_padded = ml_core::state_layout::STATE_DIM_PADDED;
@@ -28998,7 +29044,14 @@ impl GpuDqnTrainer {
let market_dim = self.config.market_dim;
let b = self.config.batch_size;
let portfolio_dim = ml_core::state_layout::STATE_DIM.saturating_sub(market_dim);
let concat_dim = bn_dim + portfolio_dim;
// SP15 Wave 4.1b: bn_d_concat_buf row stride is the full forward
// concat_dim including the dd_pct column. Same rationale as the
// aux_bottleneck_vsn_backward_dispatch site — bn_tanh_backward
// reads [0..bn_dim) and the portfolio_pad reads
// [bn_dim..bn_dim+portfolio_dim); the dd_pct column at
// bn_dim+portfolio_dim is correctly skipped because dd_pct has no
// learnable input source.
let concat_dim = bn_dim + portfolio_dim + 1;
let state_dim = ml_core::state_layout::STATE_DIM;
let state_dim_padded = ml_core::state_layout::STATE_DIM_PADDED;
@@ -29813,9 +29866,19 @@ impl GpuDqnTrainer {
let sizes = compute_param_sizes(cfg);
let total = self.total_params;
// Compute s1_input_dim matching compute_param_sizes
// Compute s1_input_dim matching compute_param_sizes — SP15 Wave 4.1b
// bumps the bottleneck-on path by +1 for the dd_pct column. Xavier
// sees the full (s1_input_dim, shared_h1) fan and initialises every
// column uniformly, including the new dd_pct column at index
// `bn_dim + portfolio_dim`. The dd_pct value at runtime is bounded
// ∈ [0, 1] (Wave 1.3 contract), matching Xavier's small-magnitude
// assumption — no special-case zero-init needed (cf. SP14 zeroed the
// aux_softmax_diff column of w_b0fc because that signal can be ±1
// which Xavier assumes is zero-mean; dd_pct is non-negative but
// similarly small-magnitude, so the gradient flows symmetrically
// from day 0).
let s1_input_dim = if cfg.bottleneck_dim > 0 {
cfg.bottleneck_dim + ml_core::state_layout::STATE_DIM.saturating_sub(cfg.market_dim)
cfg.bottleneck_dim + ml_core::state_layout::STATE_DIM.saturating_sub(cfg.market_dim) + 1
} else {
ml_core::state_layout::STATE_DIM
};
@@ -31371,7 +31434,7 @@ fn deflate_rank_one(mat: &[f32], n: usize, lambda_1: f32) -> Vec<f32> {
fn compile_training_kernels(
stream: &Arc<CudaStream>,
config: &GpuDqnTrainConfig,
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
info!(
state_dim = ml_core::state_layout::STATE_DIM,
total_params = compute_total_params(config),
@@ -31495,8 +31558,11 @@ fn compile_training_kernels(
.map_err(|e| MLError::ModelError(format!("bn_tanh_backward_kernel load: {e}")))?;
let bn_bias_grad = module.load_function("bn_bias_grad_kernel")
.map_err(|e| MLError::ModelError(format!("bn_bias_grad_kernel load: {e}")))?;
let bn_tanh_concat = module.load_function("bn_tanh_concat_kernel")
.map_err(|e| MLError::ModelError(format!("bn_tanh_concat_kernel load: {e}")))?;
// SP15 Wave 4.1b: legacy `bn_tanh_concat_kernel` no longer loaded —
// production callers migrated to `launch_sp15_bn_concat_dd`, which
// resolves the `bn_tanh_concat_dd_kernel` symbol from the same cubin.
// The legacy symbol remains in the source for the SP15 oracle parity
// tests but is not surfaced as a `CudaFunction` field on the trainer.
let her_inplace = module.load_function("her_inplace_relabel")
.map_err(|e| MLError::ModelError(format!("her_inplace_relabel load: {e}")))?;
@@ -31526,7 +31592,7 @@ fn compile_training_kernels(
.map_err(|e| MLError::ModelError(format!("popart_normalize_robust load: {e}")))?;
info!("GpuDqnTrainer: 39 utility kernels loaded from precompiled cubin (5 CUmodules)");
Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, nan_check_fused_f32, clamp_finite_f32, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust))
Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, nan_check_fused_f32, clamp_finite_f32, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust))
}
/// Load the standalone Polyak EMA kernel from precompiled cubin.

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

@@ -1143,10 +1143,11 @@ impl GpuExperienceCollector {
// expected Q via compute_expected_q kernel.
let num_atoms = num_atoms_max.max(1);
// s1_input_dim must match the trainer's layout:
// with bottleneck: compressed market → bn_dim, then concat with portfolio → bn_dim + portfolio_dim
// with bottleneck: compressed market → bn_dim, then concat with portfolio
// + 1 dd_pct column (SP15 Wave 4.1b) → bn_dim + portfolio_dim + 1
// without bottleneck: raw state_dim
let s1_input_dim = if bottleneck_dim > 0 {
bottleneck_dim + state_dim.saturating_sub(market_dim_cfg)
bottleneck_dim + state_dim.saturating_sub(market_dim_cfg) + 1
} else {
state_dim
};
@@ -1751,18 +1752,33 @@ impl GpuExperienceCollector {
.map_err(|e| MLError::ModelError(format!("alloc quantiles_scratch_buf: {e}")))?;
// #31 Bottleneck — always loaded and allocated (one production path)
// SP15 Wave 4.1b: experience collector inference forward shares the
// trunk's dd_pct-aware concat path. Loads `bn_tanh_concat_dd_kernel`
// (same cubin, multi-kernel-per-file pattern); the launcher reads
// ISV[DD_PCT_INDEX=406] via this collector's `isv_signals_dev_ptr`
// (set by `set_isv_signals_ptr` from the trainer's bus). At step 0
// the slot is sentinel 0 (per state_reset_registry); the experience
// collector's `launch_sp15_dd_state` populates it during env_step
// (which runs AFTER this forward, so step k+1 sees step k's
// dd_pct — same one-step-lag semantic as mag_concat).
use super::gpu_dqn_trainer::DQN_UTILITY_CUBIN;
let util_module = stream.context()
.load_cubin(DQN_UTILITY_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("utility cubin load for bn: {e}")))?;
let bn_tanh_concat_fn = util_module.load_function("bn_tanh_concat_f32_kernel")
.map_err(|e| MLError::ModelError(format!("bn_tanh_concat_f32_kernel load: {e}")))?;
let bn_tanh_concat_fn = util_module.load_function("bn_tanh_concat_dd_kernel")
.map_err(|e| MLError::ModelError(format!("bn_tanh_concat_dd_kernel load: {e}")))?;
let bn_alloc = bottleneck_dim.max(1);
let portfolio_dim_bn = state_dim.saturating_sub(market_dim_cfg);
let exp_bn_hidden = stream.alloc_zeros::<f32>(alloc_episodes * bn_alloc + 128)
.map_err(|e| MLError::ModelError(format!("alloc exp_bn_hidden: {e}")))?;
let exp_bn_concat = stream.alloc_zeros::<f32>(alloc_episodes * (bn_alloc + portfolio_dim_bn) + 128)
// SP15 Wave 4.1b: +1 column for dd_pct. Matches the trainer's
// bn_concat_buf / tg_bn_concat_buf / on_next_bn_concat_buf
// reallocation in `gpu_dqn_trainer.rs` (`concat_dim = bn_alloc_dim
// + portfolio_dim + 1`).
let exp_bn_concat = stream.alloc_zeros::<f32>(
alloc_episodes * (bn_alloc + portfolio_dim_bn + 1) + 128
)
.map_err(|e| MLError::ModelError(format!("alloc exp_bn_concat: {e}")))?;
// Phase 2d: per-sample, per-branch C51 support buffer [alloc_episodes, 4, 3]
@@ -3817,7 +3833,10 @@ impl GpuExperienceCollector {
let states_ptr = self.batch_states.raw_ptr();
let md = self.market_dim_bn;
let portfolio_dim = ml_core::state_layout::STATE_DIM - md;
let concat_dim = bn_dim + portfolio_dim;
// SP15 Wave 4.1b: trailing +1 column = dd_pct broadcast from
// ISV[DD_PCT_INDEX=406]. Mirrors the trainer's
// bn_tanh_concat_dd_kernel migration.
let concat_dim = bn_dim + portfolio_dim + 1;
// #30 F32 SGEMM: batch_states[N, market_dim] @ w_bn^T → h_bn[N, bn_dim]
{
@@ -3839,20 +3858,27 @@ impl GpuExperienceCollector {
)?;
}
// Tanh + concat kernel
// Tanh + concat + dd_pct kernel — SP15 Wave 4.1b. Reads
// ISV[DD_PCT_INDEX=406] via the collector's bus pointer (set
// by `set_isv_signals_ptr`); when the pointer is 0 (uninit
// case) the kernel would dereference NULL — guarded by the
// captured-graph wiring in training_loop.rs:859 which sets
// the bus pointer BEFORE the first epoch's forward graph
// capture.
{
let sdp = self.cublas_forward.state_dim_padded as i32;
let total_elems = (n * concat_dim) as i32;
let blocks = ((total_elems as u32 + 255) / 256) as u32;
let bn_i32 = bn_dim as i32;
let md_i32 = md as i32;
let cd_i32 = concat_dim as i32;
let n_i32_bn = n as i32;
let total_elems = (n * concat_dim) as i64;
let blocks = (((total_elems as u64) + 255) / 256).max(1) as u32;
unsafe {
self.stream
.launch_builder(&self.bn_tanh_concat_fn)
.arg(&bn_hidden_ptr)
.arg(&states_ptr)
.arg(&self.isv_signals_dev_ptr)
.arg(&bn_concat_ptr)
.arg(&n_i32_bn)
.arg(&bn_i32)
@@ -3865,7 +3891,7 @@ impl GpuExperienceCollector {
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(
"exp bn_tanh_concat: {e}"
"exp bn_tanh_concat_dd: {e}"
)))?;
}
}

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

File diff suppressed because one or more lines are too long