# 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
```
