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:
95
docs/superpowers/specs/2026-05-26-bypass-cudarc-raw-cuda.md
Normal file
95
docs/superpowers/specs/2026-05-26-bypass-cudarc-raw-cuda.md
Normal 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
|
||||
Reference in New Issue
Block a user