Update flashinfer CUTLASS MoE Kernel

[wenscarl]

Essential Elements of an Effective PR Description Checklist

• Update the block_scale_interleave API
• Use per expert global scaling factor for gemm1 and gemm2.

Before the change:
INFO:lm_eval.loggers.evaluation_tracker:Output path not provided, skipping saving results aggregated
vllm (pretrained=nvidia/DeepSeek-R1-FP4,quantization=modelopt_fp4,tensor_parallel_size=4,enforce_eager=True,max_model_len=2048,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8734	±	0.0092
		strict-match	5	exact_match	↑	0.8726	±	0.0092

After:
INFO:lm_eval.loggers.evaluation_tracker:Output path not provided, skipping saving results aggregated
vllm (pretrained=nvidia/DeepSeek-R1-FP4,quantization=modelopt_fp4,tensor_parallel_size=4,enforce_eager=True,max_model_len=2048,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9409	±	0.0065
		strict-match	5	exact_match	↑	0.9378	±	0.0067

Purpose

Test Plan

Test Result

(Optional) Documentation Update

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request updates the FlashInfer CUTLASS MoE kernel usage. The changes mostly involve renaming block_scale_interleave to nvfp4_block_scale_interleave to align with an API update. However, there's a potential issue in how activation scales are handled for FP4 MoE layers. A per-expert scale tensor is being incorrectly used to quantize the entire input activation tensor, which requires a single scalar scale. I've suggested a fix to restore the correct behavior.

[gemini-code-assist]

The change from torch.min(layer.w13_input_scale_quant) to layer.w13_input_scale_quant may be incorrect. The a1_gscale variable is used in extra_prepare_args and passed to FlashInferCutlassMoEPrepareAndFinalize.prepare. This method uses the scale to quantize the input activation tensor (hidden_states) before the tokens are routed to the experts. The input activation tensor has a shape of (num_tokens, hidden_dim). To quantize this tensor, a single scalar scale is required. The original code, torch.min(layer.w13_input_scale_quant), correctly produced this scalar by selecting the most conservative scale among all per-expert scales. The new code passes layer.w13_input_scale_quant, which is a tensor of per-expert scales with shape (num_experts,). Using per-expert scales to quantize the entire activation tensor before expert routing is logically incorrect and may cause a runtime error or incorrect results from flashinfer.fp4_quantize. Please revert this change to use torch.min to ensure a correct scalar scale is used for the initial input quantization.

Suggested change

	a1_gscale = layer.w13_input_scale_quant
	a2_gscale = layer.w2_input_scale_quant
	a1_gscale = torch.min(layer.w13_input_scale_quant)
	a2_gscale = torch.min(layer.w2_input_scale_quant)

[alexm-redhat]

Verified locally that works, LGTM!