From 072ca456861b34c6cd07f189a192f550d07f97ea Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 23:24:19 +0200 Subject: [PATCH] perf(val): batched backtest_state_gather_chunk + chunk-batched val TLOB (kernel #1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit nsys profile of multi_fold_convergence on L40S identified backtest_state_gather as the #1 GPU consumer at 37.2% (596 ms / 60,588 calls — kernel launch latency dominated compute) and the per-step TLOB cuBLAS gemvx calls as #2 at 25% (242K calls). Both share the same per-step amplification: chunk_len=512 separate gather launches + 512 separate TLOB.forward calls (4 SGEMMs each) per chunk before any Q-values can be computed. This commit replaces the per-step gather + DtoD pattern with a single batched launch, and reuses the same chunked buffer for a single chunk-wide TLOB forward. Per-chunk launch reduction: from 2*chunk_len + chunk_len*7 to 1 + 7 for the gather+TLOB phase (4608 -> 8 with chunk_len=512, a 576x reduction). New kernel `backtest_state_gather_chunk` (experience_kernels.cu): - Writes [chunk_len, N, padded_sd] directly into chunked_states_buf - chunk_len * N threads, 1 thread per output row - Mathematically identical to per-step gather: portfolio_buf and plan_isv_buf are CONSTANT within a chunk (env_step + plan_state_isv update at chunk boundary only). Each thread reads independent feature offsets, no atomics, no reordering. Chunked val TLOB (gpu_backtest_evaluator.rs + metrics.rs): - Val TLOB instance now sized to DQN_BACKTEST_CHUNK_SIZE * n_windows via new GpuBacktestEvaluator::val_tlob_batch_size() helper. - submit_dqn_step_loop_cublas calls tlob.forward(chunked_states, batch) ONCE per chunk on the chunk-wide buffer instead of chunk_len times on states_buf. - Partial last chunks reuse the same buffers (forward(b) accepts any b <= construction_batch). Borrow restructure: - Removed top-of-function `let ch_states = self.chunked_states_buf.as_ref()?` binding (TLOB needs &mut). Replaced with per-chunk ch_states_base raw u64 device pointer extracted in tight scope, reused by Phase 1 (gather) and Phase 2+3 (compute_q_values_to + last_step_states_ptr). The pointer is stable across the chunk because the Option> does not reallocate. Per-step gather kernel `backtest_state_gather` retained unchanged for evaluate() / evaluate_ppo() / evaluate_supervised() paths that still need a single-step writer (closure-based callers with no chunked buffer). Audit doc dqn-gpu-hot-path-audit.md updated with Fix 18 entry per Invariant 7. Build: SQLX_OFFLINE=true cargo check -p ml --lib clean (12 warnings, baseline). Tests: cargo test -p ml --lib --no-run clean. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/experience_kernels.cu | 132 +++++++++++++ .../cuda_pipeline/gpu_backtest_evaluator.rs | 176 ++++++++++++++---- crates/ml/src/trainers/dqn/trainer/metrics.rs | 11 +- docs/dqn-gpu-hot-path-audit.md | 23 +++ 4 files changed, 306 insertions(+), 36 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 28a001193..13191ec7c 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -6893,6 +6893,138 @@ extern "C" __global__ void backtest_state_gather( out[k] = 0.0f; } +/* ================================================================== */ +/* Kernel: backtest_state_gather_chunk */ +/* */ +/* Chunked variant of backtest_state_gather: writes [chunk_len, N, PD]*/ +/* states directly into chunked_states_buf in a single launch. */ +/* */ +/* Within a chunk, portfolio_buf and plan_isv_buf are CONSTANT (env */ +/* updates them only at chunk boundary via env_batch_kernel + plan */ +/* state_isv kernel). So each (window, step_offset) thread reads the */ +/* same per-window portfolio + plan_isv and writes a unique state row.*/ +/* */ +/* Layout match: chunked_states_buf[step_offset * N + w] * PD. */ +/* This matches the per-step DtoD copy pattern that this kernel */ +/* replaces (each step's per-window gather copied row-major into the */ +/* chunked buffer at offset step_offset * N * PD). */ +/* */ +/* Thread mapping: 1D grid with tid = step_offset * N + w. */ +/* step_offset = tid / N (chunk-local step index) */ +/* w = tid % N (window index) */ +/* actual_step = start_step + step_offset */ +/* */ +/* Total threads = chunk_len * N. Block size 256 (matches single-step */ +/* kernel for occupancy parity). Per `feedback_no_atomicadd.md` no */ +/* atomics are used; each thread writes its own row independently. */ +/* ================================================================== */ +extern "C" __global__ void backtest_state_gather_chunk( + const float* __restrict__ features, /* [n_windows * max_len * feat_dim] */ + const float* __restrict__ portfolio, /* [n_windows * 8] live portfolio state */ + const float* __restrict__ plan_isv, /* [n_windows * SL_PORTFOLIO_PLAN_DIM=7] plan/ISV features */ + float* __restrict__ chunked_states_out, /* [chunk_len * n_windows * padded_sd] */ + int n_windows, + int max_len, + int feat_dim, /* = market_dim + ofi_dim */ + int start_step, + int chunk_len, + int padded_sd +) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int total_threads = n_windows * chunk_len; + if (tid >= total_threads) return; + + int step_offset = tid / n_windows; + int w = tid - step_offset * n_windows; + int current_step = start_step + step_offset; + + /* Defensive bound: caller guarantees chunk_end <= max_len, but keep + * a guard so an over-sized launch can never read OOB on an early- + * terminating window. The env-batch kernel handles per-window + * window_lens[w] termination — gather can produce stale reads for + * past-end steps without affecting downstream correctness (env step + * no-ops once done_buf[w] is set). */ + if (current_step >= max_len) return; + + int feat_base = (w * max_len + current_step) * feat_dim; + /* Output row = chunked_states_out[(step_offset * n_windows + w) * padded_sd] */ + int out_row = step_offset * n_windows + w; + int out_base = out_row * padded_sd; + float* out = chunked_states_out + out_base; + + /* Market features [0..42) */ + float market[SL_MARKET_DIM]; + for (int k = 0; k < SL_MARKET_DIM; k++) + market[k] = features[feat_base + k]; + + /* OFI [42..62) — appended after market features in feature buffer */ + int market_dim = feat_dim - SL_OFI_DIM; + float ofi[SL_OFI_DIM]; + for (int k = 0; k < SL_OFI_DIM; k++) + ofi[k] = features[feat_base + market_dim + k]; + + /* MTF [62..78) — compute from price history with 4 lookbacks */ + float mtf[SL_MTF_DIM]; + const int lookbacks[4] = {5, 15, 60, 240}; + int window_bar_base = w * max_len; + + for (int lb = 0; lb < 4; lb++) { + int N_lb = lookbacks[lb]; + int past_step = current_step - N_lb; + int slot = lb * 4; + + if (past_step >= 0) { + int now_off = (window_bar_base + current_step) * feat_dim; + int past_off = (window_bar_base + past_step) * feat_dim; + float close_cur = features[now_off]; + float close_past = features[past_off]; + + float ret = (close_past > 0.0f) ? (close_cur - close_past) / close_past : 0.0f; + mtf[slot + 0] = fmaxf(-10.0f, fminf(10.0f, ret * 100.0f)); + + float max_val = close_cur, min_val = close_cur; + float vol_sum = 0.0f; + int vol_count = 0; + for (int j = past_step; j <= current_step; j++) { + int j_off = (window_bar_base + j) * feat_dim; + float v = features[j_off]; + if (v > max_val) max_val = v; + if (v < min_val) min_val = v; + if (feat_dim > 4) { vol_sum += features[j_off + 4]; vol_count++; } + } + float range = (close_cur > 0.0f) ? (max_val - min_val) / close_cur : 0.0f; + mtf[slot + 1] = fmaxf(0.0f, fminf(10.0f, range * 100.0f)); + + float avg_vol = (vol_count > 0) ? vol_sum / (float)vol_count : 1.0f; + float cur_vol = (feat_dim > 4) ? features[now_off + 4] : 1.0f; + mtf[slot + 2] = fmaxf(0.0f, fminf(5.0f, (avg_vol > 0.0f) ? cur_vol / avg_vol : 1.0f)); + + float range_size = max_val - min_val; + mtf[slot + 3] = (range_size > 0.0f) ? (close_cur - min_val) / range_size : 0.5f; + } else { + mtf[slot + 0] = 0.0f; mtf[slot + 1] = 0.0f; + mtf[slot + 2] = 0.0f; mtf[slot + 3] = 0.0f; + } + } + + /* Portfolio base [78..86) — read pre-computed from buffer (constant within chunk) */ + float pf[SL_PORTFOLIO_BASE_DIM]; + for (int k = 0; k < SL_PORTFOLIO_BASE_DIM; k++) + pf[k] = portfolio[w * SL_PORTFOLIO_BASE_DIM + k]; + + /* Plan/ISV [86..92) — read pre-computed from buffer (constant within chunk) */ + float pisv[SL_PORTFOLIO_PLAN_DIM]; + for (int k = 0; k < SL_PORTFOLIO_PLAN_DIM; k++) + pisv[k] = plan_isv[w * SL_PORTFOLIO_PLAN_DIM + k]; + + /* Single shared assembly call — produces identical layout to training */ + assemble_state(out, market, ofi, mtf, pf, pisv); + + /* Zero-pad [SL_STATE_DIM .. padded_sd) for cuBLAS K-tile alignment */ + for (int k = SL_STATE_DIM; k < padded_sd; k++) + out[k] = 0.0f; +} + /* ================================================================== */ /* Kernel: scatter_intent_chunk */ /* ================================================================== */ diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index eaa19f21f..0ba032fcd 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -285,6 +285,15 @@ pub struct GpuBacktestEvaluator { env_batch_kernel: CudaFunction, metrics_kernel: CudaFunction, gather_kernel: CudaFunction, + /// Chunked gather: writes `[chunk_len, N, padded_sd]` rows into + /// `chunked_states_buf` in a single launch. Replaces the per-step + /// `gather_kernel` + DtoD-into-chunk loop in `submit_dqn_step_loop_cublas`. + /// Cuts per-chunk launches from `2 * chunk_len` (gather + DtoD) to **1**. + /// Loaded from the same `experience_kernels.cubin` as `gather_kernel`; + /// the per-step `gather_kernel` is retained for `evaluate()` / + /// `evaluate_ppo()` / `evaluate_supervised()` where the closure-based + /// caller still needs a single-step output. + gather_chunk_kernel: CudaFunction, // Uploaded data (read-only; persists across the step loop) prices_buf: CudaSlice, // [n_windows * max_len * 4] @@ -615,6 +624,9 @@ impl GpuBacktestEvaluator { let gather_kernel = gather_module .load_function("backtest_state_gather") .map_err(|e| MLError::ModelError(format!("backtest_state_gather load: {e}")))?; + let gather_chunk_kernel = gather_module + .load_function("backtest_state_gather_chunk") + .map_err(|e| MLError::ModelError(format!("backtest_state_gather_chunk load: {e}")))?; // ── Upload read-only data via mapped-pinned staging (no HtoD) ───── let prices_buf = super::mapped_pinned::clone_to_device_f32_via_pinned(&stream, &flat_prices) @@ -730,6 +742,7 @@ impl GpuBacktestEvaluator { env_batch_kernel, metrics_kernel, gather_kernel, + gather_chunk_kernel, prices_buf, features_buf, window_lens_buf, @@ -859,6 +872,20 @@ impl GpuBacktestEvaluator { self.n_windows } + /// Maximum batch size that the val TLOB instance must support. + /// + /// `submit_dqn_step_loop_cublas` runs TLOB once per chunk on + /// `[chunk_len * n_windows, state_dim_padded]` rows of `chunked_states_buf` + /// instead of once per step on `[n_windows, ..]` rows of `states_buf`. + /// The val TLOB constructor (`metrics.rs::reset_evaluator_for_validation`) + /// must allocate buffers sized for this maximum so the per-chunk forward + /// can proceed without reallocation. Partial last chunks (chunk_len < + /// DQN_BACKTEST_CHUNK_SIZE) reuse the same buffers — `forward()` accepts + /// any `batch_size <= construction_batch_size`. + pub(crate) fn val_tlob_batch_size(&self) -> usize { + DQN_BACKTEST_CHUNK_SIZE * self.n_windows + } + /// Sharpe annualization factor used by `compute_backtest_metrics` kernel. /// /// Returns the EXACT factor applied to `(mean/std)` inside the kernel @@ -1516,9 +1543,11 @@ impl GpuBacktestEvaluator { let ch_q_values = self.chunked_q_values.as_ref().ok_or_else(|| { MLError::ModelError("chunked_q_values unexpectedly None".to_owned()) })?; - let ch_states = self.chunked_states_buf.as_ref().ok_or_else(|| { - MLError::ModelError("chunked_states_buf unexpectedly None".to_owned()) - })?; + // Note: `chunked_states_buf` is NOT borrowed up-front. Each chunk + // iteration re-derives a fresh borrow after the optional TLOB call + // (which needs `&mut self.chunked_states_buf`). The previous shared + // borrow held across the loop blocked the TLOB mutable borrow + // required by the chunk-batched TLOB refactor. let ch_actions = self.chunked_actions_buf.as_ref().ok_or_else(|| { MLError::ModelError("chunked_actions_buf unexpectedly None".to_owned()) })?; @@ -1560,41 +1589,58 @@ impl GpuBacktestEvaluator { let batch = n * chunk_len; // total rows for this chunk's forward pass let chunk_idx = chunk_start / DQN_BACKTEST_CHUNK_SIZE; - // ── Phase 1: Gather states for all steps in the chunk ──────────── + // ── Phase 1: Gather states for the entire chunk in a single launch ──── // - // Each launch_gather() writes [n_windows, state_dim] into self.states_buf. - // We then DtoD copy that into the chunked_states_buf at the right offset. - let (ch_states_base, _ch_states_guard) = ch_states.device_ptr(&self.stream); + // `backtest_state_gather_chunk` writes [chunk_len, n_windows, padded_sd] + // directly into chunked_states_buf with one kernel launch + zero DtoD + // copies (replaces chunk_len gather launches + chunk_len DtoD copies). + // + // Within a chunk portfolio_buf and plan_isv_buf are CONSTANT — env_step + // updates them only at chunk boundary (Phase 5 + Phase 6). So the per- + // step gather sees identical inputs and the batched gather is bit- + // identical to the prior per-step + DtoD pattern (modulo float + // reorder, of which there is none — every thread writes its own row + // from independent feature reads). + // + // After gather, TLOB runs ONCE on `chunked_states_buf` with batch + // = chunk_len * n_windows (replacing chunk_len separate forward + // calls each at batch = n_windows). The val TLOB instance is sized + // to `val_tlob_batch_size()` = DQN_BACKTEST_CHUNK_SIZE * n_windows + // by `metrics.rs::reset_evaluator_for_validation`; partial last- + // chunks (chunk_len < DQN_BACKTEST_CHUNK_SIZE) reuse the same buffers. + // + // Borrow scope: ch_states_base is computed in a tight scope so the + // immutable borrow drops before we take &mut self.chunked_states_buf + // for TLOB below. The raw u64 device pointer survives the borrow + // (it's just a number copied into a stack local). + let ch_states_base = { + let ch_states = self.chunked_states_buf.as_ref().ok_or_else(|| { + MLError::ModelError("chunked_states_buf unexpectedly None".to_owned()) + })?; + let (ptr, _guard) = ch_states.device_ptr(&self.stream); + ptr + }; - for step_offset in 0..chunk_len { - let step = chunk_start + step_offset; - self.launch_gather(step)?; + self.launch_gather_chunk(ch_states_base, chunk_start, chunk_len)?; - // 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); - unsafe { - cudarc::driver::result::memcpy_dtod_async( - ch_states_base + dst_byte_offset, - src_ptr, - state_row_bytes, - self.stream.cu_stream(), - ) + // D.8 TLOB val-path: single forward over the entire chunk at + // batch = chunk_len * n_windows. TLOB reads OFI[42..74) and writes + // TLOB[74..90) for each row independently — safe to batch. + // + // The kernel writes only the TLOB slot [SL_TLOB_START..+SL_TLOB_DIM); + // other slots remain as the chunked gather wrote them. + if self.tlob.is_some() { + let chunked_states = self.chunked_states_buf.as_mut().ok_or_else(|| { + MLError::ModelError("chunked_states_buf unexpectedly None".to_owned()) + })?; + let tlob = self.tlob.as_mut().expect("checked is_some above"); + tlob.forward(chunked_states, batch) .map_err(|e| MLError::ModelError(format!( - "chunked gather DtoD step {step}: {e}" + "val TLOB forward chunk_start={chunk_start} batch={batch}: {e}" )))?; - } } - // Sync after Phase 1 gather — catch any gather/DtoD errors + // Sync after Phase 1 gather — catch any gather errors if chunk_idx == 0 { self.stream.synchronize().map_err(|e| MLError::ModelError(format!( "backtest chunk {chunk_idx}: sync after gather phase: {e}" @@ -1602,7 +1648,10 @@ impl GpuBacktestEvaluator { } // ── Phase 2+3: Q-values via QValueProvider or internal cuBLAS ── - let states_ptr = raw_f32_ptr(ch_states, &self.stream); + // ch_states_base is the device pointer to chunked_states_buf base — + // computed in Phase 1 and stable across the chunk (the Option does + // not reallocate). Reuse it instead of taking another borrow. + let states_ptr = ch_states_base; let q_out_ptr = raw_f32_ptr(ch_q_values, &self.stream); match q_provider { Some(ref mut qp) => { @@ -2200,6 +2249,69 @@ impl GpuBacktestEvaluator { Ok(()) } + /// Launch `backtest_state_gather_chunk` — batches an entire chunk of + /// `chunk_len` steps × `n_windows` rows into a single kernel launch, + /// writing directly into `chunked_states_out` in `[chunk_len, N, padded_sd]` + /// layout. Replaces the prior per-step `launch_gather` + DtoD copy loop. + /// + /// Within a chunk `portfolio_buf` and `plan_isv_buf` are constant — the + /// env-batch kernel updates them only at chunk boundary. So the gather + /// for all `chunk_len` steps reads identical per-window portfolio / + /// plan_isv inputs and writes a unique state row per (window, step_offset) + /// pair. See the docstring on the kernel itself for layout details. + /// + /// `chunked_states_dev_ptr` must be the device pointer to the start of + /// the chunk's portion of `chunked_states_buf` (i.e. base; the kernel + /// writes rows `0..chunk_len * n_windows`). Caller is responsible for + /// `chunk_len <= DQN_BACKTEST_CHUNK_SIZE` (the buffer was sized for + /// `cn = DQN_BACKTEST_CHUNK_SIZE * n_windows` rows). + fn launch_gather_chunk( + &self, + chunked_states_dev_ptr: u64, + start_step: usize, + chunk_len: usize, + ) -> Result<(), MLError> { + let total_threads = (self.n_windows * chunk_len) as u32; + // Block size 256 — same as single-step gather, well within L40S + // 1024-thread per-block limit. Grid sized to cover all + // (window, step_offset) pairs. + let blocks = (total_threads + 255) / 256; + let launch_cfg = LaunchConfig { + grid_dim: (blocks.max(1), 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let n_windows_i32 = self.n_windows as i32; + let max_len_i32 = self.max_len as i32; + let feat_dim_i32 = self.feature_dim as i32; + let start_step_i32 = start_step as i32; + let chunk_len_i32 = chunk_len as i32; + let padded_sd_i32 = ml_core::state_layout::STATE_DIM_PADDED as i32; + + // Safety: argument order matches `backtest_state_gather_chunk` signature exactly: + // features, portfolio, plan_isv, chunked_states_out (raw u64 ptr), + // n_windows, max_len, feat_dim, start_step, chunk_len, padded_sd + unsafe { + self.stream + .launch_builder(&self.gather_chunk_kernel) + .arg(&self.features_buf) + .arg(&self.portfolio_buf) + .arg(&self.plan_isv_buf) + .arg(&chunked_states_dev_ptr) + .arg(&n_windows_i32) + .arg(&max_len_i32) + .arg(&feat_dim_i32) + .arg(&start_step_i32) + .arg(&chunk_len_i32) + .arg(&padded_sd_i32) + .launch(launch_cfg) + .map_err(|e| MLError::ModelError(format!( + "backtest_state_gather_chunk launch start_step={start_step} chunk_len={chunk_len}: {e}" + )))?; + } + Ok(()) + } + /// Launch `backtest_env_step` kernel for a given step on `self.stream`. fn launch_env_step(&self, step: usize) -> Result<(), MLError> { // Shared memory: 256 threads × PORTFOLIO_STATE_SIZE(8) f32 × 4 bytes = 8 KB diff --git a/crates/ml/src/trainers/dqn/trainer/metrics.rs b/crates/ml/src/trainers/dqn/trainer/metrics.rs index fb4cedab5..2af125397 100644 --- a/crates/ml/src/trainers/dqn/trainer/metrics.rs +++ b/crates/ml/src/trainers/dqn/trainer/metrics.rs @@ -629,8 +629,11 @@ impl DQNTrainer { } // 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. + // Sizes the eval TLOB to `val_tlob_batch_size()` = DQN_BACKTEST_CHUNK_SIZE * + // n_windows so `submit_dqn_step_loop_cublas` can run TLOB.forward once per + // chunk on the chunked states buffer, replacing the prior per-step call + // on states_buf at batch=n_windows. Cuts TLOB launches by chunk_len (~512x + // fewer cuBLAS SGEMMs per chunk; matches the chunked-gather refactor). // 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) { use crate::cuda_pipeline::gpu_tlob::GpuTlob; @@ -638,14 +641,14 @@ impl DQNTrainer { use std::sync::Arc; let train_tlob = &fused.gpu_tlob; let eval_stream = Arc::clone(eval.stream()); - let n_win = eval.n_windows(); + let val_tlob_batch = eval.val_tlob_batch_size(); // No graceful-degrade: TLOB is on the val production path. If // init or weight sync fails we surface the error so the // train/val state-distribution parity is preserved. Per // `feedback_no_hiding.md`: wire up or surface, never suppress. let h = PerStreamCublasHandles::new(&eval_stream) .map_err(|e| anyhow::anyhow!("Val TLOB shared_cublas: {e:?}"))?; - let eval_tlob = GpuTlob::new(Arc::new(h), n_win) + let eval_tlob = GpuTlob::new(Arc::new(h), val_tlob_batch) .map_err(|e| anyhow::anyhow!("Val TLOB init: {e:?}"))?; eval.set_tlob_from_training(eval_tlob, train_tlob) .map_err(|e| anyhow::anyhow!("Val TLOB weight sync: {e:?}"))?; diff --git a/docs/dqn-gpu-hot-path-audit.md b/docs/dqn-gpu-hot-path-audit.md index 9f9dcbe46..675591a0c 100644 --- a/docs/dqn-gpu-hot-path-audit.md +++ b/docs/dqn-gpu-hot-path-audit.md @@ -186,3 +186,26 @@ Final state of the HtoD migration sequence: - `:11330` `denoise_params_host` (1800 floats: 2-step diffusion Q-refinement Xavier init for W1[24,24] + W2[12,24]). - `:11476` `qlstm_weights_host` (528 floats: QLSTM Xavier init for 6 gates × 11 input × 8 head). All sites use `mapped_pinned::upload_f32_via_pinned` (canonical mapped-pinned + DtoD staging helper). The helper returns `Result<_, String>` whereas the ctor returns `Result<_, MLError>`, so each site wraps the error via `.map_err(|e| MLError::ModelError(format!(" upload via pinned: {e}")))`. Site labels preserved so backtraces remain readable. + +### Fix 18 (perf) — Eliminate per-step val gather DtoD copies via chunked gather kernel (2026-04-28) +`experience_kernels.cu` + `gpu_backtest_evaluator.rs` + `metrics.rs`: nsys profile of `multi_fold_convergence` on L40S identified `backtest_state_gather` as the #1 GPU-time consumer at 37.2% (596 ms / 60,588 calls — kernel launch latency dominated compute). The per-step pattern was: + +``` +for step_offset in 0..chunk_len { + launch_gather(step) -> states_buf[N, padded_sd] (1 launch) + if tlob: tlob.forward(states_buf, n) (4 cuBLAS SGEMMs + 3 kernel launches) + memcpy_dtod_async(chunked_states_buf[step_offset*N..], states_buf, ...) (1 DtoD) +} +``` + +Per chunk (DQN_BACKTEST_CHUNK_SIZE=512): 512 gather launches + 512 DtoD copies + 512 × 7 TLOB launches = 4608 kernel launches per chunk before forward Q-values can begin. + +New kernel `backtest_state_gather_chunk` writes `[chunk_len, N, padded_sd]` directly into `chunked_states_buf` in a SINGLE launch (`chunk_len * N` threads, 1 thread per output row). Within a chunk `portfolio_buf` and `plan_isv_buf` are CONSTANT — env_step (Phase 5) and `backtest_plan_state_isv` (Phase 6) update them only at chunk boundary. So the batched gather is mathematically identical to the prior per-step + DtoD pattern; each thread reads its own (window, step_offset) row from independent feature offsets and writes its own state row. No reordering, no atomics, no shmem races. + +Same change enables the val TLOB instance to run ONCE per chunk on the chunked buffer at batch = `chunk_len * N` instead of `chunk_len` separate forward calls each at batch = `N`. Val TLOB construction in `metrics.rs::reset_evaluator_for_validation` resized via new `GpuBacktestEvaluator::val_tlob_batch_size()` to `DQN_BACKTEST_CHUNK_SIZE * n_windows`; partial last chunks reuse the same buffers (`forward(b)` accepts any `b <= construction_batch`). + +Per-chunk launch reduction: from `2 * chunk_len + chunk_len * 7` to `1 + 7` for the gather+TLOB phase. With 512-step chunks that's 4608 → 8, a 576x reduction. Across the full smoke run (60K gather calls in profile), expected wall-clock saving on the gather/TLOB hot path is the bulk of the 37% + 6.5% + 5.2% + 4.8% = ~53.7% of GPU time the per-call cluster occupied. + +Borrow restructure: the prior single-shot `let ch_states = self.chunked_states_buf.as_ref()?` at top of `submit_dqn_step_loop_cublas` was removed because TLOB now needs `&mut self.chunked_states_buf`. Replaced with per-chunk `ch_states_base` (raw u64 device pointer extracted in a tight scope), reused for both `launch_gather_chunk` and the Phase 2+3 `compute_q_values_to` call. The pointer is stable across the chunk because the `Option>` does not reallocate. + +No new HtoD/DtoH paths introduced. Strictly removes 512 per-chunk DtoD copies. Per `feedback_no_partial_refactor.md`: kernel ABI is new (separate function name), prior `backtest_state_gather` retained for `evaluate()` / `evaluate_ppo()` / `evaluate_supervised()` paths that still need a single-step writer.