Skip to content

[None][perf] Use FlashInfer MXFP8 GEMM for MiniMax-M3 decode#20

Open
peihu-nv wants to merge 1 commit into
brb-nv:feat/m3_with_msafrom
peihu-nv:peihengh/m3-flashinfer-safe-default
Open

[None][perf] Use FlashInfer MXFP8 GEMM for MiniMax-M3 decode#20
peihu-nv wants to merge 1 commit into
brb-nv:feat/m3_with_msafrom
peihu-nv:peihengh/m3-flashinfer-safe-default

Conversation

@peihu-nv

Copy link
Copy Markdown
Collaborator

@coderabbitai summary

Description

MiniMax-M3 decode CUDA graphs execute 240 MXFP8 linear GEMMs per replay. This PR uses autotuned FlashInfer mm_mxfp8 for those captured GEMMs while reusing TensorRT-LLM's existing quantized activations, weights, and scales without repacking.

The automatic path is enabled only for MiniMax-M3 decode CUDA-graph capture. Eager execution, CTX/prefill, and piecewise prefill graphs remain on the native TensorRT-LLM backend; an environment override and native fallback are preserved. Matched GB200 serving improved total/output throughput by 10.4% and reduced median TPOT by 9.8%, with median TTFT unchanged.

Test Coverage

  • Focused GB200 tests covering native default, explicit override, zero-copy FlashInfer dispatch, auto eager/decode selection, invalid configuration, and autotuner warmup: 5 passed.
  • tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True]: passed on GB200 with MMLU 84.844% and GSM8K 90.637%.
  • Matched GB200 disaggregated serving A/B, TP4/EP4 with attention DP, FP8 KV cache, ISL 8192 / OSL 1024, concurrency 64, and 1,000 requests: 21,599.7 to 23,840.6 total token/s; median TPOT 25.269 to 22.798 ms.

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant