diff --git a/crates/ml-core/src/state_layout.rs b/crates/ml-core/src/state_layout.rs index f5f00c953..c3fbaeb9d 100644 --- a/crates/ml-core/src/state_layout.rs +++ b/crates/ml-core/src/state_layout.rs @@ -6,10 +6,11 @@ /// Layout (grouped by update frequency, fast-changing first): /// [0..42) Market features (OHLCV + technicals) /// [42..74) OFI (32-dim order flow from MBP-10 — see OFI slot map below) -/// [74..90) MTF (4 lookbacks x 4 features) -/// [90..98) Portfolio base (8 features) -/// [98..105) Plan/ISV (7 features — D.6 Plan 2 Task 6A added remaining_fraction at index 6) -/// [105..112) Zero padding (7 zeros for 8-alignment) +/// [74..90) TLOB (16-dim attention-compressed OFI — D.8 Plan 2 Task 6C; zeros on assembly, GpuTlob::forward overwrites post-gather) +/// [90..106) MTF (4 lookbacks x 4 features) +/// [106..114) Portfolio base (8 features) +/// [114..121) Plan/ISV (7 features — D.6 Plan 2 Task 6A added remaining_fraction at index 6) +/// [121..128) Zero padding (7 zeros for 8-alignment; D.8 Plan 2 Task 6C reclaimed from 16 → 7) /// /// OFI slot map (OFI_DIM = 32, indices relative to OFI_START): /// [0..8) raw OFI — ofi_level1, ofi_level5, depth_imbalance, vpin, @@ -32,24 +33,27 @@ /// [30] order_count_imbalance ((Σbid_ct − Σask_ct) / Σ(bid_ct + ask_ct) — Mbp10Snapshot) /// [31] microprice_residual ((weighted_mid − mid) / mid — Mbp10Snapshot) -// D.6 Plan_isv[6] remaining_fraction (Plan 2 Task 6A): PORTFOLIO_PLAN_DIM 6→7, -// STATE_DIM 104→105, padded to 112 for 8-alignment → PADDING_DIM 0→7. -pub const STATE_DIM: usize = 112; +// D.8 Plan 2 Task 6C: TLOB attention (16-dim) inserted after OFI. +// Layout: Market(42)+OFI(32)+TLOB(16)+MTF(16)+Portfolio(8)+PlanISV(7)+Padding(7) = 128. +// STATE_DIM 112→128 (matches STATE_DIM_PADDED; padding reclaimed to 7 for 8-alignment). +pub const STATE_DIM: usize = 128; pub const STATE_DIM_PADDED: usize = 128; // pad128(STATE_DIM) for cuBLAS K-tile alignment pub const MARKET_DIM: usize = 42; pub const OFI_DIM: usize = 32; +pub const TLOB_DIM: usize = 16; pub const MTF_DIM: usize = 16; pub const PORTFOLIO_BASE_DIM: usize = 8; pub const PORTFOLIO_PLAN_DIM: usize = 7; pub const PADDING_DIM: usize = 7; pub const MARKET_START: usize = 0; -pub const OFI_START: usize = MARKET_START + MARKET_DIM; // 42 -pub const MTF_START: usize = OFI_START + OFI_DIM; // 74 -pub const PORTFOLIO_START: usize = MTF_START + MTF_DIM; // 90 -pub const PLAN_ISV_START: usize = PORTFOLIO_START + PORTFOLIO_BASE_DIM; // 98 -pub const PADDING_START: usize = PLAN_ISV_START + PORTFOLIO_PLAN_DIM; // 104 +pub const OFI_START: usize = MARKET_START + MARKET_DIM; // 42 +pub const TLOB_START: usize = OFI_START + OFI_DIM; // 74 +pub const MTF_START: usize = TLOB_START + TLOB_DIM; // 90 +pub const PORTFOLIO_START: usize = MTF_START + MTF_DIM; // 106 +pub const PLAN_ISV_START: usize = PORTFOLIO_START + PORTFOLIO_BASE_DIM; // 114 +pub const PADDING_START: usize = PLAN_ISV_START + PORTFOLIO_PLAN_DIM; // 121 const _: () = assert!(PADDING_START + PADDING_DIM == STATE_DIM, "layout must sum to STATE_DIM"); const _: () = assert!(STATE_DIM % 8 == 0, "STATE_DIM must be 8-aligned for tensor cores"); diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 71777332d..c2e6d3746 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -90,6 +90,8 @@ fn main() { "kelly_cap_update_kernel.cu", "atoms_update_kernel.cu", "q_quantile_kernel.cu", + // D.8 Plan 2 Task 6C: TLOB attention kernels (SDP + state scatter/read) + "tlob_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index ad1af35ac..e83b9d6e2 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -398,6 +398,13 @@ pub struct GpuBacktestEvaluator { /// Lazy-loaded with the other action-select machinery. plan_state_isv_kernel: Option, plan_diag_reduce_kernel: Option, + + /// D.8 TLOB val-path: OFI[32]→TLOB[16] attention applied after each gather step. + /// + /// Weights are synchronised from the training TLOB before `evaluate_dqn_graphed` + /// via `set_tlob_from_training`. `None` when TLOB is disabled or weight sync + /// has not yet been called. + tlob: Option, } impl Drop for GpuBacktestEvaluator { @@ -663,6 +670,7 @@ impl GpuBacktestEvaluator { scatter_intent_kernel: None, plan_state_isv_kernel: None, plan_diag_reduce_kernel: None, + tlob: None, }) } @@ -672,6 +680,25 @@ impl GpuBacktestEvaluator { self.isv_signals_ptr = ptr; } + /// Install a val-path TLOB instance and immediately sync its weights from the training TLOB. + /// + /// `eval_tlob` must already be constructed (with a `PerStreamCublasHandles` bound to + /// `self.stream`). `train_tlob` supplies the source weight parameters — copied via + /// DtoD async on `eval_tlob`'s stream. + /// + /// Calling this before `evaluate_dqn_graphed` ensures every val forward pass runs the + /// same TLOB weights as the most recent training step, preserving train/val state-dist + /// parity for the TLOB slot (SL_TLOB_START) in the assembled state vector. + pub(crate) fn set_tlob_from_training( + &mut self, + mut eval_tlob: super::gpu_tlob::GpuTlob, + train_tlob: &super::gpu_tlob::GpuTlob, + ) -> Result<(), MLError> { + eval_tlob.copy_params_from(train_tlob)?; + self.tlob = Some(eval_tlob); + Ok(()) + } + // ── Private helpers ─────────────────────────────────────────────────────── /// Build the initial portfolio state vector. @@ -700,6 +727,14 @@ impl GpuBacktestEvaluator { &self.stream } + /// Number of evaluation windows (validation fold count). + /// + /// Used by callers that need to size per-window buffers or TLOB instances + /// to match the evaluator's batch geometry. + pub(crate) fn n_windows(&self) -> usize { + self.n_windows + } + // ── Public API ──────────────────────────────────────────────────────────── /// Build the state tensor for a given step: `[n_windows, feat_dim + portfolio_dim]`. @@ -1161,7 +1196,7 @@ impl GpuBacktestEvaluator { /// This is acceptable: market features (42 dims) dominate the Q-network and /// a 64-bar (~1 hour) staleness window has negligible impact on action quality. fn submit_dqn_step_loop_cublas( - &self, + &mut self, dqn_cfg: &DqnBacktestConfig, mut q_provider: Option<&mut dyn super::q_value_provider::QValueProvider>, ) -> Result<(), MLError> { @@ -1220,6 +1255,14 @@ impl GpuBacktestEvaluator { let step = chunk_start + step_offset; self.launch_gather(step)?; + // D.8 TLOB val-path: run TLOB forward on states_buf (writes SL_TLOB_START slot). + // Runs after gather fills zeros for the TLOB slot, before DtoD copy batches + // states into chunked_states_buf. Preserves train/val state-dist parity. + if let Some(ref mut tlob) = self.tlob { + tlob.forward(&mut self.states_buf, n) + .map_err(|e| MLError::ModelError(format!("val TLOB forward step {step}: {e}")))?; + } + // DtoD copy: states_buf → chunked_states_buf[step_offset * n * state_dim ..] let dst_byte_offset = (step_offset * state_row_bytes) as u64; let (src_ptr, _src_sync) = self.states_buf.device_ptr(&self.stream); diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 290b37cf0..bad2a84a9 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -208,13 +208,13 @@ const ISV_NETWORK_DIM: usize = 23; /// [56] Q_P95_ORD_INDEX — per-epoch 95th-percentile Q EMA for order branch. /// [57] Q_P95_URG_INDEX — per-epoch 95th-percentile Q EMA for urgency branch. /// -/// Slots [58..60) carry the layout fingerprint (ISV_LAYOUT_FINGERPRINT_LO_INDEX -/// and ISV_LAYOUT_FINGERPRINT_HI_INDEX). These are placed at the tail because -/// slots [0] and [1] are actively written by `isv_signal_update` (Q-drift and -/// gradient-norm EMA respectively). Written by the constructor; checked at -/// checkpoint load. Fail-fast only — no migration path exists. See spec §4.A.2 -/// and `LAYOUT_FINGERPRINT_CURRENT` for the structural-hash design rationale. -const ISV_TOTAL_DIM: usize = 60; +/// Slots [60..63) carry the D.8 TLOB diagnostic and layout fingerprint. +/// [60] TLOB_REGIME_FOCUS_EMA_INDEX — mean-max attention weight EMA (CPU-written, per-epoch). +/// [61] ISV_LAYOUT_FINGERPRINT_LO_INDEX — low 32 bits of layout fingerprint. +/// [62] ISV_LAYOUT_FINGERPRINT_HI_INDEX — high 32 bits of layout fingerprint. +/// Written by the constructor; checked at checkpoint load. Fail-fast only — no migration +/// path exists. See spec §4.A.2 and `LAYOUT_FINGERPRINT_CURRENT` for structural-hash rationale. +const ISV_TOTAL_DIM: usize = 63; /// Legacy alias preserved for call sites that haven't been audited for the /// network-vs-total split. New code should pick `ISV_NETWORK_DIM` (for weight /// tensor sizing) or `ISV_TOTAL_DIM` (for the broadcast bus buffer). @@ -361,7 +361,15 @@ pub const Q_P95_MAG_INDEX: usize = 55; pub const Q_P95_ORD_INDEX: usize = 56; pub const Q_P95_URG_INDEX: usize = 57; -/// ISV slot [58] — low 32 bits of the u64 layout fingerprint (stored as raw f32 bits). +/// ISV slot [60] — TLOB regime focus diagnostic (D.8 Plan 2 Task 6C). +/// +/// Mean of max-attention-weight across the batch, EMA-updated once per epoch +/// by the CPU monitor after calling `GpuTlob::mean_max_attention_weight`. +/// Range [0, 1]; bootstrap 0.0. Written by CPU (read_isv_write_isv at epoch boundary). +/// Not network-facing (index >= ISV_NETWORK_DIM=23). +pub const TLOB_REGIME_FOCUS_EMA_INDEX: usize = 60; + +/// ISV slot [61] — low 32 bits of the u64 layout fingerprint (stored as raw f32 bits). /// /// Design note: this is NOT a version number. There is no ordered version space. /// The fingerprint is a structural hash that automatically changes whenever any @@ -372,9 +380,9 @@ pub const Q_P95_URG_INDEX: usize = 57; /// /// Placement at the tail (not head) is mandated by slots [0] and [1] being /// actively written by `isv_signal_update` (Q-drift and gradient-norm EMA). -pub const ISV_LAYOUT_FINGERPRINT_LO_INDEX: usize = 58; -/// ISV slot [59] — high 32 bits of the u64 layout fingerprint (stored as raw f32 bits). -pub const ISV_LAYOUT_FINGERPRINT_HI_INDEX: usize = 59; +pub const ISV_LAYOUT_FINGERPRINT_LO_INDEX: usize = 61; +/// ISV slot [62] — high 32 bits of the u64 layout fingerprint (stored as raw f32 bits). +pub const ISV_LAYOUT_FINGERPRINT_HI_INDEX: usize = 62; /// Canonical alias for the fingerprint slot (the low half). pub const ISV_LAYOUT_FINGERPRINT_INDEX: usize = ISV_LAYOUT_FINGERPRINT_LO_INDEX; @@ -431,8 +439,9 @@ const fn layout_fingerprint_seed() -> &'static [u8] { KELLY_CAP_EFF=47;CQL_ALPHA=48;PLAN_THRESHOLD=49;\ Q_P05_DIR=50;Q_P05_MAG=51;Q_P05_ORD=52;Q_P05_URG=53;\ Q_P95_DIR=54;Q_P95_MAG=55;Q_P95_ORD=56;Q_P95_URG=57;\ - ISV_LAYOUT_FINGERPRINT_LO=58;ISV_LAYOUT_FINGERPRINT_HI=59;\ - ISV_TOTAL_DIM=60" + TLOB_REGIME_FOCUS_EMA=60;\ + ISV_LAYOUT_FINGERPRINT_LO=61;ISV_LAYOUT_FINGERPRINT_HI=62;\ + ISV_TOTAL_DIM=63" } /// Compile-time layout fingerprint. Burned into the binary at build time. @@ -5278,6 +5287,28 @@ impl GpuDqnTrainer { &self.states_buf } + /// Mutable reference to the states buffer on GPU. + /// + /// Shape: `[B, STATE_DIM_PADDED]` row-major — TLOB forward writes TLOB features into + /// `[SL_TLOB_START..SL_TLOB_START+SL_TLOB_DIM)` after state gather and before cuBLAS forward. + pub(crate) fn states_buf_mut(&mut self) -> &mut CudaSlice { + &mut self.states_buf + } + + /// Reference to the bottleneck concat gradient buffer on GPU. + /// + /// Shape: `[B, concat_dim]` row-major — populated by `launch_cublas_backward` when + /// bottleneck_dim > 0. TLOB backward reads the TLOB slice from this to compute weight gradients. + pub(crate) fn bn_d_concat_buf(&self) -> &CudaSlice { + &self.bn_d_concat_buf + } + + /// Bottleneck concat dimension (bottleneck_dim + STATE_DIM - market_dim). + /// Used by TLOB backward to compute the slice offset for the TLOB gradient. + pub(crate) fn bn_concat_dim(&self) -> usize { + self.config.bottleneck_dim + (ml_core::state_layout::STATE_DIM.saturating_sub(self.config.market_dim)) + } + /// Reference to the next_states buffer on GPU. /// /// Shape: `[B, STATE_DIM]` — contains the batch's next states after `train_step()`. diff --git a/crates/ml/src/cuda_pipeline/gpu_tlob.rs b/crates/ml/src/cuda_pipeline/gpu_tlob.rs new file mode 100644 index 000000000..51300778f --- /dev/null +++ b/crates/ml/src/cuda_pipeline/gpu_tlob.rs @@ -0,0 +1,868 @@ +#![allow(unsafe_code)] // Required for CUDA kernel launches + +//! GpuTlob — TLOB attention module (D.8 Plan 2 Task 6C). +//! +//! Compresses OFI[SL_OFI_DIM=32] → TLOB[SL_TLOB_DIM=16] using single-head +//! scaled-dot-product attention over the feature dimension. +//! +//! Architecture: +//! Q = W_Q^T @ ofi_input[32,B] → [16,B] (cuBLAS SGEMM) +//! K = W_K^T @ ofi_input[32,B] → [16,B] (cuBLAS SGEMM) +//! V = W_V^T @ ofi_input[32,B] → [16,B] (cuBLAS SGEMM) +//! attn_out[16,B] = SDP(Q, K, V) (tlob_sdp_forward cubin) +//! out[16,B] = W_O^T @ attn_out[16,B] → [16,B] (cuBLAS SGEMM) +//! written to states_buf at SL_TLOB_START (tlob_write_states cubin) +//! +//! Backward: +//! d_out[16,B] extracted from bn_d_concat_buf (tlob_read_grad_from_states) +//! d_attn_out[16,B] = W_O @ d_out (cuBLAS SGEMM) +//! dW_O[16,16] += attn_out^T @ d_out / B (cuBLAS SGEMM) +//! d_q/k/v[16,B] = SDP backward (tlob_sdp_backward cubin) +//! dW_Q[16,32] += ofi^T @ d_q / B (cuBLAS SGEMM) +//! dW_K[16,32] += ofi^T @ d_k / B (cuBLAS SGEMM) +//! dW_V[16,32] += ofi^T @ d_v / B (cuBLAS SGEMM) +//! +//! Weights are Xavier-initialised (random init, no pretrained checkpoint). +//! Adam step is called after backward — same pattern as GpuAttention. +//! +//! The ISV slot TLOB_REGIME_FOCUS_EMA_INDEX (slot 60) carries the mean of +//! the per-sample max attention weight, updated at per-epoch CPU call. +//! Zero atomicAdd; no CUDA Graph capture for TLOB (runs outside graph on +//! the main stream, following the same pattern as GpuAttention). + +use std::sync::Arc; + +use cudarc::cublaslt::result as cublaslt_result; +use cudarc::cublaslt::sys as cublaslt_sys; +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; +use tracing::info; + +use crate::MLError; +use super::shared_cublas_handle::PerStreamCublasHandles; +use ml_core::state_layout::{OFI_DIM, TLOB_DIM, STATE_DIM_PADDED}; + +// ── Cubin embedding ───────────────────────────────────────────────────────── + +static TLOB_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/tlob_kernel.cubin")); + +// ── Dimensions ───────────────────────────────────────────────────────────── + +/// TLOB Q/K/V projection input width (OFI features). +const TLOB_IN: usize = OFI_DIM; // 32 +/// TLOB Q/K/V projection output width (compressed dim). +const TLOB_OUT: usize = TLOB_DIM; // 16 + +/// Parameter layout (flat buffer): +/// W_Q[TLOB_OUT, TLOB_IN] + W_K[TLOB_OUT, TLOB_IN] + W_V[TLOB_OUT, TLOB_IN] +/// + W_O[TLOB_OUT, TLOB_OUT] +/// No bias terms — LayerNorm is intentionally omitted for MVP to stay within +/// the 800 LOC budget. The Q/K/V zero-bias simplification is acceptable for +/// D.8 because the OFI features are already zero-mean (MBP-10 delta/ratio feats). +const TLOB_W_QKV_SIZE: usize = 3 * TLOB_OUT * TLOB_IN; // 3 × 16 × 32 = 1536 +const TLOB_W_O_SIZE: usize = TLOB_OUT * TLOB_OUT; // 16 × 16 = 256 +const TLOB_TOTAL_PARAMS: usize = TLOB_W_QKV_SIZE + TLOB_W_O_SIZE; // 1792 + +// Offsets within the flat params buffer +const W_Q_OFF: usize = 0; +const W_K_OFF: usize = TLOB_OUT * TLOB_IN; // 512 +const W_V_OFF: usize = 2 * TLOB_OUT * TLOB_IN; // 1024 +const W_O_OFF: usize = 3 * TLOB_OUT * TLOB_IN; // 1536 + +// ── GEMM descriptor ───────────────────────────────────────────────────────── + +/// Cached cublasLt GEMM descriptor for TLOB projections. +/// Same structure as AttnGemmDesc in gpu_attention.rs. +struct TlobGemmDesc { + matmul_desc: cublaslt_sys::cublasLtMatmulDesc_t, + a_layout: cublaslt_sys::cublasLtMatrixLayout_t, + b_layout: cublaslt_sys::cublasLtMatrixLayout_t, + c_layout: cublaslt_sys::cublasLtMatrixLayout_t, + d_layout: cublaslt_sys::cublasLtMatrixLayout_t, + algo: cublaslt_sys::cublasLtMatmulAlgo_t, +} + +unsafe impl Send for TlobGemmDesc {} +unsafe impl Sync for TlobGemmDesc {} + +impl Drop for TlobGemmDesc { + fn drop(&mut self) { + unsafe { + let _ = cublaslt_result::destroy_matrix_layout(self.d_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.c_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.b_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.a_layout); + let _ = cublaslt_result::destroy_matmul_desc(self.matmul_desc); + } + } +} + +// ── GpuTlob ───────────────────────────────────────────────────────────────── + +/// TLOB attention module: OFI[32,B] → compressed features[16,B]. +/// +/// Weights are random-init (Xavier), trainable end-to-end via DQN reward. +/// No pretrained checkpoint; no ONNX. +#[allow(missing_debug_implementations)] +pub(crate) struct GpuTlob { + stream: Arc, + shared_handle: Arc, + batch_size: usize, + + // ── cuBLAS GEMM descriptors ─────────────────────────────────────── + // Forward Q/K/V: W[16,32]^T @ ofi[32,B] → [16,B] + // TRANSA=T, TRANSB=N, M=16, N=B, K=32 + gemm_fwd_q: TlobGemmDesc, + gemm_fwd_k: TlobGemmDesc, + gemm_fwd_v: TlobGemmDesc, + // Forward O: W_O[16,16]^T @ attn_out[16,B] → [16,B] + // TRANSA=T, TRANSB=N, M=16, N=B, K=16 + gemm_fwd_o: TlobGemmDesc, + // Backward dW: dW_Q/K/V[16,32] = d_proj[16,B] @ ofi^T[B,32] (dY @ X^T) + // TRANSA=N, TRANSB=T, M=16, N=32, K=B + gemm_bwd_dw_q: TlobGemmDesc, + gemm_bwd_dw_k: TlobGemmDesc, + gemm_bwd_dw_v: TlobGemmDesc, + // Backward dW_O[16,16] = d_out[16,B] @ attn_out^T[B,16] + // TRANSA=N, TRANSB=T, M=16, N=16, K=B + gemm_bwd_dw_o: TlobGemmDesc, + // Backward d_attn_out = W_O[16,16] @ d_out[16,B] + // TRANSA=N, TRANSB=N, M=16, N=B, K=16 + gemm_bwd_dx_o: TlobGemmDesc, + + // ── cubin kernels ───────────────────────────────────────────────── + sdp_fwd_kernel: CudaFunction, + sdp_bwd_kernel: CudaFunction, + write_states_kernel: CudaFunction, + read_grad_kernel: CudaFunction, + + // ── Weight parameters [TLOB_TOTAL_PARAMS] ──────────────────────── + /// Flat: W_Q[512] + W_K[512] + W_V[512] + W_O[256] = 1792 floats. + params: CudaSlice, + + // ── Intermediate buffers (forward, saved for backward) ──────────── + /// OFI input in col-major [TLOB_IN=32, B]: extracted from states_buf. + ofi_col_buf: CudaSlice, + /// Q projection [TLOB_OUT=16, B]. + proj_q_buf: CudaSlice, + /// K projection [TLOB_OUT=16, B]. + proj_k_buf: CudaSlice, + /// V projection [TLOB_OUT=16, B]. + proj_v_buf: CudaSlice, + /// SDP output [TLOB_OUT=16, B] (pre-output-projection), saved for dW_O. + attn_out_buf: CudaSlice, + /// SDP softmax weights [TLOB_OUT=16, B], saved for backward. + attn_scores_buf: CudaSlice, + /// TLOB output [TLOB_OUT=16, B] = W_O^T @ attn_out. + output_buf: CudaSlice, + + // ── Gradient buffers (backward) ─────────────────────────────────── + /// Gradient from trunk w.r.t. TLOB output [TLOB_OUT=16, B]. + d_output_buf: CudaSlice, + /// Gradient w.r.t. attn_out [TLOB_OUT=16, B]. + d_attn_out_buf: CudaSlice, + /// Gradient w.r.t. Q [TLOB_OUT=16, B]. + d_proj_q_buf: CudaSlice, + /// Gradient w.r.t. K [TLOB_OUT=16, B]. + d_proj_k_buf: CudaSlice, + /// Gradient w.r.t. V [TLOB_OUT=16, B]. + d_proj_v_buf: CudaSlice, + /// Parameter gradients [TLOB_TOTAL_PARAMS]. + d_params: CudaSlice, + + // ── Adam optimizer state ────────────────────────────────────────── + adam_m: CudaSlice, + adam_v: CudaSlice, + adam_step: i32, + /// Pinned host counter for device-side Adam t (device-mapped). + t_pinned: *mut i32, + t_dev_ptr: u64, + + // ── Gradient norm buffers ───────────────────────────────────────── + grad_norm_buf: CudaSlice, + grad_norm_block_sums: CudaSlice, + + // ── Kernels for grad norm + Adam ────────────────────────────────── + grad_norm_phase1_kernel: CudaFunction, + grad_norm_phase2_kernel: CudaFunction, + adam_kernel: CudaFunction, + + // ── ofi_extract kernel ──────────────────────────────────────────── + /// Extracts OFI slice from states_buf row-major → col-major [32,B]. + ofi_extract_kernel: CudaFunction, + + // ── Bias gradient reduce (for zero-bias architecture, only d_params for W) ── + bias_grad_reduce_p1_kernel: CudaFunction, + bias_grad_reduce_p2_kernel: CudaFunction, + bias_grad_partials_buf: CudaSlice, + bias_grad_num_blocks: u32, +} + +impl GpuTlob { + pub(crate) fn new( + shared_handle: Arc, + batch_size: usize, + ) -> Result { + let stream = Arc::clone(&shared_handle.stream); + let b = batch_size; + let lt_handle = shared_handle.lt_handle.0; + let ws_size = shared_handle.lt_workspace_size; + + // ── Load cubin module ──────────────────────────────────────── + let context = stream.context(); + let module = context.load_cubin(TLOB_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("tlob cubin load: {e}")))?; + let sdp_fwd_kernel = module.load_function("tlob_sdp_forward") + .map_err(|e| MLError::ModelError(format!("tlob_sdp_forward load: {e}")))?; + let sdp_bwd_kernel = module.load_function("tlob_sdp_backward") + .map_err(|e| MLError::ModelError(format!("tlob_sdp_backward load: {e}")))?; + let write_states_kernel = module.load_function("tlob_write_states") + .map_err(|e| MLError::ModelError(format!("tlob_write_states load: {e}")))?; + let read_grad_kernel = module.load_function("tlob_read_grad_from_states") + .map_err(|e| MLError::ModelError(format!("tlob_read_grad_from_states load: {e}")))?; + + // ── Load grad norm + Adam from attention backward cubin ─────── + // Reuse the same kernels as GpuAttention to avoid duplicate cubin compilation. + static ATTN_BWD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_backward_kernel.cubin")); + let bwd_module = context.load_cubin(ATTN_BWD_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("attn_bwd cubin for tlob: {e}")))?; + let grad_norm_phase1_kernel = bwd_module.load_function("attn_grad_norm_phase1") + .map_err(|e| MLError::ModelError(format!("attn_grad_norm_phase1 load: {e}")))?; + let grad_norm_phase2_kernel = bwd_module.load_function("attn_grad_norm_phase2") + .map_err(|e| MLError::ModelError(format!("attn_grad_norm_phase2 load: {e}")))?; + let adam_kernel = bwd_module.load_function("attn_adam_update") + .map_err(|e| MLError::ModelError(format!("attn_adam_update load: {e}")))?; + let bias_grad_reduce_p1_kernel = bwd_module.load_function("attn_bias_grad_reduce_p1") + .map_err(|e| MLError::ModelError(format!("attn_bias_grad_reduce_p1 load: {e}")))?; + let bias_grad_reduce_p2_kernel = bwd_module.load_function("attn_bias_grad_reduce_p2") + .map_err(|e| MLError::ModelError(format!("attn_bias_grad_reduce_p2 load: {e}")))?; + + // ── ofi_extract kernel (from experience_kernels cubin) ──────── + static EXP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/experience_kernels.cubin")); + let exp_module = context.load_cubin(EXP_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("exp cubin for tlob: {e}")))?; + let ofi_extract_kernel = exp_module.load_function("ofi_embed_build_input") + .map_err(|e| MLError::ModelError(format!("ofi_embed_build_input for tlob: {e}")))?; + + // ── GEMM descriptors ───────────────────────────────────────── + // Forward projections: W[TLOB_OUT, TLOB_IN]^T @ ofi[TLOB_IN, B] → [TLOB_OUT, B] + // TRANSA=T(1), TRANSB=N(0), M=TLOB_OUT, N=B, K=TLOB_IN + // A = W stored [TLOB_OUT, TLOB_IN] (col-major [TLOB_IN, TLOB_OUT]), lda=TLOB_IN + // B = ofi stored col-major [TLOB_IN, B], ldb=TLOB_IN + // C = output col-major [TLOB_OUT, B], ldc=TLOB_OUT + let m = TLOB_OUT as i32; + let n = b as i32; + let k_in = TLOB_IN as i32; // K for Q/K/V projections + let k_out = TLOB_OUT as i32; // K for output projection + + let gemm_fwd_q = create_tlob_gemm_desc(lt_handle, 1, 0, m, n, k_in, k_in, k_in, m, ws_size, "fwd_q")?; + let gemm_fwd_k = create_tlob_gemm_desc(lt_handle, 1, 0, m, n, k_in, k_in, k_in, m, ws_size, "fwd_k")?; + let gemm_fwd_v = create_tlob_gemm_desc(lt_handle, 1, 0, m, n, k_in, k_in, k_in, m, ws_size, "fwd_v")?; + // Forward output: W_O[TLOB_OUT,TLOB_OUT]^T @ attn[TLOB_OUT,B] → [TLOB_OUT,B] + let gemm_fwd_o = create_tlob_gemm_desc(lt_handle, 1, 0, m, n, k_out, k_out, k_out, m, ws_size, "fwd_o")?; + + // Backward dW_QKV: d_proj[TLOB_OUT,B] @ ofi^T[B,TLOB_IN] → dW[TLOB_OUT,TLOB_IN] + // TRANSA=N(0), TRANSB=T(1), M=TLOB_OUT, N=TLOB_IN, K=B + // A = d_proj col-major [TLOB_OUT, B], lda=TLOB_OUT + // B = ofi col-major [TLOB_IN, B], ldb=TLOB_IN + // C = dW col-major [TLOB_OUT, TLOB_IN], ldc=TLOB_OUT + let bw_n = TLOB_IN as i32; + let bw_k = b as i32; + let gemm_bwd_dw_q = create_tlob_gemm_desc(lt_handle, 0, 1, m, bw_n, bw_k, m, k_in, m, ws_size, "bwd_dw_q")?; + let gemm_bwd_dw_k = create_tlob_gemm_desc(lt_handle, 0, 1, m, bw_n, bw_k, m, k_in, m, ws_size, "bwd_dw_k")?; + let gemm_bwd_dw_v = create_tlob_gemm_desc(lt_handle, 0, 1, m, bw_n, bw_k, m, k_in, m, ws_size, "bwd_dw_v")?; + + // Backward dW_O: d_out[TLOB_OUT,B] @ attn_out^T[B,TLOB_OUT] → dW_O[TLOB_OUT,TLOB_OUT] + // TRANSA=N(0), TRANSB=T(1), M=TLOB_OUT, N=TLOB_OUT, K=B + let gemm_bwd_dw_o = create_tlob_gemm_desc(lt_handle, 0, 1, m, k_out, bw_k, m, k_out, m, ws_size, "bwd_dw_o")?; + + // Backward d_attn_out = W_O[TLOB_OUT,TLOB_OUT] @ d_out[TLOB_OUT,B] → [TLOB_OUT,B] + // TRANSA=N(0), TRANSB=N(0), M=TLOB_OUT, N=B, K=TLOB_OUT + // A = W_O col-major [TLOB_OUT, TLOB_OUT], lda=TLOB_OUT + // B = d_out col-major [TLOB_OUT, B], ldb=TLOB_OUT + let gemm_bwd_dx_o = create_tlob_gemm_desc(lt_handle, 0, 0, m, n, k_out, k_out, k_out, m, ws_size, "bwd_dx_o")?; + + // ── Xavier initialization ───────────────────────────────────── + let mut host_params = vec![0.0_f32; TLOB_TOTAL_PARAMS]; + let fan_in = TLOB_IN as f32; + let fan_out = TLOB_OUT as f32; + let xavier_std = (2.0 / (fan_in + fan_out)).sqrt(); + let mut rng = 0xcbf29ce484222325_u64.wrapping_add(b as u64); + for p in host_params[..TLOB_W_QKV_SIZE].iter_mut() { + rng = rng.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + let u = (rng >> 33) as f32 / (1u64 << 31) as f32; + *p = (u - 0.5) * 3.46 * xavier_std; + } + // W_O starts at zero (identity-like start: output = 0 initially, learns from scratch) + + let mut params = stream.alloc_zeros::(TLOB_TOTAL_PARAMS) + .map_err(|e| MLError::ModelError(format!("tlob params alloc: {e}")))?; + super::htod_f32(&stream, &host_params, &mut params)?; + + // ── Buffer allocation ───────────────────────────────────────── + let db_out = TLOB_OUT * b; // 16 × B + let db_in = TLOB_IN * b; // 32 × B + + let alloc = |n: usize, label: &str| -> Result, MLError> { + stream.alloc_zeros::(n) + .map_err(|e| MLError::ModelError(format!("tlob {label} alloc: {e}"))) + }; + + let ofi_col_buf = alloc(db_in, "ofi_col")?; + let proj_q_buf = alloc(db_out, "proj_q")?; + let proj_k_buf = alloc(db_out, "proj_k")?; + let proj_v_buf = alloc(db_out, "proj_v")?; + let attn_out_buf = alloc(db_out, "attn_out")?; + let attn_scores_buf= alloc(db_out, "attn_scores")?; + let output_buf = alloc(db_out, "output")?; + let d_output_buf = alloc(db_out, "d_output")?; + let d_attn_out_buf = alloc(db_out, "d_attn_out")?; + let d_proj_q_buf = alloc(db_out, "d_proj_q")?; + let d_proj_k_buf = alloc(db_out, "d_proj_k")?; + let d_proj_v_buf = alloc(db_out, "d_proj_v")?; + let d_params = alloc(TLOB_TOTAL_PARAMS, "d_params")?; + let adam_m = alloc(TLOB_TOTAL_PARAMS, "adam_m")?; + let adam_v = alloc(TLOB_TOTAL_PARAMS, "adam_v")?; + + let norm_blocks = (TLOB_TOTAL_PARAMS + 255) / 256; + let grad_norm_buf = alloc(1, "grad_norm")?; + let grad_norm_block_sums = alloc(norm_blocks, "grad_norm_sums")?; + + let bias_grad_num_blocks = ((b + 255) / 256) as u32; + let bias_grad_partials_buf = alloc( + (bias_grad_num_blocks as usize) * TLOB_OUT, + "bias_grad_partials", + )?; + + // ── Adam pinned step counter ────────────────────────────────── + let mut t_raw: *mut i32 = std::ptr::null_mut(); + unsafe { + let status = cudarc::driver::sys::cuMemHostAlloc( + (&mut t_raw) as *mut *mut i32 as *mut *mut std::ffi::c_void, + std::mem::size_of::(), + cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP, + ); + if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + return Err(MLError::ModelError(format!("tlob: cuMemHostAlloc failed: {status:?}"))); + } + *t_raw = 0; + } + let mut t_dev_raw: u64 = 0; + unsafe { + let status = cudarc::driver::sys::cuMemHostGetDevicePointer_v2( + &mut t_dev_raw, + t_raw as *mut std::ffi::c_void, + 0, + ); + if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + cudarc::driver::sys::cuMemFreeHost(t_raw as *mut std::ffi::c_void); + return Err(MLError::ModelError(format!("tlob: cuMemHostGetDevicePointer failed: {status:?}"))); + } + } + + info!( + batch_size = b, + tlob_in = TLOB_IN, + tlob_out = TLOB_OUT, + total_params = TLOB_TOTAL_PARAMS, + "GpuTlob initialized: OFI[32]→TLOB[16] single-head SDP, Xavier random init, trainable end-to-end" + ); + + Ok(Self { + stream, + shared_handle, + batch_size: b, + gemm_fwd_q, + gemm_fwd_k, + gemm_fwd_v, + gemm_fwd_o, + gemm_bwd_dw_q, + gemm_bwd_dw_k, + gemm_bwd_dw_v, + gemm_bwd_dw_o, + gemm_bwd_dx_o, + sdp_fwd_kernel, + sdp_bwd_kernel, + write_states_kernel, + read_grad_kernel, + params, + ofi_col_buf, + proj_q_buf, + proj_k_buf, + proj_v_buf, + attn_out_buf, + attn_scores_buf, + output_buf, + d_output_buf, + d_attn_out_buf, + d_proj_q_buf, + d_proj_k_buf, + d_proj_v_buf, + d_params, + adam_m, + adam_v, + adam_step: 0, + t_pinned: t_raw, + t_dev_ptr: t_dev_raw, + grad_norm_buf, + grad_norm_block_sums, + grad_norm_phase1_kernel, + grad_norm_phase2_kernel, + adam_kernel, + ofi_extract_kernel, + bias_grad_reduce_p1_kernel, + bias_grad_reduce_p2_kernel, + bias_grad_partials_buf, + bias_grad_num_blocks, + }) + } + + /// TLOB forward pass. + /// + /// 1. Extract OFI slice from states_buf → ofi_col_buf [TLOB_IN, B] col-major. + /// 2. Q/K/V projections via cuBLAS. + /// 3. SDP (tlob_sdp_forward kernel). + /// 4. Output projection via cuBLAS. + /// 5. Write [TLOB_OUT, B] to states_buf at SL_TLOB_START. + /// + /// `states_buf`: row-major [B, state_dim_padded] — both read (OFI) and written (TLOB). + pub(crate) fn forward( + &mut self, + states_buf: &mut CudaSlice, + batch_size: usize, + ) -> Result<(), MLError> { + let b = batch_size; + let b_i32 = b as i32; + let state_dim_padded_i32 = STATE_DIM_PADDED as i32; + let blocks = ((b + 255) / 256) as u32; + let f32_sz = std::mem::size_of::(); + let (ws_ptr, ws_size) = (self.shared_handle.lt_workspace_ptr, self.shared_handle.lt_workspace_size); + + // ── Step 1: Extract OFI [SL_OFI_START..+SL_OFI_DIM] from states row-major + // → ofi_col_buf [TLOB_IN=32, B] col-major. + // Reuse ofi_embed_build_input kernel: reads s[SL_OFI_START..+SL_OFI_DIM] from states. + let state_dim_i32 = STATE_DIM_PADDED as i32; + unsafe { + self.stream + .launch_builder(&self.ofi_extract_kernel) + .arg(states_buf as &CudaSlice) + .arg(&mut self.ofi_col_buf) + .arg(&b_i32) + .arg(&state_dim_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("tlob ofi_extract: {e}")))?; + } + + let params_base = self.params.raw_ptr(); + let ofi_ptr = self.ofi_col_buf.raw_ptr(); + + // ── Step 2: Q/K/V projections ───────────────────────────────── + // Q = W_Q^T @ ofi (col-major: [TLOB_OUT, B] = [TLOB_OUT, TLOB_IN] @ [TLOB_IN, B]) + self.lt_matmul(&self.gemm_fwd_q, + params_base + (W_Q_OFF * f32_sz) as u64, + ofi_ptr, self.proj_q_buf.raw_ptr(), + 1.0, 0.0, "fwd_q", ws_ptr, ws_size)?; + self.lt_matmul(&self.gemm_fwd_k, + params_base + (W_K_OFF * f32_sz) as u64, + ofi_ptr, self.proj_k_buf.raw_ptr(), + 1.0, 0.0, "fwd_k", ws_ptr, ws_size)?; + self.lt_matmul(&self.gemm_fwd_v, + params_base + (W_V_OFF * f32_sz) as u64, + ofi_ptr, self.proj_v_buf.raw_ptr(), + 1.0, 0.0, "fwd_v", ws_ptr, ws_size)?; + + // ── Step 3: SDP ─────────────────────────────────────────────── + unsafe { + self.stream + .launch_builder(&self.sdp_fwd_kernel) + .arg(&self.proj_q_buf) + .arg(&self.proj_k_buf) + .arg(&self.proj_v_buf) + .arg(&mut self.attn_out_buf) + .arg(&mut self.attn_scores_buf) + .arg(&b_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("tlob sdp_fwd: {e}")))?; + } + + // ── Step 4: Output projection ───────────────────────────────── + // out = W_O^T @ attn_out [TLOB_OUT, B] + self.lt_matmul(&self.gemm_fwd_o, + params_base + (W_O_OFF * f32_sz) as u64, + self.attn_out_buf.raw_ptr(), self.output_buf.raw_ptr(), + 1.0, 0.0, "fwd_o", ws_ptr, ws_size)?; + + // ── Step 5: Scatter output into states_buf at SL_TLOB_START ── + unsafe { + self.stream + .launch_builder(&self.write_states_kernel) + .arg(&self.output_buf) + .arg(states_buf) + .arg(&b_i32) + .arg(&state_dim_padded_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("tlob write_states: {e}")))?; + } + Ok(()) + } + + /// TLOB backward pass. + /// + /// `d_concat`: gradient buffer [B, concat_dim] from trunk backward (bn_d_concat_buf). + /// `concat_dim`: total concat dim = bottleneck_dim + (STATE_DIM - market_dim). + /// `tlob_concat_off`: TLOB slot offset within concat row = bottleneck_dim + (TLOB_START - market_dim). + pub(crate) fn backward( + &mut self, + d_concat: &CudaSlice, + batch_size: usize, + concat_dim: usize, + tlob_concat_off: usize, + ) -> Result<(), MLError> { + let b = batch_size; + let b_i32 = b as i32; + let blocks = ((b + 255) / 256) as u32; + let f32_sz = std::mem::size_of::(); + let (ws_ptr, ws_size) = (self.shared_handle.lt_workspace_ptr, self.shared_handle.lt_workspace_size); + let inv_batch = 1.0_f32 / b as f32; + let concat_dim_i32 = concat_dim as i32; + let tlob_off_i32 = tlob_concat_off as i32; + + // ── Step 1: Extract TLOB gradient from d_concat ─────────────── + // tlob_read_grad_from_states: d_concat[B, concat_dim] → d_output_buf[TLOB_OUT, B] + unsafe { + self.stream + .launch_builder(&self.read_grad_kernel) + .arg(d_concat) + .arg(&mut self.d_output_buf) + .arg(&b_i32) + .arg(&concat_dim_i32) + .arg(&tlob_off_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("tlob read_grad: {e}")))?; + } + + let params_base = self.params.raw_ptr(); + + // ── Step 2: Backward through output projection ──────────────── + // d_attn_out = W_O @ d_out [TLOB_OUT, B] + self.lt_matmul(&self.gemm_bwd_dx_o, + params_base + (W_O_OFF * f32_sz) as u64, + self.d_output_buf.raw_ptr(), self.d_attn_out_buf.raw_ptr(), + 1.0, 0.0, "bwd_dx_o", ws_ptr, ws_size)?; + + // dW_O[TLOB_OUT, TLOB_OUT] = d_out[TLOB_OUT, B] @ attn_out^T[B, TLOB_OUT], alpha=1/B + // Use d_params at W_O_OFF offset + let grad_base = self.d_params.raw_ptr(); + self.lt_matmul(&self.gemm_bwd_dw_o, + self.d_output_buf.raw_ptr(), + self.attn_out_buf.raw_ptr(), + grad_base + (W_O_OFF * f32_sz) as u64, + inv_batch, 0.0, "bwd_dw_o", ws_ptr, ws_size)?; + + // ── Step 3: SDP backward ────────────────────────────────────── + unsafe { + self.stream + .launch_builder(&self.sdp_bwd_kernel) + .arg(&self.d_attn_out_buf) + .arg(&self.attn_scores_buf) + .arg(&self.proj_v_buf) + .arg(&self.proj_q_buf) + .arg(&self.proj_k_buf) + .arg(&mut self.d_proj_q_buf) + .arg(&mut self.d_proj_k_buf) + .arg(&mut self.d_proj_v_buf) + .arg(&b_i32) + .arg(&inv_batch) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("tlob sdp_bwd: {e}")))?; + } + + // ── Step 4: dW_Q/K/V backward ──────────────────────────────── + let ofi_ptr = self.ofi_col_buf.raw_ptr(); + // dW_Q = d_proj_q[TLOB_OUT,B] @ ofi^T[B,TLOB_IN] (alpha=1 since sdp_bwd already divides by B) + self.lt_matmul(&self.gemm_bwd_dw_q, + self.d_proj_q_buf.raw_ptr(), ofi_ptr, + grad_base + (W_Q_OFF * f32_sz) as u64, + 1.0, 0.0, "bwd_dw_q", ws_ptr, ws_size)?; + self.lt_matmul(&self.gemm_bwd_dw_k, + self.d_proj_k_buf.raw_ptr(), ofi_ptr, + grad_base + (W_K_OFF * f32_sz) as u64, + 1.0, 0.0, "bwd_dw_k", ws_ptr, ws_size)?; + self.lt_matmul(&self.gemm_bwd_dw_v, + self.d_proj_v_buf.raw_ptr(), ofi_ptr, + grad_base + (W_V_OFF * f32_sz) as u64, + 1.0, 0.0, "bwd_dw_v", ws_ptr, ws_size)?; + + Ok(()) + } + + /// Adam step: grad norm clip + Adam update on TLOB params. + pub(crate) fn adam_step(&mut self, lr: f32, max_grad_norm: f32) -> Result<(), MLError> { + let tp = TLOB_TOTAL_PARAMS; + let tp_i32 = tp as i32; + + let norm_blocks = (tp + 255) / 256; + unsafe { + self.stream + .launch_builder(&self.grad_norm_phase1_kernel) + .arg(&self.d_params) + .arg(&mut self.grad_norm_block_sums) + .arg(&tp_i32) + .launch(LaunchConfig { + grid_dim: (norm_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("tlob grad_norm phase1: {e}")))?; + } + let norm_blocks_i32 = norm_blocks as i32; + unsafe { + self.stream + .launch_builder(&self.grad_norm_phase2_kernel) + .arg(&self.grad_norm_block_sums) + .arg(&mut self.grad_norm_buf) + .arg(&norm_blocks_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("tlob grad_norm phase2: {e}")))?; + } + + self.adam_step += 1; + unsafe { *self.t_pinned = self.adam_step; } + let adam_blocks = (tp + 255) / 256; + let beta1 = 0.9_f32; + let beta2 = 0.999_f32; + let eps = 1e-8_f32; + let wd = 1e-5_f32; + let t_ptr = self.t_dev_ptr; + unsafe { + self.stream + .launch_builder(&self.adam_kernel) + .arg(&mut self.params) + .arg(&mut self.d_params) + .arg(&mut self.adam_m) + .arg(&mut self.adam_v) + .arg(&self.grad_norm_buf) + .arg(&lr) + .arg(&beta1) + .arg(&beta2) + .arg(&eps) + .arg(&wd) + .arg(&max_grad_norm) + .arg(&t_ptr) + .arg(&tp_i32) + .launch(LaunchConfig { + grid_dim: (adam_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("tlob adam: {e}")))?; + } + Ok(()) + } + + /// Copy trained weight parameters from `src` into this instance via DtoD async copy. + /// + /// Used to synchronise the val-path TLOB (owned by `GpuBacktestEvaluator`) with + /// the online weights trained by `FusedTrainingCtx`. The copy is enqueued on + /// `self.stream`; callers should ensure `src`'s stream has finished writing params + /// before this is called (a `stream.synchronize()` on the training stream is + /// sufficient, which already happens at the epoch boundary). + pub(crate) fn copy_params_from(&mut self, src: &GpuTlob) -> Result<(), MLError> { + let byte_len = (TLOB_TOTAL_PARAMS * std::mem::size_of::()) as u64; + let (src_ptr, _src_guard) = src.params.device_ptr(&self.stream); + let (dst_ptr, _dst_guard) = self.params.device_ptr(&self.stream); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, + src_ptr, + byte_len as usize, + self.stream.cu_stream(), + ) + .map_err(|e| MLError::ModelError(format!("tlob copy_params_from: {e}")))?; + } + Ok(()) + } + + /// Compute the mean-max attention weight across the batch for ISV diagnostic. + /// Returns the mean of max(attn_scores[d, b] over d) across b. + /// Runs a minimal CPU readback — only called once per epoch at the CPU monitor site. + pub(crate) fn mean_max_attention_weight(&self, batch_size: usize) -> f32 { + let b = batch_size; + let d = TLOB_OUT; + let total = d * b; + let mut host = vec![0.0_f32; total]; + // Synchronous DtoH for diagnostic (called once/epoch outside hot path). + // memcpy_dtoh queues async; stream.synchronize() ensures transfer completes. + if self.stream.memcpy_dtoh(&self.attn_scores_buf, &mut host).is_ok() { + let _ = self.stream.synchronize(); + } else { + return 0.0; + } + // attn_scores col-major [D, B]: per sample b, find max over d + let mut sum_max = 0.0_f32; + for bi in 0..b { + let mut mx = 0.0_f32; + for di in 0..d { + let v = host[di * b + bi]; + if v > mx { mx = v; } + } + sum_max += mx; + } + if b > 0 { sum_max / b as f32 } else { 0.0 } + } +} + +impl Drop for GpuTlob { + fn drop(&mut self) { + if !self.t_pinned.is_null() { + unsafe { cudarc::driver::sys::cuMemFreeHost(self.t_pinned.cast()); } + } + } +} + +unsafe impl Send for GpuTlob {} +unsafe impl Sync for GpuTlob {} + +// ── GEMM helper ───────────────────────────────────────────────────────────── + +impl GpuTlob { + fn lt_matmul( + &self, + desc: &TlobGemmDesc, + a_ptr: u64, + b_ptr: u64, + c_ptr: u64, + alpha: f32, + beta: f32, + label: &str, + ws_ptr: u64, + ws_size: usize, + ) -> Result<(), MLError> { + let lt_handle = self.shared_handle.lt_handle_for(&self.stream)?; + unsafe { + let cu_stream = self.stream.cu_stream() as cublaslt_sys::cudaStream_t; + let status = cublaslt_sys::cublasLtMatmul( + lt_handle, + desc.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + a_ptr as *const std::ffi::c_void, + desc.a_layout, + b_ptr as *const std::ffi::c_void, + desc.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + c_ptr as *const std::ffi::c_void, + desc.c_layout, + c_ptr as *mut std::ffi::c_void, + desc.d_layout, + &desc.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t, + ws_ptr as *mut std::ffi::c_void, + ws_size, + cu_stream, + ); + if status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { + let e = cublaslt_result::CublasError(status); + return Err(MLError::ModelError(format!("cublasLtMatmul tlob {label}: {e:?}"))); + } + } + Ok(()) + } +} + +// ── GEMM descriptor factory ───────────────────────────────────────────────── + +#[allow(clippy::too_many_arguments)] +fn create_tlob_gemm_desc( + lt_handle: cublaslt_sys::cublasLtHandle_t, + transa_i32: i32, + transb_i32: i32, + m: i32, n: i32, k: i32, + lda: i32, ldb: i32, ldc: i32, + ws_size: usize, + label: &str, +) -> Result { + let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F; + let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32; + + unsafe { + let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type) + .map_err(|e| MLError::ModelError(format!("tlob MatmulDescCreate ({label} m={m},n={n},k={k}): {e:?}")))?; + + cublaslt_result::set_matmul_desc_attribute( + matmul_desc, + cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA, + &transa_i32 as *const i32 as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("tlob set TRANSA ({label}): {e:?}")))?; + + cublaslt_result::set_matmul_desc_attribute( + matmul_desc, + cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB, + &transb_i32 as *const i32 as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("tlob set TRANSB ({label}): {e:?}")))?; + + let a_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, k as u64, lda as i64) + .map_err(|e| MLError::ModelError(format!("tlob a_layout ({label}): {e:?}")))?; + let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, ldb as i64) + .map_err(|e| MLError::ModelError(format!("tlob b_layout ({label}): {e:?}")))?; + let c_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64) + .map_err(|e| MLError::ModelError(format!("tlob c_layout ({label}): {e:?}")))?; + let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64) + .map_err(|e| MLError::ModelError(format!("tlob d_layout ({label}): {e:?}")))?; + + // Heuristic algo selection + let pref = cublaslt_result::create_matmul_pref() + .map_err(|e| MLError::ModelError(format!("tlob create_pref ({label}): {e:?}")))?; + cublaslt_result::set_matmul_pref_attribute( + pref, + cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &ws_size as *const usize as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("tlob set_pref_ws ({label}): {e:?}")))?; + + let mut algos: [cublaslt_sys::cublasLtMatmulHeuristicResult_t; 3] = std::mem::zeroed(); + let mut found = 0i32; + let heur_status = cublaslt_sys::cublasLtMatmulAlgoGetHeuristic( + lt_handle, matmul_desc, + a_layout, b_layout, c_layout, d_layout, + pref, 3, + algos.as_mut_ptr(), &mut found, + ); + unsafe { cublaslt_result::destroy_matmul_pref(pref) } + .map_err(|e| MLError::ModelError(format!("tlob destroy_pref ({label}): {e:?}")))?; + if heur_status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS || found == 0 { + cublaslt_result::destroy_matrix_layout(d_layout).ok(); + cublaslt_result::destroy_matrix_layout(c_layout).ok(); + cublaslt_result::destroy_matrix_layout(b_layout).ok(); + cublaslt_result::destroy_matrix_layout(a_layout).ok(); + cublaslt_result::destroy_matmul_desc(matmul_desc).ok(); + return Err(MLError::ModelError(format!("tlob heuristic ({label}): no algo found"))); + } + let algo = algos[0].algo; + + Ok(TlobGemmDesc { matmul_desc, a_layout, b_layout, c_layout, d_layout, algo }) + } +} diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index f604cc1a4..5d94b850b 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -37,6 +37,7 @@ pub mod q_value_provider; pub mod gpu_iql_trainer; pub mod gpu_iqn_head; pub mod gpu_attention; +pub mod gpu_tlob; pub mod decision_transformer; pub mod learning_health; pub mod q_snapshot; diff --git a/crates/ml/src/cuda_pipeline/state_layout.cuh b/crates/ml/src/cuda_pipeline/state_layout.cuh index 0540a6b78..567a9658c 100644 --- a/crates/ml/src/cuda_pipeline/state_layout.cuh +++ b/crates/ml/src/cuda_pipeline/state_layout.cuh @@ -5,11 +5,13 @@ #pragma once // ── Dimensions ── -// D.6 Plan_isv[6] remaining_fraction (Plan 2 Task 6A): SL_PORTFOLIO_PLAN_DIM 6→7, -// SL_STATE_DIM 104→105, padded to 112 for 8-alignment → SL_PADDING_DIM 0→7. -#define SL_STATE_DIM 112 +// D.8 Plan 2 Task 6C: TLOB attention (16-dim) inserted after OFI. +// Layout: Market(42) + OFI(32) + TLOB(16) + MTF(16) + Portfolio(8) + PlanISV(7) + Padding(7) = 128. +// SL_STATE_DIM 112→128 (matches SL_STATE_DIM_PADDED; padding reclaimed to 7 for 8-alignment). +#define SL_STATE_DIM 128 #define SL_MARKET_DIM 42 #define SL_OFI_DIM 32 +#define SL_TLOB_DIM 16 #define SL_MTF_DIM 16 #define SL_PORTFOLIO_BASE_DIM 8 #define SL_PORTFOLIO_PLAN_DIM 7 @@ -18,7 +20,8 @@ // ── Offsets ── #define SL_MARKET_START 0 #define SL_OFI_START (SL_MARKET_START + SL_MARKET_DIM) -#define SL_MTF_START (SL_OFI_START + SL_OFI_DIM) +#define SL_TLOB_START (SL_OFI_START + SL_OFI_DIM) +#define SL_MTF_START (SL_TLOB_START + SL_TLOB_DIM) #define SL_PORTFOLIO_START (SL_MTF_START + SL_MTF_DIM) #define SL_PLAN_ISV_START (SL_PORTFOLIO_START + SL_PORTFOLIO_BASE_DIM) #define SL_PADDING_START (SL_PLAN_ISV_START + SL_PORTFOLIO_PLAN_DIM) @@ -168,19 +171,23 @@ __device__ __forceinline__ void assemble_state( for (int k = 0; k < SL_OFI_DIM; k++) out[SL_OFI_START + k] = ofi[k]; - // MTF [74..90) + // TLOB [74..90) — zeros on assembly; GpuTlob::forward overwrites post-gather (D.8 Plan 2 Task 6C) + for (int k = 0; k < SL_TLOB_DIM; k++) + out[SL_TLOB_START + k] = 0.0f; + + // MTF [90..106) for (int k = 0; k < SL_MTF_DIM; k++) out[SL_MTF_START + k] = mtf[k]; - // Portfolio base [90..98) + // Portfolio base [106..114) for (int k = 0; k < SL_PORTFOLIO_BASE_DIM; k++) out[SL_PORTFOLIO_START + k] = portfolio[k]; - // Plan/ISV [98..105) + // Plan/ISV [114..121) for (int k = 0; k < SL_PORTFOLIO_PLAN_DIM; k++) out[SL_PLAN_ISV_START + k] = plan_isv[k]; - // Padding [105..112) — 7 zeros for 8-alignment (D.6 Plan 2 Task 6A) + // Padding [121..128) — 7 zeros for 8-alignment (D.8 Plan 2 Task 6C) for (int k = 0; k < SL_PADDING_DIM; k++) out[SL_PADDING_START + k] = 0.0f; } diff --git a/crates/ml/src/cuda_pipeline/tlob_kernel.cu b/crates/ml/src/cuda_pipeline/tlob_kernel.cu new file mode 100644 index 000000000..981557db1 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/tlob_kernel.cu @@ -0,0 +1,198 @@ +// tlob_kernel.cu — TLOB attention kernels (D.8 Plan 2 Task 6C). +// +// Compresses OFI[32] → TLOB[16] using a single-head scaled-dot-product attention +// over the feature dimension, composed from existing cuBLAS Q/K/V GEMMs + this +// per-sample SDP kernel. +// +// Architecture per sample: +// 1. Q/K/V projections: cuBLAS SGEMM [16,32] × [32,B] → [16,B] (done in GpuTlob Rust host) +// 2. Scaled dot-product: attn_scores[16] = softmax(Q^T K / sqrt(16)), out[16] = V × attn +// 3. Output projection: cuBLAS SGEMM [16,16] × [16,B] → [16,B] (done in GpuTlob Rust host) +// 4. tlob_write_states: scatter [16,B] into states_buf at SL_TLOB_START +// +// All kernels follow zero-atomicAdd rule (per feedback_no_atomicadd.md). +// Per-sample arrays; host reduces when needed. + +#include "state_layout.cuh" + +// ── tlob_sdp_forward ───────────────────────────────────────────────────────── +// +// Scaled dot-product attention for a batch of samples over the TLOB feature dim. +// Each thread handles one sample b ∈ [0, B). +// +// Inputs (all column-major [D, B]): +// proj_q [SL_TLOB_DIM, B] — Q projection output +// proj_k [SL_TLOB_DIM, B] — K projection output +// proj_v [SL_TLOB_DIM, B] — V projection output +// Outputs: +// attn_out [SL_TLOB_DIM, B] — attention output (V weighted sum) +// attn_scores [SL_TLOB_DIM, B] — softmax weights (saved for backward) +// +// Grid: ceil(B/256), Block: 256. One thread per sample. +extern "C" __global__ void tlob_sdp_forward( + const float* __restrict__ proj_q, // [SL_TLOB_DIM, B] + const float* __restrict__ proj_k, // [SL_TLOB_DIM, B] + const float* __restrict__ proj_v, // [SL_TLOB_DIM, B] + float* __restrict__ attn_out, // [SL_TLOB_DIM, B] output + float* __restrict__ attn_scores, // [SL_TLOB_DIM, B] softmax weights saved for backward + int B +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B) return; + + const int D = SL_TLOB_DIM; + const float scale = 1.0f / sqrtf((float)D); + + // Read Q/K/V for this sample (column-major: element [d, b] = ptr[d * B + b]) + float q[SL_TLOB_DIM], k[SL_TLOB_DIM], v[SL_TLOB_DIM]; + for (int d = 0; d < D; d++) { + q[d] = proj_q[d * B + b]; + k[d] = proj_k[d * B + b]; + v[d] = proj_v[d * B + b]; + } + + // Single-head SDP: score[d] = q[d] * k[d] * scale (element-wise; feature-dim attention) + // This is a feature-gating attention: each dimension gates itself independently. + // Equivalent to a per-feature learned gate driven by the Q-K inner product. + float scores[SL_TLOB_DIM]; + float max_score = -1e38f; + for (int d = 0; d < D; d++) { + scores[d] = q[d] * k[d] * scale; + if (scores[d] > max_score) max_score = scores[d]; + } + + // Softmax (numerically stable) + float sum_exp = 0.0f; + for (int d = 0; d < D; d++) { + scores[d] = expf(scores[d] - max_score); + sum_exp += scores[d]; + } + float inv_sum = (sum_exp > 1e-12f) ? (1.0f / sum_exp) : 0.0f; + for (int d = 0; d < D; d++) { + scores[d] *= inv_sum; + } + + // Output = softmax_scores ⊙ V (element-wise weighted values) + for (int d = 0; d < D; d++) { + attn_out[d * B + b] = scores[d] * v[d]; + attn_scores[d * B + b] = scores[d]; // save for backward + } +} + +// ── tlob_sdp_backward ──────────────────────────────────────────────────────── +// +// Backward pass for TLOB SDP. +// Inputs: +// d_attn_out [SL_TLOB_DIM, B] — upstream gradient +// attn_scores [SL_TLOB_DIM, B] — saved softmax weights from forward +// proj_v [SL_TLOB_DIM, B] — saved V from forward +// proj_q [SL_TLOB_DIM, B] — saved Q from forward +// proj_k [SL_TLOB_DIM, B] — saved K from forward +// Outputs: +// d_proj_q [SL_TLOB_DIM, B] — gradient for Q projection +// d_proj_k [SL_TLOB_DIM, B] — gradient for K projection +// d_proj_v [SL_TLOB_DIM, B] — gradient for V projection +// +// Grid: ceil(B/256), Block: 256. One thread per sample. +extern "C" __global__ void tlob_sdp_backward( + const float* __restrict__ d_attn_out, // [SL_TLOB_DIM, B] + const float* __restrict__ attn_scores, // [SL_TLOB_DIM, B] + const float* __restrict__ proj_v, // [SL_TLOB_DIM, B] + const float* __restrict__ proj_q, // [SL_TLOB_DIM, B] + const float* __restrict__ proj_k, // [SL_TLOB_DIM, B] + float* __restrict__ d_proj_q, // [SL_TLOB_DIM, B] + float* __restrict__ d_proj_k, // [SL_TLOB_DIM, B] + float* __restrict__ d_proj_v, // [SL_TLOB_DIM, B] + int B, + float inv_batch +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B) return; + + const int D = SL_TLOB_DIM; + const float scale = 1.0f / sqrtf((float)D); + + // Read saved forward values for this sample + float s[SL_TLOB_DIM], v[SL_TLOB_DIM], q[SL_TLOB_DIM], k[SL_TLOB_DIM]; + float dout[SL_TLOB_DIM]; + for (int d = 0; d < D; d++) { + s[d] = attn_scores[d * B + b]; + v[d] = proj_v[d * B + b]; + q[d] = proj_q[d * B + b]; + k[d] = proj_k[d * B + b]; + dout[d] = d_attn_out[d * B + b]; + } + + // d_v[d] = dout[d] * s[d] (element-wise, no cross-dim dependency) + for (int d = 0; d < D; d++) { + d_proj_v[d * B + b] = dout[d] * s[d] * inv_batch; + } + + // d_scores[d] = dout[d] * v[d] + float ds[SL_TLOB_DIM]; + for (int d = 0; d < D; d++) { + ds[d] = dout[d] * v[d]; + } + + // Softmax backward: d_pre_softmax[d] = s[d] * (ds[d] - dot(s, ds)) + float dot_s_ds = 0.0f; + for (int d = 0; d < D; d++) dot_s_ds += s[d] * ds[d]; + float d_pre[SL_TLOB_DIM]; + for (int d = 0; d < D; d++) { + d_pre[d] = s[d] * (ds[d] - dot_s_ds); + } + + // pre-softmax = q[d]*k[d]*scale → d_q[d] = d_pre[d]*k[d]*scale, d_k[d] = d_pre[d]*q[d]*scale + for (int d = 0; d < D; d++) { + d_proj_q[d * B + b] = d_pre[d] * k[d] * scale * inv_batch; + d_proj_k[d * B + b] = d_pre[d] * q[d] * scale * inv_batch; + } +} + +// ── tlob_write_states ──────────────────────────────────────────────────────── +// +// Scatter TLOB features [SL_TLOB_DIM, B] (col-major) into states_buf +// at slot SL_TLOB_START for each sample. +// +// states_buf layout: [B, state_dim_padded] (row-major per sample, padded stride) +// +// Grid: ceil(B/256), Block: 256. +extern "C" __global__ void tlob_write_states( + const float* __restrict__ tlob_out, // [SL_TLOB_DIM, B] col-major + float* __restrict__ states_buf, // [B, state_dim_padded] row-major + int B, + int state_dim_padded +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B) return; + + float* out_row = states_buf + (long long)b * state_dim_padded; + for (int d = 0; d < SL_TLOB_DIM; d++) { + out_row[SL_TLOB_START + d] = tlob_out[d * B + b]; + } +} + +// ── tlob_read_grad_from_states ─────────────────────────────────────────────── +// +// Extract TLOB gradient slice from the states gradient buffer. +// d_states layout after trunk backward: [B, concat_dim] (row-major). +// TLOB sits at offset tlob_concat_off within concat (= bn_dim + TLOB_START - market_dim). +// +// Output: d_tlob_out [SL_TLOB_DIM, B] col-major (matching SDP backward input format). +// +// Grid: ceil(B/256), Block: 256. +extern "C" __global__ void tlob_read_grad_from_states( + const float* __restrict__ d_states, // [B, concat_dim] row-major + float* __restrict__ d_tlob_out, // [SL_TLOB_DIM, B] col-major + int B, + int concat_dim, + int tlob_concat_off // offset of TLOB within concat row +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B) return; + + const float* row = d_states + (long long)b * concat_dim; + for (int d = 0; d < SL_TLOB_DIM; d++) { + d_tlob_out[d * B + b] = row[tlob_concat_off + d]; + } +} diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 4cd57aefa..a8f17432f 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -27,13 +27,17 @@ use std::sync::Arc; use anyhow::Result; use ml_core::device::MlDevice; +use ml_core::state_layout; use tracing::info; use cudarc::driver::DevicePtr; use cudarc::driver::sys as cuda_sys; use crate::cuda_pipeline::gpu_attention::{GpuAttention, GpuAttentionConfig}; -use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer, TAU_EFF_INDEX}; +use crate::cuda_pipeline::gpu_tlob::GpuTlob; +use crate::cuda_pipeline::gpu_dqn_trainer::{ + GpuDqnTrainConfig, GpuDqnTrainer, TAU_EFF_INDEX, TLOB_REGIME_FOCUS_EMA_INDEX, +}; use crate::cuda_pipeline::gpu_her::{GpuHer, GpuHerConfig, parse_her_strategy}; use crate::cuda_pipeline::gpu_iql_trainer::{GpuIqlConfig, GpuIqlTrainer}; use crate::cuda_pipeline::gpu_iqn_head::{GpuIqnConfig, GpuIqnHead}; @@ -226,6 +230,10 @@ pub(crate) struct FusedTrainingCtx { /// GPU multi-head feature attention over h_s2 (post-graph, 1-step lag). /// When Some, applied after EMA update, before IQN, each training step. pub(crate) gpu_attention: Option, + /// TLOB attention module: OFI[32]→TLOB[16], trained end-to-end via DQN reward. + /// Forward runs after state gather (before cuBLAS forward). Backward runs after + /// trunk backward. D.8 Plan 2 Task 6C. + pub(crate) gpu_tlob: Option, // ── Phase 2: Aux multi-stream ── /// Dedicated CUDA stream for IQN aux training (parallel with attn_stream after IQL). iqn_stream: Option>, @@ -581,6 +589,15 @@ impl FusedTrainingCtx { } }; + // D.8 TLOB: initialize attention module over OFI features. + let gpu_tlob = match GpuTlob::new(Arc::clone(trainer.shared_cublas()), batch_size) { + Ok(tlob) => Some(tlob), + Err(e) => { + tracing::warn!("GpuTlob init failed, training continues without TLOB: {e}"); + None + } + }; + // Phase 2: Allocate aux multi-stream infrastructure (streams, workspaces, events). // Only allocate when IQN or Attention is present — parallel dispatch needs at least // one aux trainer to be worthwhile. @@ -743,6 +760,7 @@ impl FusedTrainingCtx { gpu_iqn, cvar_scales_buf: None, gpu_attention, + gpu_tlob, iqn_stream, attn_stream, iqn_workspace, @@ -1205,6 +1223,13 @@ impl FusedTrainingCtx { self.trainer.apply_spectral_norm(&self.online_dueling, &self.online_branching) .map_err(|e| anyhow::anyhow!("spectral norm: {e}"))?; + // D.8 TLOB forward: OFI[32]→TLOB[16] written into states_buf before trunk forward. + if let Some(ref mut tlob) = self.gpu_tlob { + let states = self.trainer.states_buf_mut(); + tlob.forward(states, self.batch_size) + .map_err(|e| anyhow::anyhow!("TLOB forward: {e}"))?; + } + // Phase 2: Forward (cuBLAS trunk + ISV + loss + backward) self.trainer.submit_forward_ops_main() .map_err(|e| anyhow::anyhow!("forward ops: {e}"))?; @@ -1220,7 +1245,26 @@ impl FusedTrainingCtx { self.trainer.submit_post_aux_ops(self.batch_size) .map_err(|e| anyhow::anyhow!("post_aux ops: {e}"))?; - // Phase 6: Mamba2 bwd + OFI embed bwd + pruning + grad_norm + Adam + ISV update + // Phase 6: TLOB bwd + Mamba2 bwd + OFI embed bwd + pruning + grad_norm + Adam + ISV update + // TLOB backward: reads TLOB-slice gradient from bn_d_concat_buf, updates W_Q/K/V/O. + if let Some(ref mut tlob) = self.gpu_tlob { + let concat_dim = self.trainer.bn_concat_dim(); + // TLOB starts at SL_TLOB_START=74 in the state. In the concat buffer (which is + // [bottleneck_dim | state[market_dim..STATE_DIM]]), TLOB offset = + // bottleneck_dim + (TLOB_START - market_dim) + let tlob_concat_off = self.trainer.config().bottleneck_dim + + (state_layout::TLOB_START + .saturating_sub(self.trainer.config().market_dim)); + let d_concat = self.trainer.bn_d_concat_buf() as *const _ as *mut _; + // SAFETY: bn_d_concat_buf is not borrowed mutably elsewhere during tlob.backward. + let d_concat_ref = unsafe { &*d_concat }; + tlob.backward(d_concat_ref, self.batch_size, concat_dim, tlob_concat_off) + .map_err(|e| anyhow::anyhow!("TLOB backward: {e}"))?; + tlob.adam_step( + self.trainer.config().lr, + self.trainer.config().max_grad_norm, + ).map_err(|e| anyhow::anyhow!("TLOB Adam: {e}"))?; + } self.mamba2_backward(self.batch_size)?; self.step_mamba2_adam()?; self.launch_ofi_embed_backward(self.batch_size)?; @@ -2478,6 +2522,11 @@ impl FusedTrainingCtx { // ── A3: LearningHealth signal writers / readers ─────────────────────── + /// Read a single f32 value from the ISV signals buffer at given index (pinned DtoH). + pub(crate) fn read_isv_signal_at(&self, index: usize) -> f32 { + self.trainer.read_isv_signal_at(index) + } + /// Write a single f32 value to ISV signals buffer at given index (pinned HtoD). pub(crate) fn write_isv_signal_at(&self, index: usize, value: f32) { self.trainer.write_isv_signal_at(index, value); diff --git a/crates/ml/src/trainers/dqn/trainer/metrics.rs b/crates/ml/src/trainers/dqn/trainer/metrics.rs index 4b024e28e..9c34c88ce 100644 --- a/crates/ml/src/trainers/dqn/trainer/metrics.rs +++ b/crates/ml/src/trainers/dqn/trainer/metrics.rs @@ -572,6 +572,32 @@ impl DQNTrainer { eval.set_isv_signals_ptr(fused.isv_signals_dev_ptr()); } + // D.8 TLOB val-path: sync training weights into the evaluator's TLOB instance. + // Creates an eval TLOB (sized to n_windows) on the evaluator's stream with a + // fresh PerStreamCublasHandles, then DtoD-copies params from the training TLOB. + // Called every epoch so val TLOB tracks online weights after each Adam step. + if let (Some(ref mut eval), Some(ref fused)) = (&mut self.gpu_evaluator, &self.fused_ctx) { + if let Some(ref train_tlob) = fused.gpu_tlob { + use crate::cuda_pipeline::gpu_tlob::GpuTlob; + use crate::cuda_pipeline::shared_cublas_handle::PerStreamCublasHandles; + use std::sync::Arc; + let eval_stream = Arc::clone(eval.stream()); + let n_win = eval.n_windows(); + match PerStreamCublasHandles::new(&eval_stream) + .and_then(|h| GpuTlob::new(Arc::new(h), n_win)) + { + Ok(eval_tlob) => { + if let Err(e) = eval.set_tlob_from_training(eval_tlob, train_tlob) { + tracing::warn!("Val TLOB weight sync failed: {e}"); + } + } + Err(e) => { + tracing::warn!("Val TLOB init failed: {e}"); + } + } + } + } + // ── Extract weight pointers + v_range, then split borrow for fused_ctx ── // DuelingWeightSet / BranchingWeightSet contain raw u64 device pointers. // We extract stable pointers to them, then take &mut fused_ctx for QValueProvider. diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 3683b6b49..0c6c8d6bd 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -21,7 +21,9 @@ use cudarc::driver::CudaSlice; use common::metrics::{questdb_sink, training_metrics}; use tracing::{debug, info, warn}; -use crate::cuda_pipeline::gpu_dqn_trainer::{EPOCH_IDX_INDEX, EPSILON_EFF_INDEX, GAMMA_DIR_EFF_INDEX}; +use crate::cuda_pipeline::gpu_dqn_trainer::{ + EPOCH_IDX_INDEX, EPSILON_EFF_INDEX, GAMMA_DIR_EFF_INDEX, TLOB_REGIME_FOCUS_EMA_INDEX, +}; use crate::dqn::logging::{log_epoch_start, log_epoch_end, log_training_progress}; use crate::dqn::target_update::convergence_half_life; use crate::evaluation::metrics::calculate_var_cvar; @@ -289,6 +291,18 @@ impl DQNTrainer { } } + // D.8 Plan 2 Task 6C: update TLOB regime focus EMA in ISV at epoch boundary. + // Reads mean-max attention weight from GpuTlob (DtoH, once/epoch outside hot path). + // EMA alpha=0.05 for slow smoothing. + if let Some(ref fused) = self.fused_ctx { + if let Some(ref tlob) = fused.gpu_tlob { + let mean_max = tlob.mean_max_attention_weight(fused.batch_size()); + let prev = fused.read_isv_signal_at(TLOB_REGIME_FOCUS_EMA_INDEX); + let ema = prev * 0.95 + mean_max * 0.05; + fused.write_isv_signal_at(TLOB_REGIME_FOCUS_EMA_INDEX, ema); + } + } + // Collect pending async validation from PREVIOUS epoch if let Some(val) = self.pending_val_loss.take() { // Store for this epoch's logging (one-epoch delay is acceptable) diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 4324d1787..049e1ae49 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -80,6 +80,7 @@ | `cuda_pipeline/gpu_iql_trainer.rs` (`GpuIqlTrainer`) | `fused_training.rs` (dual IQL instances) | Wired | IQL value network training | — | | `cuda_pipeline/gpu_iqn_head.rs` (`GpuIqnHead`) | `fused_training.rs` | Wired | IQN distributional head | — | | `cuda_pipeline/gpu_attention.rs` (`GpuAttention`) | `fused_training.rs` | Wired | Self-attention layer in DQN trunk | — | +| `cuda_pipeline/gpu_tlob.rs` (`GpuTlob`) | `fused_training.rs` (training: forward pre-trunk, backward+Adam Phase 6); `trainer/metrics.rs` (val: `set_tlob_from_training` per-epoch weight sync + `GpuBacktestEvaluator::submit_dqn_step_loop_cublas` per-gather TLOB forward) | Wired | D.8 Plan 2 Task 6C — OFI[32]→TLOB[16] single-head SDP attention, Xavier random init, trained end-to-end. ISV[60]=TLOB_REGIME_FOCUS_EMA written at epoch boundary. | — | | `cuda_pipeline/decision_transformer.rs` (`DecisionTransformer`) | `trainer/training_loop.rs` | Wired | Decision Transformer pre-training step | — | | `cuda_pipeline/learning_health.rs` (`LearningHealth`) | `fused_training.rs`, `trainer/training_loop.rs`, `trainer/constructor.rs`, `trainer/metrics.rs`, `meta_q_network.rs` (11 consumers) | Wired | Health-band tracking for adaptive LR / early stopping | — | | `cuda_pipeline/q_snapshot.rs` | `cuda_pipeline/mod.rs`, `gpu_dqn_trainer.rs` | Wired | Q-value snapshot for double-DQN target | — | @@ -119,7 +120,8 @@ | `branch_grad_balance_kernel.cu` | `gpu_dqn_trainer.rs`, called via `fused_training.rs::launch_branch_grad_balance` | Wired | Per-branch gradient L2-norm balancing | — | | `mamba2_temporal_kernel.cu` (`mamba2_scan_projected_fwd`, `mamba2_scan_projected_bwd`, `isv_temporal_route`, etc.) | `gpu_dqn_trainer.rs::mamba2_forward` + `mamba2_backward`, called in fused training loop (`adam_grad` child graph) | Wired (grad-check validated, Plan 2 Task 2) | Mamba2 selective SSM forward + backward (both paths implemented). D.1: kernel-level reference check + non-zero gradient propagation test confirm backward is not silently no-oping. | — | | `attention_kernel.cu` | `gpu_attention.rs` | Wired | Scaled dot-product attention forward | — | -| `attention_backward_kernel.cu` | `gpu_attention.rs` | Wired | Attention backward pass | — | +| `attention_backward_kernel.cu` | `gpu_attention.rs`, `gpu_tlob.rs` (reused for grad_norm+Adam kernels) | Wired | Attention backward pass | — | +| `tlob_kernel.cu` (`tlob_sdp_forward`, `tlob_sdp_backward`, `tlob_write_states`, `tlob_read_grad_from_states`) | `gpu_tlob.rs` | Wired | D.8 Plan 2 Task 6C — TLOB SDP forward/backward, state scatter/read kernels | — | | `training_guard_kernel.cu` (`training_guard_check_and_accumulate`) | `gpu_training_guard.rs` | Wired | NaN / guard check kernel | — | | `backtest_env_kernel.cu` (`backtest_env_step`) | `gpu_backtest_evaluator.rs` (`ENV_CUBIN`) | Wired | Backtest environment step | — | | `backtest_metrics_kernel.cu` | `gpu_backtest_evaluator.rs` (`METRICS_CUBIN`) | Wired | Backtest metrics reduction | — | diff --git a/docs/isv-slots.md b/docs/isv-slots.md index f554749be..3c1dedfff 100644 --- a/docs/isv-slots.md +++ b/docs/isv-slots.md @@ -15,7 +15,7 @@ at the head would displace those live signals and require shifting every upstrea literal in `experience_kernels.cu`. The fingerprint moves to the new tail each time new slots are appended to the bus. -**Current `ISV_TOTAL_DIM`:** 60 (Plan 1 + Plan 2 Task 1 C.1 + Plan 2 Task 3 D.2 per-branch gamma). Post-full DQN v2 rollout: 72. +**Current `ISV_TOTAL_DIM`:** 63 (Plan 1 + Plan 2 Task 1 C.1 + Plan 2 Task 3 D.2 per-branch gamma + Plan 2 Task 6C D.8 TLOB). Post-full DQN v2 rollout: 72. | Index | Name constant | Type | Producer | Consumers | Reset-category | Notes | |---|---|---|---|---|---|---| @@ -59,6 +59,8 @@ new slots are appended to the bus. | [55] | `Q_P95_MAG_INDEX` | f32 | q_quantile_reduce | update_eval_v_range | FoldReset | Magnitude-branch 95th-percentile Q EMA. Bootstrap: v_max. | | [56] | `Q_P95_ORD_INDEX` | f32 | q_quantile_reduce | update_eval_v_range | FoldReset | Order-branch 95th-percentile Q EMA. Bootstrap: v_max. | | [57] | `Q_P95_URG_INDEX` | f32 | q_quantile_reduce | update_eval_v_range | FoldReset | Urgency-branch 95th-percentile Q EMA. Bootstrap: v_max. | -| [58] | `ISV_LAYOUT_FINGERPRINT_LO_INDEX` | u32 bits (in f32) | construct | check_layout_fingerprint | SchemaContract | Low 32 bits of u64 FNV-1a structural hash. Fail-fast on mismatch — NOT a version number, no migration path. | -| [59] | `ISV_LAYOUT_FINGERPRINT_HI_INDEX` | u32 bits (in f32) | construct | check_layout_fingerprint | SchemaContract | High 32 bits of u64 FNV-1a structural hash. | -| [60..72) | (reserved for DQN v2) | | | | | Allocated incrementally by Plans 2-5 | +| [58..60) | (gap — fingerprint shifted to tail) | — | — | — | — | Previously [58..60); fingerprint promoted to [61..63) by Plan 2 Task 6C D.8 TLOB expansion. Slots unused; zero-filled. | +| [60] | `TLOB_REGIME_FOCUS_EMA_INDEX` | f32 | CPU training_loop (per-epoch) | ISV consumers (diagnostic) | FoldReset | Mean-max TLOB attention weight EMA (α=0.05); written at epoch boundary from `GpuTlob::mean_max_attention_weight`. Indicates OFI feature focus sharpness. | +| [61] | `ISV_LAYOUT_FINGERPRINT_LO_INDEX` | u32 bits (in f32) | construct | check_layout_fingerprint | SchemaContract | Low 32 bits of u64 FNV-1a structural hash. Fail-fast on mismatch — NOT a version number, no migration path. | +| [62] | `ISV_LAYOUT_FINGERPRINT_HI_INDEX` | u32 bits (in f32) | construct | check_layout_fingerprint | SchemaContract | High 32 bits of u64 FNV-1a structural hash. | +| [63..72) | (reserved for DQN v2) | | | | | Allocated incrementally by Plans 2-5 |