foxhunt/docs/superpowers/plans/2026-04-04-cuda-kernel-safety.md

CUDA Kernel Safety + Integration Test Fixes

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Fix the __restrict__ aliasing NaN bug, audit all CUDA kernels, tune smoketest profile, get all integration tests passing on RTX 3050 in <2 minutes.

Architecture: Add a dedicated dqn_scale_f32_kernel to the CUDA utility kernels, replace the 2 aliased SAXPY calls, audit all other kernels, tune the smoketest TOML for fast validation.

Tech Stack: CUDA C (.cu kernels), Rust (cudarc), cargo test

---

File Map

File	Change
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu	Add dqn_scale_f32_kernel
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs	Load scale kernel, fix 2 aliased calls
config/training/dqn-smoketest.toml	Tune gpu_n_episodes, buffer_size, etc.
crates/ml/src/training_profile.rs	Update test assertions

---

Task 1: Add dqn_scale_f32_kernel to CUDA utility kernels

Files:

• Modify: crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu

• Step 1: Add the in-place scale kernel

In crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu, add after the existing dqn_saxpy_f32_kernel (after line 166):

/* ══════════════════════════════════════════════════════════════════════
 * F32 IN-PLACE SCALE KERNEL
 *
 * y[i] *= alpha   for i = 0..n-1
 *
 * Single-pointer variant — no aliasing possible, no __restrict__ needed.
 * Used for gradient scaling where SAXPY y+=α*x with x==y violates
 * __restrict__ and causes undefined behavior (NaN from register
 * corruption under NVCC's no-alias optimizations).
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_scale_f32_kernel(
    float* y,
    float alpha,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) y[i] *= alpha;
}

• Step 2: Rebuild the cubin

The project precompiles CUDA kernels into a .cubin or .ptx file. Check how the existing kernels are compiled:

grep -r "dqn_utility_kernels" crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | head -5

Follow the same compilation pattern. The kernel will be available via module.load_function("dqn_scale_f32_kernel").

• Step 3: Verify the kernel compiles

SQLX_OFFLINE=true cargo check -p ml

• Step 4: Commit

git add crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
git commit -m "feat: add dqn_scale_f32_kernel (in-place scale, no __restrict__ aliasing)"

---

Task 2: Load scale kernel + fix aliased SAXPY calls

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

• Step 1: Add scale_f32_kernel field to GpuDqnTrainer

Find the struct definition (search for saxpy_f32_kernel: CudaFunction). Add after it:

    scale_f32_kernel: CudaFunction,

• Step 2: Load the kernel in the constructor

Find where saxpy_f32 is loaded (around line 5594). Add after it:

    let scale_f32 = module.load_function("dqn_scale_f32_kernel")
        .map_err(|e| MLError::ModelError(format!("dqn_scale_f32_kernel load: {e}")))?;

Find where the struct is constructed (search for saxpy_f32_kernel: saxpy_f32). Add:

    scale_f32_kernel: scale_f32,

• Step 3: Fix the 2 aliased SAXPY calls

In the C51/MSE blending code (around line 4090-4130), replace the aliased calls.

Current code (lines 4096, 4109-4130):

let scale_c51 = alpha - 1.0; // SAXPY: y += (α-1)*y → y *= α

Replace the entire unsafe block (lines 4109-4130) with:

            unsafe {
                // d_value *= α  (in-place scale — no pointer aliasing)
                self.stream.launch_builder(&self.scale_f32_kernel)
                    .arg(&val_ptr)
                    .arg(&alpha).arg(&n_val)
                    .launch(cfg_val).map_err(|e| MLError::ModelError(format!("scale c51 val: {e}")))?;
                // d_value += (1-α) * mse_scratch  (SAXPY with distinct pointers — safe)
                self.stream.launch_builder(&self.saxpy_f32_kernel)
                    .arg(&val_ptr).arg(&val_mse_ptr)
                    .arg(&scale_mse).arg(&n_val)
                    .launch(cfg_val).map_err(|e| MLError::ModelError(format!("blend mse val: {e}")))?;
                // d_adv *= α  (in-place scale — no pointer aliasing)
                self.stream.launch_builder(&self.scale_f32_kernel)
                    .arg(&adv_ptr)
                    .arg(&alpha).arg(&n_adv)
                    .launch(cfg_adv).map_err(|e| MLError::ModelError(format!("scale c51 adv: {e}")))?;
                // d_adv += (1-α) * mse_scratch  (SAXPY with distinct pointers — safe)
                self.stream.launch_builder(&self.saxpy_f32_kernel)
                    .arg(&adv_ptr).arg(&adv_mse_ptr)
                    .arg(&scale_mse).arg(&n_adv)
                    .launch(cfg_adv).map_err(|e| MLError::ModelError(format!("blend mse adv: {e}")))?;
            }

Key change: The scale kernel takes alpha directly (not alpha - 1.0). The old SAXPY did y += (α-1)*y which equals y*α. The new scale kernel does y *= α directly. So remove the let scale_c51 = alpha - 1.0; line (line 4096) — it's no longer needed.

• Step 4: Verify compilation

SQLX_OFFLINE=true cargo check -p ml

• Step 5: Commit

git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "fix: replace aliased SAXPY with in-place scale kernel (fixes __restrict__ NaN bug)"

---

Task 3: Audit all CUDA kernels for __restrict__ aliasing

Files:

• All .cu files in crates/ml/src/cuda_pipeline/

• All .rs files in crates/ml/src/cuda_pipeline/

• Step 1: Find all kernels with multiple __restrict__ pointer params

grep -n "__restrict__" crates/ml/src/cuda_pipeline/*.cu | grep "extern" 

For each kernel that has 2+ __restrict__ pointer parameters, note the kernel name.

• Step 2: Find all Rust call sites for those kernels

For each kernel found in step 1, search for its call sites in the .rs files:

grep -rn "launch_builder.*kernel_name\|kernel_name.*launch" crates/ml/src/cuda_pipeline/*.rs

• Step 3: Check each call site for aliased pointers

For each call site, verify that the pointer arguments are distinct variables. The already-confirmed safe patterns from our earlier analysis:

Location	Pointers	Status
gpu_dqn_trainer.rs:4112	val_ptr, val_ptr	FIXED in Task 2
gpu_dqn_trainer.rs:4122	adv_ptr, adv_ptr	FIXED in Task 2
gpu_dqn_trainer.rs:4117	val_ptr, val_mse_ptr	SAFE — distinct
gpu_dqn_trainer.rs:4127	adv_ptr, adv_mse_ptr	SAFE — distinct
gpu_dqn_trainer.rs:1822	w_ptr, u_ptr, v_ptr	SAFE — distinct
All decision_transformer.rs calls	distinct pointers	SAFE

• Step 4: Add safety comments to any ambiguous call sites

For any SAXPY call where aliasing isn't immediately obvious from variable names, add a comment:

// SAFETY: val_ptr and val_mse_ptr are distinct buffers (d_value_logits_buf vs d_value_logits_mse)

• Step 5: Document findings and commit

git add crates/ml/src/cuda_pipeline/
git commit -m "audit: all __restrict__ CUDA kernels verified for pointer aliasing safety"

---

Task 4: Tune smoketest profile for fast validation

Files:

• Modify: config/training/dqn-smoketest.toml

• Modify: crates/ml/src/training_profile.rs

• Step 1: Update smoketest TOML

In config/training/dqn-smoketest.toml, change:

# [training] section:
buffer_size = 256

# [replay_buffer] section:
buffer_size = 256
min_replay_size = 64

# [experience] section:
gpu_timesteps_per_episode = 50
gpu_n_episodes = 4

Read the file first to find the exact lines. The buffer_size may appear in both [training] and [replay_buffer] sections.

• Step 2: Update training_profile test assertions

Search crates/ml/src/training_profile.rs for assertions on gpu_n_episodes, gpu_timesteps_per_episode, buffer_size, or min_replay_size that reference old smoketest values. Update to match the new TOML.

Also check any test that asserts buffer_size = 10000 or gpu_n_episodes = 32.

• Step 3: Verify profile tests pass

SQLX_OFFLINE=true cargo test -p ml --lib -- training_profile

• Step 4: Commit

git add config/training/dqn-smoketest.toml crates/ml/src/training_profile.rs
git commit -m "perf: tune smoketest profile for fast validation (buffer=256, episodes=4)"

---

Task 5: Run integration tests and verify all pass

Files: None (verification only)

• Step 1: Run early stopping tests

SQLX_OFFLINE=true cargo test -p ml --release --test dqn_early_stopping_termination_test -- --nocapture

Expected: all 4 pass, no NaN, <60s.

• Step 2: Run pipeline tests

SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --release --test dqn_training_pipeline_test -- --nocapture

Expected: all 5 pass (or gracefully skip if data not found), no NaN, <60s.

• Step 3: Run CQL test

SQLX_OFFLINE=true cargo test -p ml --release --test dqn_action_collapse_fix_test -- --nocapture

Expected: all pass.

• Step 4: Run GPU smoke tests

SQLX_OFFLINE=true cargo test -p ml --release --lib -- trainers::dqn::smoke_tests -- --nocapture

Expected: all pass, no CUDA OOB.

• Step 5: Run GPU profile + unit tests

SQLX_OFFLINE=true cargo test -p ml-core --lib -- gpu::profile
SQLX_OFFLINE=true cargo test -p ml --lib -- fxcache feature_cache training_profile

Expected: all pass.

• Step 6: Push

git push origin main