fix(sp15-wave4.1b-OOB-followup): pre-load bn_tanh_concat_dd_kernel — fixes forward-capture SEGV
Root cause: SP15 Wave 4.1b (eb9515e41) migrated the trainer's 3 production bottleneck-concat call sites (online forward, target forward, DDQN argmax) from the pre-loaded `bn_tanh_concat_kernel: CudaFunction` field to a `launch_sp15_bn_concat_dd` free function that resolved the symbol via `load_cubin` + `load_function` ON EVERY CALL. That pattern is incompatible with CUDA Graph capture: `cuModuleLoadData` and `cuModuleGetFunction` are HOST-side driver API calls that allocate memory and mutate driver state, and they are NOT capturable inside a `cuStreamBeginCapture` region. Calling the launcher from `submit_forward_ops_main` (graph-captured forward child) caused a SEGV at exit 139 — the loader raced with the capture-mode driver state, the host-side corruption surfaced as a segfault before `CAPTURE_PHASE_FORWARD_DONE` could print. Diagnostic evidence (L40S smoke `train-vg5f7` on commit `bfc3ffa9d`): ALL 16 step-0 ungraphed checkpoints printed clean. Capture begins: CAPTURE_PHASE_BEGIN CAPTURE_PHASE_PER_SAMPLE_DONE CAPTURE_PHASE_COUNTERS_DONE CAPTURE_PHASE_SPECTRAL_DONE Then: SEGV. The next checkpoint that didn't print was CAPTURE_PHASE_FORWARD_DONE -> SEGV is inside `submit_forward_ops_main`'s `forward` child capture. The same function ran cleanly ungraphed in step 0 because the loader-host-branch was harmless without active capture. The experience collector's equivalent caller (gpu_experience_collector.rs:4127) was already doing this correctly — pre-loaded `bn_tanh_concat_fn` field populated at construction. Wave 4.1b's mistake was asymmetric: it kept the collector's pre-load pattern but introduced an on-demand loader for the trainer's call sites. Fix (atomic, restores graph-safe contract): - Add `bn_tanh_concat_dd_kernel: CudaFunction` field on `GpuDqnTrainer` (back-fills what Wave 4.1b removed, with updated docstring naming the new dd_pct-aware kernel). - Pre-load the symbol in `compile_training_kernels` from the same `module` as `bn_tanh_bw` / `bn_bias_grad` (forward-child captured replay group). Tuple grows 43 -> 44 CudaFunctions; struct ctor wires the field. - Change `launch_sp15_bn_concat_dd` signature to take `&CudaFunction` as parameter; remove the per-call `load_cubin` + `load_function`. All 3 trainer call sites pass `&self.bn_tanh_concat_dd_kernel`. - Promote `DQN_UTILITY_CUBIN` from `pub(crate)` to `pub` so the oracle tests in `crates/ml/tests/sp15_phase1_oracle_tests.rs` can pre-load the kernel handle (the launcher no longer hides this for them). - Update the 2 test call sites (kernel-level oracle + Wave 4.1c behavioral KL test) to load the kernel handle once up-front and pass it through. Verification: - `cargo check -p ml --tests` clean (release + dev profile). - `cargo test -p ml --test sp15_phase1_oracle_tests -- --ignored`: 36/36 pass, including the Wave 4.1c behavioral KL test that exercises two end-to-end forward passes through the modified launcher. Refs: SP15 Wave 4.1b (eb9515e41), pearl_no_host_branches_in_captured_graph, feedback_no_partial_refactor, feedback_wire_everything_up. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -76,7 +76,7 @@ use super::sp4_isv_slots::{
|
||||
};
|
||||
|
||||
// ── Precompiled cubins (build.rs → include_bytes! → ZERO runtime nvcc) ──────
|
||||
pub(crate) static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin"));
|
||||
pub static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin"));
|
||||
static EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ema_kernel.cubin"));
|
||||
static RELU_MASK_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/relu_mask_kernel.cubin"));
|
||||
static C51_LOSS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/c51_loss_kernel.cubin"));
|
||||
@@ -1070,8 +1070,26 @@ pub fn launch_sp15_dd_state_reduce(
|
||||
/// buffer and computing the dim.
|
||||
///
|
||||
/// Grid: ceil(B * concat_dd_dim / 256), Block: 256.
|
||||
///
|
||||
/// SP15 Wave 4.1b OOB fix (2026-05-07): the kernel handle is now passed in
|
||||
/// by the caller as a pre-loaded `&CudaFunction`, NOT resolved via
|
||||
/// `load_cubin` + `load_function` on every call. The on-demand resolution
|
||||
/// pattern is incompatible with CUDA Graph capture: `cuModuleLoadData` and
|
||||
/// `cuModuleGetFunction` are HOST-side driver API calls that allocate
|
||||
/// memory and mutate driver state, and they are NOT capturable inside a
|
||||
/// `cuStreamBeginCapture` region. Calling them from `submit_forward_ops_main`
|
||||
/// (which runs inside the `forward` child capture) caused a SEGV at exit
|
||||
/// 139 — the captured-step path entered the loader while the stream was
|
||||
/// already in capture mode, the cudarc wrapper raced with the driver's
|
||||
/// internal capture tracking, and the host-side state corruption surfaced
|
||||
/// as a segfault before `CAPTURE_PHASE_FORWARD_DONE` could print. Pre-
|
||||
/// loading the handle in `compile_training_kernels` (mirroring the legacy
|
||||
/// `bn_tanh_concat_kernel` field that Wave 4.1b removed) restores the
|
||||
/// graph-safe contract: each child capture sees only kernel launches, not
|
||||
/// driver-state mutations. See `pearl_no_host_branches_in_captured_graph`.
|
||||
pub fn launch_sp15_bn_concat_dd(
|
||||
stream: &Arc<CudaStream>,
|
||||
kernel: &CudaFunction,
|
||||
bn_hidden: cudarc::driver::sys::CUdeviceptr,
|
||||
states: cudarc::driver::sys::CUdeviceptr,
|
||||
isv: cudarc::driver::sys::CUdeviceptr,
|
||||
@@ -1082,22 +1100,11 @@ pub fn launch_sp15_bn_concat_dd(
|
||||
state_dim_padded: i32,
|
||||
concat_dd_dim: i32,
|
||||
) -> Result<(), MLError> {
|
||||
let module = stream
|
||||
.context()
|
||||
.load_cubin(DQN_UTILITY_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"load dqn_utility_kernels cubin for bn_tanh_concat_dd: {e}"
|
||||
)))?;
|
||||
let kernel = module
|
||||
.load_function("bn_tanh_concat_dd_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"load bn_tanh_concat_dd_kernel function: {e}"
|
||||
)))?;
|
||||
let total_elems = (batch_size as i64).saturating_mul(concat_dd_dim as i64);
|
||||
let blocks = (((total_elems as u64) + 255) / 256).max(1) as u32;
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&kernel)
|
||||
.launch_builder(kernel)
|
||||
.arg(&bn_hidden)
|
||||
.arg(&states)
|
||||
.arg(&isv)
|
||||
@@ -5633,17 +5640,30 @@ pub struct GpuDqnTrainer {
|
||||
on_next_bn_hidden_buf: CudaSlice<f32>,
|
||||
/// DDQN argmax-pass mirror of `bn_concat_buf` (online-on-next_states).
|
||||
on_next_bn_concat_buf: CudaSlice<f32>,
|
||||
// SP15 Wave 4.1b removed the legacy `bn_tanh_concat_kernel` field —
|
||||
// production callers (3 forward sites here + 1 in `gpu_experience_collector`)
|
||||
// migrated to `launch_sp15_bn_concat_dd` (the bottleneck-aware
|
||||
// dd_pct-appending launcher), which loads the
|
||||
// `bn_tanh_concat_dd_kernel` symbol on-demand from the shared
|
||||
// `dqn_utility_kernels.cubin`. The legacy `bn_tanh_concat_kernel`
|
||||
// and `bn_tanh_concat_f32_kernel` symbols stay in the cubin source
|
||||
// for the SP15 oracle tests in
|
||||
// `crates/ml/tests/sp15_phase1_oracle_tests.rs` (parity check
|
||||
// between old and new kernels), but no production code path resolves
|
||||
// them anymore.
|
||||
/// SP15 Wave 4.1b OOB fix (2026-05-07): pre-loaded handle for the
|
||||
/// bottleneck-aware `bn_tanh_concat_dd_kernel` (replaces the legacy
|
||||
/// `bn_tanh_concat_kernel` + `bn_tanh_concat_f32_kernel` fields with a
|
||||
/// single dd_pct-appending kernel). Wave 4.1b originally migrated the 3
|
||||
/// production trainer call sites (online forward, target forward, DDQN
|
||||
/// argmax) to a `launch_sp15_bn_concat_dd` free function that resolved
|
||||
/// the symbol via `load_cubin` + `load_function` ON EVERY CALL — that
|
||||
/// pattern is incompatible with CUDA Graph capture: those CUDA driver
|
||||
/// API calls allocate memory and mutate driver state, which is NOT
|
||||
/// capturable inside `cuStreamBeginCapture`. Calling the launcher from
|
||||
/// `submit_forward_ops_main` (graph-captured forward child) caused a
|
||||
/// SEGV at exit 139 — the loader raced with the capture-mode driver
|
||||
/// state and the host-side corruption surfaced as a segfault before
|
||||
/// `CAPTURE_PHASE_FORWARD_DONE` could print. The experience collector's
|
||||
/// equivalent caller (line 4127 of gpu_experience_collector.rs) was
|
||||
/// already pre-loading correctly via `bn_tanh_concat_fn` field — this
|
||||
/// field restores the same graph-safe pattern on the trainer side. All
|
||||
/// 3 trainer call sites consume this single handle (Hopper's
|
||||
/// "no-CUfunction-sharing-across-children" comment elsewhere refers to
|
||||
/// `dqn_utility_kernels` workhorse kernels like `dqn_saxpy_f32_kernel`
|
||||
/// that are invoked from MULTIPLE child captures — this kernel only
|
||||
/// runs in the forward + ddqn children, mirroring `bn_tanh_backward_kernel`'s
|
||||
/// single-handle pattern for the forward+backward chain).
|
||||
bn_tanh_concat_dd_kernel: CudaFunction,
|
||||
/// #31 Backward scratch: d_loss/d_bn_concat [B, concat_dim] f32.
|
||||
bn_d_concat_buf: CudaSlice<f32>,
|
||||
/// #31 Backward scratch: d_loss/d_bn [B, bn_dim] f32 (after tanh derivative).
|
||||
@@ -17743,7 +17763,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, 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) =
|
||||
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, 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, bn_tanh_concat_dd_kernel) =
|
||||
compile_training_kernels(&stream, &config)?;
|
||||
|
||||
// Separate grad_norm instance for non-graph launches (outside CUDA graph).
|
||||
@@ -22584,8 +22604,9 @@ impl GpuDqnTrainer {
|
||||
on_next_bn_concat_buf,
|
||||
bn_d_concat_buf,
|
||||
bn_d_hidden_buf,
|
||||
// bn_tanh_concat_kernel field removed in SP15 Wave 4.1b — see
|
||||
// the field-removal comment block in the struct definition.
|
||||
// SP15 Wave 4.1b OOB fix (2026-05-07): pre-loaded handle, see
|
||||
// the field docstring in the struct definition.
|
||||
bn_tanh_concat_dd_kernel,
|
||||
bn_tanh_backward_kernel,
|
||||
bn_bias_grad_kernel,
|
||||
// ── Plan 4 Task 1B-iii: VSN feature-selection buffers ──────
|
||||
@@ -27417,6 +27438,7 @@ impl GpuDqnTrainer {
|
||||
// so by the time this kernel reads slot 406 the value is fresh.
|
||||
launch_sp15_bn_concat_dd(
|
||||
&self.stream,
|
||||
&self.bn_tanh_concat_dd_kernel,
|
||||
bn_hidden_ptr,
|
||||
on_next_states_for_bn,
|
||||
self.isv_signals_dev_ptr,
|
||||
@@ -27754,6 +27776,7 @@ impl GpuDqnTrainer {
|
||||
|
||||
launch_sp15_bn_concat_dd(
|
||||
&self.stream,
|
||||
&self.bn_tanh_concat_dd_kernel,
|
||||
bn_hidden_ptr,
|
||||
states_for_bn,
|
||||
self.isv_signals_dev_ptr,
|
||||
@@ -27948,6 +27971,7 @@ impl GpuDqnTrainer {
|
||||
let b_i32 = b as i32;
|
||||
launch_sp15_bn_concat_dd(
|
||||
&self.stream,
|
||||
&self.bn_tanh_concat_dd_kernel,
|
||||
tg_bn_hidden_ptr,
|
||||
tg_states_for_bn,
|
||||
self.isv_signals_dev_ptr,
|
||||
@@ -31678,7 +31702,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),
|
||||
@@ -31802,11 +31826,15 @@ fn compile_training_kernels(
|
||||
.map_err(|e| MLError::ModelError(format!("bn_tanh_backward_kernel load: {e}")))?;
|
||||
let bn_bias_grad = module.load_function("bn_bias_grad_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("bn_bias_grad_kernel load: {e}")))?;
|
||||
// SP15 Wave 4.1b: legacy `bn_tanh_concat_kernel` no longer loaded —
|
||||
// production callers migrated to `launch_sp15_bn_concat_dd`, which
|
||||
// resolves the `bn_tanh_concat_dd_kernel` symbol from the same cubin.
|
||||
// The legacy symbol remains in the source for the SP15 oracle parity
|
||||
// tests but is not surfaced as a `CudaFunction` field on the trainer.
|
||||
// SP15 Wave 4.1b OOB fix (2026-05-07): pre-load the dd_pct-appending
|
||||
// bottleneck concat kernel from the same `module` as `bn_tanh_bw` /
|
||||
// `bn_bias_grad` (forward-child captured replay group). The 3 trainer
|
||||
// production callers (online forward, target forward, DDQN argmax)
|
||||
// consume this handle via `&self.bn_tanh_concat_dd_kernel`; resolving
|
||||
// the symbol on-demand inside the captured forward child caused a
|
||||
// SEGV (cf. the field-level docstring in the struct definition).
|
||||
let bn_tanh_concat_dd = module.load_function("bn_tanh_concat_dd_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("bn_tanh_concat_dd_kernel load: {e}")))?;
|
||||
|
||||
let her_inplace = module.load_function("her_inplace_relabel")
|
||||
.map_err(|e| MLError::ModelError(format!("her_inplace_relabel load: {e}")))?;
|
||||
@@ -31835,8 +31863,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: 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, 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))
|
||||
info!("GpuDqnTrainer: 40 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, 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, bn_tanh_concat_dd))
|
||||
}
|
||||
|
||||
/// Load the standalone Polyak EMA kernel from precompiled cubin.
|
||||
|
||||
@@ -880,9 +880,16 @@ mod gpu {
|
||||
#[test]
|
||||
#[ignore = "requires GPU"]
|
||||
fn bn_tanh_concat_dd_kernel_writes_dd_pct_column() {
|
||||
use ml::cuda_pipeline::gpu_dqn_trainer::launch_sp15_bn_concat_dd;
|
||||
use ml::cuda_pipeline::gpu_dqn_trainer::{launch_sp15_bn_concat_dd, DQN_UTILITY_CUBIN};
|
||||
|
||||
let stream = make_test_stream();
|
||||
// SP15 Wave 4.1b OOB fix (2026-05-07): the launcher now consumes a
|
||||
// pre-loaded `&CudaFunction` (graph-safe), so tests resolve the
|
||||
// symbol once up-front rather than relying on the launcher itself.
|
||||
let module = stream.context().load_cubin(DQN_UTILITY_CUBIN.to_vec())
|
||||
.expect("load dqn_utility_kernels cubin");
|
||||
let kernel = module.load_function("bn_tanh_concat_dd_kernel")
|
||||
.expect("load bn_tanh_concat_dd_kernel symbol");
|
||||
|
||||
// Synthetic batch matches production layout: bn_dim=16,
|
||||
// market_dim=42, STATE_DIM=128 → state_dim_padded=128 (no
|
||||
@@ -940,6 +947,7 @@ mod gpu {
|
||||
|
||||
launch_sp15_bn_concat_dd(
|
||||
&stream,
|
||||
&kernel,
|
||||
bn_hidden_buf.dev_ptr,
|
||||
states_buf.dev_ptr,
|
||||
isv_buf.dev_ptr,
|
||||
@@ -3670,7 +3678,8 @@ mod sp15_wave_4_1c_behavioral {
|
||||
|
||||
use cudarc::cublas::sys as cublas_sys;
|
||||
use cudarc::driver::{CudaContext, CudaStream};
|
||||
use ml::cuda_pipeline::gpu_dqn_trainer::launch_sp15_bn_concat_dd;
|
||||
use cudarc::driver::CudaFunction;
|
||||
use ml::cuda_pipeline::gpu_dqn_trainer::{launch_sp15_bn_concat_dd, DQN_UTILITY_CUBIN};
|
||||
use ml::cuda_pipeline::mapped_pinned::MappedF32Buffer;
|
||||
use ml::cuda_pipeline::shared_cublas_handle::PerStreamCublasHandles;
|
||||
use ml::cuda_pipeline::sp15_isv_slots::{DD_PCT_INDEX, SP15_SLOT_END};
|
||||
@@ -3729,6 +3738,7 @@ mod sp15_wave_4_1c_behavioral {
|
||||
fn synthetic_forward(
|
||||
stream: &Arc<CudaStream>,
|
||||
cublas_handle: cublas_sys::cublasHandle_t,
|
||||
bn_concat_kernel: &CudaFunction,
|
||||
bn_hidden_dev: u64,
|
||||
states_dev: u64,
|
||||
isv_buf: &MappedF32Buffer,
|
||||
@@ -3764,6 +3774,7 @@ mod sp15_wave_4_1c_behavioral {
|
||||
// write returns).
|
||||
launch_sp15_bn_concat_dd(
|
||||
stream,
|
||||
bn_concat_kernel,
|
||||
bn_hidden_dev,
|
||||
states_dev,
|
||||
isv_buf.dev_ptr,
|
||||
@@ -3940,10 +3951,18 @@ mod sp15_wave_4_1c_behavioral {
|
||||
.classic_handle_for(&stream)
|
||||
.expect("classic_handle_for default stream");
|
||||
|
||||
// ── Pre-load the bn_tanh_concat_dd_kernel handle (graph-safe pattern;
|
||||
// SP15 Wave 4.1b OOB fix 2026-05-07: launcher consumes &CudaFunction). ──
|
||||
let util_module = stream.context().load_cubin(DQN_UTILITY_CUBIN.to_vec())
|
||||
.expect("load dqn_utility_kernels cubin");
|
||||
let bn_concat_kernel = util_module.load_function("bn_tanh_concat_dd_kernel")
|
||||
.expect("load bn_tanh_concat_dd_kernel symbol");
|
||||
|
||||
// ── At-ATH pass (dd_pct = 0.0) ──
|
||||
let logits_ath = synthetic_forward(
|
||||
&stream,
|
||||
cublas_handle,
|
||||
&bn_concat_kernel,
|
||||
bn_hidden_buf.dev_ptr,
|
||||
states_buf.dev_ptr,
|
||||
&isv_buf,
|
||||
@@ -3962,6 +3981,7 @@ mod sp15_wave_4_1c_behavioral {
|
||||
let logits_dd = synthetic_forward(
|
||||
&stream,
|
||||
cublas_handle,
|
||||
&bn_concat_kernel,
|
||||
bn_hidden_buf.dev_ptr,
|
||||
states_buf.dev_ptr,
|
||||
&isv_buf,
|
||||
|
||||
Reference in New Issue
Block a user