# 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
