Files
foxhunt/docs/superpowers/specs/2026-04-04-cuda-kernel-safety-design.md
2026-04-04 11:07:47 +02:00

5.7 KiB
Raw Blame History

CUDA Kernel Safety + Integration Test Fixes

Goal: Fix the __restrict__ aliasing violation that causes NaN in training, audit all CUDA kernels for similar bugs, tune smoketest profile for fast validation, and get all integration tests passing on RTX 3050.

Root cause: gpu_dqn_trainer.rs lines 4111-4124 pass the same pointer as both y (output) and x (input) to dqn_saxpy_f32_kernel. The kernel declares both parameters __restrict__, which promises the compiler they don't alias. When they DO alias, the compiler's register caching and reordering produces undefined behavior → NaN. This affects ALL network sizes — small networks trigger NaN faster, large networks may silently produce wrong gradients.


Part 1: In-place scale kernel

New CUDA kernel

Add dqn_scale_f32_kernel to crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu:

extern "C" __global__ void dqn_scale_f32_kernel(
    float* y,     // No __restrict__ needed — single pointer, no aliasing
    float alpha,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) y[i] *= alpha;
}

This kernel does y *= alpha in-place. Safe, simple, no aliasing possible.

Load kernel in GpuDqnTrainer

Add scale_f32_kernel: CudaFunction field to GpuDqnTrainer. Load from the precompiled cubin alongside the existing SAXPY kernel.

Fix the aliased calls

Replace gpu_dqn_trainer.rs lines 4111-4114 and 4121-4124:

Before (buggy):

// d_value += (α-1) * d_value → d_value *= α  (ALIASED — UB!)
self.stream.launch_builder(&self.saxpy_f32_kernel)
    .arg(&val_ptr).arg(&val_ptr)   // SAME POINTER — __restrict__ violation
    .arg(&scale_c51).arg(&n_val)
    .launch(cfg_val)?;

After (correct):

// d_value *= α  (in-place scale, no aliasing)
self.stream.launch_builder(&self.scale_f32_kernel)
    .arg(&val_ptr)
    .arg(&scale_c51_plus_one).arg(&n_val)  // alpha = scale_c51 + 1.0 (SAXPY was y += (α-1)*y → y *= α)
    .launch(cfg_val)?;

Note: the SAXPY call did y += (alpha-1) * y which equals y *= alpha. The scale kernel takes alpha directly (not alpha-1). So scale_c51_plus_one = scale_c51 + 1.0.

Same fix for d_adv (lines 4121-4124).

The two non-aliased SAXPY calls (lines 4116-4119, 4126-4129) stay as-is — they use distinct pointers (val_ptr + val_mse_ptr) correctly.


Part 2: Full __restrict__ audit

Audit every CUDA kernel in crates/ml/src/cuda_pipeline/ for aliasing violations:

Scope

  1. All .cu kernel files — find every __restrict__ parameter
  2. All Rust launch_builder calls — find every call site that passes the same pointer for multiple args
  3. Cross-reference — for each __restrict__ kernel with 2+ pointer args, verify no Rust call site aliases them

Files to audit

File Kernel count Risk
dqn_utility_kernels.cu ~27 kernels HIGH — SAXPY bug found here
experience_collector_kernels.cu ~3 kernels MEDIUM
gpu_per_kernels.cu ~5 kernels LOW — PER uses distinct buffers
c51_loss_kernels.cu ~2 kernels HIGH — loss computation
gpu_attention.cu ~4 kernels MEDIUM
gpu_iqn_head.cu ~9 kernels HIGH — IQN uses quantile buffers
gpu_iql_trainer.cu ~5 kernels MEDIUM
gpu_her.cu ~2 kernels LOW

Deliverable

For each violation found:

  1. Document which kernel, which call site, which pointers alias
  2. Fix with appropriate kernel (in-place scale, or remove __restrict__ if aliasing is intentional and performance-neutral)
  3. Add a comment at each call site documenting pointer aliasing safety

Part 3: Smoketest profile tuning

config/training/dqn-smoketest.toml changes

Field Old New Reason
gpu_n_episodes 32 4 8x fewer episodes per collection
gpu_timesteps_per_episode 100 50 2x shorter episodes
buffer_size 10000 256 Fills fast, trains immediately
min_replay_size 100 64 Start training after 1 collection round
epochs 3 3 (already correct)

With these settings: 4 episodes × 50 timesteps = 200 experiences per collection. Buffer fills in 2 rounds. Training starts almost immediately. 3 epochs should complete in <30s on RTX 3050.

Update test assertions

Any training_profile test that asserts these values needs updating.


Part 4: Integration test verification

After Parts 1-3, run all integration tests on RTX 3050 (release mode, one at a time):

SQLX_OFFLINE=true cargo test -p ml --release --test dqn_early_stopping_termination_test -- --nocapture
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --release --test dqn_training_pipeline_test -- --nocapture
SQLX_OFFLINE=true cargo test -p ml --release --test dqn_action_collapse_fix_test -- --nocapture
SQLX_OFFLINE=true cargo test -p ml --lib -- trainers::dqn::smoke_tests

Success criteria: All tests pass. No NaN. Total time <2 minutes on RTX 3050.


Files touched

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 SAXPY calls
crates/ml/src/cuda_pipeline/*.cu Audit all __restrict__ kernels
crates/ml/src/cuda_pipeline/*.rs Fix any other aliased calls found in audit
config/training/dqn-smoketest.toml Tune profile for fast validation
crates/ml/src/training_profile.rs Update test assertions for new profile values

What this does NOT change

  • Production training profiles or hyperparameters
  • The fxcache pipeline (fixed in previous spec)
  • Argo workflow templates
  • Any CUDA kernel logic beyond the aliasing fix