[NVIDIA] Explicitly disable shuffled weights for flashinfer blockscale moe fp8 kernels

[kaixih]

The latest Flashinfer (PR) introduces a new flag to the trtllm_fp8_block_scale_moe API, which defaults to True. This PR explicitly disables it to restore the previous behavior.

I have verified the perf and accuracy with the tot and we recommend to use flashinfer v0.2.9.

cc. @kushanam @mgoin

[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 aims to restore previous behavior for FlashInfer MoE kernels by explicitly disabling a new use_shuffled_weight flag. While the change is correct in its intent, it introduces a critical backward compatibility issue for users with older versions of FlashInfer. I've provided a comment with a suggested fix to address this.

[gemini-code-assist]

This change explicitly adds the use_shuffled_weight argument, which was introduced in FlashInfer v0.2.9. This will cause a TypeError for users with older versions of FlashInfer, breaking backward compatibility.

To fix this, we should only pass the argument if the installed FlashInfer version supports it. This can be done with a version check, which requires refactoring the function call to use a kwargs dictionary.

Here is a suggested implementation to replace lines 1097-1131:

    from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe

    a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1])
    # NOTE: scales of hidden states have to be transposed!
    a_sf_t = a_sf.t().contiguous()

    kwargs = dict(
        routing_logits=routing_logits,
        routing_bias=routing_bias,
        hidden_states=a_q,
        hidden_states_scale=a_sf_t,
        gemm1_weights=w13_weight,
        gemm1_weights_scale=w13_weight_scale_inv,
        gemm2_weights=w2_weight,
        gemm2_weights_scale=w2_weight_scale_inv,
        num_experts=global_num_experts,
        top_k=top_k,
        n_group=num_expert_group,
        topk_group=topk_group,
        intermediate_size=intermediate_size,
        local_expert_offset=expert_offset,
        local_num_experts=local_num_experts,
        routed_scaling_factor=routed_scaling,
        tile_tokens_dim=_get_tile_tokens_dim(x.shape[0], top_k,
                                             global_num_experts),
        routing_method_type=2,  # DeepSeek-styled routing method
    )

    try:
        import flashinfer
        from packaging.version import Version
        # The use_shuffled_weight argument was added in flashinfer v0.2.9
        if Version(flashinfer.__version__) >= Version("0.2.9"):
            kwargs["use_shuffled_weight"] = False
    except (ImportError, AttributeError):
        # Older flashinfer version or flashinfer not installed.
        # The lazy loader will handle the ImportError later if it's missing.
        pass

    return flashinfer_trtllm_fp8_block_scale_moe(**kwargs)

[mgoin]

Thanks Kaixi, should we wait for the 0.2.9 release to land this?

[kaixih]

@mgoin Sure.

[kaixih]

@mgoin the flashinfer has released the 0.2.9rc1. I took a quick look at the vLLM codebase and noticed that only the Dockerfile explicitly references the Flashinfer version. I'm not sure how vLLM determines or enforces the Flashinfer version elsewhere. Can you advise?

[mgoin]

@kaixih Let's wait on landing this PR until after #21485 then, as it already updates the dockerfile. We will enforce the version in the future by adding it to requirements/cuda.txt as well

[kaixih]

@mgoin can we merge this PR since the flashinfer 0.2.9rc1 is in.