diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index aad6791e..4455ff1d 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -182,8 +182,9 @@ class CuTeDSLMoERunner: # Scatter x_sf into padded_scales padded_scales[dst_rows, :K_sf] = x_sf - # Apply swizzle to the whole padded tensor - return pad_and_swizzle_single(padded_scales) + # Apply swizzle to the whole padded tensor, return 2D for 2D-side scale_a + swizzled_flat = pad_and_swizzle_single(padded_scales) + return swizzled_flat.reshape(total_padded_rows.item(), -1) def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None): num_tokens, hidden_size = hidden_states.shape