From 4300775bfec557d514e52ed00459bef6fdff161a Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 18:29:12 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20remove=20.item()=20sync=20in=20scale=20r?= =?UTF-8?q?eshape=20=E2=80=94=20use=20padded=5Fscales.shape[0]=20instead?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- 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 4455ff1d..78b51590 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -183,8 +183,9 @@ class CuTeDSLMoERunner: padded_scales[dst_rows, :K_sf] = x_sf # 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) + # to_blocked preserves element count, so reshape to match padded shape + swizzled = pad_and_swizzle_single(padded_scales) + return swizzled.reshape(padded_scales.shape[0], -1) def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None): num_tokens, hidden_size = hidden_states.shape