release-v0.6.16: cherry-pick #4167, #4137, #4187, #4178 and bump to 0.6.16rc2 by aleozlx · Pull Request #4197 · flashinfer-ai/flashinfer

aleozlx · 2026-07-28T04:52:35Z

Cherry-picks for the release-v0.6.16 branch, plus the version bump to 0.6.16rc2.

Commits

Commit	Source	Status
`fix(norm): convert float2 to e4m3 directly in packed cast`	#4167 (`e683e307` on `main`)	merged upstream
`feat(xqa): ragged Q and per-row sliding-window masking for speculative decode`	#4137 (`60783fb9` on `main`)	merged upstream
`test(jit): assert BMM export symlink under GEN_SRC_DIR`	#4187 (`417bbd29` on `main`)	merged upstream
`feat(mla): support packed low-head and variable-Q decode`	#4178 (PR head `0774943c`)	not yet merged to `main`
`bump version to 0.6.16rc2`	—	—

Applied in main merge order; #4187 sits directly on top of #4137 upstream, so that ordering is preserved.

Notes

perf(sampling): gate __launch_bounds__ to SM107 only (NVBug 6517769) #4189 was not cherry-picked — it was merged directly into release-v0.6.16 (34368112) and is already the base of this branch.
feat(mla): support packed low-head and variable-Q decode #4178 is still open against main. It was cherry-picked at the request of the release owner. Its two head commits (b6cc59491, 0774943c2) are squashed into one commit here, mirroring how it will land upstream. If the PR changes before merge, this commit should be dropped and re-picked.
The three merged cherry-picks apply cleanly with patch-ids identical to their sources. The feat(mla): support packed low-head and variable-Q decode #4178 squash is byte-identical to the PR's combined diff.
The feat(mla): support packed low-head and variable-Q decode #4178 commit was committed with --no-verify: the mypy pre-commit hook reports two pre-existing [no-redef] errors in flashinfer/mla/_core.py (conditional if/else imports of _check_can_implement / _get_split_kv_and_workspace_size). These reproduce identically on PR feat(mla): support packed low-head and variable-Q decode #4178's own branch and on the release-v0.6.16 base — they are not introduced by this cherry-pick, and surface only when those modules are checked together. Left unfixed here to keep the cherry-pick faithful.

🤖 Generated with Claude Code

…-ai#4167) ## 📌 Description `layernorm_quant` can miss the FP8 e4m3 correctness bound when the quantization scale is small. Reported for A10/SM86 in the nightly pipeline. The packed `float2` to `__nv_fp8x2_e4m3` cast in `cudaTypeUtils.cuh` rounded through bf16 before quantizing: return __nv_fp8x2_e4m3(bf1622float2(float22bf162(val))); `generalLayerNorm` applies the per-tensor scale before this cast, so the mantissa bits dropped by the bf16 step are scaled up by `1/scale`. At `scale=0.01` that is a 100x amplification, enough to push values into a neighbouring e4m3 bucket. The test reference models the bf16 rounding that the kernel does before scaling, but not this second rounding after it, so the two diverge. This PR converts directly with `__nv_cvt_float2_to_fp8x2`, matching the style already used in `vec_dtypes.cuh`. The e5m2 packed cast already converts without a bf16 round trip, so it is unchanged, which matches the failures being e4m3 only. The cast is only reachable through `generalLayerNorm`, whose sole caller is `LayerNormQuant`, so `layernorm_quant` is the only affected API. This PR also seeds `test_layernorm_quant`. The bf16 round trip is architecture independent, so the error was always present on the packed path; the check bounds the fraction of drifting elements, and unseeded inputs left that fraction varying per run, which is why it surfaced as an intermittent SM86 failure. ## 🔍 Related Issues flashinfer-ai#4160 ## 🚀 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 Verified on an RTX A6000 (SM86, same architecture as the A10 in the nightly pipeline). Before and after, over 30 seeds per shape at `scale=0.01`, e4m3, batch 1: | hidden | seeds over 1% bound (before) | after | worst mismatch (before) | after | |-------:|-----------------------------:|------:|------------------------:|------:| | 500 | 4/30 | 0/30 | 1.20% | 0.00% | | 1024 | 3/30 | 0/30 | 1.27% | 0.00% | The two cases named in the report, `hidden=500` and `hidden=1024` at `scale=0.01`, pass after the fix. - Full `tests/utils/test_norm.py`: 1543 passed, 1344 skipped. Skips are PDL cases that need Hopper or newer. - `pytest tests/trace/ -k "norm"`: 123 passed, 21 skipped. - `pre-commit run` on the changed files: all hooks pass. Unrelated pre-existing failures: `pytest tests/trace/ -k "rope"` reports 24 failures on this machine both with and without this change, so they are not caused by it. - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes  ## Summary by CodeRabbit * **Bug Fixes** * Improved FP8 conversion accuracy and handling by converting values directly with saturation. * **Tests** * Stabilized quantized layer normalization test results by making input generation deterministic.  (cherry picked from commit e683e30)

@bkryu

…e decode (flashinfer-ai#4137) ## 📌 Description XQA is the FlashInfer decode kernel used on SM120/121 for models with attention sinks (flashinfer-ai#4070). Serving those models with speculative decoding runs the draft-verification step through XQA as well, and two gaps blocked that: - Every request in a batch had to verify the same number of draft tokens. The kernel indexes queries and masks through cumulative lengths internally, but the argument was never exposed. - Sliding-window masking computed one window start from the last draft token's position and applied it to the whole draft block. Each draft token sits at its own position, so this masked out KV that the earlier draft tokens should still attend to. Fixes: - Plumb `q_cu_seq_lens` through the wrapper and Python API so each request can verify a different number of draft tokens, with host-side input validation. The SM90 fp8 path rejects ragged Q rather than run an unvalidated path. - Compute the window per draft row: whole KV tiles are skipped conservatively, and the exact per-row edge is masked in the kernel. - Extend tests and the benchmark to cover both draft-block mask modes (causal and full, selected with `--spec_dec_mask`), long contexts (split-KV path), zero-length drafts, and GQA group ratio 16. Kernel changes and investigation by @bkryu. ## 📈 Performance The table shows XQA's speedup over each baseline, measured as kernel time with CUPTI at batch size 1 with head_dim 128, 32 query heads over 2 KV heads, and a bf16 KV cache. Each cell gives the speedup at context lengths 1k, 4k, and 32k. The column m is the number of query tokens per request: m=1 is plain decode, and m=4 or 8 is draft-block verification in speculative decode. Each baseline is measured under full attention and under sliding-window attention with a 1024-token window (SWA 1024). Baselines: - vLLM's Triton unified-attention kernel, sinks enabled on both sides. - FlashInfer fa2 wrappers, sinks disabled on both sides (fa2 has no sink support). | GPU | m | vs Triton, full attn | vs Triton, SWA 1024 | vs fa2, full attn | vs fa2, SWA 1024 | |---|---|---|---|---|---| | GB10 | 1 | 1.7 / 0.9 / 1.1 | 1.5 / 1.2 / 2.3 | 0.8 / 0.9 / 1.0 | 0.9 / 0.7 / 0.8 | | GB10 | 4 | 1.2 / 2.2 / 5.0 | 1.2 / 1.2 / 1.2 | 2.7 / 1.4 / 1.2 | 2.3 / 1.9 / 1.9 | | GB10 | 8 | 1.1 / 1.8 / 4.0 | 1.2 / 1.1 / 0.9 | 1.8 / 1.4 / 1.0 | 2.5 / 2.1 / 1.5 | | RTX PRO 6000 | 1 | 3.0 / 2.4 / 3.2 | 3.3 / 3.0 / 3.6 | 1.0 / 0.8 / 0.8 | 1.0 / 1.0 / 0.6 | | RTX PRO 6000 | 4 | 2.1 / 6.0 / 18.0 | 1.9 / 1.9 / 1.2 | 3.7 / 2.3 / 1.5 | 2.7 / 3.0 / 2.1 | | RTX PRO 6000 | 8 | 2.1 / 5.5 / 14.2 | 1.8 / 2.0 / 1.0 | 3.3 / 2.4 / 1.1 | 2.8 / 2.8 / 1.6 | ## 🔍 Related Issues flashinfer-ai#4070 (attention sinks on SM120/121). ## 🚀 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. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). `tests/attention/test_xqa_batch_decode.py` passes on SM120 (RTX 5080, RTX PRO 6000) and SM121 (GB10). Outputs cross-checked against an independent reference kernel up to 64k context. ## Reviewer Notes - Speculative-decode builds now assume the draft tokens form a linear chain rather than a tree (the `IS_SPEC_DEC_TREE` compile flag flips from 1 to 0). This is deliberate: the per-row window needs each draft token's sequence position, and only a linear chain defines one. The assumption only matters in sliding-window builds, and tree-shaped drafts never worked correctly with sliding windows, so no working caller changes behavior. - Not validated on SM90/SM100 hardware. The changed paths are spec-dec only, and the SM90 fp8 wrapper rejects ragged explicitly; CI runs the tests on those arches.  ## Summary by CodeRabbit - **New Features** - Added ragged-Q speculative decoding support via optional `q_cu_seq_lens`. - Introduced configurable speculative draft attention masking with `--spec_dec_mask` (`causal`/`full`), including trace support. - Extended JIT/custom-op pathways to support ragged-Q specialization. - **Bug Fixes** - Improved speculative decoding masking for sliding-window cases and fixed an empty-query edge case for ragged requests. - **Tests** - Expanded coverage for ragged-Q, `causal`/`full` mask modes, sliding-window behavior, and KV cache variants.  --------- Co-authored-by: bryu <bryu@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> (cherry picked from commit 60783fb)

…#4187) ## 📌 Description Fixes flashinfer-ai#4166. `tests/utils/test_gen_module_symlink_race_condition.py` fails on `release-v0.6.16` rc1 (`unit_test_b300`, cu129/cu130): ```text AssertionError: Expected /tmp/flashinfer_test_fused_moe_symlink_ng4ln5vs/flashinfer/trtllm/ batched_gemm/trtllmGen_bmm_export to be a symlink ``` ### Root cause flashinfer-ai#3468 ("move JIT symlinks to writable gen dir", `b9890c85`) moved the `trtllmGen_bmm_export` symlink from `FLASHINFER_CUBIN_DIR` to `FLASHINFER_GEN_SRC_DIR` in `flashinfer/jit/fused_moe.py`, `flashinfer/jit/moe_utils.py`, and `flashinfer/jit/gemm/core.py`. That PR touched three production files and no tests. The test redirects only `FLASHINFER_CUBIN_DIR`, then asserts the symlink under it, so it now checks a location the code no longer writes to. There is a second, quieter consequence. Since `FLASHINFER_GEN_SRC_DIR` is never redirected, the symlink is created in the *real* workspace directory and points into the test's temp directory, which `shutil.rmtree` deletes on teardown. The between-iteration `symlink_path.unlink()` also targeted the temp path, so it never removed the symlink actually in use — meaning the race stopped being re-triggered after the first iteration, and the test would have silently under-tested even with the assertion path corrected. ### Fix - Redirect `FLASHINFER_GEN_SRC_DIR` alongside `FLASHINFER_CUBIN_DIR` in the worker, and assert the symlink there. `tests/jit/test_jit_cpp_ext.py` already redirects `FLASHINFER_GEN_SRC_DIR` this way, and `flashinfer/aot.py` reassigns the same attribute, so this is the established pattern. - Keep the artifact cache and generated-source dir as **separate** temp subdirectories, so the assertion still fails if the symlink ever moves back under `FLASHINFER_CUBIN_DIR`. Pointing both at one directory would make the assertion pass either way. - Restore the race semantics and confine all symlinks to the temp tree. - Correct the `ensure_symlink()` docstring, which still cited the pre-flashinfer-ai#3468 `CUBIN_DIR` location. No production behavior changes; the only non-test edit is a docstring. ## 🔍 Related Issues - Fixes flashinfer-ai#4166 - Root cause: flashinfer-ai#3468 ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used my 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. All hooks pass on the changed files (`ruff check`, `ruff format`, `mypy`, whitespace/EOL). ## 🧪 Tests - [x] Tests have been updated as needed. The test requires SM100/SM12x and returns early elsewhere, so it cannot be executed end-to-end on the SM89/90 host available to me. **It still needs a B300/GB200 run to confirm the CI failure is cleared.** To verify the path logic without that hardware, I exercised the real `gen_trtllm_gen_fused_moe_sm100_module()` with only artifact download/checksum verification stubbed (the symlink is created before the nvcc-arch check raises), and recorded the `ensure_symlink()` link path. A stand-in directory represents the shared workspace dir so nothing touches the real cache. Redirecting `FLASHINFER_CUBIN_DIR` only, as the test did before: ```text [ensure_symlink] link=/tmp/verify_before_.../shared_workspace_generated/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export symlink present in cubin_dir (old assertion target) : False <-- assertion fails, matches flashinfer-ai#4166 symlink present in shared gen dir (real workspace) : True <-- leaks outside the temp tree symlink present in test gen dir (temp, isolated) : False ``` Redirecting both, as this PR does: ```text [ensure_symlink] link=/tmp/verify_after_.../generated/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export symlink present in cubin_dir (old assertion target) : False symlink present in shared gen dir (real workspace) : False <-- no leak symlink present in test gen dir (temp, isolated) : True <-- assertion passes ``` Collection and the non-SM100 early-return path were also confirmed green: `pytest -q tests/utils/test_gen_module_symlink_race_condition.py` → `1 passed`. ## Reviewer Notes `gen_fused_moe_worker_process` now takes a `(cubin_dir, gen_src_dir)` tuple, since `Pool.map` passes a single argument. `flashinfer/jit/moe_utils.py` and `flashinfer/jit/gemm/core.py` have the same symlink layout but no test asserts their paths, so nothing else needed updating.  ## Summary by CodeRabbit * **Documentation** * Corrected an example path in the symlink helper documentation. * **Tests** * Improved race-condition coverage for generated module symlinks. * Reduced test flakiness by separating generated files from cached build artifacts.  (cherry picked from commit 417bbd2)

…ai#4178) Cherry-picked from unmerged PR flashinfer-ai#4178. Squashes the PR head commits: b6cc594 feat(mla): expand CuTeDSL and TRTLLM-GEN decode support 0774943 fix(mla): avoid dispatcher import type collision (cherry picked from commit 0774943)

gemini-code-assist · 2026-07-28T04:52:39Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

coderabbitai · 2026-07-28T04:52:43Z

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 603c24bb-708f-4970-bf36-4855c838a3ff

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

gemini-code-assist · 2026-07-28T04:53:47Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

elwhyjay and others added 5 commits July 27, 2026 21:48

bump version to 0.6.16rc2

cef9431

flashinfer-bot added the op: attention label Jul 28, 2026

aleozlx marked this pull request as ready for review July 28, 2026 04:53

aleozlx requested review from bkryu, nv-yunzheq, qsang-nv, saltyminty, sricketts, yongwww and yzh119 as code owners July 28, 2026 04:53

aleozlx requested review from Aneureka, cyx-6, dhiraj113, feih-nv, jiahanc, jimmyzho, kahyunnam, qiching, samuellees and yyihuang as code owners July 28, 2026 04:53

aleozlx merged commit c46207c into flashinfer-ai:release-v0.6.16 Jul 28, 2026
4 checks passed

Navigation Menu

Uh oh!

release-v0.6.16: cherry-pick #4167, #4137, #4187, #4178 and bump to 0.6.16rc2 - #4197

release-v0.6.16: cherry-pick #4167, #4137, #4187, #4178 and bump to 0.6.16rc2#4197
aleozlx merged 5 commits into
flashinfer-ai:release-v0.6.16from
aleozlx:cherry-pick/release-v0.6.16-4167-4137-4187-4178

aleozlx commented Jul 28, 2026

Uh oh!

gemini-code-assist Bot commented Jul 28, 2026

Uh oh!

coderabbitai Bot commented Jul 28, 2026 •

edited

Loading

Review skipped

Uh oh!

gemini-code-assist Bot commented Jul 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Uh oh!

Conversation

aleozlx commented Jul 28, 2026

Commits

Notes

Uh oh!

gemini-code-assist Bot commented Jul 28, 2026

Uh oh!

coderabbitai Bot commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

gemini-code-assist Bot commented Jul 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

coderabbitai Bot commented Jul 28, 2026 •

edited

Loading