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

Fix cudaErrorStreamCaptureUnsupported: pre-allocate all tensors used during capture

Browse Source 
torch.full(), torch.zeros(), torch.arange() allocate new tensors during
cudagraph capture, which triggers cudaErrorStreamCaptureUnsupported.

Pre-allocate:
- _l1_gsa_buf / _l2_gsa_buf (use .fill_() instead of torch.full)
- _output_buf (use .zero_() on pre-allocated slice)
- _row_indices_buf (pre-allocated arange, sliced during use)
This commit is contained in:
biondizzle
2026-05-17 12:31:25 +00:00
parent b0221662e7
commit 7e692c3aec
1 changed files with 23 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
vllm/nvfp4_cutedsl.py
View File

@@ -76,6 +76,10 @@ class CuTeDSLMoERunner:
self._per_expert_scale_bufs_l2 = None
self._padded_x_sf_buf_l1 = None
self._padded_x_sf_buf_l2 = None
self._l1_gsa_buf = None
self._l2_gsa_buf = None
self._output_buf = None
self._row_indices_buf = None
self._buffers_allocated = False
def _fill_token_indices(self):
@@ -114,6 +118,20 @@ class CuTeDSLMoERunner:
self.num_experts * 128, padded_cols_l2, dtype=torch.float16, device=self.device
).to(torch.float8_e4m3fn)
# Pre-allocated global_scale_a buffers (filled via .fill_(), no torch.full during capture)
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 output buffer
self._output_buf = torch.zeros(
self.max_num_tokens, self.hidden_size, dtype=torch.bfloat16, 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
)
self._buffers_allocated = True
def _ensure_stacked(self):
@@ -204,7 +222,7 @@ class CuTeDSLMoERunner:
# Tokens beyond 128 per expert wrap to zero rows (harmless — zero scale
# means zero contribution in GEMM). All indexing is fixed-shape.
total_rows = x_sf.shape[0]
row_indices = torch.arange(total_rows, device=x_sf.device)
row_indices = self._row_indices_buf[:total_rows]
expert_assign = torch.searchsorted(
expert_offsets[1:], row_indices, right=True
).clamp(max=num_experts - 1)
@@ -360,10 +378,7 @@ class CuTeDSLMoERunner:
x_sf, expert_offsets[:self.num_experts + 1],
self._padded_x_sf_buf_l1, self._per_expert_scale_bufs_l1
)
l1_gsa = torch.full(
(self.num_experts,), self._l1_activation_global_scale,
dtype=torch.float32, device=device
)
l1_gsa = self._l1_gsa_buf.fill_(self._l1_activation_global_scale)
l1_out = run_nvfp4_grouped_gemm(
mat_a=x_fp4, mat_b=self._l1_mat_b,
@@ -386,10 +401,7 @@ class CuTeDSLMoERunner:
l2_x_sf, expert_offsets[:self.num_experts + 1],
self._padded_x_sf_buf_l2, self._per_expert_scale_bufs_l2
)
l2_gsa = torch.full(
(self.num_experts,), self._l2_activation_global_scale,
dtype=torch.float32, device=device
)
l2_gsa = self._l2_gsa_buf.fill_(self._l2_activation_global_scale)
l2_out = run_nvfp4_grouped_gemm(
mat_a=l2_x_fp4, mat_b=self._l2_mat_b,
@@ -399,7 +411,8 @@ class CuTeDSLMoERunner:
)
# === Scatter -> final output ===
y = torch.zeros(num_tokens, self.hidden_size, dtype=torch.bfloat16, device=device)
y = self._output_buf[:num_tokens]
y.zero_()
weighted_out = l2_out * sorted_weights.unsqueeze(1).to(l2_out.dtype)
y.scatter_add_(
0,