diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 63235048..3fb19285 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -238,12 +238,6 @@ class CuTeDSLMoERunner: padded_x_sf = padded_x_sf_buf padded_x_sf.zero_() - # Compute padded expert offsets (each expert padded to 128 rows) - 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 fixed-layout per-expert sections # Each expert gets max_chunks * 128 rows at offset e * max_chunks * 128. # This matches Phase 2's fixed reading pattern. @@ -276,7 +270,8 @@ class CuTeDSLMoERunner: all_flat = torch.cat([p.view(torch.uint8) for p in swizzled_parts], dim=0) all_flat = all_flat.view(torch.float8_e4m3fn) - total_padded = padded_expert_offsets[num_experts] + # Total rows = num_experts * max_chunks * 128 (fixed) + total_padded = num_experts * max_chunks * 128 return all_flat.reshape(total_padded, -1) def compute_activation_global_scales(self, hidden_states_sample, topk_weights, topk_ids):