5.7 KiB
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:
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):
// 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):
// 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
- All
.cukernel files — find every__restrict__parameter - All Rust
launch_buildercalls — find every call site that passes the same pointer for multiple args - 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:
- Document which kernel, which call site, which pointers alias
- Fix with appropriate kernel (in-place scale, or remove
__restrict__if aliasing is intentional and performance-neutral) - 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):
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