From 37d78e02c8f91fd93809d1a9ef6fc2312d004a42 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 4 Apr 2026 11:07:47 +0200 Subject: [PATCH] docs: spec for CUDA kernel safety fix (__restrict__ aliasing NaN bug) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../2026-04-04-cuda-kernel-safety-design.md | 145 ++++++++++++++++++ 1 file changed, 145 insertions(+) create mode 100644 docs/superpowers/specs/2026-04-04-cuda-kernel-safety-design.md diff --git a/docs/superpowers/specs/2026-04-04-cuda-kernel-safety-design.md b/docs/superpowers/specs/2026-04-04-cuda-kernel-safety-design.md new file mode 100644 index 000000000..900d03dba --- /dev/null +++ b/docs/superpowers/specs/2026-04-04-cuda-kernel-safety-design.md @@ -0,0 +1,145 @@ +# 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`: + +```cuda +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):** +```rust +// 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):** +```rust +// 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): + +```bash +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