Skip to content

[Kernel] adopt mxfp8 grouped_gemm and grouped_quant kernel#34381

Open
EdalatiAli wants to merge 2 commits intovllm-project:mainfrom
EdalatiAli:main
Open

[Kernel] adopt mxfp8 grouped_gemm and grouped_quant kernel#34381
EdalatiAli wants to merge 2 commits intovllm-project:mainfrom
EdalatiAli:main

Conversation

@EdalatiAli
Copy link

@EdalatiAli EdalatiAli commented Feb 11, 2026

Purpose

In order to enable serving mxfp8 MoE models, this PR Integrates SGLang’s SM100+ expert-specialization MXFP8 blockscaled grouped kernels into vLLM so they are built, registered, importable, and test-covered in the vLLM codebase.
Here is source PR for the adopted kernels.

This PR:

  • Adds es_sm100_mxfp8_blockscaled_grouped_mm and es_sm100_mxfp8_blockscaled_grouped_quant kernel sources into vLLM’s _C build path (CUDA-gated for SM100-compatible targets).
  • Registers op schemas in csrc/torch_bindings.cpp and wires CUDA impls already present in copied kernel files.
  • Adds Python wrappers in vllm/_custom_ops.py, including guarded register_fake hooks for fake/meta tracing compatibility.
  • Integrates copied test into vLLM test conventions:
    • tests/kernels/moe/test_es_mxfp8_blockscaled_moe.py
    • switches from sgl_kernel import to vllm._custom_ops
    • fixes argument ordering to match vLLM wrapper signatures
    • applies robust SM100+/CUDA skip gating.
  • Adds source provenance/citation headers to copied kernel and test files (adapted from SGLang, commit snapshot ded068a76).

Test Plan

  1. Static/syntax checks:
    • python3 -m py_compile vllm/_custom_ops.py
    • python3 -m py_compile tests/kernels/moe/test_es_mxfp8_blockscaled_moe.py
  2. Lint/diagnostic check on modified files.
  3. (On SM100 + CUDA>=12.8 environment) run:
    • pytest -q tests/kernels/moe/test_es_mxfp8_blockscaled_moe.py

Test Result

  • python3 -m py_compile vllm/_custom_ops.py passed.
  • python3 -m py_compile tests/kernels/moe/test_es_mxfp8_blockscaled_moe.py passed.
  • ✅ Lint diagnostics for modified files: no issues found.
  • pytest -q tests/kernels/moe/test_es_mxfp8_blockscaled_moe.py passed.

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.

@github-actions
Copy link

👋 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 ask your reviewers to trigger select CI tests on top of fastcheck CI.

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.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

@mergify mergify bot added the ci/build label Feb 11, 2026
Copy link
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 integrates SGLang's SM100+ MXFP8 blockscaled grouped kernels into vLLM. The changes include adding new CUDA source files for the kernels, updating CMakeLists.txt to handle their compilation, and adding corresponding Python wrappers and tests. The implementation appears solid and well-integrated. I have one suggestion to improve the clarity of an error message in one of the new CUDA files to prevent potential confusion during debugging.

@EdalatiAli EdalatiAli marked this pull request as ready for review February 12, 2026 16:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant