feat: trtllm-gen fp8 moe kernels #1212

aleozlx · 2025-07-03T16:35:26Z

📌 Description

🔍 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 @aleozlx, 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 introduces a significant new capability: fused Mixture-of-Experts operations with FP8 precision, specifically tailored for NVIDIA Blackwell GPUs. It integrates highly optimized batched GEMM kernels and custom CUDA logic for key MoE stages, exposed through PyTorch custom operators. The changes aim to enhance performance and memory efficiency for large-scale language models utilizing MoE architectures.

Highlights

FP8 Mixture-of-Experts (MoE) Implementation: Introduced a comprehensive implementation of fused Mixture-of-Experts (MoE) operations with support for FP8 (Float 8) precision, including both per-tensor and block-scale quantization. This enables more memory-efficient and potentially faster inference for large language models.
NVIDIA Blackwell (SM100a) Optimization: The entire MoE pipeline is heavily optimized for NVIDIA Blackwell (SM100a) architecture, leveraging advanced hardware features like TMA (Tensor Memory Access) for efficient data movement and programmatic stream serialization (PDL) for improved kernel launch efficiency.
Modular C++/CUDA Architecture: The MoE pipeline is designed with a modular C++/CUDA architecture, breaking down the complex operation into distinct stages: batched GEMM, expert routing, activation, permutation, and finalization. This modularity allows for fine-grained optimization and easier maintenance.
PyTorch Integration: The optimized MoE operations are exposed as PyTorch custom operators, providing a seamless and high-performance interface for integration into existing PyTorch-based deep learning frameworks.
Extensive Validation and Testing: Comprehensive unit tests are included, featuring detailed Python reference implementations for various MoE stages and FP8 quantization schemes. These tests ensure the correctness and numerical accuracy of the new CUDA kernels.

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 significant new functionality for FP8 Mixture of Experts (MoE) by integrating kernels from TensorRT-LLM. The changes are extensive and add considerable value.

My review focused on ensuring the correctness and maintainability of the new code. I have identified a few issues, including a critical bug in the kernel selection logic that could prevent the system from functioning correctly, as well as some high-severity issues related to performance and logical correctness. Additionally, there are opportunities to improve code quality by reducing duplication and replacing estimated values with precise calculations.

I've provided detailed comments and suggestions for each of these points. Addressing them will be important for the stability and performance of this new feature.

csrc/trtllm_batched_gemm_runner.cu

flashinfer/fused_moe.py

include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h

csrc/trtllm_fused_moe_dev_kernel.cu

csrc/trtllm_fused_moe_kernel_launcher.cu

yzh119 · 2025-07-08T06:50:36Z

include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaKernelLauncher.h

@@ -0,0 +1,92 @@
+/*


I'm confused about the code structure here, seems these functions are general and why we place them under a "trtllm/gen" subfolder of "trtllmGen_bmm_export" which seems to be operator-specific?

yzh119 · 2025-07-08T06:56:51Z

include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CommonUtils.h

+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>


These functions seems to have same functionalities of ceil_div and round_up in flashinfer/utils.cuh, can we just rely on flashinfer/utils.cuh?

addressed in upcoming commit

yzh119 · 2025-07-08T06:57:32Z

tests/test_trtllm_gen_fused_moe.py

+import torch
+from torch.nn import functional as F
+
+sys.path.append(os.path.join(os.path.dirname(__file__), ".."))


Why do we need this?

addressed in upcoming commit

yzh119 · 2025-07-08T06:59:22Z

tests/test_trtllm_gen_fused_moe.py

+    return permute_info, scores
+
+
+def dequant_reference_dsfp8(input, scale, transpose_scale, block_m, block_n):


Can we use matrix operations instead of for-loop instead? The unittests could be accelerated with that.
You can see my comments in #1211 (comment)

addressed in upcoming commit

include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/config.json

yzh119 · 2025-07-08T07:03:57Z

csrc/trtllm_fused_moe_kernel_launcher.cu

+#include <cuda_runtime.h>
+#include <flashinfer/exception.h>
+#include <nvrtc.h>
+#include <torch/all.h>


Avoid using <torch/all.h> as it will greatly increse compilation speed.

addressed in upcoming commit

yzh119 · 2025-07-08T07:04:50Z

csrc/trtllm_fused_moe_dev_kernel.cu

+#include <cuda/std/functional>
+
+#include "flashinfer/trtllm/fused_moe/DevKernel.h"
+


Kernel files should be put under include/, the include/ folder should be treated as a header-only library and self-contained, so kernel definitions can also be put under flashinfer/trtllm/fused_moe/DevKernel.h.

discussed offline

yzh119 · 2025-07-08T07:21:58Z

csrc/trtllm_fused_moe_routing_kernel.cu

+#include <type_traits>
+
+#include "flashinfer/trtllm/fused_moe/DevKernel.h"
+#include "flashinfer/trtllm/fused_moe/RoutingKernel.h"


We can move kernel definitions to RoutingKernels.h under include/.
csrc/ is the directory for operator registration and torch-specific.

discussed offline

yzh119

Thanks @aleozlx and @azhurkevich for the great work.
The current form is acceptable to me and let's merge it now.
Will Postpone refactor to later PRs.

azhurkevich · 2025-07-10T17:09:52Z

@yzh119 small commit coming to fix most issues once I get access). I will do some refactorings we've discussed in nvfp4 PR.

Patch fp8 cubin availability (followup of #1212 ) --------- Co-authored-by: Zihao Ye <expye@outlook.com>

.

d46bf02

gemini-code-assist bot reviewed Jul 3, 2025

View reviewed changes

aleozlx added 11 commits July 3, 2025 17:30

fix for bot reviewer

593a0a8

run pre-commit

add8210

confirmed pre-commit totally harmless on metadata

72316b4

checksums

bfadadd

..

5e21da5

..

3ea682b

try cuda::std

20597c2

refresh

41b68d2

revert sampling.cuh

e30c97d

config fp8_nvfp4

b074c9d

cubin path

a2bcd76

yzh119 reviewed Jul 8, 2025

View reviewed changes

refresh

f7d0973

kaixih mentioned this pull request Jul 8, 2025

[NVIDIA] Add SM100 Flashinfer MoE blockscale fp8 backend for low latency vllm-project/vllm#20645

Merged

yzh119 and others added 2 commits July 8, 2025 23:57

Merge branch 'main' into trtllmgen-moe-fp8

fc9f0ef

missing module cache

3dfda05

aleozlx force-pushed the trtllmgen-moe-fp8 branch from 7d74feb to 3dfda05 Compare July 9, 2025 20:08

Fixes and feedback

b54aeb9

yzh119 changed the title ~~trtllmgen-moe-fp8~~ feat: trtllm-gen fp8 moe kernels Jul 10, 2025

yzh119 approved these changes Jul 10, 2025

View reviewed changes

precommit fixes

161de1e

yzh119 merged commit bd74e15 into flashinfer-ai:main Jul 10, 2025
2 checks passed

yzh119 mentioned this pull request Jul 11, 2025

Patch fp8 cubin availability #1240

Merged

yzh119 added a commit that referenced this pull request Jul 11, 2025

hotfix: Patch fp8 cubin availability (#1240)

d0c104d

Patch fp8 cubin availability (followup of #1212 ) --------- Co-authored-by: Zihao Ye <expye@outlook.com>


		////////////////////////////////////////////////////////////////////////////////////////////////////

		template <typename T>

		return permute_info, scores


		def dequant_reference_dsfp8(input, scale, transpose_scale, block_m, block_n):

		#include <cuda/std/functional>

		#include "flashinfer/trtllm/fused_moe/DevKernel.h"

feat: trtllm-gen fp8 moe kernels #1212

feat: trtllm-gen fp8 moe kernels #1212

Conversation

aleozlx commented Jul 3, 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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yzh119 left a comment

Choose a reason for hiding this comment

Uh oh!

azhurkevich commented Jul 10, 2025

Uh oh!

Uh oh!

Uh oh!

aleozlx commented Jul 3, 2025 •

edited

Loading