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

perf(htod): migrate gpu_dqn_trainer.rs residual COLD ctor sites to mapped-pinned (8 sites)

Browse Source 
Residual COLD ctor sites missed by the original audit are now migrated
off explicit HtoD per `feedback_no_htod_htoh_only_mapped_pinned.md`:

- `:9922` sel_clip_buf 1-element init (sigmoid head clip-norm seed)
- `:10075-10078` spectral norm init_u_s1/v_s1/u_s2/v_s2 (4 calls)
- `:10095-10096` `alloc_spec_pair!` macro body — both u/v init uploads
  inside the macro now go through `upload_f32_via_pinned`; the macro's
  `$lbl_u`/`$lbl_v` are reused in the error-message format so per-pair
  failures stay diagnostically distinct
- `:11276` graph_params_host (60 floats: cross-branch graph message
  passing weights)
- `:11330` denoise_params_host (1800 floats: 2-step diffusion Q-refinement)
- `:11476` qlstm_weights_host (528 floats: QLSTM Xavier init)

All sites use `mapped_pinned::upload_f32_via_pinned` (the canonical
mapped-pinned + DtoD staging helper). The helper returns
`Result<_, String>` whereas this constructor 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.

docs/dqn-gpu-hot-path-audit.md updated with Fix 14 entry.

cargo check -p ml --lib clean at 12 warnings.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-28 21:22:57 +02:00
parent 45879f3901
commit 24afb2baf0
2 changed files with 30 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -9919,7 +9919,8 @@ impl GpuDqnTrainer {
// Fixed clip norm of 1.0 — selectivity gate is a small sigmoid head, clip at 1.0 norm.
let mut sel_clip_buf = stream.alloc_zeros::<f32>(1)
.map_err(|e| MLError::ModelError(format!("alloc sel_clip_buf: {e}")))?;
super::htod_f32(&stream, &[1.0_f32], &mut sel_clip_buf)?;
super::mapped_pinned::upload_f32_via_pinned(&stream, &[1.0_f32], &mut sel_clip_buf)
.map_err(|e| MLError::ModelError(format!("sel_clip_buf upload via pinned: {e}")))?;
// Selectivity Adam step counter — pinned device-mapped (no per-step HtoD copy)
let sel_t_pinned: *mut i32 = unsafe {
let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
@@ -10072,10 +10073,14 @@ impl GpuDqnTrainer {
let mut spec_v_s1 = alloc_f32(&stream, ml_core::state_layout::STATE_DIM, "spec_v_s1")?;
let mut spec_u_s2 = alloc_f32(&stream, config.shared_h2, "spec_u_s2")?;
let mut spec_v_s2 = alloc_f32(&stream, config.shared_h1, "spec_v_s2")?;
super::htod_f32(&stream, &init_u_s1, &mut spec_u_s1)?;
super::htod_f32(&stream, &init_v_s1, &mut spec_v_s1)?;
super::htod_f32(&stream, &init_u_s2, &mut spec_u_s2)?;
super::htod_f32(&stream, &init_v_s2, &mut spec_v_s2)?;
super::mapped_pinned::upload_f32_via_pinned(&stream, &init_u_s1, &mut spec_u_s1)
.map_err(|e| MLError::ModelError(format!("spec_u_s1 upload via pinned: {e}")))?;
super::mapped_pinned::upload_f32_via_pinned(&stream, &init_v_s1, &mut spec_v_s1)
.map_err(|e| MLError::ModelError(format!("spec_v_s1 upload via pinned: {e}")))?;
super::mapped_pinned::upload_f32_via_pinned(&stream, &init_u_s2, &mut spec_u_s2)
.map_err(|e| MLError::ModelError(format!("spec_u_s2 upload via pinned: {e}")))?;
super::mapped_pinned::upload_f32_via_pinned(&stream, &init_v_s2, &mut spec_v_s2)
.map_err(|e| MLError::ModelError(format!("spec_v_s2 upload via pinned: {e}")))?;
// Head spectral norm vectors: value head, adv/branch heads
let vh = config.value_h;
@@ -10092,8 +10097,10 @@ impl GpuDqnTrainer {
let mut $v_name = alloc_f32(&stream, $v_n, $lbl_v)?;
let init_u = rand_unit($u_n, &mut rng_state);
let init_v = rand_unit($v_n, &mut rng_state);
super::htod_f32(&stream, &init_u, &mut $u_name)?;
super::htod_f32(&stream, &init_v, &mut $v_name)?;
super::mapped_pinned::upload_f32_via_pinned(&stream, &init_u, &mut $u_name)
.map_err(|e| MLError::ModelError(format!("{} upload via pinned: {e}", $lbl_u)))?;
super::mapped_pinned::upload_f32_via_pinned(&stream, &init_v, &mut $v_name)
.map_err(|e| MLError::ModelError(format!("{} upload via pinned: {e}", $lbl_v)))?;
};
}
@@ -11273,7 +11280,8 @@ impl GpuDqnTrainer {
}
let mut graph_params = stream.alloc_zeros::<f32>(60)
.map_err(|e| MLError::ModelError(format!("alloc graph_params: {e}")))?;
super::htod_f32(&stream, &graph_params_host, &mut graph_params)?;
super::mapped_pinned::upload_f32_via_pinned(&stream, &graph_params_host, &mut graph_params)
.map_err(|e| MLError::ModelError(format!("graph_params upload via pinned: {e}")))?;
info!("GpuDqnTrainer: branch_graph_message_pass kernel loaded (60 params, near-identity W_msg init)");
// ── Diffusion Q-refinement ────────────────────────────────────────────────
@@ -11327,7 +11335,8 @@ impl GpuDqnTrainer {
}
let mut denoise_params = stream.alloc_zeros::<f32>(DENOISE_TOTAL_PARAMS)
.map_err(|e| MLError::ModelError(format!("alloc denoise_params: {e}")))?;
super::htod_f32(&stream, &denoise_params_host, &mut denoise_params)?;
super::mapped_pinned::upload_f32_via_pinned(&stream, &denoise_params_host, &mut denoise_params)
.map_err(|e| MLError::ModelError(format!("denoise_params upload via pinned: {e}")))?;
let cpbi_module_denoise = stream.context().load_cubin(EXPECTED_Q_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("cpbi cubin (denoise): {e}")))?;
let q_denoise_kernel = cpbi_module_denoise.load_function("q_denoise_step")
@@ -11473,7 +11482,8 @@ impl GpuDqnTrainer {
}
let mut qlstm_weights = stream.alloc_zeros::<f32>(QLSTM_WEIGHT_COUNT)
.map_err(|e| MLError::ModelError(format!("alloc qlstm_weights: {e}")))?;
super::htod_f32(&stream, &qlstm_weights_host, &mut qlstm_weights)?;
super::mapped_pinned::upload_f32_via_pinned(&stream, &qlstm_weights_host, &mut qlstm_weights)
.map_err(|e| MLError::ModelError(format!("qlstm_weights upload via pinned: {e}")))?;
let qlstm_c = stream.alloc_zeros::<f32>(QLSTM_HEAD_DIM * QLSTM_HEAD_DIM) // [64]
.map_err(|e| MLError::ModelError(format!("alloc qlstm_c: {e}")))?;
let qlstm_n = stream.alloc_zeros::<f32>(QLSTM_HEAD_DIM) // [8]

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

@@ -159,3 +159,13 @@ Out-of-scope (remaining host-side stream syncs in eval, not part of this fix): `
- Constructor (cold): `portfolio_states` init (htod_f32) → `upload_f32_via_pinned`; `rng_states` init (memcpy_htod) → `upload_u32_via_pinned`.
- `collect_experiences` (warm but per-collect): `episode_starts_buf` (memcpy_htod) → `upload_i32_via_pinned`; `barrier_config` (htod_f32) → `upload_f32_via_pinned`.
- `reset_episodes` (warm, per-epoch): replaced the `Vec<f32>` / `Vec<u32>` host staging fields with persistent `MappedF32Buffer` / `MappedU32Buffer` allocated once at construction. Resets fill via `host_slice_mut()` (zero-allocation, zero-HtoD); a single async DtoD seeds the persistent device-resident `portfolio_states` / `rng_states` (which the kernel mutates each step). Adds `host_slice_mut()` accessor on both mapped-pinned types.
### Fix 14 — Migrate `gpu_dqn_trainer.rs` residual COLD ctor sites (8 sites, 2026-04-28)
`gpu_dqn_trainer.rs`: residual COLD ctor sites missed by the original Fix-1..13 audit are now migrated off explicit `super::htod_f32` per `feedback_no_htod_htoh_only_mapped_pinned.md`.
- `:9922` `sel_clip_buf` 1-element init — selectivity gate clip-norm seed (1.0).
- `:10075-10078` spectral-norm `init_u_s1` / `init_v_s1` / `init_u_s2` / `init_v_s2` (4 separate uploads of trunk-encoder spectral u/v vectors).
- `:10095-10096` `alloc_spec_pair!` macro body — both u/v init uploads inside the macro now go through `upload_f32_via_pinned`. The macro's `$lbl_u` / `$lbl_v` are reused in the error-message format so per-pair failures (`spec_u_v1`, `spec_v_v1`, `spec_u_a1`, …) stay diagnostically distinct.
- `:11276` `graph_params_host` (60 floats: cross-branch graph message passing weights, 4 edges × 15 params each, near-identity W_msg init).
- `: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.