Skip to content

[Kernel] Add MXFP8 to Marlin GEMM/MoE and refactor Mxfp8LinearOp#34664

Open
mgoin wants to merge 15 commits intovllm-project:mainfrom
neuralmagic:mxfp8-marlin
Open

[Kernel] Add MXFP8 to Marlin GEMM/MoE and refactor Mxfp8LinearOp#34664
mgoin wants to merge 15 commits intovllm-project:mainfrom
neuralmagic:mxfp8-marlin

Conversation

@mgoin
Copy link
Copy Markdown
Member

@mgoin mgoin commented Feb 17, 2026 •
edited
Loading

Purpose

The Marlin kernel already supports FP8 (per-channel/group scales) and MXFP4 (per-32-element e8m0 scales). MXFP8 is a natural combination: FP8 weights (like existing FP8 Marlin) with e8m0 microscaling block scales (like existing MXFP4 Marlin). We just have to wire the kernel building blocks together.

This PR also consolidates gemm kernel backend specific logic more into the Mxfp8LinearOp class for modelopt.py and mxfp8.py

Test Plan

Existing online quant test for mxfp8 will now run on L4 in CI tests/models/quantization/test_mxfp8.py

Test Result

vllm serve mgoin/Qwen3-0.6B-MXFP8
vllm serve Qwen/Qwen3-0.6B --quantization mxfp8
vllm serve Qwen/Qwen3-30B-A3B --quantization mxfp8
Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
  • (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

gemini-code-assist[bot]
gemini-code-assist bot reviewed Feb 17, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request adds support for MXFP8 quantization in the Marlin kernel, providing a faster alternative to the existing emulation path. The changes span across kernel generation, C++ dispatch logic, and Python-level integration. The implementation introduces new utility functions for handling MXFP8-specific weight and scale preparation for Marlin. My review identifies a critical issue in the hardware capability check that could lead to runtime errors on unsupported GPUs.

@mgoin mgoin added performance Performance-related issues quantization ready ONLY add when PR is ready to merge/full CI is needed labels Feb 20, 2026
@danisereb
Copy link
Copy Markdown
Contributor

danisereb commented Feb 22, 2026 •
edited
Loading

Hey @mgoin,
please also see my PR for the Flashinfer cutlass MXFP8 GEMM:
#35053

The GEMM is available in flashinfer 0.6.4 (recently bumped in vLLM).

danisereb
danisereb reviewed Feb 22, 2026
View reviewed changes
@danisereb danisereb mentioned this pull request Feb 23, 2026
Merged
5 tasks
@mergify
Copy link
Copy Markdown

mergify bot commented Feb 24, 2026

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

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 Feb 24, 2026
mgoin added 4 commits March 19, 2026 17:49
@mgoin
Add MXFP8 to Marlin dense kernel
382b416 
Signed-off-by: mgoin <mgoin64@gmail.com>
@mgoin
Fix scale dequant
00a5d9b 
Signed-off-by: mgoin <mgoin64@gmail.com>
@mgoin
Refactor: integrate Marlin backend into Mxfp8LinearOp
07354dd 
Move backend selection, weight processing, and apply logic from
ModelOptMxFp8LinearMethod into Mxfp8LinearOp so all MXFP8 linear
backends (emulation, flashinfer CUTLASS, Marlin) are managed in
one place.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Signed-off-by: mgoin <mgoin64@gmail.com>
@mgoin
Update Mxfp8OnlineLinearMethod to use shared backend selection
4126c85 
Use select_mxfp8_linear_backend() and delegate weight processing
to Mxfp8LinearOp.process_weights(), enabling Marlin backend support
for online MXFP8 quantization on SM80+ GPUs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Signed-off-by: mgoin <mgoin64@gmail.com>
@mergify mergify bot removed the needs-rebase label Mar 19, 2026
mgoin added 5 commits March 19, 2026 18:07
@mgoin
Remove comment on weight processing delegation
72ca1b2 
Removed comment about backend-specific weight processing.
@mgoin
Remove redundant layer attributes for Marlin MXFP8
b6c9e87 
Use torch.get_default_dtype() and layer.output_size_per_partition /
layer.input_size_per_partition directly instead of stashing copies
as layer.orig_dtype, layer.marlin_size_n, layer.marlin_size_k.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Signed-off-by: mgoin <mgoin64@gmail.com>
@mgoin
Cleanup
b8b4b09 
Signed-off-by: mgoin <mgoin64@gmail.com>
@mgoin
Fix scale logic
ef27d9b 
Signed-off-by: mgoin <mgoin64@gmail.com>
@mgoin
Fix
bf8ade8 
Signed-off-by: mgoin <mgoin64@gmail.com>
@mgoin mgoin changed the title Add MXFP8 to Marlin dense kernel [Kernel] Add MXFP8 to Marlin dense kernel Mar 19, 2026
@mgoin mgoin added the nvidia label Mar 19, 2026
@mgoin mgoin changed the title [Kernel] Add MXFP8 to Marlin dense kernel [Kernel] Add MXFP8 to Marlin dense kernel and refactor Mxfp8LinearOp Mar 19, 2026
@mgoin mgoin mentioned this pull request Mar 24, 2026
Open
mgoin added 2 commits March 30, 2026 16:13
@mgoin mgoin requested a review from WoosukKwon as a code owner March 30, 2026 21:15
@mgoin mgoin changed the title [Kernel] Add MXFP8 to Marlin dense kernel and refactor Mxfp8LinearOp [Kernel] Add MXFP8 to Marlin GEMM/MoE and refactor Mxfp8LinearOp Mar 30, 2026
@mgoin
Fix moe kernel
c57f65e 
Signed-off-by: mgoin <mgoin64@gmail.com>
@mergify
Copy link
Copy Markdown

mergify bot commented Mar 31, 2026

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

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 Mar 31, 2026
@mgoin
Merge branch 'main' into mxfp8-marlin
d736dd6 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Signed-off-by: mgoin <mgoin64@gmail.com>
@mergify mergify bot removed the needs-rebase label Mar 31, 2026
@mgoin
Fix backend selection
f939b89 
Signed-off-by: mgoin <mgoin64@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat is a code owner

@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth is a code owner

@yewentao256 yewentao256 Awaiting requested review from yewentao256 yewentao256 is a code owner

@pavanimajety pavanimajety Awaiting requested review from pavanimajety pavanimajety is a code owner

@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner

@DarkLight1337 DarkLight1337 Awaiting requested review from DarkLight1337 DarkLight1337 is a code owner

@ywang96 ywang96 Awaiting requested review from ywang96 ywang96 is a code owner

2 more reviewers

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

@danisereb danisereb danisereb left review comments

Reviewers whose approvals may not affect merge requirements

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

nvidia performance Performance-related issues quantization ready ONLY add when PR is ready to merge/full CI is needed

Projects

NVIDIA
Status: No status

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@mgoin @danisereb