fix(sp6): IQN τ buffers — MappedF32Buffer per feedback_no_htod_htoh_only_mapped_pinned
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:
@@ -318,21 +318,20 @@ pub struct GpuIqnHead {
|
||||
bias_grad_num_blocks: usize,
|
||||
|
||||
// ── Per-step buffers ─────────────────────────────────────────────
|
||||
/// Pre-sampled τ values for online network [B, N]
|
||||
online_taus: CudaSlice<f32>,
|
||||
/// Pre-sampled τ values for target network [B, N]
|
||||
target_taus: CudaSlice<f32>,
|
||||
/// Precomputed cosine features [D, N] — cos(π·(d+1)·τ_i)
|
||||
/// Fixed at construction (τ are midpoints). Column-major.
|
||||
cos_features: CudaSlice<f32>,
|
||||
/// SP6 Pearl 5: per-branch τ tensors [B, N], one slab per branch (4 total).
|
||||
/// Populated by `refresh_taus_for_branch`; activated via `activate_branch_taus`
|
||||
/// before each per-branch IQN forward pass. Same allocation size as `online_taus`.
|
||||
online_taus_branch: [CudaSlice<f32>; 4],
|
||||
/// SP6 Pearl 5: per-branch target τ tensors [B, N], one slab per branch.
|
||||
target_taus_branch: [CudaSlice<f32>; 4],
|
||||
/// SP6 Pearl 5: per-branch cosine features [D, N], one slab per branch.
|
||||
cos_features_branch: [CudaSlice<f32>; 4],
|
||||
/// Per-branch online τ values [4][B, N] — mapped pinned (cuMemHostAlloc DEVICEMAP).
|
||||
/// Host writes via host_ptr; IQN kernel reads via dev_ptr. No DtoD copy.
|
||||
/// Indexed by `active_branch_idx` at each kernel launch site.
|
||||
online_taus_per_branch: [super::mapped_pinned::MappedF32Buffer; 4],
|
||||
/// Per-branch target τ values [4][B, N] — mapped pinned.
|
||||
target_taus_per_branch: [super::mapped_pinned::MappedF32Buffer; 4],
|
||||
/// Per-branch precomputed cosine features [4][D, N] — mapped pinned.
|
||||
/// cos_feat[dim + q*D] = cos(π·(d+1)·τ_q). Column-major.
|
||||
cos_features_per_branch: [super::mapped_pinned::MappedF32Buffer; 4],
|
||||
/// Index of the branch whose τ buffers are currently active for the IQN
|
||||
/// forward/backward pass. Set by `activate_branch_taus(branch_idx)`,
|
||||
/// reset to 0 by `deactivate_branch_taus`. Kernel launch sites index
|
||||
/// the per-branch arrays with this value.
|
||||
active_branch_idx: usize,
|
||||
/// Branch actions decoded from flat actions [B, 4]
|
||||
branch_actions: CudaSlice<i32>,
|
||||
/// Target h_s2 computed from next_states + target trunk weights [B, H]
|
||||
@@ -512,50 +511,57 @@ impl GpuIqnHead {
|
||||
|
||||
// ── Per-step buffers ────────────────────────────────────────────
|
||||
// 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
|
||||
// quantiles, which is the same uniform-grid mean estimator but
|
||||
// randomized via Philox at run time. Fixing the τ values lets the
|
||||
// 5 quantile heads be interpretable / consistent across batches and
|
||||
// enables the median-Q action-ranking switch in `iqn_forward_kernel`.
|
||||
// Both online and target IQN read this same static buffer — the
|
||||
// off-median heads are diagnostic only (ISV[99..103) producers) and
|
||||
// must agree on the fixed τ between online and target so the
|
||||
// quantile-Huber regression learns a stable distributional surface.
|
||||
let online_taus;
|
||||
let target_taus;
|
||||
let cos_features;
|
||||
{
|
||||
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;
|
||||
// All 4 per-branch slabs are initialised from FIXED_TAUS at construction.
|
||||
// refresh_taus_for_branch / refresh_taus_from_isv write updated values
|
||||
// directly to host_ptr (no copy, no DtoD) per
|
||||
// feedback_no_htod_htoh_only_mapped_pinned.
|
||||
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]:
|
||||
// cos_features[d + q*D] = cos(π·(d+1)·τ_q)
|
||||
let mut cos_feat_host = vec![0.0_f32; d * n];
|
||||
for q in 0..n {
|
||||
let tau_q = fixed[q];
|
||||
for dim in 0..d {
|
||||
cos_feat_host[dim + q * d] =
|
||||
(std::f32::consts::PI * ((dim + 1) as f32) * tau_q).cos();
|
||||
}
|
||||
// Precompute cosine features [D, N] from FIXED_TAUS (col-major):
|
||||
// cos_feat[dim + q*D] = cos(π·(d+1)·τ_q)
|
||||
let mut cos_feat_host = vec![0.0_f32; d * n];
|
||||
for q in 0..n {
|
||||
let tau_q = fixed[q];
|
||||
for dim in 0..d {
|
||||
cos_feat_host[dim + q * d] =
|
||||
(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| {
|
||||
MLError::ModelError(format!("IQN alloc branch_actions: {e}"))
|
||||
})?;
|
||||
@@ -591,20 +597,6 @@ impl GpuIqnHead {
|
||||
let iqr_buf = alloc_f32(&stream, tba, "iqn_iqr")?;
|
||||
|
||||
// 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
|
||||
+ d * bq + b * h * 2 + b * 4 + tba + b * 2 + 2) * 4;
|
||||
|
||||
@@ -703,12 +695,10 @@ impl GpuIqnHead {
|
||||
d_h_s2_tiled_buf,
|
||||
bias_grad_partials_buf,
|
||||
bias_grad_num_blocks,
|
||||
online_taus,
|
||||
target_taus,
|
||||
cos_features,
|
||||
online_taus_branch: [online_taus_b0, online_taus_b1, online_taus_b2, online_taus_b3],
|
||||
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],
|
||||
online_taus_per_branch: [online_taus_b0, online_taus_b1, online_taus_b2, online_taus_b3],
|
||||
target_taus_per_branch: [target_taus_b0, target_taus_b1, target_taus_b2, target_taus_b3],
|
||||
cos_features_per_branch: [cos_features_b0, cos_features_b1, cos_features_b2, cos_features_b3],
|
||||
active_branch_idx: 0,
|
||||
branch_actions,
|
||||
target_h_s2,
|
||||
cached_target_h_s2_ptr: None,
|
||||
@@ -1009,7 +999,7 @@ impl GpuIqnHead {
|
||||
unsafe {
|
||||
effective_stream
|
||||
.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(&d_i32)
|
||||
.arg(&b_i32)
|
||||
@@ -1242,7 +1232,7 @@ impl GpuIqnHead {
|
||||
.launch_builder(&self.quantile_huber_loss_kernel)
|
||||
.arg(&self.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(&mut self.per_sample_loss)
|
||||
.arg(&mut self.d_branch_logits_buf)
|
||||
@@ -1979,12 +1969,12 @@ impl GpuIqnHead {
|
||||
///
|
||||
/// Reads 4 branches × 5 quantiles from ISV, averages across branches per
|
||||
/// quantile slot, applies Invariant-1 cold-start floors (if ISV reads 0,
|
||||
/// fall back to `FIXED_TAUS[q]`), then re-uploads `online_taus`,
|
||||
/// `target_taus`, and `cos_features` in-place via mapped-pinned staging.
|
||||
/// fall back to `FIXED_TAUS[q]`), then writes the averaged τ schedule and
|
||||
/// 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
|
||||
/// training pipeline runs. Non-fatal on mapped-pinned upload error
|
||||
/// (logs a warning).
|
||||
/// training pipeline runs.
|
||||
pub fn refresh_taus_from_isv(
|
||||
&mut self,
|
||||
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]);
|
||||
}
|
||||
|
||||
// Upload all three buffers in-place via mapped-pinned (no new allocations).
|
||||
if let Err(e) = super::mapped_pinned::upload_f32_via_pinned(
|
||||
&self.stream, &tiled, &mut self.online_taus,
|
||||
) {
|
||||
tracing::warn!("SP5 Pearl 5 online_taus upload failed (non-fatal): {e}");
|
||||
return;
|
||||
}
|
||||
if let Err(e) = super::mapped_pinned::upload_f32_via_pinned(
|
||||
&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}");
|
||||
// Write the averaged tau schedule to all 4 per-branch mapped-pinned slabs.
|
||||
// Direct host_ptr writes — no DtoD copy, no upload. The kernel reads via
|
||||
// dev_ptr of the same mapped-pinned allocation after the next stream sync
|
||||
// (feedback_no_htod_htoh_only_mapped_pinned).
|
||||
for b_idx in 0..4 {
|
||||
self.online_taus_per_branch[b_idx].write_from_slice(&tiled);
|
||||
self.target_taus_per_branch[b_idx].write_from_slice(&tiled);
|
||||
self.cos_features_per_branch[b_idx].write_from_slice(&cos_feat_host);
|
||||
}
|
||||
}
|
||||
|
||||
/// 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]).
|
||||
///
|
||||
/// Uploads to `online_taus_branch[branch_idx]`, `target_taus_branch[branch_idx]`,
|
||||
/// and `cos_features_branch[branch_idx]`. Called 4 times per epoch refresh (once
|
||||
/// per branch) from `FusedTrainingCtx::run_full_step`.
|
||||
/// Writes directly to `online_taus_per_branch[branch_idx]`,
|
||||
/// `target_taus_per_branch[branch_idx]`, and `cos_features_per_branch[branch_idx]`
|
||||
/// 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),
|
||||
/// uses `FIXED_TAUS[q]` to preserve the interpretable quantile distribution.
|
||||
@@ -2084,45 +2067,30 @@ impl GpuIqnHead {
|
||||
tiled.extend_from_slice(&taus[..n]);
|
||||
}
|
||||
|
||||
if let Err(e) = super::mapped_pinned::upload_f32_via_pinned(
|
||||
&self.stream, &tiled, &mut self.online_taus_branch[branch_idx],
|
||||
) {
|
||||
tracing::warn!("SP6 Pearl 5 online_taus_branch[{branch_idx}] upload failed: {e}");
|
||||
return;
|
||||
}
|
||||
if let Err(e) = super::mapped_pinned::upload_f32_via_pinned(
|
||||
&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}");
|
||||
}
|
||||
// Write directly to mapped-pinned host_ptr — no DtoD copy, no upload.
|
||||
// The IQN kernel reads via dev_ptr of the same mapped-pinned allocation
|
||||
// after the stream sync that precedes kernel launch
|
||||
// (feedback_no_htod_htoh_only_mapped_pinned).
|
||||
self.online_taus_per_branch[branch_idx].write_from_slice(&tiled);
|
||||
self.target_taus_per_branch[branch_idx].write_from_slice(&tiled);
|
||||
self.cos_features_per_branch[branch_idx].write_from_slice(&cos_feat_host);
|
||||
}
|
||||
|
||||
/// SP6 Pearl 5: Swap the main tau/cos buffers with branch `branch_idx`'s per-branch
|
||||
/// buffers so the IQN forward pass uses that branch's τ schedule.
|
||||
/// SP6 Pearl 5: Set `active_branch_idx` so all subsequent IQN kernel launches
|
||||
/// use `branch_idx`'s per-branch mapped-pinned τ / cos_features slabs.
|
||||
///
|
||||
/// 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) {
|
||||
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]);
|
||||
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]);
|
||||
self.active_branch_idx = 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)`.
|
||||
pub fn deactivate_branch_taus(&mut self, branch_idx: usize) {
|
||||
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]);
|
||||
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]);
|
||||
self.active_branch_idx = 0;
|
||||
}
|
||||
|
||||
/// Compute CVaR-based position scaling from IQN quantiles.
|
||||
@@ -2147,12 +2115,12 @@ impl GpuIqnHead {
|
||||
let embed_dim_i32 = self.config.embed_dim as i32;
|
||||
|
||||
// Plan 4 Task 3 (E.3): per-step Philox τ sampling deleted. The
|
||||
// constructor already populated `online_taus` (and `target_taus`)
|
||||
// with FIXED_TAUS broadcast across all B rows; CVaR's quantile
|
||||
// estimates are therefore drawn from the same fixed quantile
|
||||
// levels {0.05, 0.25, 0.50, 0.75, 0.95} as the training pipeline,
|
||||
// making the CVaR estimate stable across steps and consistent
|
||||
// with the spec's interpretable-quantile contract.
|
||||
// constructor populated all per-branch mapped-pinned τ slabs with
|
||||
// FIXED_TAUS broadcast across all B rows; CVaR's quantile estimates
|
||||
// are drawn from the same fixed quantile levels {0.05, 0.25, 0.50,
|
||||
// 0.75, 0.95} as the training pipeline, making the CVaR estimate
|
||||
// stable across steps and consistent with the spec's
|
||||
// interpretable-quantile contract.
|
||||
let b0_i32 = self.config.branch_0_size as i32;
|
||||
let b1_i32 = self.config.branch_1_size as i32;
|
||||
let b2_i32 = self.config.branch_2_size as i32;
|
||||
@@ -2170,9 +2138,9 @@ impl GpuIqnHead {
|
||||
self.stream
|
||||
.launch_builder(&self.forward_kernel)
|
||||
.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.cos_features)
|
||||
.arg(&self.cos_features_per_branch[self.active_branch_idx].dev_ptr)
|
||||
.arg(&mut self.save_q_online)
|
||||
.arg(&batch_i32)
|
||||
.arg(&shared_h1_i32)
|
||||
|
||||
@@ -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: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_iqn_head.rs:469` | `stream.clone_htod` → `cos_features` | COLD-PATH | Constructor: precompute cosine embedding table; one-shot | OK |
|
||||
| `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` (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: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: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 |
|
||||
|
||||
Reference in New Issue
Block a user