SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

feat: Add 8x4 swizzle layout support to MXFP4 and MXFP8 CuTe-DSL kernels (#3357)

<!-- .github/pull_request_template.md -->

## ๐Ÿ“Œ Description
Brings the CuTe-DSL backend to parity with CUDA: `mxfp8_quantize` and
the MXFP4 path of `fp4_quantize` previously raised ValueError when asked
for 8x4. NVFP4 `cute-dsl` already supported 8x4.

### Changes ###
  
Kernels (`flashinfer/quantization/kernels/mxfp{4,8}_quantize.py`) --
Same pattern previously used for `NVFP4QuantizeSwizzledKernel`.
- Added `SF_LAYOUT_8x4` constant and threaded sf_layout through
`MXFP{4,8}QuantizeSwizzledKernel`.
- Added a `_compute_sf_offset` JIT helper that dispatches between
`compute_sf_index_swizzled_{128x4,8x4}_gpu`.
- Added `sf_layout` to the `@functools.cache` key on the compile
helpers.
  
Notes
- No public-API breakage (the new mxfp8 kwarg is keyword-only with False
default).
- High-level mxfp4_quantize still hardcodes 128x4 on both backends โ€”
parity preserved; 8x4 reachable via fp4_quantize.

In all cases tested, the CUDA and CuTe-DSL backends have 100% bitwise
matches in quantized output and scale factor tensors. Perf numbers from
B200:
<details>
<summary>Backend performance comparisons from
bench_nvfp4_quantize_backend_comparison.py</summary>

<img width="1644" height="1477"
alt="nvfp4_quantize_backend_comparison_linear_bfloat16"
src="https://github.com/user-attachments/assets/c2da9699-f9df-4a0b-ba30-3991cc1a8fa6"
/>
<img width="1644" height="1478"
alt="nvfp4_quantize_backend_comparison_swizzled_128x4_bfloat16"
src="https://github.com/user-attachments/assets/320d4cae-dc8a-4dab-af90-82684704170d"
/>
<img width="1644" height="1478"
alt="nvfp4_quantize_backend_comparison_swizzled_8x4_bfloat16"
src="https://github.com/user-attachments/assets/b6e16502-6314-4d7b-afdb-f3af5a6e3685"
/>


</details>
<details>
<summary>Backend performance comparisons from
bench_mxfp4_quantize_backend_comparison.py</summary>

<img width="1644" height="1477"
alt="mxfp4_quantize_backend_comparison_linear_bfloat16"
src="https://github.com/user-attachments/assets/64841ea7-d023-491b-84f1-3fa65fa18721"
/>
<img width="1644" height="1478"
alt="mxfp4_quantize_backend_comparison_swizzled_128x4_bfloat16"
src="https://github.com/user-attachments/assets/5586914b-f698-4fc1-880f-68892440637a"
/>
<img width="1644" height="1478"
alt="mxfp4_quantize_backend_comparison_swizzled_8x4_bfloat16"
src="https://github.com/user-attachments/assets/15734f58-3ba9-47be-9653-76eafb98b3b9"
/>


</details>
<details>
<summary>Backend performance comparisons from
bench_mxfp8_quantize_backend_comparison.py</summary>
<img width="1644" height="1477"
alt="mxfp8_backend_comparison_linear_bfloat16"
src="https://github.com/user-attachments/assets/9af95a9c-fd66-4e2d-8277-3c1e8c1f7883"
/>
<img width="1644" height="1478"
alt="mxfp8_backend_comparison_swizzled_128x4_bfloat16"
src="https://github.com/user-attachments/assets/f32233e3-9dd4-4a11-8a8f-f17d8fd73185"
/>
<img width="1644" height="1478"
alt="mxfp8_backend_comparison_swizzled_8x4_bfloat16"
src="https://github.com/user-attachments/assets/b00f6ec8-1ce7-4215-999a-36b4ecddd455"
/>



</details>

<!-- What does this PR do? Briefly describe the changes and why theyโ€™re
needed. -->

## ๐Ÿ” Related Issues

<!-- Link any related issues here -->

#3356

## ๐Ÿš€ 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

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] 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](https://pre-commit.com/).

## ๐Ÿงช Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->
B
Brian K. Ryu committed
f925953cec0cbad87fc185fd38d996477b98c715
Parent: fc9b2f6
Committed by GitHub <noreply@github.com> on 5/20/2026, 5:16:38 PM