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

CUDA graph: Fix MoE bincount and per-call allocations (Hazard #4)

Browse Source 
1. Replace torch.bincount with scatter_add_ into pre-allocated buffer
   - bincount produces data-dependent shapes → breaks graph capture
   - scatter_add_ with pre-allocated _tokens_per_expert_buf (fixed shape)
   - Pre-allocated _ones_buf to avoid per-call torch.ones()

2. Replace torch.full for l1_gsa with pre-allocated buffer + fill_
   - torch.full allocates every call → breaks graph capture
   - Use self._l1_gsa_buf.fill_(l1_gs) instead
This commit is contained in:
biondizzle
2026-06-03 17:37:03 +00:00
parent df05289d6f
commit 84655d066a
1 changed files with 17 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
dsv4/layers/moe.py
View File

@@ -164,6 +164,10 @@ class Nvfp4MoE:
self._l1_gsa_buf = torch.zeros(self.num_experts, dtype=torch.float32, device=self.device)
self._l2_gsa_buf = torch.zeros(self.num_experts, dtype=torch.float32, device=self.device)
# Pre-allocated tokens-per-expert buffer — replaces torch.bincount
# (bincount produces data-dependent shapes, breaks CUDA graph capture)
self._tokens_per_expert_buf = torch.zeros(self.num_experts, dtype=torch.int32, device=self.device)
# Row indices for scale assembly (max_num_tokens * top_k slots)
self._row_indices_buf = torch.arange(
self.max_num_tokens * self.top_k, device=self.device
@@ -466,7 +470,16 @@ class Nvfp4MoE:
# Quantize slot_hidden for GEMM
slot_x_fp4, slot_x_sf = quantize_activation_nvfp4(slot_hidden, l1_gs)
tokens_per_expert = torch.bincount(sorted_ids, minlength=self.num_experts)[:self.num_experts].int()
# Compute tokens_per_expert — CUDA-graph-safe alternative to torch.bincount.
# torch.bincount produces data-dependent shapes (violates graph capture).
# Instead, use scatter_add_ into a pre-allocated buffer (fixed shape, GPU-only).
self._tokens_per_expert_buf.zero_()
# Pre-allocated ones buffer — avoids per-call torch.ones() allocation
n_slots = sorted_ids.shape[0]
if not hasattr(self, '_ones_buf') or self._ones_buf.shape[0] < n_slots:
self._ones_buf = torch.ones(self.max_num_tokens * self.top_k, dtype=torch.int32, device=sorted_ids.device)
self._tokens_per_expert_buf.scatter_add_(0, sorted_ids, self._ones_buf[:n_slots])
tokens_per_expert = self._tokens_per_expert_buf[:self.num_experts]
expert_offsets = self._expert_offsets_buf
expert_offsets.zero_()
expert_offsets[1:self.num_experts + 1] = tokens_per_expert.cumsum(0)
@@ -494,7 +507,9 @@ class Nvfp4MoE:
padded_expert_offsets,
self._padded_x_sf_buf_l1, self._per_expert_scale_bufs_l1
)
l1_gsa = torch.full((self.num_experts,), l1_gs, dtype=torch.float32, device=device)
# l1_gsa: pre-allocated buffer, no per-call allocation
self._l1_gsa_buf.fill_(l1_gs)
l1_gsa = self._l1_gsa_buf
l1_out = run_nvfp4_grouped_gemm(
mat_a=padded_x_fp4, mat_b=self._l1_mat_b,