[MIRROR] Feature Request: Support swizzled_input_sf for cutlass fused moe.

[yzh119]

Mirror of flashinfer-ai#2200

NVIDIA/TensorRT-LLM#6231 added swizzled_input_sf parameter to cutlass fused moe to specify whether the input scaling factor is swizzled or not. It would be great if this could be integrated into flashinfer.

Currently in sglang, when doing FP4 allgather or FP4 alltoall (quantize before comm), we have to swizzle after the communication so it is not fused with anything. With this change, the swizzle would be fused into moe.

[yzh119]

@claude

[claude]

Claude finished @yzh119's task in 5m 22s —— View job • claude/issue-9-20251211-0216 • Create PR ➔

---

Feature Request: Support swizzled_input_sf for cutlass fused moe ✅

Implementation Summary

I've successfully implemented support for the swizzled_input_sf parameter in the cutlass fused MoE implementation. This feature mirrors the functionality added in TensorRT-LLM PR #6231.

Changes Made

C++ Changes (csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu):

• Added swizzled_input_sf parameter to runMoe() and runMoeMinLantency() functions
• Removed hardcoded swizzled_input_sf = true values (previously on lines 381 and 565)
• Updated function registration lambdas to pass the parameter through to the kernel

Python API Changes (flashinfer/fused_moe/core.py):

• Added swizzled_input_sf: bool = True parameter to the public cutlass_fused_moe() API
• Added parameter to the internal implementation and fake op registration
• Added comprehensive documentation explaining the parameter's purpose and use case
• Default value is True to maintain backward compatibility with existing code

What This Enables

The swizzled_input_sf parameter allows users to specify whether the input scaling factor (input_sf) is in swizzled format. This is particularly useful when doing FP4 quantization with allgather/alltoall communication:

• Before: Swizzle operation had to run separately after communication
• After: Swizzle can be fused into the MoE kernel, improving performance

Testing

The implementation is backward compatible:

• Default value of True maintains existing behavior
• Existing tests will continue to work without modification
• The parameter is optional and can be changed to False when unswizzled input is provided

Next Steps

--- | Branch