biondizzle 0ca7bed0e1 CUDA graph: Fix sync violations found by B200 detector ...

Fixes from running Section A detector on B200:

1. single_shot_inference.py: Use pinned CPU buffers for token/position transfer
   - dec_tid_buf[0] = python_int causes CPU→GPU sync
   - Fixed: write to pinned CPU buffer, then copy_ (async, graph-capturable)

2. grouped_linear.py: Fix expert_offsets Python loop
   - expert_offsets[g] = python_int * padded_rows → CPU→GPU sync per iteration
   - Fixed: element-wise multiply with pre-allocated range tensor (GPU-only)

3. grouped_linear.py: Vectorized output extraction for T=1 decode
   - Python loop z[:, g, :] = out[...] → CPU sync for each slice
   - Fixed: GPU gather with pre-computed indices for T=1

4. grouped_linear.py: Pre-allocate output buffer
   - torch.empty() per call → allocation inside graph
   - Fixed: use self._output_buf (pre-allocated at max size)

5. grouped_linear.py: Pre-allocate expert_offsets_range_buf
   - torch.arange() per call → allocation inside graph
   - Fixed: compute once at init, reuse via element-wise multiply

2026-06-03 16:52:19 +00:00

95 KiB

Raw Blame History

View Raw