From 95a1345b92ffa7699536216b5cb905f3b5009476 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 18:26:57 +0000 Subject: [PATCH] fix: return 2D scale tensor from _assemble_scales_cudagraph_safe --- vllm/nvfp4_cutedsl.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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