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

fix(sp6): IQN τ buffers — MappedF32Buffer per feedback_no_htod_htoh_only_mapped_pinned

Browse Source 
Pearl 5's online_taus/target_taus/cos_features were declared as
CudaSlice<f32> (device-only), populated via upload_f32_via_pinned
which does a DtoD copy from a separate mapped-pinned staging buffer.
The DtoD inside CUDA Graph capture triggers
CUDA_ERROR_STREAM_CAPTURE_INVALIDATED and the 'continuing ungraphed'
fallback observed in smoke-test-hhr5q.

This violates feedback_no_htod_htoh_only_mapped_pinned: the rule is
mapped-pinned (cuMemHostAlloc DEVICEMAP) for ALL CPU↔GPU paths. No
DtoD copies, no HtoD copies, no exceptions.

Fix: convert all 3 buffers (online_taus, target_taus, cos_features)
to MappedF32Buffer per-branch [MappedF32Buffer; 4] arrays. Host writes
go directly to host_ptr; IQN kernel reads dev_ptr of the same memory
— no copy step at all. The mem::swap pattern is replaced with pure
selection: activate_branch_taus sets active_branch_idx; kernel launch
sites index online_taus_per_branch[active_branch_idx].dev_ptr.

Eliminates upload_f32_via_pinned calls for these buffers entirely.

Refresh becomes a host write to mapped-pinned host_ptr at fold
boundary; subsequent kernel launches see the write through the
mapped-pinned coherence guarantee after stream sync.

cargo check + cargo build --release + cargo test --lib (sp4 sp5
state_reset_registry: 13/13) all clean. Sanity grep for
upload_f32_via_pinned in gpu_iqn_head.rs returns zero.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-02 09:39:13 +02:00
parent b861567890
commit facbf76eb5
2 changed files with 106 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files

242
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -318,21 +318,20 @@ pub struct GpuIqnHead {
bias_grad_num_blocks: usize, bias_grad_num_blocks: usize,
// ── Per-step buffers ───────────────────────────────────────────── // ── Per-step buffers ─────────────────────────────────────────────
/// Pre-sampled τ values for online network [B, N] /// Per-branch online τ values [4][B, N] — mapped pinned (cuMemHostAlloc DEVICEMAP).
online_taus: CudaSlice<f32>, /// Host writes via host_ptr; IQN kernel reads via dev_ptr. No DtoD copy.
/// Pre-sampled τ values for target network [B, N] /// Indexed by `active_branch_idx` at each kernel launch site.
target_taus: CudaSlice<f32>, online_taus_per_branch: [super::mapped_pinned::MappedF32Buffer; 4],
/// Precomputed cosine features [D, N] — cos(π·(d+1)·τ_i) /// Per-branch target τ values [4][B, N] — mapped pinned.
/// Fixed at construction (τ are midpoints). Column-major. target_taus_per_branch: [super::mapped_pinned::MappedF32Buffer; 4],
cos_features: CudaSlice<f32>, /// Per-branch precomputed cosine features [4][D, N] — mapped pinned.
/// SP6 Pearl 5: per-branch τ tensors [B, N], one slab per branch (4 total). /// cos_feat[dim + q*D] = cos(π·(d+1)·τ_q). Column-major.
/// Populated by `refresh_taus_for_branch`; activated via `activate_branch_taus` cos_features_per_branch: [super::mapped_pinned::MappedF32Buffer; 4],
/// before each per-branch IQN forward pass. Same allocation size as `online_taus`. /// Index of the branch whose τ buffers are currently active for the IQN
online_taus_branch: [CudaSlice<f32>; 4], /// forward/backward pass. Set by `activate_branch_taus(branch_idx)`,
/// SP6 Pearl 5: per-branch target τ tensors [B, N], one slab per branch. /// reset to 0 by `deactivate_branch_taus`. Kernel launch sites index
target_taus_branch: [CudaSlice<f32>; 4], /// the per-branch arrays with this value.
/// SP6 Pearl 5: per-branch cosine features [D, N], one slab per branch. active_branch_idx: usize,
cos_features_branch: [CudaSlice<f32>; 4],
/// Branch actions decoded from flat actions [B, 4] /// Branch actions decoded from flat actions [B, 4]
branch_actions: CudaSlice<i32>, branch_actions: CudaSlice<i32>,
/// Target h_s2 computed from next_states + target trunk weights [B, H] /// Target h_s2 computed from next_states + target trunk weights [B, H]
@@ -512,50 +511,57 @@ impl GpuIqnHead {
// ── Per-step buffers ──────────────────────────────────────────── // ── Per-step buffers ────────────────────────────────────────────
// Plan 4 Task 3 (E.3): fixed-τ {0.05, 0.25, 0.50, 0.75, 0.95}. // Plan 4 Task 3 (E.3): fixed-τ {0.05, 0.25, 0.50, 0.75, 0.95}.
// Pre-edit used the QR-DQN midpoints τ_i = (2i+1)/(2N) over 32 // All 4 per-branch slabs are initialised from FIXED_TAUS at construction.
// quantiles, which is the same uniform-grid mean estimator but // refresh_taus_for_branch / refresh_taus_from_isv write updated values
// randomized via Philox at run time. Fixing the τ values lets the // directly to host_ptr (no copy, no DtoD) per
// 5 quantile heads be interpretable / consistent across batches and // feedback_no_htod_htoh_only_mapped_pinned.
// enables the median-Q action-ranking switch in `iqn_forward_kernel`. assert_eq!(
// Both online and target IQN read this same static buffer — the n, FIXED_TAUS.len(),
// off-median heads are diagnostic only (ISV[99..103) producers) and "GpuIqnConfig::num_quantiles ({n}) must equal FIXED_TAUS.len() \
// must agree on the fixed τ between online and target so the ({}) — kernel-side IQN_NUM_QUANTILES is sized to FIXED_TAUS \
// quantile-Huber regression learns a stable distributional surface. in lockstep (Plan 4 Task 3, E.3)",
let online_taus; FIXED_TAUS.len(),
let target_taus; );
let cos_features; let fixed: &[f32] = &FIXED_TAUS;
{
assert_eq!(
n, FIXED_TAUS.len(),
"GpuIqnConfig::num_quantiles ({n}) must equal FIXED_TAUS.len() \
({}) — kernel-side IQN_NUM_QUANTILES is sized to FIXED_TAUS \
in lockstep (Plan 4 Task 3, E.3)",
FIXED_TAUS.len(),
);
let fixed: &[f32] = &FIXED_TAUS;
// Precompute cosine features in col-major [D, N]: // Precompute cosine features [D, N] from FIXED_TAUS (col-major):
// cos_features[d + q*D] = cos(π·(d+1)·τ_q) // cos_feat[dim + q*D] = cos(π·(d+1)·τ_q)
let mut cos_feat_host = vec![0.0_f32; d * n]; let mut cos_feat_host = vec![0.0_f32; d * n];
for q in 0..n { for q in 0..n {
let tau_q = fixed[q]; let tau_q = fixed[q];
for dim in 0..d { for dim in 0..d {
cos_feat_host[dim + q * d] = cos_feat_host[dim + q * d] =
(std::f32::consts::PI * ((dim + 1) as f32) * tau_q).cos(); (std::f32::consts::PI * ((dim + 1) as f32) * tau_q).cos();
}
} }
cos_features = super::mapped_pinned::clone_to_device_f32_via_pinned(&stream, &cos_feat_host)
.map_err(|e| MLError::ModelError(format!("IQN cos_features upload via pinned ({} f32): {e}", d * n)))?;
let mut tiled = Vec::with_capacity(b * n);
for _ in 0..b {
tiled.extend_from_slice(fixed);
}
online_taus = super::mapped_pinned::clone_to_device_f32_via_pinned(&stream, &tiled)
.map_err(|e| MLError::ModelError(format!("IQN online_taus upload via pinned ({} f32): {e}", b * n)))?;
target_taus = super::mapped_pinned::clone_to_device_f32_via_pinned(&stream, &tiled)
.map_err(|e| MLError::ModelError(format!("IQN target_taus upload via pinned ({} f32): {e}", b * n)))?;
} }
// Tile taus [B, N]: each row is FIXED_TAUS.
let mut tiled_fixed = Vec::with_capacity(b * n);
for _ in 0..b {
tiled_fixed.extend_from_slice(fixed);
}
// Allocate and initialise 4 per-branch mapped-pinned slab triplets.
// MappedF32Buffer::new zero-inits; write_from_slice populates with FIXED_TAUS data.
let make_taus_buf = |len: usize| -> Result<super::mapped_pinned::MappedF32Buffer, MLError> {
unsafe { super::mapped_pinned::MappedF32Buffer::new(len) }
.map_err(|e| MLError::ModelError(format!("IQN tau buf alloc ({len} f32): {e}")))
};
let tau_buf_len = b * n;
let cos_buf_len = d * n;
let online_taus_b0 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
let online_taus_b1 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
let online_taus_b2 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
let online_taus_b3 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
let target_taus_b0 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
let target_taus_b1 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
let target_taus_b2 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
let target_taus_b3 = { let buf = make_taus_buf(tau_buf_len)?; buf.write_from_slice(&tiled_fixed); buf };
let cos_features_b0 = { let buf = make_taus_buf(cos_buf_len)?; buf.write_from_slice(&cos_feat_host); buf };
let cos_features_b1 = { let buf = make_taus_buf(cos_buf_len)?; buf.write_from_slice(&cos_feat_host); buf };
let cos_features_b2 = { let buf = make_taus_buf(cos_buf_len)?; buf.write_from_slice(&cos_feat_host); buf };
let cos_features_b3 = { let buf = make_taus_buf(cos_buf_len)?; buf.write_from_slice(&cos_feat_host); buf };
let branch_actions = stream.alloc_zeros::<i32>(b * 4).map_err(|e| { let branch_actions = stream.alloc_zeros::<i32>(b * 4).map_err(|e| {
MLError::ModelError(format!("IQN alloc branch_actions: {e}")) MLError::ModelError(format!("IQN alloc branch_actions: {e}"))
})?; })?;
@@ -591,20 +597,6 @@ impl GpuIqnHead {
let iqr_buf = alloc_f32(&stream, tba, "iqn_iqr")?; let iqr_buf = alloc_f32(&stream, tba, "iqn_iqr")?;
// SP6 Pearl 5: allocate per-branch tau/cos buffer triplets (same size as main buffers). // SP6 Pearl 5: allocate per-branch tau/cos buffer triplets (same size as main buffers).
// 4 branches × ([B*N] online + [B*N] target + [D*N] cos).
let online_taus_b0 = alloc_f32(&stream, b * n, "iqn_online_taus_b0")?;
let online_taus_b1 = alloc_f32(&stream, b * n, "iqn_online_taus_b1")?;
let online_taus_b2 = alloc_f32(&stream, b * n, "iqn_online_taus_b2")?;
let online_taus_b3 = alloc_f32(&stream, b * n, "iqn_online_taus_b3")?;
let target_taus_b0 = alloc_f32(&stream, b * n, "iqn_target_taus_b0")?;
let target_taus_b1 = alloc_f32(&stream, b * n, "iqn_target_taus_b1")?;
let target_taus_b2 = alloc_f32(&stream, b * n, "iqn_target_taus_b2")?;
let target_taus_b3 = alloc_f32(&stream, b * n, "iqn_target_taus_b3")?;
let cos_features_b0 = alloc_f32(&stream, d * n, "iqn_cos_features_b0")?;
let cos_features_b1 = alloc_f32(&stream, d * n, "iqn_cos_features_b1")?;
let cos_features_b2 = alloc_f32(&stream, d * n, "iqn_cos_features_b2")?;
let cos_features_b3 = alloc_f32(&stream, d * n, "iqn_cos_features_b3")?;
let vram_bytes = (total_params * 6 + h * bq * 8 + tba * bq * 4 let vram_bytes = (total_params * 6 + h * bq * 8 + tba * bq * 4
+ d * bq + b * h * 2 + b * 4 + tba + b * 2 + 2) * 4; + d * bq + b * h * 2 + b * 4 + tba + b * 2 + 2) * 4;
@@ -703,12 +695,10 @@ impl GpuIqnHead {
d_h_s2_tiled_buf, d_h_s2_tiled_buf,
bias_grad_partials_buf, bias_grad_partials_buf,
bias_grad_num_blocks, bias_grad_num_blocks,
online_taus, online_taus_per_branch: [online_taus_b0, online_taus_b1, online_taus_b2, online_taus_b3],
target_taus, target_taus_per_branch: [target_taus_b0, target_taus_b1, target_taus_b2, target_taus_b3],
cos_features, cos_features_per_branch: [cos_features_b0, cos_features_b1, cos_features_b2, cos_features_b3],
online_taus_branch: [online_taus_b0, online_taus_b1, online_taus_b2, online_taus_b3], active_branch_idx: 0,
target_taus_branch: [target_taus_b0, target_taus_b1, target_taus_b2, target_taus_b3],
cos_features_branch: [cos_features_b0, cos_features_b1, cos_features_b2, cos_features_b3],
branch_actions, branch_actions,
target_h_s2, target_h_s2,
cached_target_h_s2_ptr: None, cached_target_h_s2_ptr: None,
@@ -1009,7 +999,7 @@ impl GpuIqnHead {
unsafe { unsafe {
effective_stream effective_stream
.launch_builder(&self.cos_tile_kernel) .launch_builder(&self.cos_tile_kernel)
.arg(&self.cos_features) .arg(&self.cos_features_per_branch[self.active_branch_idx].dev_ptr)
.arg(&mut self.cos_features_tiled_buf) .arg(&mut self.cos_features_tiled_buf)
.arg(&d_i32) .arg(&d_i32)
.arg(&b_i32) .arg(&b_i32)
@@ -1242,7 +1232,7 @@ impl GpuIqnHead {
.launch_builder(&self.quantile_huber_loss_kernel) .launch_builder(&self.quantile_huber_loss_kernel)
.arg(&self.branch_logits_buf) .arg(&self.branch_logits_buf)
.arg(&self.target_branch_logits_buf) .arg(&self.target_branch_logits_buf)
.arg(&self.online_taus) .arg(&self.online_taus_per_branch[self.active_branch_idx].dev_ptr)
.arg(&self.branch_actions) .arg(&self.branch_actions)
.arg(&mut self.per_sample_loss) .arg(&mut self.per_sample_loss)
.arg(&mut self.d_branch_logits_buf) .arg(&mut self.d_branch_logits_buf)
@@ -1979,12 +1969,12 @@ impl GpuIqnHead {
/// ///
/// Reads 4 branches × 5 quantiles from ISV, averages across branches per /// Reads 4 branches × 5 quantiles from ISV, averages across branches per
/// quantile slot, applies Invariant-1 cold-start floors (if ISV reads 0, /// quantile slot, applies Invariant-1 cold-start floors (if ISV reads 0,
/// fall back to `FIXED_TAUS[q]`), then re-uploads `online_taus`, /// fall back to `FIXED_TAUS[q]`), then writes the averaged τ schedule and
/// `target_taus`, and `cos_features` in-place via mapped-pinned staging. /// recomputed cosine features to all 4 per-branch mapped-pinned slabs via
/// `write_from_slice` (direct host_ptr write — no DtoD copy).
/// ///
/// Must be called once per epoch from `FusedTrainingCtx` before the IQN /// Must be called once per epoch from `FusedTrainingCtx` before the IQN
/// training pipeline runs. Non-fatal on mapped-pinned upload error /// training pipeline runs.
/// (logs a warning).
pub fn refresh_taus_from_isv( pub fn refresh_taus_from_isv(
&mut self, &mut self,
isv_per_branch_taus: &[f32; 20], // ISV[250..270): branch b, quantile q at [b*5+q] isv_per_branch_taus: &[f32; 20], // ISV[250..270): branch b, quantile q at [b*5+q]
@@ -2022,32 +2012,25 @@ impl GpuIqnHead {
tiled.extend_from_slice(&taus[..n]); tiled.extend_from_slice(&taus[..n]);
} }
// Upload all three buffers in-place via mapped-pinned (no new allocations). // Write the averaged tau schedule to all 4 per-branch mapped-pinned slabs.
if let Err(e) = super::mapped_pinned::upload_f32_via_pinned( // Direct host_ptr writes — no DtoD copy, no upload. The kernel reads via
&self.stream, &tiled, &mut self.online_taus, // dev_ptr of the same mapped-pinned allocation after the next stream sync
) { // (feedback_no_htod_htoh_only_mapped_pinned).
tracing::warn!("SP5 Pearl 5 online_taus upload failed (non-fatal): {e}"); for b_idx in 0..4 {
return; self.online_taus_per_branch[b_idx].write_from_slice(&tiled);
} self.target_taus_per_branch[b_idx].write_from_slice(&tiled);
if let Err(e) = super::mapped_pinned::upload_f32_via_pinned( self.cos_features_per_branch[b_idx].write_from_slice(&cos_feat_host);
&self.stream, &tiled, &mut self.target_taus,
) {
tracing::warn!("SP5 Pearl 5 target_taus upload failed (non-fatal): {e}");
return;
}
if let Err(e) = super::mapped_pinned::upload_f32_via_pinned(
&self.stream, &cos_feat_host, &mut self.cos_features,
) {
tracing::warn!("SP5 Pearl 5 cos_features upload failed (non-fatal): {e}");
} }
} }
/// SP6 Pearl 5: Refresh the τ tensors for a single branch from its 5-quantile /// SP6 Pearl 5: Refresh the τ tensors for a single branch from its 5-quantile
/// schedule `tau5` (from ISV[IQN_TAU_BASE + branch_idx*5 .. +5]). /// schedule `tau5` (from ISV[IQN_TAU_BASE + branch_idx*5 .. +5]).
/// ///
/// Uploads to `online_taus_branch[branch_idx]`, `target_taus_branch[branch_idx]`, /// Writes directly to `online_taus_per_branch[branch_idx]`,
/// and `cos_features_branch[branch_idx]`. Called 4 times per epoch refresh (once /// `target_taus_per_branch[branch_idx]`, and `cos_features_per_branch[branch_idx]`
/// per branch) from `FusedTrainingCtx::run_full_step`. /// via `write_from_slice` (mapped-pinned host_ptr — no DtoD copy).
/// Called 4 times per epoch refresh (once per branch) from
/// `FusedTrainingCtx::run_full_step`.
/// ///
/// Cold-start floor: if a tau value is below 1e-6 (ISV not yet populated), /// Cold-start floor: if a tau value is below 1e-6 (ISV not yet populated),
/// uses `FIXED_TAUS[q]` to preserve the interpretable quantile distribution. /// uses `FIXED_TAUS[q]` to preserve the interpretable quantile distribution.
@@ -2084,45 +2067,30 @@ impl GpuIqnHead {
tiled.extend_from_slice(&taus[..n]); tiled.extend_from_slice(&taus[..n]);
} }
if let Err(e) = super::mapped_pinned::upload_f32_via_pinned( // Write directly to mapped-pinned host_ptr — no DtoD copy, no upload.
&self.stream, &tiled, &mut self.online_taus_branch[branch_idx], // The IQN kernel reads via dev_ptr of the same mapped-pinned allocation
) { // after the stream sync that precedes kernel launch
tracing::warn!("SP6 Pearl 5 online_taus_branch[{branch_idx}] upload failed: {e}"); // (feedback_no_htod_htoh_only_mapped_pinned).
return; self.online_taus_per_branch[branch_idx].write_from_slice(&tiled);
} self.target_taus_per_branch[branch_idx].write_from_slice(&tiled);
if let Err(e) = super::mapped_pinned::upload_f32_via_pinned( self.cos_features_per_branch[branch_idx].write_from_slice(&cos_feat_host);
&self.stream, &tiled, &mut self.target_taus_branch[branch_idx],
) {
tracing::warn!("SP6 Pearl 5 target_taus_branch[{branch_idx}] upload failed: {e}");
return;
}
if let Err(e) = super::mapped_pinned::upload_f32_via_pinned(
&self.stream, &cos_feat_host, &mut self.cos_features_branch[branch_idx],
) {
tracing::warn!("SP6 Pearl 5 cos_features_branch[{branch_idx}] upload failed: {e}");
}
} }
/// SP6 Pearl 5: Swap the main tau/cos buffers with branch `branch_idx`'s per-branch /// SP6 Pearl 5: Set `active_branch_idx` so all subsequent IQN kernel launches
/// buffers so the IQN forward pass uses that branch's τ schedule. /// use `branch_idx`'s per-branch mapped-pinned τ / cos_features slabs.
/// ///
/// Must be paired with `deactivate_branch_taus(branch_idx)` after the forward pass. /// Must be paired with `deactivate_branch_taus(branch_idx)` after the forward pass.
/// `activate` followed by `deactivate` with the same index is its own inverse.
pub fn activate_branch_taus(&mut self, branch_idx: usize) { pub fn activate_branch_taus(&mut self, branch_idx: usize) {
debug_assert!(branch_idx < 4, "activate_branch_taus: branch_idx {branch_idx} >= 4"); debug_assert!(branch_idx < 4, "activate_branch_taus: branch_idx {branch_idx} >= 4");
std::mem::swap(&mut self.online_taus, &mut self.online_taus_branch[branch_idx]); self.active_branch_idx = branch_idx;
std::mem::swap(&mut self.target_taus, &mut self.target_taus_branch[branch_idx]);
std::mem::swap(&mut self.cos_features, &mut self.cos_features_branch[branch_idx]);
} }
/// SP6 Pearl 5: Restore the main tau/cos buffers after a per-branch IQN forward pass. /// SP6 Pearl 5: Reset `active_branch_idx` to 0 after a per-branch IQN forward pass.
/// ///
/// Symmetric inverse of `activate_branch_taus(branch_idx)`. /// Symmetric inverse of `activate_branch_taus(branch_idx)`.
pub fn deactivate_branch_taus(&mut self, branch_idx: usize) { pub fn deactivate_branch_taus(&mut self, branch_idx: usize) {
debug_assert!(branch_idx < 4, "deactivate_branch_taus: branch_idx {branch_idx} >= 4"); debug_assert!(branch_idx < 4, "deactivate_branch_taus: branch_idx {branch_idx} >= 4");
std::mem::swap(&mut self.online_taus, &mut self.online_taus_branch[branch_idx]); self.active_branch_idx = 0;
std::mem::swap(&mut self.target_taus, &mut self.target_taus_branch[branch_idx]);
std::mem::swap(&mut self.cos_features, &mut self.cos_features_branch[branch_idx]);
} }
/// Compute CVaR-based position scaling from IQN quantiles. /// Compute CVaR-based position scaling from IQN quantiles.
@@ -2147,12 +2115,12 @@ impl GpuIqnHead {
let embed_dim_i32 = self.config.embed_dim as i32; let embed_dim_i32 = self.config.embed_dim as i32;
// Plan 4 Task 3 (E.3): per-step Philox τ sampling deleted. The // Plan 4 Task 3 (E.3): per-step Philox τ sampling deleted. The
// constructor already populated `online_taus` (and `target_taus`) // constructor populated all per-branch mapped-pinned τ slabs with
// with FIXED_TAUS broadcast across all B rows; CVaR's quantile // FIXED_TAUS broadcast across all B rows; CVaR's quantile estimates
// estimates are therefore drawn from the same fixed quantile // are drawn from the same fixed quantile levels {0.05, 0.25, 0.50,
// levels {0.05, 0.25, 0.50, 0.75, 0.95} as the training pipeline, // 0.75, 0.95} as the training pipeline, making the CVaR estimate
// making the CVaR estimate stable across steps and consistent // stable across steps and consistent with the spec's
// with the spec's interpretable-quantile contract. // interpretable-quantile contract.
let b0_i32 = self.config.branch_0_size as i32; let b0_i32 = self.config.branch_0_size as i32;
let b1_i32 = self.config.branch_1_size as i32; let b1_i32 = self.config.branch_1_size as i32;
let b2_i32 = self.config.branch_2_size as i32; let b2_i32 = self.config.branch_2_size as i32;
@@ -2170,9 +2138,9 @@ impl GpuIqnHead {
self.stream self.stream
.launch_builder(&self.forward_kernel) .launch_builder(&self.forward_kernel)
.arg(h_s2) .arg(h_s2)
.arg(&self.online_taus) .arg(&self.online_taus_per_branch[self.active_branch_idx].dev_ptr)
.arg(&self.online_params) .arg(&self.online_params)
.arg(&self.cos_features) .arg(&self.cos_features_per_branch[self.active_branch_idx].dev_ptr)
.arg(&mut self.save_q_online) .arg(&mut self.save_q_online)
.arg(&batch_i32) .arg(&batch_i32)
.arg(&shared_h1_i32) .arg(&shared_h1_i32)

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

@@ -44,9 +44,7 @@
| `gpu_dqn_trainer.rs:7680` | `stream.memcpy_htod` → stochastic depth scale | COLD-PATH | Constructor: one-shot HtoD for stochastic depth kernel seed | OK | | `gpu_dqn_trainer.rs:7680` | `stream.memcpy_htod` → stochastic depth scale | COLD-PATH | Constructor: one-shot HtoD for stochastic depth kernel seed | OK |
| `gpu_dqn_trainer.rs:7686` | `stream.memcpy_htod` → stochastic depth RNG | COLD-PATH | Constructor: one-shot HtoD for stochastic depth RNG | OK | | `gpu_dqn_trainer.rs:7686` | `stream.memcpy_htod` → stochastic depth RNG | COLD-PATH | Constructor: one-shot HtoD for stochastic depth RNG | OK |
| `gpu_dqn_trainer.rs` (removed) | `cuMemcpyDtoHAsync_v2` in `run_causal_intervention_unconditional` | **MIGRATED** | **Fix 1**: removed dead copy — result `causal_mean_scratch` was never consumed by any caller; result stays on GPU | FIXED | | `gpu_dqn_trainer.rs` (removed) | `cuMemcpyDtoHAsync_v2` in `run_causal_intervention_unconditional` | **MIGRATED** | **Fix 1**: removed dead copy — result `causal_mean_scratch` was never consumed by any caller; result stays on GPU | FIXED |
| `gpu_iqn_head.rs:469` | `stream.clone_htod` → `cos_features` | COLD-PATH | Constructor: precompute cosine embedding table; one-shot | OK | | `gpu_iqn_head.rs` (removed) | `clone_to_device_f32_via_pinned` → `cos_features` / `online_taus` / `target_taus` (3 CudaSlice fields) | **MIGRATED** | **SP6 fix**: converted to `[MappedF32Buffer; 4]` per-branch arrays (`online_taus_per_branch`, `target_taus_per_branch`, `cos_features_per_branch`). Constructor writes via `write_from_slice` (host_ptr direct); IQN kernels read via `dev_ptr`. `mem::swap` activate/deactivate replaced by `active_branch_idx` selection. `upload_f32_via_pinned` (which did an intermediate DtoD inside graph capture → `STREAM_CAPTURE_INVALIDATED`) fully eliminated. | FIXED |
| `gpu_iqn_head.rs:477` | `stream.clone_htod` → `online_taus` | COLD-PATH | Constructor: tau quantile tiling; one-shot | OK |
| `gpu_iqn_head.rs:480` | `stream.clone_htod` → `target_taus` | COLD-PATH | Constructor: tau quantile tiling; one-shot | OK |
| `gpu_iqn_head.rs:498` | `malloc_host` → `total_loss_pinned` | OK-pinned | Constructor: GPU writes IQN loss; host reads via `read_total_loss()` | OK | | `gpu_iqn_head.rs:498` | `malloc_host` → `total_loss_pinned` | OK-pinned | Constructor: GPU writes IQN loss; host reads via `read_total_loss()` | OK |
| `gpu_iqn_head.rs:534` | `malloc_host` → `t_pinned` | OK-pinned | Constructor: Adam step counter — CPU increments, GPU reads via dev ptr | OK | | `gpu_iqn_head.rs:534` | `malloc_host` → `t_pinned` | OK-pinned | Constructor: Adam step counter — CPU increments, GPU reads via dev ptr | OK |
| `gpu_iqn_head.rs:548` | `malloc_host` → `tau_pinned` | OK-pinned | Constructor (Fix 3): tau scalar — CPU writes, GPU reads via `tau_dev_ptr`; replaces `CudaSlice<f32>` + `cuMemcpyHtoDAsync_v2` | OK | | `gpu_iqn_head.rs:548` | `malloc_host` → `tau_pinned` | OK-pinned | Constructor (Fix 3): tau scalar — CPU writes, GPU reads via `tau_dev_ptr`; replaces `CudaSlice<f32>` + `cuMemcpyHtoDAsync_v2` | OK |