Skip to content

[kernel][perf][minimax m3] adds flashinfer MSA backend for sm100 - #56483

Closed
NVShreyas wants to merge 1 commit into
vllm-project:mainfrom
NVShreyas:shreyasm/minimax-m3-flashinfer-msa
Closed

NVShreyas wants to merge 1 commit into
vllm-project:mainfrom
NVShreyas:shreyasm/minimax-m3-flashinfer-msa

Conversation

@NVShreyas

@NVShreyas NVShreyas commented Sep 11, 2026 •

Copy link
Copy Markdown
Contributor

Purpose

Add a FlashInfer MSA decode backend for MiniMax M3 on SM100/GB200. The new
backend consumes vLLM's packed FP8 K/V cache and compact token-major top-k
indices directly, writes into vLLM-owned output and LSE buffers, and supports
CUDA graph capture and replay without per-step K/V or top-k materialization.

The backend is selectable through minimax_m3_msa_decode_backend="flashinfer"
or the FLASHINFER_MSA attention-backend alias. Unsupported geometry and cache
dtypes retain the existing backend behavior.

This backend is also a prerequisite for enabling DCP for MiniMax M3 on GB200.
The DCP integration itself remains separate from this PR.

Dependency

Do not merge this PR before
flashinfer-ai/flashinfer#5094.

That PR provides the packed-FP8 paged MSA kernel, token-major top-k support, and
caller-owned out/lse_out API used here.

Duplicate-work check

This does not duplicate
vllm-project/vllm#48994.
That PR targets the SM120/SM121 consumer-Blackwell path and couples a
FlashInfer indexer with its attention implementation. This PR targets the
existing SM100/GB200 NVIDIA MSA path, preserves the current indexer, and adds a
direct packed-FP8 decode integration for the kernel introduced by FlashInfer
#5094.

AI assistance was used to implement, test, and prepare this change. The human
submitter reviewed the changed code and validation results.

Test Plan

  • Run vLLM's staged-file pre-commit suite.
  • Validate backend selection and unsupported-geometry fallback cases.
  • Compare Q1 and Q4 FlashInfer output/LSE against both the existing Triton
    implementation and a direct PyTorch sparse-attention reference on GB200.
  • Capture and replay the FlashInfer decode path with caller-owned output and
    LSE buffers.
  • Run matched 1,000-request end-to-end Mooncake benchmarks for FlashInfer and
    CUTLASS MSA.

Test Result

Correctness and CUDA graphs

  • pre-commit run: passed all staged-file hooks, including ruff, mypy, SPDX,
    forbidden-import, configuration-default, and sign-off checks.
  • Fresh-source FlashInfer GB200 validation: 4 passed.
  • vLLM FlashInfer-focused GB200 suite: 9 passed, 32 deselected.
  • Q1 and Q4 output/LSE parity passed against Triton and PyTorch references.
  • CUDA graph capture and replay passed with caller-owned out and lse_out
    buffers.

GB200 performance

Both runs used the same immutable image, MiniMax-M3-NVFP4 model, TP4/K=3,
FULL_AND_PIECEWISE CUDA graphs, local argmax, four GB200 GPUs, and the same
1,000-request Mooncake trace with 64K-class inputs, 400 OSL, and 90% prefix
reuse.

Metric FlashInfer MSA (E18) CUTLASS MSA (E21) FlashInfer improvement
Input throughput 49,475 tok/s 46,220 tok/s +7.04%
Input throughput/GPU 12,368.81 tok/s 11,554.90 tok/s +7.04%
Output throughput 1,967.71 tok/s 1,838.23 tok/s +7.04%
Output throughput/GPU 491.93 tok/s 459.56 tok/s +7.04%
Request throughput 0.746 req/s 0.697 req/s +7.04%
p50 output tok/s/user 57.09 55.15 +3.51%
p50 TPOT 17.52 ms 18.13 ms 3.39% lower
p90 TPOT 32.02 ms 35.22 ms 9.09% lower
p50 TTFT 373.09 ms 485.54 ms 23.16% lower
p90 TTFT 36.23 s 38.87 s 6.79% lower
p50 request latency 26.72 s 27.52 s 2.90% lower
p90 request latency 160.94 s 174.49 s 7.76% lower
Benchmark duration 1,338.68 s 1,432.98 s 6.58% lower

Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, including dependencies and duplicate-work analysis.
  • The test plan and commands/checks used.
  • Correctness, CUDA-graph, and end-to-end performance results.
  • Documentation impact considered; this adds an internal model backend and
    requires no supported-model-list change.

Assisted-by: OpenAI Codex
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@NVShreyas NVShreyas changed the title flashinfer MSA backend for minimax m3 [kernel][perf][minimax m3] adds flashinfer MSA backend Sep 11, 2026
@NVShreyas NVShreyas changed the title [kernel][perf][minimax m3] adds flashinfer MSA backend [kernel][perf][minimax m3] adds flashinfer MSA backend for sm100 Sep 11, 2026
@mergify

mergify Bot commented Oct 1, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @NVShreyas.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Oct 1, 2026
@NVShreyas NVShreyas closed this Oct 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

1 participant