Mistral Large 3 NVFP4 TRTLLM MoE support

[elvischenv]

Motivation

Support Mistral Large 3 NVFP4 TRTLLM MoE.
Tested with cmd from #14485 + --moe-runner-backend flashinfer_trtllm:

Accuracy

TRTLLM MoE:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     8|exact_match|↑  |0.9249|±  |0.0073|
|     |       |strict-match    |     8|exact_match|↑  |0.7180|±  |0.0124|

Default cutlass MoE:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     8|exact_match|↑  |0.9272|±  |0.0072|
|     |       |strict-match    |     8|exact_match|↑  |0.7104|±  |0.0125|

Perf

TRTLLM MoE:

Output token throughput (tok/s):         15138.92
Median TPOT (ms):                        58.37

Default cutlass MoE:

Output token throughput (tok/s):         8677.97
Median TPOT (ms):                        80.68

Modifications

Support Compressed Tensors W4A4 NVFP4 TRTLLM MoE method.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Summary of Changes

Hello @elvischenv, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request integrates TensorRT-LLM (TRTLLM) support for Mixture-of-Experts (MoE) layers, specifically targeting Mistral Large 3 models with NVFP4 quantization. The primary goal is to enhance inference performance, which has been validated by benchmark results showing a substantial increase in output token throughput. The changes involve adapting the quantization pipeline to prepare weights for TRTLLM kernels and modifying the MoE layer to leverage these optimized operations, ensuring both speed and accuracy are maintained.

Highlights

• TRTLLM MoE Support: Introduced support for Mistral Large 3 NVFP4 quantization using the TensorRT-LLM (TRTLLM) backend for Mixture-of-Experts (MoE) layers.
• Performance Improvement: Achieved significant throughput gains, with TRTLLM MoE reaching 15138.92 tokens/second compared to 8677.97 tokens/second for the default Cutlass MoE.
• Accuracy Maintenance: Demonstrated comparable or slightly improved accuracy on benchmarks like gsm8k, ensuring performance enhancements do not compromise model quality.
• Compressed Tensors Integration: Enabled the FlashInferFusedMoE layer to support compressed_tensors quantization, specifically for W4A4 NVFP4.
• Optimized Weight Handling: Implemented TRTLLM-specific weight reordering and shuffling (w1w3 to w3w1) and scaling factor adjustments for efficient kernel execution.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request adds support for Mistral Large 3 with NVFP4 quantization on the TRTLLM backend, which is a valuable enhancement. The changes are well-structured, primarily introducing a new execution path for this specific configuration. The refactoring of prepare_static_weights_for_kernel into a reusable utility function is a notable improvement in code organization. The performance and accuracy benchmarks provided in the pull request description are very helpful and demonstrate significant gains. My review includes one suggestion to improve code clarity and maintainability regarding a hardcoded value.

[ispobock]

/tag-and-rerun-ci

[ispobock]

LGTM