Conversation
lesj0610
changed the title
lesj0610
force-pushed
the
lesj/flashinfer-nvfp4-gemma-integration-20260626
branch
2 times, most recently
from
July 1, 2026 09:28
db5278e to
58530e6
Compare
…nfer-ai#3250) <!-- .github/pull_request_template.md --> ## 📌 Description <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 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 new block-sparse attention support for Blackwell GPUs, including a 64-block path and broader sparse pattern handling. * Added support for per-head block masks, variable block counts, and GQA/MQA layouts. * Added new sparse attention test coverage and reference utilities. * **Bug Fixes** * Improved validation and tensor shape handling for sparse attention runs. * Added more robust output and log-sum-exp handling for attention results. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: hsr <shao@nvidia.com> Co-authored-by: root <root@umbriel-b200-028.ipp4a1.colossus.nvidia.com> Co-authored-by: root <root@umbriel-b200-092.ipp4a1.colossus.nvidia.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: root <root@umbriel-b200-068.ipp4a1.colossus.nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: root <root@umbriel-b200-094.ipp4a1.colossus.nvidia.com> Co-authored-by: root <root@umbriel-b200-091.ipp4a1.colossus.nvidia.com> Co-authored-by: root <root@umbriel-b200-093.ipp4a1.colossus.nvidia.com> Co-authored-by: root <root@umbriel-b200-022.ipp4a1.colossus.nvidia.com> Co-authored-by: root <root@umbriel-b200-019.ipp4a1.colossus.nvidia.com>
## 📌 Description Add FlashInfer-native NVFP4 paged KV append APIs for runtimes that write paged KV cache entries directly. This PR: - Adds the base NVFP4 paged KV append API. - Adds a slot-mapping variant that interprets slot mappings as flat paged-cache locations and skips negative padding slots. - Exposes FFI and Python bindings that accept device scalar scale tensors. - Adds coverage for layouts, append tensor contiguity, slot index dtypes, boundary page sizes, and padding slots. ## 🔍 Related Issues N/A ## 🚀 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. - [ ] All tests are passing (`unittest`, etc.). Validation performed: - `pre-commit run --all-files` - `python -m pytest tests/attention/test_page.py::test_nvfp4_quantize_append_paged_kv_cache_with_slot_mapping -q -s -x --tb=short` - `python -m pytest tests/attention/test_page.py::test_nvfp4_quantize_append_paged_kv_cache -q -s -x --tb=short` - Direct packed-cache view smoke check on SM86 for NHD and HND layouts - `git diff --check` - `python -m py_compile flashinfer/page.py tests/attention/test_page.py flashinfer/__init__.py` Not yet verified: - Full CUDA/JIT regression matrix across all dtype, page-size, layout, and slot-mapping combinations. ## Reviewer Notes This PR is independent against `main`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for FP4-quantized paged KV-cache append operations, including a variant that works with slot mappings. * Exposed the new functionality in the public API and low-level interface. * **Bug Fixes** * Added stronger input validation for tensor types, shapes, layouts, and scale values. * Improved handling of empty inputs, invalid slots, and no-op cases so caches remain unchanged when expected. * **Tests** * Expanded GPU test coverage across layouts, dtypes, head sizes, slot-mapping cases, and scale validation. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
## 📌 Description Add a paged NVFP4 KV dequantization helper for cases where the caller already has paged NVFP4 K/V cache and needs contiguous K/V tensors again. This PR: - Adds `nvfp4_kv_dequantize_paged` for paged NVFP4 K/V cache input. - Gathers physical pages by caller-provided block tables and sequence lengths. - Writes dequantized K and V into caller-owned contiguous output tensors. - Supports tuple or stacked paged cache input, NHD and HND layouts, int32/int64 block tables, fp16/bf16 output, non-contiguous outer strides, and different K/V head dimensions. - Rejects block tables that cannot cover the requested output length before launching the CUDA kernel. - Keeps the existing flat `nvfp4_kv_dequantize` API and existing paged attention paths unchanged. The main point is simple: when K/V is stored in paged NVFP4 cache, caller should not need to manually gather pages and unpack E2M1 values on the Python side. This helper keeps that logic in FlashInfer. Scope note: the explicit FP4 scale-stride handling in this PR is applied to the paged prefill path. Ragged params carry the same fields for parameter/JIT layout consistency, but ragged prefill still uses the existing compact scale layout because there is no current caller passing strided NVFP4 scale views through the ragged path. Ragged strided-scale support can be handled separately if a caller needs it. ## 🔍 Related Issues N/A ## 🚀 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. - [ ] All tests are passing (`unittest`, etc.). Validation performed: - `pre-commit run --files csrc/fp4_kv_dequantization.cu tests/utils/test_fp4_kv_quantization.py` - `git diff --check` - `python -m py_compile flashinfer/quantization/fp4_quantization.py tests/utils/test_fp4_kv_quantization.py` - `FLASHINFER_DISABLE_VERSION_CHECK=1 FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-jit-test-paged-dequant-resume PYTHONPATH=. python -m pytest -q tests/utils/test_fp4_kv_quantization.py::test_nvfp4_kv_dequantize_paged tests/utils/test_fp4_kv_quantization.py::test_nvfp4_kv_dequantize_paged_rejects_short_block_tables` Not yet verified: - Full CUDA/JIT regression matrix across all page sizes, layouts, output dtypes, and larger batch shapes. - Performance comparison with a consumer-side gather/dequant implementation. ## Reviewer Notes This PR is independent against `main`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added paged NVFP4 KV-cache dequantization that writes dequantized K and V outputs. * Supports both NHD and HND cache layouts with per-page FP8 E4M3 scaling, including functional trace support. * Extended paged KV prefill to handle strided FP4 scale-factor tensors. * **Bug Fixes** * Improved validation for kv layout, block-table coverage, scalar scale tensors, and required head/dimension alignment. * **Tests** * Expanded coverage for layout/dtype combinations, non-contiguous and stacked inputs, long contexts, and rejection cases; added coverage for strided scale views. * **Chores** * Exposed the new paged dequantization API at the package top level and related modules. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
<!-- .github/pull_request_template.md --> ## 📌 Description Fix CuteDSL paged decode for `is_causal=False` when `seq_len` is not aligned to the 256-token sequence tile. The old path skipped the boundary mask loop for non-causal decode. Scores for KV lanes after `seq_len` stayed finite, so padding lanes entered the softmax denominator. This made short or unaligned sequences return scaled-down outputs. This PR: - masks non-causal boundary tiles with `key_pos >= seq_len`; - keeps the causal mask condition unchanged; - adds a paged non-causal decode regression test with `q_len_per_req > 1`. Validation summary on B200: - `q = 0` repro: before the fix, `ref/out` followed padding scale (`16.0` at `seq_len=16`, `4.0` at `seq_len=64`); after the fix, all checked lengths are `1.000`. - Random GQA + multi-query checks match fp32 reference and prefill within `2.7e-4` max absolute error. ## 🔍 Related Issues Fixes flashinfer-ai#3716 ## 🚀 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 non-causal paged attention decode with correct boundary handling when sequence lengths don’t align. * Expanded non-causal decoding to support variable query lengths per request. * **Bug Fixes** * Corrected non-causal masking so out-of-range KV lanes are excluded consistently (using proper negative-infinity masking) for split/boundary cases. * **Tests** * Added regression coverage for non-causal paged decode, including variable-length queries and split/boundary scenarios, validated against a reference implementation. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
SM120 measurements showed a small regression from the vectorized V-scale path while SM86 benefits from it. Keep the vector path on SM8x, where scalar scale loads were measured as the bottleneck, and leave SM90+ on the existing scalar path.
lesj0610
force-pushed
the
lesj/flashinfer-nvfp4-gemma-integration-20260626
branch
from
July 6, 2026 02:39
7f6e0ac to
cce3532
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Create a fork-only integration branch that combines the independent FlashInfer NVFP4 support branches so local testing can track upstream main without stacking those branches.
Changes
large-head-fa2-sm8.lesj/nvfp4-slot-mapping-kv-cache-20260624.lesj/flashinfer-workspace-sizing-helper-20260626.Validation
origin/mainat647c52dawithout conflicts.pre-commit run --all-fileson the integration branch.FLASHINFER_DISABLE_VERSION_CHECK=1 MAX_JOBS=4 FLASHINFER_NVCC_THREADS=2 python -m pytest tests/attention/test_workspace_size.py -q -s -x --tb=shorton the integration branch.Notes