Feature/sm100 low latency nvfp4 kernels #1214

azhurkevich · 2025-07-04T04:59:09Z

📌 Description

Enable Blackwell with speed of light low latency kernels. Collaboration with @nekorobov. Supporting: @aleozlx, Kaixi Hou, Dongfeng Yu.

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

I have installed pre-commit by running pip install pre-commit (or used your preferred method).
I have installed the hooks with pre-commit install.
I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

Tests have been added or updated as needed.
All tests are passing (unittest, etc.).

Reviewer Notes

gemini-code-assist

Summary of Changes

Hello @azhurkevich, 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 significantly enhances FlashInfer's Mixture-of-Experts (MoE) capabilities by introducing a new, highly optimized backend for Blackwell GPUs based on NVIDIA's TensorRT-LLM Gen framework. It brings support for advanced FP4/FP8 quantization schemes, flexible routing algorithms, and specialized weight transformations, aiming to deliver speed-of-light low-latency kernels for large language models.

Highlights

New TRT-LLM Gen MoE Backend: Introduces a new fused Mixture-of-Experts (MoE) kernel implementation specifically for Blackwell (SM100) GPUs, leveraging NVIDIA's TensorRT-LLM Gen framework for highly optimized low-latency operations.
Advanced Quantization Support: Enables FP4 and FP8 block-scaled quantization for MoE layers, including support for different scale factor layouts (swizzled, linear) and quantization types (FP16_TO_FP4, FP8_TO_FP4, FP16_TO_MXFP8).
Flexible Routing Methods: Integrates multiple routing algorithms, such as DeepSeekV3, Renormalize (TopK -> Softmax), and RenormalizeNaive (Softmax -> TopK -> Renormalize), providing diverse options for MoE expert selection.
Optimized Weight Transformations: Implements specialized preprocessing steps for weight matrices, including row reordering for gated activations and shuffling for transposed MMA outputs, to maximize performance on Blackwell architecture.
Code Refactoring: Renames the existing FusedMoeRunner class to FusedMoeCutlassRunner and updates its references, improving code clarity and distinguishing it from the new TRT-LLM Gen backend.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 is currently in preview and 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 to provide feedback.

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.

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 introduces enhancements for FP4 quantization on SM100 architecture, including a new trtllmgen backend for fused MoE operations. The changes involve refactoring of existing quantization kernels for better generality and the addition of new functionalities and APIs. The new tests are comprehensive and cover the new features well. I've provided a few comments to improve code clarity and robustness.

gemini-code-assist · 2025-07-04T05:03:41Z

csrc/fused_moe/trtllmgen_backend/fp4BlockScaleMoe.cpp

+    at::Tensor num_non_exiting_ctas
+        = at::empty({1}, at::TensorOptions().device(routing_logits.device()).dtype(at::ScalarType::Int));
+
+    // FIXME: check shape


There's a FIXME: check shape comment here. Please clarify what needs to be checked or resolve this before merging. It's important to address FIXMEs to maintain code quality and prevent potential issues.

gemini-code-assist · 2025-07-04T05:03:41Z

csrc/nv_internal/tensorrt_llm/thop/fp4Quantize.cpp

  CHECK_CONTIGUOUS(self);
-  CHECK_INPUT_TYPE(globalScale, c10::ScalarType::Float);
-  TORCH_CHECK(sfVecSize == 16, "sfVecSize can only be 16");
+  if (sfUseUE8M0) {
+    TORCH_CHECK(globalScale.has_value(), "globalScale is required for UE8M0");


The sfVecSize argument is checked to be 16 or 32, but it's not used to select the kernel implementation. The selection is based on sfUseUE8M0, which implicitly determines the vector size (16 for false, 32 for true). This can be confusing and error-prone. To make the code more robust and self-documenting, add a check to enforce this relationship. For example:

TORCH_CHECK((sfUseUE8M0 && sfVecSize == 32) || (!sfUseUE8M0 && sfVecSize == 16), "sfVecSize must be 32 for mxfp4 (sfUseUE8M0=true) and 16 for nvfp4 (sfUseUE8M0=false)");

This will ensure that the arguments are consistent with the kernel's expectations.

gemini-code-assist · 2025-07-04T05:03:41Z

csrc/nv_internal/tensorrt_llm/thop/fp4Quantize.cpp

+        TORCH_CHECK(false, "Other layout not implemented yet.");
+    }


This TORCH_CHECK(false, ...) will always fail if an unsupported layout is provided. It's better to provide a more informative error message by checking the layout value directly. This makes debugging easier for users.

Suggested change

TORCH_CHECK(false, "Other layout not implemented yet.");

}

else

{

TORCH_CHECK(false, "Unsupported QuantizationSFLayout. Only SWIZZLED and LINEAR are supported.");

}

azhurkevich · 2025-07-04T05:16:50Z

(cherry picked from commit efc0d82)

…4feb)

flashinfer/fp4_quantization.py

fzyzcjy · 2025-07-16T00:05:33Z

"ETA for kernel: 7/15" (src: sgl-project/sglang#8037)

Hi, may I know whether this will be merged today? Thanks!

azhurkevich · 2025-07-16T00:07:25Z

thats the plan yes

fzyzcjy · 2025-07-16T00:48:00Z

Looking forward to it!

azhurkevich · 2025-07-16T01:07:14Z

its gonna be long night)

gemini-code-assist bot reviewed Jul 4, 2025

View reviewed changes

azhurkevich and others added 4 commits July 4, 2025 06:17

Fast nvfp4 decode

055780e

Fast nvfp4 decode

a0bb250

precommit

70da651

cherrypick cubin loader

79191e0

(cherry picked from commit efc0d82)

yzh119 marked this pull request as ready for review July 6, 2025 00:20

zhyncs assigned yzh119 and zhyncs Jul 6, 2025

zhyncs added the priority: high label Jul 6, 2025

azhurkevich mentioned this pull request Jul 8, 2025

[POC] enable trtllm fp4 from trtllm wheel sgl-project/sglang#7711

Open

6 tasks

refresh

7d1d265

aleozlx force-pushed the feature/sm100_low_latency_nvfp4_kernels branch from a65a04f to 7d1d265 Compare July 8, 2025 21:38

aleozlx added 2 commits July 8, 2025 15:35

more cubin loader plumbing

9f7939f

Update KernelRunner.cpp (bringing in cuda graph fix flashinfer-ai@7d7…

581998b

…4feb)

kaixih reviewed Jul 10, 2025

View reviewed changes

flashinfer/fp4_quantization.py Show resolved Hide resolved

azhurkevich mentioned this pull request Jul 10, 2025

feat: trtllm-gen fp8 moe kernels #1212

Merged

5 tasks

kushanam mentioned this pull request Jul 15, 2025

[Feature] Integrating FlashInfer FP4/FP8 Low-Latency MoE Kernels for DSR1 sgl-project/sglang#8037

Open

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Feature/sm100 low latency nvfp4 kernels #1214

Feature/sm100 low latency nvfp4 kernels #1214

azhurkevich commented Jul 4, 2025 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jul 4, 2025

Uh oh!

gemini-code-assist bot Jul 4, 2025

Uh oh!

gemini-code-assist bot Jul 4, 2025

Uh oh!

azhurkevich commented Jul 4, 2025

Uh oh!

Uh oh!

fzyzcjy commented Jul 16, 2025

Uh oh!

azhurkevich commented Jul 16, 2025

Uh oh!

fzyzcjy commented Jul 16, 2025

Uh oh!

azhurkevich commented Jul 16, 2025

Uh oh!

Uh oh!

Feature/sm100 low latency nvfp4 kernels #1214

Are you sure you want to change the base?

Feature/sm100 low latency nvfp4 kernels #1214

Conversation

azhurkevich commented Jul 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 4, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 4, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 4, 2025

Choose a reason for hiding this comment

Uh oh!

azhurkevich commented Jul 4, 2025

Uh oh!

Uh oh!

fzyzcjy commented Jul 16, 2025

Uh oh!

azhurkevich commented Jul 16, 2025

Uh oh!

fzyzcjy commented Jul 16, 2025

Uh oh!

azhurkevich commented Jul 16, 2025

Uh oh!

Uh oh!

azhurkevich commented Jul 4, 2025 •

edited

Loading