spec: bypass cudarc for hot path — raw CUDA driver API + mega-kernel fusion

cudarc adds 18.8ms/step overhead at b=256 (GPU kernels take 0.16ms).
Phase 1: raw cuLaunchKernel wrapper. Phase 2: pre-extracted raw ptrs.
Phase 3: fused mega-kernels (20 launches → 3).
Target: 1ms/step → 500+ sps.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 09:39:17 +02:00
parent 62adf13187
commit dac9cfef5c

View File

@@ -0,0 +1,95 @@
# Bypass cudarc — Raw CUDA Driver API Hot Path
**Problem:** cudarc adds 18.8ms/step of host overhead at b=256. GPU kernels complete in 160μs. The framework overhead is 117× the actual compute. Sources: `PushKernelArg` trait dispatch, `device_ptr()` event guards, `launch_builder` config assembly, internal Arc bookkeeping per CudaSlice operation.
**Goal:** Replace the entire step_with_lobsim + step_synthetic hot path with raw CUDA driver API calls. Zero cudarc in the inner loop. Target: 1ms/step host overhead → 500+ sps at b=256.
---
## Architecture: Pre-baked Kernel Launch Configs
At trainer init (cold path, cudarc OK):
1. Load all cubins → extract `CUfunction` handles (raw pointers)
2. Pre-compute all `LaunchConfig` params (grid/block/shmem) for each kernel
3. Pre-extract all device pointers from `CudaSlice` → raw `CUdeviceptr` values
4. Store everything in a `HotPathState` struct of plain `u64` / raw pointers
At step time (hot path, zero cudarc):
1. For each kernel: `cuLaunchKernel(fn, grid, block, shmem, stream, args, 0)`
2. Args are pre-packed `*mut c_void` arrays pointing to the `CUdeviceptr` values
3. Zero trait dispatch, zero Arc increment, zero event check
### Fused Mega-Kernels
While bypassing cudarc, also fuse the remaining small kernels:
**Fused step_synthetic kernel:** The training step has ~20 kernel launches (Q fwd, IQN fwd, Bellman, C51 loss, IQN loss, PPO surrogate, V loss, 4× backward, 4× reduce_axis0, 4× Adam). At b=256, each is <10μs but the host overhead per launch is ~200μs. Fuse into 3 mega-kernels:
1. `rl_fused_fwd`: All forward passes (DQN + IQN + policy + value + outcome)
2. `rl_fused_loss_bwd`: All losses + all backward passes + grad_h accumulate
3. `rl_fused_adam`: All reduce_axis0 + all Adam steps + spectral norm + per-branch LR
**Fused env-step kernel:** Graph A/A2 already capture 27 kernels. But graph launch itself has overhead. Consider a single mega-kernel for the pre-fill:
1. All head forwards (DQN, IQN, value, policy, FRD)
2. Ensemble + noisy exploration
3. Q→π agreement + distill grad
4. Action selection + gates + trails + heat + market targets
This is effectively a CUDA graph but as a single persistent kernel — zero launch overhead.
---
## Implementation Approach
### Phase 1: Raw launch wrapper (minimal change)
Replace cudarc's `launch_builder` with a thin raw wrapper:
```rust
unsafe fn raw_launch(
func: cudarc::driver::sys::CUfunction,
grid: (u32, u32, u32),
block: (u32, u32, u32),
shmem: u32,
stream: cudarc::driver::sys::CUstream,
args: &mut [*mut std::ffi::c_void],
) -> Result<()> {
let r = cudarc::driver::sys::cuLaunchKernel(
func, grid.0, grid.1, grid.2, block.0, block.1, block.2,
shmem, stream, args.as_mut_ptr(), std::ptr::null_mut(),
);
if r != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
anyhow::bail!("cuLaunchKernel failed: {:?}", r);
}
Ok(())
}
```
Replace every `self.stream.launch_builder(&self.fn).arg(...).launch(cfg)` with `raw_launch(fn_raw, grid, block, shmem, stream_raw, &mut [&mut ptr1, &mut ptr2, ...])`.
### Phase 2: Pre-extract raw pointers at init
```rust
struct HotPathPtrs {
isv: u64, // CUdeviceptr
rewards: u64,
dones: u64,
actions: u64,
// ... all device buffers as raw u64
stream: CUstream, // raw stream handle
}
```
Extract once at construction from the CudaSlice fields. The pointers are stable (pre-allocated).
### Phase 3: Mega-kernel fusion
Write the 3 fused training kernels in CUDA, replacing 20+ individual launches.
---
## Constraints
- Per `feedback_no_nvrtc`: mega-kernels are still pre-compiled cubins
- Per `feedback_no_atomicadd`: fused kernels use same tree-reduce patterns
- cudarc stays for init/alloc/graph capture — only the hot path bypasses it
- Raw CUDA driver API is `unsafe` — acceptable for the inner training loop