perf: optimize per-token nvfp4 quantization kernel. (#3237)
<!-- .github/pull_request_template.md --> ## ๐ Description Optimize the performance of the per-token nvfp4 quantization kernel introduced by #3027. 1. default block size to 128. 2. default to fast math path. rename `TE_EXACT_FP4` to `TRTLLM_DISABLE_FP4_QUANT_FAST_MATH` and controlled by environmental variable. 3. change argument list of `get_sf_out_offset_128x4` and `get_sf_out_offset_8x4`. TODOs: 1. optimize low latency cases. ## ๐ Related Issues <!-- Link any related issues here --> ## ๐ 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. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added environment variable configuration to disable fast-math optimization in FP4 quantization, enabling behavior alignment with alternative implementations. * **Tests** * Added test fixture to validate FP4 quantization functionality with fast-math mode disabled. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Siyuan Fu <siyuanf@nvidia.com> Co-authored-by: Ziang Li <ziangli@umich.edu>
L
Lain committed
202af00dff85c1d5b50d15eab59d936365042af3
Parent: 2fe28ca
Committed by GitHub <noreply@github.com>
on 5/8/2026, 3:58:39 AM