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

Fix shape mismatch: shared padded buffers, revert max_num_tokens cap

Browse Source 
Root cause: capping max_num_tokens to 512 made buffers too small for the
actual 8192-token warmup. slot_hidden had 49152 rows but padded_hidden
only had 6144.

Fix: Revert the 512 cap. Use SHARED padded buffers (not per-layer) to
avoid OOM. Only 72 MB total (not 4.3 GB) since layers run sequentially
and reuse the same buffer. Cudagraph-safe since capture and replay both
run layers sequentially on the same tensor.
This commit is contained in:
biondizzle
2026-05-17 15:47:10 +00:00
parent e2f33596a2
commit b1ac74bb4d
2 changed files with 19 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
vllm/nvfp4_cutedsl.py
View File

@@ -142,15 +142,22 @@ class CuTeDSLMoERunner:
self.max_num_tokens * self.top_k, device=self.device
)
# Padded hidden/activated: sized for num_experts * max_chunks * 128 (the actual padded size)
max_rows_per_expert = self._max_chunks_per_expert * 128
padded_max_slots = self.num_experts * max_rows_per_expert
self._padded_hidden_buf = torch.zeros(
padded_max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device
)
self._padded_activated_buf = torch.zeros(
padded_max_slots, self.intermediate_size, dtype=torch.bfloat16, device=self.device
)
# Padded hidden/activated: SHARED across all runners (not per-layer)
# These are only used during run() which is sequential across layers.
# Per-layer allocation would be 72 MB × 60 layers = 4.3 GB → OOM.
device_key = str(self.device)
if not hasattr(CuTeDSLMoERunner, '_shared_padded_bufs'):
CuTeDSLMoERunner._shared_padded_bufs = {}
if device_key not in CuTeDSLMoERunner._shared_padded_bufs:
CuTeDSLMoERunner._shared_padded_bufs[device_key] = {
'hidden': torch.zeros(
padded_max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device
),
'activated': torch.zeros(
padded_max_slots, self.intermediate_size, dtype=torch.bfloat16, device=self.device
),
}
self._shared_bufs = CuTeDSLMoERunner._shared_padded_bufs[device_key]
# Padded expert offsets buffer: [0, max_rows, 2*max_rows, ...] (fixed)
self._padded_expert_offsets_buf = torch.zeros(
@@ -389,7 +396,7 @@ class CuTeDSLMoERunner:
# -- Gather hidden states into slot order, pad per expert --
slot_hidden = hidden_states[sorted_token_ids]
total_padded_slots = self.num_experts * max_rows_per_expert
padded_hidden = self._padded_hidden_buf[:total_padded_slots]
padded_hidden = self._shared_bufs['hidden'][:total_padded_slots]
padded_hidden.zero_()
row_indices = self._row_indices_buf[:num_slots]
expert_assign = torch.searchsorted(
@@ -427,7 +434,7 @@ class CuTeDSLMoERunner:
activated = torch.nn.functional.silu(gate) * up
# === L2: down ===
padded_activated = self._padded_activated_buf[:total_padded_slots]
padded_activated = self._shared_bufs['activated'][:total_padded_slots]
padded_activated.zero_()
padded_activated[padded_dst] = activated

5
vllm/patches/deepseek_v4.py
View File

@@ -499,14 +499,11 @@ class DeepseekV4MegaMoEExperts(nn.Module):
l2_gs.append(down_gs)
# Create CuTeDSL runner with directly-cast weights
# Max tokens for buffer pre-allocation: use cudagraph max capture size
# (not scheduler max which can be 8192, causing OOM with 60 layers)
max_cg_size = getattr(self, '_cudagraph_max_capture_size', 512)
self._cutedsl_runner = CuTeDSLMoERunner(
num_experts=self.num_local_experts,
hidden_size=self.hidden_size,
intermediate_size=self.intermediate_size,
max_num_tokens=min(self.max_num_tokens, max_cg_size),
max_num_tokens=self.max_num_tokens,
top_k=self.top_k,
device=l1_fp4[0].device,
experts_start_idx=self.experts_start_idx,