docs: implementation plan for CUDA kernel safety fix

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-04 11:12:47 +02:00
parent 37d78e02c8
commit 938d75fdac
1 changed files with 310 additions and 0 deletions
--- a/docs/superpowers/plans/2026-04-04-cuda-kernel-safety.md
+++ b/docs/superpowers/plans/2026-04-04-cuda-kernel-safety.md
@@ -0,0 +1,310 @@
+# 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):
+
+```cuda
+/* ══════════════════════════════════════════════════════════════════════
+ * 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:
+
+```bash
+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**
+
+```bash
+SQLX_OFFLINE=true cargo check -p ml
+```
+
+- [ ] **Step 4: Commit**
+
+```bash
+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:
+
+```rust
+    scale_f32_kernel: CudaFunction,
+```
+
+- [ ] **Step 2: Load the kernel in the constructor**
+
+Find where `saxpy_f32` is loaded (around line 5594). Add after it:
+
+```rust
+    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:
+
+```rust
+    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):**
+```rust
+let scale_c51 = alpha - 1.0; // SAXPY: y += (α-1)*y → y *= α
+```
+
+Replace the entire `unsafe` block (lines 4109-4130) with:
+
+```rust
+            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**
+
+```bash
+SQLX_OFFLINE=true cargo check -p ml
+```
+
+- [ ] **Step 5: Commit**
+
+```bash
+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**
+
+```bash
+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:
+
+```bash
+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:
+
+```rust
+// 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**
+
+```bash
+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:
+
+```toml
+# [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**
+
+```bash
+SQLX_OFFLINE=true cargo test -p ml --lib -- training_profile
+```
+
+- [ ] **Step 4: Commit**
+
+```bash
+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**
+
+```bash
+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**
+
+```bash
+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**
+
+```bash
+SQLX_OFFLINE=true cargo test -p ml --release --test dqn_action_collapse_fix_test -- --nocapture
+```
+
+Expected: all pass.
+
+- [ ] **Step 4: Run GPU smoke tests**
+
+```bash
+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**
+
+```bash
+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**
+
+```bash
+git push origin main
+```