feat(dqn): SP2 A3 — fused NaN check populate + launch wrapper
Replaces A2's CudaSlice<u64>/<i32> field types with MappedU64Buffer/ MappedI32Buffer per feedback_no_htod_htoh_only_mapped_pinned. Mapped- pinned eliminates the HtoD copy entirely — the kernel reads via the device-mapped pointer (cuMemHostGetDevicePointer_v2) while the trainer writes through the same mapped pages on the host side. Adds populate_nan_check_meta on GpuDqnTrainer (one-shot construction- time write of 12 (ptr, len) tuples for slots 24-35). Slot 31 deferred (null entry); slots 27/28 nullable on Option<u64> (None when IQN inactive); slots 33-35 null (inline checks fire separately at backward orchestration phases — kept individual for entry-point localization). Adds launch_nan_check_fused_f32 (per-step kernel launch wrapper with grid_dim=12, block_dim=256, base_flag_idx=24). Registers dqn_nan_check_fused_f32_kernel in compile_training_kernels (tuple 43→44, info log 38→39 utility kernels) — same module as the per-buffer dqn_nan_check_f32 to share the captured replay group. Constructor-time wire-up lands in FusedTrainingCtx::new after gpu_iqn construction (gpu_iqn is owned by FusedTrainingCtx, not GpuDqnTrainer — mirrors the same Option<u64> arg pattern used by apply_iqn_trunk_gradient and run_nan_checks_post_backward). Wrapper unused yet — call-site replacement (8 individual check_nan_f32 calls in run_nan_checks_post_backward → single fused launch) lands in A4. Audit doc updated (Invariant 7). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -2655,6 +2655,14 @@ pub struct GpuDqnTrainer {
|
||||
pub(crate) nan_check_f32_kernel: CudaFunction,
|
||||
/// Second NaN check kernel handle (formerly f32, now f32 — both check f32 buffers).
|
||||
pub(crate) nan_check_f32_kernel_b: CudaFunction,
|
||||
/// SP2: fused multi-buffer NaN check kernel handle
|
||||
/// (`dqn_nan_check_fused_f32_kernel` in `dqn_utility_kernels.cu`). One launch
|
||||
/// processes 12 backward-path slots (24-35) in 12 blocks. Reads `(ptr, len)`
|
||||
/// tables from `nan_check_buf_ptrs` / `nan_check_buf_lens` via the mapped
|
||||
/// device pointer; null-pointer entries (deferred slot 31, optional IQN
|
||||
/// slots when inactive, slots 33-35 inline elsewhere) no-op via the
|
||||
/// kernel's null-pointer guard.
|
||||
pub(crate) nan_check_fused_f32_kernel: CudaFunction,
|
||||
/// SP1 Phase C surgical fix: in-place finite clamp kernel
|
||||
/// (`dqn_clamp_finite_f32_kernel` in `dqn_utility_kernels.cu`). Replaces
|
||||
/// NaN/Inf with 0 and clamps finite |v| to ±max_abs. max_abs is sourced
|
||||
@@ -2673,14 +2681,18 @@ pub struct GpuDqnTrainer {
|
||||
pub(crate) nan_flags_buf: CudaSlice<i32>,
|
||||
|
||||
/// SP2: fused NaN-check buffer pointer table. 12 device pointers (u64) for
|
||||
/// slots 24-35. Populated ONCE in the constructor (slot pointers are
|
||||
/// stable across the trainer's lifetime). Slot 31 (deferred ensemble)
|
||||
/// holds 0 — the fused kernel skips that block.
|
||||
pub(crate) nan_check_buf_ptrs: CudaSlice<u64>,
|
||||
/// slots 24-35. Host+device-visible mapped-pinned buffer; populated via
|
||||
/// host-side write (no HtoD copy) per
|
||||
/// `feedback_no_htod_htoh_only_mapped_pinned`. Populated ONCE in
|
||||
/// `populate_nan_check_meta` (slot pointers are stable across the
|
||||
/// trainer's lifetime). Slot 31 (deferred ensemble) holds 0 — the fused
|
||||
/// kernel skips that block.
|
||||
pub(crate) nan_check_buf_ptrs: MappedU64Buffer,
|
||||
|
||||
/// SP2: fused NaN-check buffer length table. 12 i32 lengths for slots
|
||||
/// 24-35. Populated alongside nan_check_buf_ptrs.
|
||||
pub(crate) nan_check_buf_lens: CudaSlice<i32>,
|
||||
/// 24-35. Host+device-visible mapped-pinned buffer; populated via
|
||||
/// host-side write (no HtoD copy) alongside `nan_check_buf_ptrs`.
|
||||
pub(crate) nan_check_buf_lens: MappedI32Buffer,
|
||||
|
||||
/// #20 Pruning epoch (epoch at which to compute the mask).
|
||||
pruning_epoch: usize,
|
||||
@@ -9448,7 +9460,7 @@ impl GpuDqnTrainer {
|
||||
// per array. Stack is set once in DQNTrainer::new() (64KB for all kernels).
|
||||
|
||||
// ── Compile utility kernels (grad_norm, adam_update, etc.) ─
|
||||
let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, clamp_finite_f32_kernel, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel) =
|
||||
let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, nan_check_fused_f32_kernel, clamp_finite_f32_kernel, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel) =
|
||||
compile_training_kernels(&stream, &config)?;
|
||||
|
||||
// Separate grad_norm instance for non-graph launches (outside CUDA graph).
|
||||
@@ -11311,12 +11323,20 @@ impl GpuDqnTrainer {
|
||||
.map_err(|e| MLError::ModelError(format!("nan_flags alloc: {e}")))?;
|
||||
|
||||
// SP2: fused NaN-check (ptr, len) tables. 12 entries each. Allocated
|
||||
// as zeros; populated once (next commit) after all backward-path
|
||||
// buffers are constructed. The kernel reads `buf_ptrs[blockIdx.x]`
|
||||
// and `buf_lens[blockIdx.x]` per block.
|
||||
let nan_check_buf_ptrs = stream.alloc_zeros::<u64>(12)
|
||||
// as mapped-pinned (`cuMemHostAlloc(DEVICEMAP|PORTABLE)`) per
|
||||
// `feedback_no_htod_htoh_only_mapped_pinned` — host-side writes are
|
||||
// visible to the kernel through the device-mapped pointer with zero
|
||||
// HtoD copies. Populated once via `populate_nan_check_meta` after all
|
||||
// backward-path buffers are constructed AND `gpu_iqn` is known. The
|
||||
// kernel reads `buf_ptrs[blockIdx.x]` and `buf_lens[blockIdx.x]` per
|
||||
// block.
|
||||
//
|
||||
// Safety: a CUDA context is active on this thread (the GpuDqnTrainer
|
||||
// constructor runs within `FusedTrainingCtx::new` which has already
|
||||
// initialised CUDA via the shared cublas handle).
|
||||
let nan_check_buf_ptrs = unsafe { MappedU64Buffer::new(12) }
|
||||
.map_err(|e| MLError::ModelError(format!("nan_check_buf_ptrs alloc: {e}")))?;
|
||||
let nan_check_buf_lens = stream.alloc_zeros::<i32>(12)
|
||||
let nan_check_buf_lens = unsafe { MappedI32Buffer::new(12) }
|
||||
.map_err(|e| MLError::ModelError(format!("nan_check_buf_lens alloc: {e}")))?;
|
||||
|
||||
// v8: PopArt running statistics buffers
|
||||
@@ -12701,6 +12721,7 @@ impl GpuDqnTrainer {
|
||||
her_inplace_kernel,
|
||||
nan_check_f32_kernel,
|
||||
nan_check_f32_kernel_b,
|
||||
nan_check_fused_f32_kernel,
|
||||
clamp_finite_f32_kernel,
|
||||
nan_flags_buf,
|
||||
nan_check_buf_ptrs,
|
||||
@@ -15227,6 +15248,132 @@ impl GpuDqnTrainer {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// SP2: populate the fused NaN-check `(ptr, len)` tables. Called once
|
||||
/// during construction after all backward-path buffers are allocated AND
|
||||
/// `gpu_iqn` is known (its accessors feed slot 27/28 entries when IQN is
|
||||
/// active; null/0 when None — the kernel's null-pointer guard skips them).
|
||||
///
|
||||
/// Slot index → buffer mapping (mirrors the per-slot accessor table in
|
||||
/// `docs/dqn-backward-nan-audit.md` and the `run_nan_checks_post_backward`
|
||||
/// per-slot launches):
|
||||
/// 0 (slot 24) d_value_logits_buf — len = `d_value_logits_buf.len()`
|
||||
/// 1 (slot 25) d_adv_logits_buf — len = `d_adv_logits_buf.len()`
|
||||
/// 2 (slot 26) iqn_trunk_m — len = `trunk_param_count`
|
||||
/// 3 (slot 27) iqn d_h_s2 buf — len = b * sh2 (or 0 if IQN inactive)
|
||||
/// 4 (slot 28) iqn d_branch_logits_buf — len = total_branch_actions * b * q (or 0)
|
||||
/// 5 (slot 29) cql_d_value_logits — len = `cql_d_value_logits.len()`
|
||||
/// 6 (slot 30) aux_dh_s2_nb_buf — len = `aux_dh_s2_nb_buf.len()`
|
||||
/// 7 (slot 31) ensemble (deferred) — ptr=0, len=0 → kernel skips block
|
||||
/// 8 (slot 32) bn_d_concat_buf — len = `bn_d_concat_buf().len()`
|
||||
/// 9-11 (slots 33-35) bw_d_h_s2 inline — ptr=0, len=0 → fused kernel skips
|
||||
/// (inline checks fire at backward orchestration phases — see
|
||||
/// `launch_cublas_backward_to` slots 33/34 and
|
||||
/// `apply_iqn_trunk_gradient` slot 35).
|
||||
///
|
||||
/// Mapped-pinned host write — no HtoD copy. The kernel reads the same
|
||||
/// memory through the device-mapped pointer (`dev_ptr`) once the launch
|
||||
/// stream barrier ensures coherence.
|
||||
pub(crate) fn populate_nan_check_meta(
|
||||
&mut self,
|
||||
b: usize,
|
||||
sh2: usize,
|
||||
iqn_d_h_s2_ptr: Option<u64>,
|
||||
iqn_d_branch_logits_buf_ptr: Option<u64>,
|
||||
iqn_q: Option<usize>,
|
||||
) -> Result<(), MLError> {
|
||||
// Slot-28 length expression: `tba × b × q` per
|
||||
// `run_nan_checks_post_backward` slot 28. tba = b0+b1+b2+b3.
|
||||
let tba = self.config.branch_0_size
|
||||
+ self.config.branch_1_size
|
||||
+ self.config.branch_2_size
|
||||
+ self.config.branch_3_size;
|
||||
let bn_d_concat_len = self.bn_d_concat_buf().len() as i32;
|
||||
let trunk_param_count = self.trunk_param_count as i32;
|
||||
|
||||
let entries: [(u64, i32); 12] = [
|
||||
// slot 24 — post-c51_grad value gradient
|
||||
(self.d_value_logits_buf_ptr(), self.d_value_logits_buf.len() as i32),
|
||||
// slot 25 — post-c51_grad branch advantage
|
||||
(self.d_adv_logits_buf_ptr(), self.d_adv_logits_buf.len() as i32),
|
||||
// slot 26 — IQN trunk cuBLAS bwd output
|
||||
(self.ptrs.iqn_trunk_m, trunk_param_count),
|
||||
// slot 27 — IQN dh_s2 (caller-supplied; 0 when IQN inactive)
|
||||
(
|
||||
iqn_d_h_s2_ptr.unwrap_or(0),
|
||||
if iqn_d_h_s2_ptr.is_some() { (b * sh2) as i32 } else { 0 },
|
||||
),
|
||||
// slot 28 — IQN d_branch_logits_buf
|
||||
(
|
||||
iqn_d_branch_logits_buf_ptr.unwrap_or(0),
|
||||
match (iqn_d_branch_logits_buf_ptr, iqn_q) {
|
||||
(Some(_), Some(q)) => (tba * b * q.max(1)) as i32,
|
||||
_ => 0,
|
||||
},
|
||||
),
|
||||
// slot 29 — CQL gradient buffer
|
||||
(self.cql_d_value_logits_ptr(), self.cql_d_value_logits.len() as i32),
|
||||
// slot 30 — aux next-bar dh_s2
|
||||
(self.aux_dh_s2_nb_buf_ptr(), self.aux_dh_s2_nb_buf.len() as i32),
|
||||
// slot 31 — deferred ensemble (owned by FusedDqnTraining); null entry → kernel skip
|
||||
(0, 0),
|
||||
// slot 32 — bottleneck Linear backward dy
|
||||
(self.bn_d_concat_buf().raw_ptr(), bn_d_concat_len),
|
||||
// slots 33-35 — bw_d_h_s2 inline checks fire separately at backward
|
||||
// orchestration phases (post-main / post-aux / post-iqn). Fused
|
||||
// kernel skips these via null-pointer guard.
|
||||
(0, 0),
|
||||
(0, 0),
|
||||
(0, 0),
|
||||
];
|
||||
|
||||
// Host-side write through the mapped-pinned pages — kernel sees the
|
||||
// values via dev_ptr after the next stream-sync barrier (no HtoD copy
|
||||
// required, per `feedback_no_htod_htoh_only_mapped_pinned`).
|
||||
let mut ptrs_view = [0u64; 12];
|
||||
let mut lens_view = [0i32; 12];
|
||||
for (i, (ptr, len)) in entries.iter().enumerate() {
|
||||
ptrs_view[i] = *ptr;
|
||||
lens_view[i] = *len;
|
||||
}
|
||||
self.nan_check_buf_ptrs.write_from_slice(&ptrs_view);
|
||||
self.nan_check_buf_lens.write_from_slice(&lens_view);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// SP2: launch the fused NaN-check kernel. Single launch processes all 12
|
||||
/// backward-path slots (24-35) in 12 blocks, in parallel. Slots with null
|
||||
/// pointer (deferred slot 31, optional IQN slots 27/28 when inactive,
|
||||
/// slots 33-35 inline elsewhere) no-op via the kernel's null-pointer
|
||||
/// guard. Sticky-flag semantics preserved (kernel only writes 1, never
|
||||
/// clears).
|
||||
///
|
||||
/// Replaces the 8-launch sequence in `run_nan_checks_post_backward`.
|
||||
/// Wrapper unused at A3 commit time — call-site replacement lands in A4.
|
||||
pub(crate) fn launch_nan_check_fused_f32(&mut self) -> Result<(), MLError> {
|
||||
let nan_flags_ptr = self.nan_flags_buf.raw_ptr();
|
||||
let buf_ptrs_dev = self.nan_check_buf_ptrs.dev_ptr;
|
||||
let buf_lens_dev = self.nan_check_buf_lens.dev_ptr;
|
||||
const BLOCKS: u32 = 12;
|
||||
const THREADS: u32 = 256;
|
||||
const BASE_FLAG_IDX: i32 = 24;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (BLOCKS, 1, 1),
|
||||
block_dim: (THREADS, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.nan_check_fused_f32_kernel)
|
||||
.arg(&buf_ptrs_dev)
|
||||
.arg(&buf_lens_dev)
|
||||
.arg(&BASE_FLAG_IDX)
|
||||
.arg(&nan_flags_ptr)
|
||||
.launch(cfg)
|
||||
.map_err(|e| MLError::ModelError(format!("nan_check_fused_f32: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Zero the NaN flags buffer (call before a batch of check_nan calls).
|
||||
pub fn reset_nan_flags(&mut self) -> Result<(), MLError> {
|
||||
self.stream.memset_zeros(&mut self.nan_flags_buf)
|
||||
@@ -20529,7 +20676,7 @@ fn deflate_rank_one(mat: &[f32], n: usize, lambda_1: f32) -> Vec<f32> {
|
||||
fn compile_training_kernels(
|
||||
stream: &Arc<CudaStream>,
|
||||
config: &GpuDqnTrainConfig,
|
||||
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
|
||||
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
|
||||
info!(
|
||||
state_dim = ml_core::state_layout::STATE_DIM,
|
||||
total_params = compute_total_params(config),
|
||||
@@ -20662,6 +20809,13 @@ fn compile_training_kernels(
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_nan_check_f32 load: {e}")))?;
|
||||
let nan_check_f32_b = module.load_function("dqn_nan_check_f32_b")
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_nan_check_f32_b load: {e}")))?;
|
||||
// SP2: fused multi-buffer NaN check kernel. Single launch processes 12
|
||||
// backward-path slots (24-35) in 12 blocks. Replaces 8 individual
|
||||
// dqn_nan_check_f32 launches in `run_nan_checks_post_backward` (Task A4
|
||||
// wires the call site). Loaded from the same `module` as the per-buffer
|
||||
// variant — both belong to the post_aux_child captured replay group.
|
||||
let nan_check_fused_f32 = module.load_function("dqn_nan_check_fused_f32_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_nan_check_fused_f32_kernel load: {e}")))?;
|
||||
// SP1 Phase C surgical fix: in-place finite clamp kernel (replaces NaN/Inf
|
||||
// with 0; clamps finite |v| to ±max_abs from caller). Loaded from the same
|
||||
// `module` as the NaN check kernels because both are launched inside
|
||||
@@ -20676,8 +20830,8 @@ fn compile_training_kernels(
|
||||
let popart_robust = ungraphed_module.load_function("popart_normalize_robust")
|
||||
.map_err(|e| MLError::ModelError(format!("popart_normalize_robust load: {e}")))?;
|
||||
|
||||
info!("GpuDqnTrainer: 38 utility kernels loaded from precompiled cubin (5 CUmodules)");
|
||||
Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, clamp_finite_f32, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust))
|
||||
info!("GpuDqnTrainer: 39 utility kernels loaded from precompiled cubin (5 CUmodules)");
|
||||
Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, nan_check_fused_f32, clamp_finite_f32, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust))
|
||||
}
|
||||
|
||||
/// Load the standalone Polyak EMA kernel from precompiled cubin.
|
||||
|
||||
@@ -653,6 +653,26 @@ impl FusedTrainingCtx {
|
||||
None
|
||||
};
|
||||
|
||||
// SP2: populate fused NaN-check (ptr, len) tables via host-side write
|
||||
// through the trainer's mapped-pinned `nan_check_buf_ptrs` /
|
||||
// `nan_check_buf_lens` (no HtoD copy per
|
||||
// `feedback_no_htod_htoh_only_mapped_pinned`). Slot pointers are
|
||||
// stable across the trainer's lifetime — populate ONCE here, after
|
||||
// `gpu_iqn` has been constructed so its slot 27/28 accessors resolve
|
||||
// correctly. When IQN is inactive (`iqn_lambda == 0.0` →
|
||||
// `gpu_iqn = None`), slots 27/28 receive null entries which the fused
|
||||
// kernel skips via its null-pointer guard.
|
||||
let iqn_q_for_meta = gpu_iqn
|
||||
.as_ref()
|
||||
.map(|_| crate::cuda_pipeline::gpu_iqn_head::FIXED_TAUS.len());
|
||||
trainer.populate_nan_check_meta(
|
||||
batch_size,
|
||||
shared_h2,
|
||||
gpu_iqn.as_ref().map(|h| h.d_h_s2_raw_ptr()),
|
||||
gpu_iqn.as_ref().map(|h| h.d_branch_logits_buf_ptr()),
|
||||
iqn_q_for_meta,
|
||||
).map_err(|e| anyhow::anyhow!("populate_nan_check_meta: {e}"))?;
|
||||
|
||||
// GPU attention — always active. 4-head self-attention on h_s2.
|
||||
let attn_config = GpuAttentionConfig {
|
||||
state_dim: shared_h2,
|
||||
|
||||
@@ -2268,3 +2268,5 @@ SP1 closure (2026-04-29): F1 cold-start clamp pathology fixed at commit `ab21334
|
||||
SP2 Phase A1 — fused NaN-check kernel (2026-04-29): foundational kernel-only commit appending `dqn_nan_check_fused_f32_kernel` to `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu`. Single-launch replacement target for the 8 per-step `dqn_nan_check_f32` calls in `run_nan_checks_post_backward` (slots 24-30 + 32). Each block (`blockIdx.x = 0..N_SLOTS`) scans its assigned buffer end-to-end via grid-strided loop, block-local `__syncthreads_or` reduce (no atomicAdd per `feedback_no_atomicadd`), sticky-flag write `nan_flags_buf[base_flag_idx + slot] = 1` only when `has_nan && threadIdx.x == 0`. Nullptr bounds check `if (buf == nullptr || n == 0) return;` makes deferred slots (e.g. slot 31 ensemble) and unused entries no-op cleanly — same call site can pass null for cross-struct slots without spurious flag writes. Invariants preserved from `dqn_nan_check_f32`: sticky-flag (writes 1, never clears), graph-capture safe (pure kernel launch), no DtoD/HtoD/HtoH copies. `extern "C"` linkage matches existing `dqn_*` symbol convention for cudarc `module.load_function` lookup. Unused yet — Rust wrapper, metadata buffer (per-slot ptr/len arrays), and call-site replacement land in subsequent A2/A3/A4 commits. Drives F0 regression remediation: 8 per-step kernel launches collapse to 1 fused launch (~7× fewer graph nodes for backward NaN coverage) while preserving identical diagnostic semantics. Pattern reference for future fused diagnostic launches if SP3 expands the slot range.
|
||||
|
||||
SP2 Phase A2 — nan-check (ptr, len) device tables (2026-04-29): added two device buffers on `GpuDqnTrainer` to feed the fused NaN-check kernel from A1: `nan_check_buf_ptrs: CudaSlice<u64>` (12 entries, one per slot 24-35) and `nan_check_buf_lens: CudaSlice<i32>` (12 entries). Both are `stream.alloc_zeros` allocations adjacent to `nan_flags_buf`. Two separate buffers chosen over a packed-stride layout for direct ABI match with the kernel signature `(const float* const* buf_ptrs, const int* buf_lens, ...)` — no struct-stride alignment concerns. Sized at 12 (slots 24-35); slots 36-47 headroom is reserved for SP3 Mech 5 extension which will resize both buffers in lockstep. Slot 31 (deferred ensemble, cross-struct on `FusedDqnTraining` per A1's null-skip pattern) entry will hold 0 — kernel skips that block on `buf == nullptr`. Allocated as zeros; populated once in A3 (slot pointers are stable across the trainer's lifetime — no per-step host→device traffic, satisfying `feedback_no_htod_htoh_only_mapped_pinned` once A3 selects the population path: mapped-pinned helper from `mapped_pinned.rs` for the one-shot construction-time write). Unused yet — A3 wires the population + Rust launch wrapper, A4 replaces the 8 individual `check_nan_f32` calls in `run_nan_checks_post_backward` with a single fused launch. Field-level only — no behavioral change.
|
||||
|
||||
SP2 Phase A3 — populate + launch wrapper (2026-04-29): refactored A2's `nan_check_buf_ptrs: CudaSlice<u64>` / `nan_check_buf_lens: CudaSlice<i32>` to mapped-pinned (`MappedU64Buffer` / `MappedI32Buffer` from `mapped_pinned.rs`) per `feedback_no_htod_htoh_only_mapped_pinned` — host-side write through the mapped pages is visible to the kernel via the `dev_ptr` returned by `cuMemHostGetDevicePointer_v2`, eliminating the HtoD copy entirely. Added `populate_nan_check_meta(b, sh2, iqn_d_h_s2_ptr, iqn_d_branch_logits_buf_ptr, iqn_q)` on `GpuDqnTrainer` — one-shot construction-time writer of 12 `(ptr, len)` tuples covering slots 24-35 (mirrors the per-slot accessor table at the end of `docs/dqn-backward-nan-audit.md` and the per-slot launches in `run_nan_checks_post_backward`). Slot 31 deferred ensemble entry: `(0, 0)`. Slots 27/28 IQN entries: `Option<u64>` ptrs gated on `gpu_iqn.is_some()` — when IQN is inactive the entries stay zero and the fused kernel skips them via its null-pointer guard. Slots 33-35 (`bw_d_h_s2` inline checks) hold null entries — fused kernel skips; the inline `check_nan_f32` calls at the three backward orchestration phases (`launch_cublas_backward_to` post-main / post-aux / `apply_iqn_trunk_gradient` post-iqn) continue to fire individually for entry-point localization. Added `launch_nan_check_fused_f32` Rust wrapper — single launch with `grid_dim=12, block_dim=256, BASE_FLAG_IDX=24`, mirrors the kernel signature `(buf_ptrs_dev, buf_lens_dev, base_flag_idx, nan_flags_ptr)`. Kernel registered in the precompiled-cubin loader (`compile_training_kernels` tuple 43→44, 38→39 utility kernels logged) and stored on `GpuDqnTrainer` as `nan_check_fused_f32_kernel: CudaFunction` — same `module` (forward_child / aux_child / post_aux_child captured replay group) as the per-buffer `dqn_nan_check_f32` to preserve graph-capture compatibility for A4's call-site replacement. Constructor-time wire-up lives in `FusedTrainingCtx::new` (after `gpu_iqn` is constructed) — chosen over the GpuDqnTrainer constructor because `gpu_iqn` is owned by `FusedTrainingCtx`, mirroring the same `Option<u64>` argument pattern used by `apply_iqn_trunk_gradient(iqn_d_h_s2_ptr, ...)` and `run_nan_checks_post_backward`. Wrapper unused at A3 commit time — A4 replaces the 8-launch sequence in `run_nan_checks_post_backward` with the single fused launch. Zero new HtoD copies; pre-commit Invariant 7 satisfied; no ISV slots added; sticky-flag semantics preserved at the kernel.
|
||||
|
||||
Reference in New Issue
Block a user