Masked batch nvfp4 quantization

[wenscarl]

📌 Description

This PR quantizes the tensor to NVFP4 along with the associated scales, making it directly usable by FlashInfer’s grouped_gemm_nt_masked.

Note: Leave the permutation to framework .

cc @kaixih

🔍 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 @wenscarl, 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 extends FlashInfer's NVFP4 quantization capabilities by introducing support for masked batch quantization and a combined SiLU activation with element-wise multiplication prior to quantization. These enhancements are designed to make quantized tensors directly compatible with FlashInfer's grouped_gemm_nt_masked function, thereby improving efficiency and adaptability for various model architectures. The changes span across CUDA kernels, C++ bindings, and Python interfaces to seamlessly integrate these new, more flexible quantization options.

Highlights

• Masked Batch NVFP4 Quantization: Introduced the capability to perform NVFP4 quantization on batched tensors with an optional mask, allowing selective quantization based on a provided mask tensor. This enhances flexibility for operations like FlashInfer's grouped_gemm_nt_masked.
• SiLU and Multiplication with NVFP4 Quantization: Added a new silu_and_mul_fp4_batched_quantize operation. This function performs a SiLU activation and element-wise multiplication on the input tensor before quantizing it to NVFP4, streamlining common computational patterns.
• Kernel Refactoring and Policy-Based Execution: The core CUDA quantization kernel (quantize_with_block_size) was refactored into a generic quantize_with_block_size_impl that accepts a Policy template parameter. This allows for flexible execution paths (e.g., with or without SiLU/multiplication) and integrates mask handling efficiently.
• Python API and Testing Updates: The Python bindings (flashinfer/__init__.py, flashinfer/fp4_quantization.py) were updated to expose the new masked and SiLU-enabled quantization functionalities. Comprehensive unit tests (tests/test_fp4_quantize.py) were added or modified to ensure the correctness and robustness of these new features.

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.

Footnotes

1. 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 masked batch quantization for NVFP4 and adds a new fused silu_and_mul quantization kernel. The changes are well-structured, particularly the use of a policy-based design in the CUDA kernel to handle different quantization paths. My review focuses on improving correctness, maintainability, and testing. I've identified a missing validation check in the C++ code, incorrect dtypes in Python fake ops that could affect memory profiling, and some inconsistencies in the new Python API and its documentation. Additionally, I've found a couple of issues in the new tests that should be addressed to ensure correctness and full coverage of the new functionality.

[gemini-code-assist]

The silu_and_mul operation requires splitting the last dimension of the input tensor in half. To ensure this is always possible and to prevent potential runtime errors, you should add a check to verify that the last dimension (k_by_2) is even.

  TORCH_CHECK(k_by_2 % 2 == 0, "The last dimension of the input tensor for silu_and_mul must be even.");
  int64_t k = k_by_2 / 2;

[gemini-code-assist]

The global_scale tensor has a shape of (b,), containing per-batch item scales. However, you are passing the entire tensor to fp4_quantize inside a loop that iterates through batches. The fp4_quantize function expects a scalar-like tensor (shape [1]) for the scale and will only use the first element, leading to incorrect test results for batches other than the first one. You should pass the scale corresponding to the current batch item i.

Suggested change

	single_out, single_scale = fp4_quantize(
	x_silu_mul, global_scale, 16, False, True
	)
	single_out, single_scale = fp4_quantize(
	x_silu_mul, global_scale[i:i+1], 16, False, True
	)

[gemini-code-assist]

The fake op _fp4_batched_quantize_sm100 allocates tensors with dtype=torch.int64 and torch.int32, which use significantly more memory than the actual torch.uint8 output tensors. This can be misleading for memory profiling and might cause out-of-memory issues during model tracing or exporting. Please use the correct dtype torch.uint8 for both output tensors, as this is what the underlying C++ extension returns.

Suggested change

	input.new_empty([b, m, k // 2], dtype=torch.int64), # float4_e2m1_x2
	input.new_empty(
	[b, m * k // sf_vec_size], dtype=torch.int32
	), # Scale factors
	input.new_empty([b, m, k // 2], dtype=torch.uint8), # float4_e2m1_x2
	input.new_empty(
	[b, m * k // sf_vec_size], dtype=torch.uint8
	), # Scale factors

[gemini-code-assist]

The fake op _silu_and_mul_fp4_batched_quantize_sm100 allocates tensors with dtype=torch.int64 and torch.int32, which use significantly more memory than the actual torch.uint8 output tensors. This can be misleading for memory profiling and might cause out-of-memory issues during model tracing. Please use the correct dtype torch.uint8 for both output tensors to match the behavior of the C++ kernel.

Suggested change

	input.new_empty([b, m, k // 4], dtype=torch.int64), # float4_e2m1_x2
	input.new_empty(
	[b, m * k // sf_vec_size], dtype=torch.int32
	), # Scale factors
	input.new_empty([b, m, k // 4], dtype=torch.uint8), # float4_e2m1_x2
	input.new_empty(
	[b, m * k // (sf_vec_size * 2)], dtype=torch.uint8
	), # Scale factors

[gemini-code-assist]

The docstring for silu_and_mul_fp4_batched_quantize appears to be copied from nvfp4_batched_quantize and is inaccurate. It doesn't mention the fused silu_and_mul operation, which is a key part of this function's behavior. Please update the docstring to accurately describe what the function does. You can use the more detailed docstring from silu_and_mul_fp4_batched_quantize_sm100 as a reference.

[gemini-code-assist]

The function signature for silu_and_mul_fp4_batched_quantize is inconsistent with nvfp4_batched_quantize. Here, mask is a required positional argument, whereas in nvfp4_batched_quantize it's an optional keyword argument. For better API consistency and user experience, consider making mask an optional keyword argument. This would likely require changes in the C++ backend to accept an optional mask, similar to how fp4_batched_quantize is implemented.

Suggested change

	def silu_and_mul_fp4_batched_quantize(
	a,
	mask,
	a_global_sf,
	sf_vec_size=16,
	):
	def silu_and_mul_fp4_batched_quantize(
	a,
	a_global_sf,
	sf_vec_size=16,
	mask=None,
	):

[gemini-code-assist]

The mask for this test is generated with low=m and high=m+1, which means all values in the mask tensor will be m. This doesn't effectively test the masking functionality, as no rows are actually masked out. To properly test the masking logic, you should generate masks with varying lengths.

Suggested change

	mask = torch.randint(low=m, high=m + 1, size=(b,), dtype=torch.int32, device=device)
	mask = torch.randint(low=1, high=m + 1, size=(b,), dtype=torch.int32, device=device)

[kaixih]

Thx for the PR.

Do you think we should put the silu_and_mul_fp4_batched_quantize into the activation module like here.

The `nvfp4_batched_quantize) is natually here.

The main reason to separate their location is that the former assumes the ..., 2*hidden->..., hidden mapping but the latter assumes ..., hidden->..., hidden.

[kaixih]

nit: ..., /*mask=*/nullptr); for clarity.

[kaixih]

So, IIUC, the main motivation of this PR is to support masked FP4 quantization and to better maintain the different variants related to quantization methods:

1. Add masking support to nvfp4_batched_quantize.
2. Add a new function (or extend the existing silu_and_mul) to support NVFP4 output with masking.

The intended use case is to prepare inputs for the CuteDSL masked GEMM.

fyi. @zihaoye

[yzh119]

Hi @kaixih I think you at another @zihaoye who is also a flashinfer contributor :)

[yzh119]

As frameworks starts using torch 2.8 where torch.float4_e2m1_x2 is available, we should considering moving to native fp4x2 data type at some point (not necessarily in this PR)

[kaixih]

@wenscarl seems there is some conflicts? can you rebase?

[kaixih]

Approved this for now.
Note, there are more following cleanup works to do. @yzh119 let us know if you need us to do that.

[yzh119]

LGTM