fix: return 2D scale tensor from _assemble_scales_cudagraph_safe

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