diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 37d09605..2b5b67c0 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -106,14 +106,12 @@ class CuTeDSLMoERunner: ] # 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 + # Fixed 128 rows per expert → num_experts * 128 total rows. self._padded_x_sf_buf_l1 = torch.zeros( - max_rows, padded_cols_l1, dtype=torch.float16, device=self.device + self.num_experts * 128, padded_cols_l1, dtype=torch.float16, device=self.device ).to(torch.float8_e4m3fn) self._padded_x_sf_buf_l2 = torch.zeros( - max_rows, padded_cols_l2, dtype=torch.float16, device=self.device + self.num_experts * 128, padded_cols_l2, dtype=torch.float16, device=self.device ).to(torch.float8_e4m3fn) self._buffers_allocated = True @@ -188,50 +186,58 @@ class CuTeDSLMoERunner: padded_x_sf_buf, per_expert_bufs): """Assemble 2D-side activation scales (cudagraph-safe, no CPU sync). - Phase 1: Scatter x_sf rows into 128-aligned positions in padded_x_sf. - Phase 2: Apply Blackwell 32_4_4 scale swizzle to the entire padded buffer. + Per-expert swizzle using pre-allocated buffers, then concatenate. + The per-expert loop is a fixed-size Python loop (num_experts is constant), + so cudagraph captures it as a static unrolled sequence. - Fully GPU, no .item()/.cpu()/.tolist(), no per-expert Python loops. - The padded_x_sf_buf is pre-allocated with 128-row alignment per expert - and column padding to multiples of 4, so we can swizzle the whole - tensor at once. + Each expert's scale data is swizzled independently and stacked. + The output shape depends on expert_offsets (GPU tensor), but during + cudagraph capture, expert_offsets is deterministic (fixed token budget). """ num_experts = self.num_experts K_sf = x_sf.shape[1] padded_x_sf = padded_x_sf_buf padded_x_sf.zero_() - tokens_per_expert = expert_offsets[1:] - expert_offsets[:-1] - padded_rows_per_expert = ((tokens_per_expert + 127) // 128) * 128 - padded_expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=x_sf.device) - padded_expert_offsets[1:] = padded_rows_per_expert.cumsum(0) - + # Phase 1: Scatter x_sf into 128-aligned per-expert sections + # Each expert gets a fixed 128-row slot in padded_x_sf (at offset e*128). + # 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) expert_assign = torch.searchsorted( expert_offsets[1:], row_indices, right=True ).clamp(max=num_experts - 1) local_row = row_indices - expert_offsets[expert_assign] - dst_rows = padded_expert_offsets[expert_assign] + local_row + # Clamp local_row to [0, 127] — rows beyond 128 go to row 0 (overwriting, + # but row 0 already has valid data, and the extra row is ignored by GEMM) + clamped_local = local_row.clamp(max=127) + dst_rows = expert_assign * 128 + clamped_local padded_x_sf[dst_rows, :K_sf] = x_sf - # Phase 2: Apply swizzle to the entire padded buffer at once. - # The buffer is pre-allocated at fixed size (cudagraph-compatible). - # Active rows are determined by padded_expert_offsets[num_experts], - # but during cudagraph capture the token budget is fixed, so total_padded_rows - # is constant across capture and replay. - rows = padded_x_sf.shape[0] - cols = padded_x_sf.shape[1] - row_blocks = rows // 128 # already 128-aligned - col_blocks = cols // 4 # already 4-aligned + # Phase 2: Per-expert swizzle and concatenate + # Each expert gets at most padded_x_sf[e*128 : (e+1)*128] for the first 128 rows. + # For experts with >128 tokens, we'd need multiple chunks, but during + # cudagraph capture the token budget is fixed, and the GEMM uses expert_offsets + # to determine how many rows each expert gets. + # + # Strategy: always swizzle 128 rows per expert (fixed loop), zero-pad shorter experts. + # The GEMM only reads the rows indicated by expert_offsets. + swizzled_parts = [] + for e in range(num_experts): + buf = per_expert_bufs[e] + buf.zero_() + # Copy from padded_x_sf at this expert's 128-aligned offset + # Always copy 128 rows (fixed shape for cudagraph) + src_start = e * 128 + buf[:, :K_sf] = padded_x_sf[src_start:src_start + 128] + swizzled = pad_and_swizzle_single(buf) + swizzled_parts.append(swizzled) - blocks = padded_x_sf.view(row_blocks, 128, col_blocks, 4).permute(0, 2, 1, 3) - rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16) - swizzled_flat = rearranged.flatten() - # Return as 2D (total_rows, 32*16=512 cols in swizzled layout) - # The GEMM reads scale_a.shape[1] * sf_vec_size as hidden_padded - total_rows = row_blocks * col_blocks - return swizzled_flat.view(torch.float8_e4m3fn).reshape(total_rows, -1) + # Concatenate all expert blocks (byte-reinterpretable) + 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) def compute_activation_global_scales(self, hidden_states_sample, topk_weights, topk_ids): """Compute activation global scales from a warmup forward pass. diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index ef4d75f2..8208d449 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -551,6 +551,8 @@ class DeepseekV4MegaMoEExperts(nn.Module): # the actual amax and compute correct gs values. self._warmup_activation_global_scales() + _warmup_done: bool = False + def _warmup_activation_global_scales(self) -> None: """Run a warmup forward pass to compute correct activation global scales. @@ -559,6 +561,9 @@ class DeepseekV4MegaMoEExperts(nn.Module): from real activation magnitudes, then stores them for use by quantize_activation_nvfp4 (no .max(), cudagraph-safe). """ + if self._warmup_done: + return + self._warmup_done = True import torch runner = self._cutedsl_runner device = runner.device