Three optimizations to IQN forward+loss and backward kernels:
1. Block size 32→256 (8 warps): occupancy 3.1%→25% on H100.
Inner hidden_dim loops now stride by 256 instead of 32.
Block-level reduction via shared memory replaces warp-only shuffle.
2. Precomputed cosine features [N, embed_dim]: eliminates 16.7M
cosf() calls per step. cos(π·(d+1)·τ_i) computed once at
construction (τ are fixed QR-DQN midpoints).
3. Quantile Huber loss distributed across 256 threads: was single-lane
serial (32×32=1024 iterations on lane 0). Now each thread handles
~4 pairs.
Expected impact: IQN from ~80ms to ~10-15ms per step (occupancy +
cosine elimination + parallel loss).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
6ff45466f9