From 938d75fdac0271c965951f9bf36418b9204494be Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 4 Apr 2026 11:12:47 +0200 Subject: [PATCH] docs: implementation plan for CUDA kernel safety fix Co-Authored-By: Claude Opus 4.6 (1M context) --- .../plans/2026-04-04-cuda-kernel-safety.md | 310 ++++++++++++++++++ 1 file changed, 310 insertions(+) create mode 100644 docs/superpowers/plans/2026-04-04-cuda-kernel-safety.md diff --git a/docs/superpowers/plans/2026-04-04-cuda-kernel-safety.md b/docs/superpowers/plans/2026-04-04-cuda-kernel-safety.md new file mode 100644 index 000000000..2ed049ed4 --- /dev/null +++ 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 +```