From 59a734efb0048cd3f7cedb46503500b5f759d9d5 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 15 Apr 2026 19:57:16 +0200 Subject: [PATCH] feat: regime-conditioned branch gating -- 20 learned params for per-branch importance W_regime[4,4] + b_regime[4] produce softmax importance weights from [ADX, CUSUM, Q_gap, atom_utilization]. Scales per-branch Q-values so direction matters more in trends, magnitude matters more in ranges. NUM_WEIGHT_TENSORS: 50 -> 52. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/batched_backward.rs | 2 +- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 105 +++++++++++++++++- 2 files changed, 103 insertions(+), 4 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index 46ab10290..0a4473058 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -178,7 +178,7 @@ pub struct CublasBackwardSet { branch_2_size: usize, branch_3_size: usize, - /// Pre-computed parameter sizes for all 50 weight tensors. + /// Pre-computed parameter sizes for all weight tensors. /// Used to compute `padded_byte_offset` for gradient buffer offsets. param_sizes: [usize; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS], diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index a99d77137..34a0a5287 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -351,8 +351,9 @@ impl Default for CausalInterventionConfig { /// Number of weight tensors in the flat parameter buffer. /// 20 original (DQN network) + 4 branch 3 (magnitude) + 2 bottleneck (w_bn, b_bn) -/// + 8 VSN bottleneck (R=16) + 8 GLU gate + 8 KAN spline (coeff+resid per branch) = 50. -pub(crate) const NUM_WEIGHT_TENSORS: usize = 50; +/// + 8 VSN bottleneck (R=16) + 8 GLU gate + 8 KAN spline (coeff+resid per branch) +/// + 2 regime branch gate (w_regime + b_regime) = 52. +pub(crate) const NUM_WEIGHT_TENSORS: usize = 52; /// Compute the size (element count) of each weight tensor. /// @@ -362,6 +363,7 @@ pub(crate) const NUM_WEIGHT_TENSORS: usize = 50; /// Tensors 26-33: Variable Selection Network bottleneck (R=16). /// Tensors 34-41: GLU gate weights and biases. /// Tensors 42-49: KAN spline coefficients and residual weights (per branch). +/// Tensors 50-51: Regime branch gate (W_regime[4,4] + b_regime[4]). /// /// When bottleneck is active, w_s1 input dimension changes from state_dim to /// (bottleneck_dim + portfolio_dim) where portfolio_dim = state_dim - market_dim. @@ -431,6 +433,9 @@ pub(crate) fn compute_param_sizes(cfg: &GpuDqnTrainConfig) -> [usize; NUM_WEIGHT cfg.adv_h, // [47] kan_resid_2 [AH] cfg.adv_h * 8, // [48] kan_coeff_3 [AH, 8] cfg.adv_h, // [49] kan_resid_3 [AH] + // ── Regime branch gate ── + 4 * 4, // [50] w_regime [4, 4] + 4, // [51] b_regime [4] ] } @@ -690,6 +695,14 @@ pub struct GpuDqnTrainer { kan_gate_combine_kernel: CudaFunction, kan_gate_backward_kernel: CudaFunction, + // ── Regime branch gate ── + regime_gate_kernel: CudaFunction, + /// Per-sample Q-gap for regime gate input [B]. + regime_q_gap_buf: CudaSlice, + /// Atom utilization scalar for regime gate [1] — pinned device-mapped. + regime_util_pinned: *mut f32, + regime_util_dev_ptr: u64, + save_current_lp: CudaSlice, // [B, NUM_BRANCHES(4), NUM_ATOMS] save_projected: CudaSlice, // [B, NUM_BRANCHES(4), NUM_ATOMS] @@ -1624,6 +1637,57 @@ impl GpuDqnTrainer { Ok(()) } + /// Apply regime branch gate to Q-values (after Q-mean centering, before Q-attention). + /// + /// Uses W_regime[4,4] + b_regime[4] to produce softmax importance weights from + /// [ADX, CUSUM, Q_gap, atom_utilization]. Scales per-branch Q-values so direction + /// matters more in trends, magnitude matters more in ranges. + pub(crate) fn apply_regime_gate(&self, batch_size: usize) -> Result<(), MLError> { + let param_sizes = compute_param_sizes(&self.config); + let w_regime_ptr = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 50); + let b_regime_ptr = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 51); + + let blocks = ((batch_size as u32 + 255) / 256).max(1); + let q_out_ptr = self.q_out_buf.raw_ptr(); + let states_ptr = self.ptrs.states_buf; + let q_gap_ptr = self.regime_q_gap_buf.raw_ptr(); + let util_ptr = self.regime_util_dev_ptr; + let b_i32 = batch_size as i32; + let sd = self.config.state_dim as i32; + let b0 = self.config.branch_0_size as i32; + let b1 = self.config.branch_1_size as i32; + let b2 = self.config.branch_2_size as i32; + let b3 = self.config.branch_3_size as i32; + + unsafe { + self.stream.launch_builder(&self.regime_gate_kernel) + .arg(&q_out_ptr) + .arg(&states_ptr) + .arg(&q_gap_ptr) + .arg(&util_ptr) + .arg(&w_regime_ptr) + .arg(&b_regime_ptr) + .arg(&b_i32) + .arg(&sd) + .arg(&b0) + .arg(&b1) + .arg(&b2) + .arg(&b3) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("regime_branch_gate launch: {e}")))?; + } + Ok(()) + } + + /// Update regime gate atom utilization from latest Q-stats. + pub fn set_regime_util(&mut self, util: f32) { + unsafe { *self.regime_util_pinned = util; } + } + /// Accumulate first SH2 columns of d_mag_concat [B, SH2+3] into d_h_s2 [B, SH2]. pub(crate) fn accumulate_d_h_s2_from_concat( &self, @@ -3218,7 +3282,33 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("strided_scatter load: {e}")))?; let concat_ofi_kernel = exp_module_for_mag.load_function("concat_ofi_features") .map_err(|e| MLError::ModelError(format!("concat_ofi_features load: {e}")))?; - info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter + concat_ofi kernels loaded"); + let regime_gate_kernel = exp_module_for_mag.load_function("regime_branch_gate") + .map_err(|e| MLError::ModelError(format!("regime_branch_gate load: {e}")))?; + info!("GpuDqnTrainer: mag_concat + strided_accumulate/scatter + concat_ofi + regime_gate kernels loaded"); + + // ── Regime branch gate buffers ─────────────────────────────── + let regime_q_gap_buf = stream.alloc_zeros::(b) + .map_err(|e| MLError::ModelError(format!("regime_q_gap alloc: {e}")))?; + + let (regime_util_pinned, regime_util_dev_ptr) = { + let mut host_ptr: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut dev_ptr_out: u64 = 0; + unsafe { + let rc = cudarc::driver::sys::cuMemAllocHost_v2( + &mut host_ptr, + std::mem::size_of::(), + ); + assert_eq!(rc, cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS, "cuMemAllocHost for regime_util"); + let rc2 = cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + &mut dev_ptr_out, + host_ptr, + 0, + ); + assert_eq!(rc2, cudarc::driver::sys::cudaError_enum::CUDA_SUCCESS, "cuMemHostGetDevicePointer for regime_util"); + *(host_ptr as *mut f32) = 1.0; // Initialize to healthy utilization + } + (host_ptr as *mut f32, dev_ptr_out) + }; // ── Cross-Branch Q Attention buffers ────────────────────────── let q_attn_params = stream.alloc_zeros::(624) @@ -4045,6 +4135,10 @@ impl GpuDqnTrainer { glu_backward_kernel, kan_gate_combine_kernel, kan_gate_backward_kernel, + regime_gate_kernel, + regime_q_gap_buf, + regime_util_pinned, + regime_util_dev_ptr, save_current_lp, save_projected, per_sample_loss_buf, @@ -5492,6 +5586,9 @@ impl GpuDqnTrainer { // Q-mean centering: removes bootstrapping drift before attention and stats. self.launch_q_mean_center(batch_size)?; + // Regime branch gate: scale Q-values by learned per-branch importance + self.apply_regime_gate(batch_size)?; + // Cross-branch Q attention: q_out_buf → q_coord_buf [B, 12]. // Runs after compute_expected_q so attention weights train on real Q-values. self.launch_q_attention(batch_size)?; @@ -7577,6 +7674,8 @@ impl GpuDqnTrainer { (0, 0), // [47] kan_resid_2 (0, 0), // [48] kan_coeff_3 (0, 0), // [49] kan_resid_3 + (4, 4), // [50] w_regime (Xavier) + (0, 0), // [51] b_regime (zero) ]; // Build flat host buffer: Xavier init for weights, zeros for biases + padding.