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

perf(val): batched backtest_state_gather_chunk + chunk-batched val TLOB (kernel #1)

Browse Source 
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<CudaSlice<f32>> 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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-28 23:24:19 +02:00
parent 0d630c799e
commit 072ca45686
4 changed files with 306 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
docs/dqn-gpu-hot-path-audit.md
View File

@@ -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!("<site> 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<CudaSlice<f32>>` 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.