feat:enable fp8 blockscale moe for fused cultass for sm90 #1819

djmmoss · 2025-09-30T15:30:24Z

📌 Description

Adds FP8 Block Scaling Fused Cutlass MoE for SM90.

✅ 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

Signed-off-by: Duncan Moss <[email protected]>

gemini-code-assist · 2025-09-30T15:31:20Z

Summary of Changes

Hello @djmmoss, 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 the Fused Cutlass Mixture-of-Experts (MoE) functionality by introducing FP8 Block Scaling capabilities, primarily for SM90 architectures. It integrates a sophisticated DeepGEMM JIT compilation system that dynamically generates and optimizes CUDA kernels for efficient FP8 GEMM operations. This allows for more flexible and performant FP8 computations in MoE layers, moving from a stub implementation to a full-fledged feature with comprehensive kernel support and updated build configurations.

Highlights

FP8 Block Scaling for MoE: Introduced support for FP8 Block Scaling in Fused Cutlass Mixture-of-Experts (MoE) operations, specifically targeting SM90 (Hopper) architectures.
DeepGEMM JIT Integration: Integrated a DeepGEMM JIT compilation framework, enabling dynamic generation and optimization of CUDA kernels for FP8 GEMM operations.
New CUDA Kernels: Added several new CUDA kernel files (fp8_gemm_impl.cuh, sm89_fp8_gemm_1d1d.cuh, fp8_blockscale_gemm_kernel.cuh, etc.) to handle the specialized FP8 GEMM computations and block-wise scaling.
Build System Updates: Modified the build process in flashinfer/fused_moe/core.py to include the new FP8 blockscale GEMM implementation and added -lnvrtc to linker flags for NVRTC support.
Test Case Enhancements: Updated existing test cases in tests/test_trtllm_cutlass_fused_moe.py to validate the new FP8 block scaling functionality, including changes to scale initialization and output handling.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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 support for FP8 block scaling in fused Mixture-of-Experts (MoE) kernels for SM90 architectures, which is a significant feature enhancement. The implementation includes a new JIT compilation framework for CUDA kernels, the FP8 GEMM kernel implementations, and updates to the build system and tests. The code is generally well-structured, but I've identified a few issues that should be addressed to improve correctness and maintainability, including a bug in kernel name generation, an inconsistent function declaration, and a redundant loop.

gemini-code-assist · 2025-09-30T15:34:22Z

csrc/nv_internal/tensorrt_llm/deep_gemm/compiler.cuh

+                       std::to_string(num_groups) + "_" + std::to_string(num_stages) +
+                       std::to_string(num_groups) + "_" + std::to_string(num_stages) + "_" +


There is a duplication in the construction of the kernel name. The num_groups and num_stages parameters are appended twice. This appears to be a copy-paste error and will result in an incorrect and unnecessarily long kernel name, which could lead to caching issues.

std::to_string(num_groups) + "_" + std::to_string(num_stages) + "_" +

gemini-code-assist · 2025-09-30T15:34:22Z

csrc/nv_internal/tensorrt_llm/deep_gemm/fp8_gemm.cuh

+
+namespace deep_gemm {
+template <typename T>
+static CUtensorMap make_2d_tma_a_desc(T* global_address, uint32_t shape_m, uint32_t shape_k,


The function make_2d_tma_a_desc is declared static, while other similar helper functions in this file (e.g., make_2d_tma_b_desc) are not. This is inconsistent. For function templates in header files, the static keyword is generally not necessary and can be removed for consistency and adherence to modern C++ practices.

CUtensorMap make_2d_tma_a_desc(T* global_address, uint32_t shape_m, uint32_t shape_k,

gemini-code-assist · 2025-09-30T15:34:22Z

.../nv_internal/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.cu

+    for (int i = 0; i < num_problems; i++) {
+      fp8_mat_b = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<void*>(mat_b));
+      per_block_scales = const_cast<float*>(scales_b);
+    }


This loop is redundant. It repeatedly assigns the same pointer values to fp8_mat_b and per_block_scales in each iteration. The loop can be removed and replaced with a single assignment to improve clarity and remove dead code.

fp8_mat_b = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<void*>(mat_b)); per_block_scales = const_cast<float*>(scales_b);

Signed-off-by: Duncan Moss <[email protected]>

yzh119 · 2025-09-30T22:40:32Z

flashinfer/jit/core.py

+        "-O3",
+        "-std=c++17",
+        "-Wno-switch-bool",
+        "-D__CUDACC_VER_MAJOR__=" + str(torch.version.cuda.split(".")[0]),


Why do we need this? They are nvcc native macros.

When I was compiling, they weren't enabled in the JIT mode, I'm not sure it this was an environment issue, but I just installed the package with pip install -e . in a fresh venv. I believe it's because this is c++ instead of nvcc?

E [2/3] c++ -MMD -MF fp4_quantization_90/fp4Op.o.d -DTORCH_EXTENSION_NAME=fp4_quantization_90 -DTORCH_API_INCLUDE_EXTENSION_H -DPy_LIMITED_API=0x03090000 -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1018\" -D_GLIBCXX_USE_CXX11_ABI=1 -I/home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc/nv_internal -I/home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc/nv_internal/include -isystem /usr/include/python3.10 -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/venv/lib/python3.10/site-packages/torch/include -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /home/scratch.svc_compute_arch/release/cuda_toolkit/internal/latest_r12.9/include -isystem /home/scratch.svc_compute_arch/release/cuda_toolkit/internal/latest_r12.9/include/cccl -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/include -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/3rdparty/cutlass/include -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/3rdparty/cutlass/tools/util/include -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/3rdparty/spdlog/include -fPIC -O3 -std=c++17 -Wno-switch-bool -DENABLE_BF16 -DENABLE_FP8 -DENABLE_FP4 -c /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc/nv_internal/tensorrt_llm/thop/fp4Op.cpp -o fp4_quantization_90/fp4Op.o E FAILED: [code=1] fp4_quantization_90/fp4Op.o E c++ -MMD -MF fp4_quantization_90/fp4Op.o.d -DTORCH_EXTENSION_NAME=fp4_quantization_90 -DTORCH_API_INCLUDE_EXTENSION_H -DPy_LIMITED_API=0x03090000 -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1018\" -D_GLIBCXX_USE_CXX11_ABI=1 -I/home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc/nv_internal -I/home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc/nv_internal/include -isystem /usr/include/python3.10 -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/venv/lib/python3.10/site-packages/torch/include -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /home/scratch.svc_compute_arch/release/cuda_toolkit/internal/latest_r12.9/include -isystem /home/scratch.svc_compute_arch/release/cuda_toolkit/internal/latest_r12.9/include/cccl -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/include -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/3rdparty/cutlass/include -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/3rdparty/cutlass/tools/util/include -isystem /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/3rdparty/spdlog/include -fPIC -O3 -std=c++17 -Wno-switch-bool -DENABLE_BF16 -DENABLE_FP8 -DENABLE_FP4 -c /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc/nv_internal/tensorrt_llm/thop/fp4Op.cpp -o fp4_quantization_90/fp4Op.o E /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc/nv_internal/tensorrt_llm/thop/fp4Op.cpp: In function ‘at::Tensor torch_ext::mxfp4_dequantize_host(at::Tensor, at::Tensor, int64_t)’: E /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc/nv_internal/tensorrt_llm/thop/fp4Op.cpp:405:44: error: ‘__CUDACC_VER_MAJOR__’ was not declared in this scope E 405 | std::cout << "Current CUDA version: " << __CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 << std::endl; E | ^~~~~~~~~~~~~~~~~~~~ E /home/scratch.dmoss_gpu_1/repos/squirtle/flashinfer/csrc/nv_internal/tensorrt_llm/thop/fp4Op.cpp:405:75: error: ‘__CUDACC_VER_MINOR__’ was not declared in this scope E 405 | std::cout << "Current CUDA version: " << __CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 << std::endl; E | ^~~~~~~~~~~~~~~~~~~~ E ninja: build stopped: subcommand failed.

Yes c++ files will be compiled with gcc instead of nvcc.
Considering these file includes cuda headers, can we rename it to fp4Op.cu instead?

ah, this was actually fixed in different PR. I've removed this change

…cale

Signed-off-by: Duncan Moss <[email protected]>

djmmoss · 2025-10-01T18:10:12Z

@yzh119 ready for final review

yzh119 · 2025-10-02T01:30:30Z

csrc/nv_internal/tensorrt_llm/deep_gemm/compiler.cuh

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


How are these implementations from https://github.com/deepseek-ai/DeepGEMM

We already have deepgemm kernels

https://docs.flashinfer.ai/generated/flashinfer.gemm.batch_deepgemm_fp8_nt_groupwise.html

https://docs.flashinfer.ai/generated/flashinfer.gemm.group_gemm_mxfp4_nt_groupwise.html

If there are not significant changes, we should unify them (maybe not in this PR).

The tensorrt_llm/deep_gemm files are generated here: https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/deep_gemm

I checked the deepgemm kernels that have already been integrated into flashinfer, currently they only seem to work on blackwell (this PR is for hopper). I can look in a follow-up PR into merging the two.

Signed-off-by: Duncan Moss <[email protected]>

…cale

Signed-off-by: Duncan Moss <[email protected]>

yzh119 · 2025-10-07T18:12:37Z

/run bot

yongwww · 2025-10-07T18:15:16Z

/bot run

flashinfer-bot · 2025-10-07T18:16:11Z

GitLab MR !69 has been created, and the CI pipeline #36196253 is currently running. I'll report back once the pipeline job completes.

yongwww · 2025-10-07T18:18:20Z

/run bot

it is bot run @yzh119 :)

djmmoss added 2 commits September 29, 2025 13:18

feat:enable fp8 blockscale moe for fused cultass

4097760

Signed-off-by: Duncan Moss <[email protected]>

pre-commit

7e507ac

Signed-off-by: Duncan Moss <[email protected]>

gemini-code-assist bot reviewed Sep 30, 2025

View reviewed changes

djmmoss added 2 commits September 30, 2025 12:30

disable for blackwell

853f465

Signed-off-by: Duncan Moss <[email protected]>

proper support

601305d

Signed-off-by: Duncan Moss <[email protected]>

yzh119 reviewed Sep 30, 2025

View reviewed changes

djmmoss changed the title ~~feat:enable fp8 blockscale moe for fused cultass~~ feat:enable fp8 blockscale moe for fused cultass for sm90 Sep 30, 2025

djmmoss added 3 commits September 30, 2025 17:06

Merge remote-tracking branch 'origin/main' into dmoss/sm90_fp8_blocks…

d4ae8b2

…cale

updates

72b49cc

Signed-off-by: Duncan Moss <[email protected]>

cuda 13 fix

db1b960

Signed-off-by: Duncan Moss <[email protected]>

yzh119 reviewed Oct 2, 2025

View reviewed changes

djmmoss added 5 commits October 2, 2025 10:58

remove unneeded

2c7e541

Signed-off-by: Duncan Moss <[email protected]>

Merge remote-tracking branch 'origin/main' into dmoss/sm90_fp8_blocks…

ac3a771

…cale

pre-commit

9fd76ec

Signed-off-by: Duncan Moss <[email protected]>

remove unneeded functions

20082ac

Signed-off-by: Duncan Moss <[email protected]>

revert jit core.py

8b04389

Signed-off-by: Duncan Moss <[email protected]>

		std::to_string(num_groups) + "_" + std::to_string(num_stages) +
		std::to_string(num_groups) + "_" + std::to_string(num_stages) + "_" +

feat:enable fp8 blockscale moe for fused cultass for sm90 #1819

Are you sure you want to change the base?

feat:enable fp8 blockscale moe for fused cultass for sm90 #1819

Conversation

djmmoss commented Sep 30, 2025

📌 Description

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Uh oh!

gemini-code-assist bot commented Sep 30, 2025

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 Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

yzh119 Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

djmmoss Sep 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yzh119 Oct 2, 2025

Choose a reason for hiding this comment

Uh oh!

djmmoss Oct 2, 2025

Choose a reason for hiding this comment

Uh oh!

djmmoss commented Oct 1, 2025

Uh oh!

yzh119 Oct 2, 2025

Choose a reason for hiding this comment

Uh oh!

djmmoss Oct 2, 2025

Choose a reason for hiding this comment

Uh oh!

yzh119 commented Oct 7, 2025

Uh oh!

yongwww commented Oct 7, 2025

Uh oh!

flashinfer-bot commented Oct 7, 2025

Uh oh!

yongwww commented Oct 7, 2025

Uh oh!

Uh oh!

djmmoss Sep 30, 2025 •

edited

Loading