biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

CUDA graph: Fix remaining sync violations from B200 detector run 2

Browse Source 
1. grouped_linear.py: Remove conditional host read of GPU tensor
   - 'if group_offsets[0] != 0' reads GPU value on host → sync
   - Fix: unconditionally update offsets every call (GPU-only multiply)

2. test_cuda_graph_readiness.py: Use pinned CPU buffers for token transfer
   - dec_tid_buf[0] = python_int → CPU→GPU sync
   - Fix: write to pinned CPU buffer, then copy_ (async, graph-capturable)

3. Add dsv4/decode/cuda_graph_decoder.py (skeleton)
This commit is contained in:
biondizzle
2026-06-03 17:20:34 +00:00
parent e07d79868f
commit df05289d6f
3 changed files with 199 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
tests/unit/test_cuda_graph_readiness.py
View File

@@ -374,12 +374,23 @@ def run_sync_debug_mode():
dec_tid_buf = torch.zeros(1, dtype=torch.long, device='cuda:0')
dec_pos_buf = torch.zeros(1, dtype=torch.long, device='cuda:0')
dec_tid32_buf = torch.zeros(1, dtype=torch.int32, device='cuda:0')
# Pinned CPU buffers for graph-capturable token/position transfer
dec_tid_pinned = torch.zeros(1, dtype=torch.long, device='cpu').pin_memory()
dec_tid32_pinned = torch.zeros(1, dtype=torch.int32, device='cpu').pin_memory()
dec_pos_pinned = torch.zeros(1, dtype=torch.long, device='cpu').pin_memory()
def write_token_to_gpu(token_id, position):
"""Write token/position to GPU buffers via pinned CPU (no CPU→GPU sync)."""
dec_tid_pinned[0] = token_id
dec_tid_buf.copy_(dec_tid_pinned)
dec_tid32_pinned[0] = token_id
dec_tid32_buf.copy_(dec_tid32_pinned)
dec_pos_pinned[0] = position
dec_pos_buf.copy_(dec_pos_pinned)
# Warmup step first (so CuTeDSL kernels are compiled)
print(" Warmup decode step (compiling CuTeDSL kernels)...", flush=True)
dec_tid_buf[0] = all_tokens[-1]
dec_tid32_buf[0] = all_tokens[-1]
dec_pos_buf[0] = len(all_tokens) - 1
write_token_to_gpu(all_tokens[-1], len(all_tokens) - 1)
X = mHCLayer.init_state(embed(dec_tid_buf))
for li in range(n_layers):
gpu = li % NUM_GPUS
@@ -408,9 +419,7 @@ def run_sync_debug_mode():
sync_errors = []
try:
detector.phase = "decode_forward"
dec_tid_buf[0] = all_tokens[-1]
dec_tid32_buf[0] = all_tokens[-1]
dec_pos_buf[0] = len(all_tokens) - 1
write_token_to_gpu(all_tokens[-1], len(all_tokens) - 1)
X = mHCLayer.init_state(embed(dec_tid_buf))
for li in range(n_layers):
@@ -473,10 +482,13 @@ def run_sync_debug_mode():
g = torch.cuda.CUDAGraph()
torch.cuda.set_device(0)
# Fill static buffers with current decode state
static_token[0] = all_tokens[-1]
static_token32[0] = all_tokens[-1]
static_pos[0] = len(all_tokens) - 1
# Fill static buffers with current decode state (via pinned CPU — no sync)
dec_tid_pinned[0] = all_tokens[-1]
static_token.copy_(dec_tid_pinned)
dec_tid32_pinned[0] = all_tokens[-1]
static_token32.copy_(dec_tid32_pinned)
dec_pos_pinned[0] = len(all_tokens) - 1
static_pos.copy_(dec_pos_pinned)
with torch.cuda.graph(g):
X = mHCLayer.init_state(embed(static_token))