diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index a049398e..00c1260d 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -90,9 +90,6 @@ class CuTeDSLMoERunner: def _allocate_buffers(self): """Pre-allocate scale buffers at max size for cudagraph compatibility.""" - K_sf = cutedsl_ceil_div(self.hidden_size, 16) - padded_cols = cutedsl_ceil_div(K_sf, 4) * 4 - # Per-expert scale buffers: separate L1/L2 since K_sf differs K_sf_l1 = cutedsl_ceil_div(self.hidden_size, 16) padded_cols_l1 = cutedsl_ceil_div(K_sf_l1, 4) * 4 @@ -108,12 +105,15 @@ class CuTeDSLMoERunner: for _ in range(self.num_experts) ] - # Padded x_sf buffers: num_experts * 128 rows, separate L1/L2 + # Padded x_sf buffers for Phase 1 scatter. + # Sized for max_num_tokens * top_k rows (worst case: all tokens in one expert, + # padded to 128). This is larger than num_experts * 128 when tokens >> experts. + max_rows = cutedsl_ceil_div(self.max_num_tokens * self.top_k, 128) * 128 self._padded_x_sf_buf_l1 = torch.zeros( - self.num_experts * 128, padded_cols_l1, dtype=torch.float16, device=self.device + max_rows, padded_cols_l1, dtype=torch.float16, device=self.device ).to(torch.float8_e4m3fn) self._padded_x_sf_buf_l2 = torch.zeros( - self.num_experts * 128, padded_cols_l2, dtype=torch.float16, device=self.device + max_rows, padded_cols_l2, dtype=torch.float16, device=self.device ).to(torch.float8_e4m3fn) self._buffers_allocated = True @@ -149,10 +149,6 @@ class CuTeDSLMoERunner: self.max_num_tokens * self.top_k, dtype=torch.int32, device=self.device ) self._fill_token_indices() - # Verify the fill worked (CuTeDSL JIT may have corrupted GPU memory) - torch.cuda.synchronize() - assert self._token_indices.max().item() == self.max_num_tokens - 1, \ - f"token_indices corrupted after fill! max={self._token_indices.max().item()}" self._needs_token_refill = True # GEMM JIT may corrupt; refill after first run self._expert_id_range = torch.arange( @@ -192,18 +188,14 @@ class CuTeDSLMoERunner: padded_x_sf_buf, per_expert_bufs): """Assemble 2D-side activation scales (cudagraph-safe, no CPU sync). - Two-phase approach: - 1. Scatter x_sf rows into 128-aligned positions in padded_x_sf (GPU-only) - 2. Per-expert: copy 128 rows from padded_x_sf, swizzle, accumulate + Phase 1: Scatter x_sf rows into 128-aligned positions in padded_x_sf. + Phase 2: Per-expert, copy from padded_x_sf at the right offset, + pad_and_swizzle, then concatenate. - All operations are fixed-shape. No .item(), no .tolist(), no variable - slicing with GPU scalars. + The output has sum(padded_rows_per_expert) rows (variable per expert). """ num_experts = self.num_experts K_sf = x_sf.shape[1] - - # ── Phase 1: Scatter x_sf into 128-aligned padded buffer ── - # Use GPU-computed indices (same as original _assemble_scales_cudagraph_safe) padded_x_sf = padded_x_sf_buf padded_x_sf.zero_() @@ -221,21 +213,29 @@ class CuTeDSLMoERunner: dst_rows = padded_expert_offsets[expert_assign] + local_row padded_x_sf[dst_rows, :K_sf] = x_sf - # ── Phase 2: Per-expert swizzle ── - # Each expert gets its own 128-row buf, copied from padded_x_sf at - # the 128-aligned offset, then swizzled independently + # Phase 2: Per-expert swizzle swizzled_parts = [] for e in range(num_experts): + n_padded = padded_rows_per_expert[e].item() + if n_padded == 0: + continue + start = padded_expert_offsets[e].item() buf = per_expert_bufs[e] - buf.zero_() - buf[:, :K_sf] = padded_x_sf[e * 128:e * 128 + 128] - swizzled = pad_and_swizzle_single(buf) - swizzled_parts.append(swizzled) + # buf is only 128 rows; process in 128-row chunks + offset = start + remaining = n_padded + while remaining > 0: + chunk = min(remaining, 128) + buf.zero_() + buf[:chunk, :K_sf] = padded_x_sf[offset:offset + chunk] + swizzled = pad_and_swizzle_single(buf) + swizzled_parts.append(swizzled) + offset += chunk + remaining -= chunk - # Concatenate all expert blocks all_flat = torch.cat([p.view(torch.uint8) for p in swizzled_parts], dim=0) all_flat = all_flat.view(torch.float8_e4m3fn) - return all_flat.reshape(num_experts * 128, -1) + return all_flat def compute_activation_global_scales(self, hidden_states_sample, topk_weights, topk_ids): """Compute activation global scales from a warmup forward pass. @@ -347,19 +347,6 @@ class CuTeDSLMoERunner: expert_offsets[1:self.num_experts + 1] = tokens_per_expert.cumsum(0) # -- Gather hidden states into slot order -- - # DEBUG: Check for OOB indices before indexing - ti_max = self._token_indices[:num_slots].max().item() - ti_min = self._token_indices[:num_slots].min().item() - si_max = sorted_token_ids.max().item() - si_min = sorted_token_ids.min().item() - print(f"[CLAWMINE] num_tokens={num_tokens} top_k={top_k} num_slots={num_slots} " - f"token_indices max={ti_max} min={ti_min} " - f"sorted_token_ids max={si_max} min={si_min} " - f"experts_start_idx={self.experts_start_idx} num_experts={self.num_experts} " - f"topk_ids max={topk_ids.max().item()} min={topk_ids.min().item()}", flush=True) - if si_max >= num_tokens or si_min < 0: - print(f"[CLAWMINE BUG] sorted_token_ids OOB! " - f"max={si_max} min={si_min} num_tokens={num_tokens}", flush=True) slot_hidden = hidden_states[sorted_token_ids] # === L1: gate + up === diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 7cef41d9..6b787da5 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -503,6 +503,8 @@ class DeepseekV4MegaMoEExperts(nn.Module): num_experts=self.num_local_experts, hidden_size=self.hidden_size, intermediate_size=self.intermediate_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, )