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

feat: regime-conditioned branch gating -- 20 learned params for per-branch importance

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 19:57:16 +02:00
parent db9d99709a
commit 59a734efb0
2 changed files with 103 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

@@ -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<f32>,
/// Atom utilization scalar for regime gate [1] — pinned device-mapped.
regime_util_pinned: *mut f32,
regime_util_dev_ptr: u64,
save_current_lp: CudaSlice<f32>, // [B, NUM_BRANCHES(4), NUM_ATOMS]
save_projected: CudaSlice<f32>, // [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(&param_sizes, 50);
let b_regime_ptr = self.ptrs.params_ptr + padded_byte_offset(&param_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::<f32>(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::<f32>(),
);
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::<f32>(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.