- gemm_runner.py: Add out= parameter to run_nvfp4_grouped_gemm and
run_fused_swiglu_grouped_gemm to accept pre-allocated output buffers
- quantize.py: Replace torch.zeros_like/torch.zeros with scalar 0.0 in
torch.where() calls (graph-capturable, no memory allocation)
- Both fixes prevent 'Disallowed operation during CUDA stream capture'
errors during graph capture
Patch torch.cuda.current_device to return the tensor's device index
during from_dlpack calls inside CUDA graph capture. This bypasses the
device check in __dlpack__ without changing the CUDA stream (which
caused 'Capture must end on the same stream' in v1) and without
triggering a cross-device copy (which caused 'Cannot copy between
CPU and CUDA tensors' in v2).
Previous fix (set_device) caused 'Capture must end on the same stream'.
New fix: wrap tensor in _DLPatchTensor during graph capture, which forces
dl_device in __dlpack__ to bypass the device check without changing the stream.
This enables CUDA graph capture on all 8 GPUs, not just cuda:0.
When capturing CUDA graphs on non-default GPUs, torch.cuda.current_device()
may not match the tensor's device. from_dlpack() checks this and fails.
Fix: set the current device to match the tensor's device before from_dlpack.
This enables graph capture on all 8 GPUs, not just cuda:0.
- run_nvfp4_grouped_gemm: use_2cta = tokens_sum >= 256 and cluster_m even
- run_fused_swiglu_grouped_gemm: same conditional
- Auto-warms up on first use via lazy compilation cache
- 1.7-1.9× throughput at prefill shapes (M>=256)
- Decode (M<256) stays 1-CTA (correct, no waste)
