perf(gemm): Improve mm_fp4 cute-dsl autotune time via disk-cache and parallel compilation - #4029
Merged
Merged
Conversation
bkryu
requested review from
aleozlx,
cyx-6,
dhiraj113,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 17, 2026 17:09
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughCuTe-DSL FP4 GEMM now uses shared compilation and cache helpers, deterministic disk names, and parallel tactic precompilation. The runner integrates these utilities, reuses the default alpha tensor, updates dtype/device handling, and tolerates precompilation failures. Cache-name coverage was added. ChangesCuTe-DSL FP4 GEMM caching
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/gemm/gemm_mm_fp4_cute_dsl.py`:
- Around line 78-115: Update the cache identity in the surrounding GEMM
compilation function so it includes the device-specific compile architecture and
max_active_clusters. Compute max_active_clusters before consulting the in-memory
cache, then construct and use the architecture-scoped key consistently for cache
lookup and storage, while preserving the existing disk artifact naming and
compilation flow.
- Around line 296-327: Update _available_host_memory_bytes to account for cgroup
memory limits before using /proc/meminfo: read cgroup v2 memory.max and
memory.current, or the equivalent v1 limit and usage files, and return the
remaining available bytes when a finite limit is present. Fall back to
MemAvailable when cgroup data is unavailable or unlimited, preserving
_get_mm_fp4_cute_dsl_compile_workers’ existing worker-cap behavior.
- Around line 228-289: Propagate the caller’s CUDA device through
precompile_mm_fp4_tactics by adding a.device.index to every worker payload. In
_mm_fp4_precompile_worker, select the payload device before constructing
JitSpecCuteDsl or invoking compilation, so _get_compile_arch() and
cute.compile() use the caller’s GPU.
In `@flashinfer/jit/cute_dsl_core.py`:
- Around line 208-215: Add an inline Ruff suppression with the specified
“persistence is best-effort” justification to the broad exception handler in
flashinfer/jit/cute_dsl_core.py at lines 208-215. Also add an inline BLE001
suppression with the “serial fallback is intentional” justification to the
corresponding broad exception handler in flashinfer/gemm/gemm_base.py at lines
5975-5993; no other behavior changes are needed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e7984e1c-c865-434c-8b90-1a8be101714f
📒 Files selected for processing (4)
flashinfer/gemm/gemm_base.pyflashinfer/gemm/gemm_mm_fp4_cute_dsl.pyflashinfer/jit/cute_dsl_core.pytests/jit/test_cute_dsl_cache.py
Comment on lines
+78
to
+115
| if cache_key in cache: | ||
| return cache[cache_key] | ||
|
|
||
| from flashinfer.cute_dsl.utils import get_max_active_clusters | ||
|
|
||
| gemm = make_gemm_kernel() | ||
|
|
||
| launch_cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] * cluster_shape_k | ||
| max_active_clusters = get_max_active_clusters(launch_cluster_size) | ||
|
|
||
| compile_kernel = _make_blockscaled_gemm_compile_fn( | ||
| gemm, | ||
| ab_cutlass_dtype=ab_cutlass_dtype, | ||
| sf_dtype=sf_dtype, | ||
| c_cutlass_dtype=c_cutlass_dtype, | ||
| ab_assumed_align=ab_assumed_align, | ||
| swap_ab=swap_ab, | ||
| sf_m=sf_m, | ||
| sf_n=sf_n, | ||
| sf_k=sf_k, | ||
| batch_size=batch_size, | ||
| max_active_clusters=max_active_clusters, | ||
| ) | ||
|
|
||
| if cache_module_name is None: | ||
| compiled_gemm = compile_kernel() | ||
| else: | ||
| from ..jit.cute_dsl_core import build_and_load_cute_dsl_kernel | ||
|
|
||
| compiled_gemm = build_and_load_cute_dsl_kernel( | ||
| cache_module_name, | ||
| _blockscaled_kernel_disk_name(cache_key, batch_size, max_active_clusters), | ||
| compile_kernel, | ||
| extra_key_files=_blockscaled_gemm_cache_key_files(), | ||
| ) | ||
|
|
||
| result = (compiled_gemm, max_active_clusters) | ||
| cache[cache_key] = result |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Include device-specific compilation state in the memory-cache identity.
The global cache is consulted before recomputing max_active_clusters, although both that value and the compilation architecture are device-dependent and embedded in the compiled kernel. Running the same tactic on heterogeneous GPUs can therefore reuse the first device's kernel instead of loading the architecture-scoped disk artifact.
Compute the device-dependent values first and include the compile architecture and max_active_clusters in the in-memory key.
Also applies to: 351-357
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/gemm/gemm_mm_fp4_cute_dsl.py` around lines 78 - 115, Update the
cache identity in the surrounding GEMM compilation function so it includes the
device-specific compile architecture and max_active_clusters. Compute
max_active_clusters before consulting the in-memory cache, then construct and
use the architecture-scoped key consistently for cache lookup and storage, while
preserving the existing disk artifact naming and compilation flow.
Collaborator
Author
There was a problem hiding this comment.
Agreed the in-memory key was missing device identity (a pre-existing gap on main)
Implemented with one deviation from the suggested mechanics in next commit. The in-memory key is now (device_index, cache_key).
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline flashinfer/gemm/gemm_mm_fp4_cute_dsl.py --items all --type function --match '_compile_block_scaled_gemm|_mm_fp4_cache_key|precompile_mm_fp4_tactics'
rg -n -A55 -B8 'def _compile_block_scaled_gemm\b|def _mm_fp4_cache_key\b|_compile_block_scaled_gemm\s*\(' flashinfer/gemm/gemm_mm_fp4_cute_dsl.pyLength of output: 5267
🏁 Script executed:
#!/bin/bash
set -euo pipefail
repo="flashinfer-ai/flashinfer"
pr="4029"
gh api "repos/${repo}/pulls/${pr}" --jq '{head_sha: .head.sha, head_ref: .head.ref, updated_at: .updated_at}'
gh api "repos/${repo}/pulls/${pr}/commits?per_page=10" \
--jq '.[] | {sha: .sha, message: (.commit.message | split("\n")[0])}'
head_sha="$(gh api "repos/${repo}/pulls/${pr}" --jq '.head.sha')"
gh api \
-H "Accept: application/vnd.github.raw+json" \
"repos/${repo}/contents/flashinfer/gemm/gemm_mm_fp4_cute_dsl.py?ref=${head_sha}" \
| nl -ba | sed -n '70,125p;345,415p'Length of output: 1266
@bkryu, (device_index, cache_key) is a valid minimal fix for the heterogeneous-device in-memory reuse scenario, provided it is used consistently for both lookup and insertion.
The current #4029 head is still 59490c96d033911683b8efcd5aefdcabd89b502c, which retains the original cache[cache_key] lookup. I’ll consider this addressed once the stated follow-up commit is pushed.
Collaborator
Author
|
/bot run tests/jit |
Collaborator
Collaborator
Author
|
/bot run tests/gemm |
Collaborator
qiching
requested changes
Jul 17, 2026
Collaborator
|
[FAILED] Pipeline #58480642: 14/20 passed |
dhiraj113
reviewed
Jul 20, 2026
Collaborator
Author
|
/bot run tests/gemm |
Collaborator
dhiraj113
reviewed
Jul 20, 2026
| enable_pdl, | ||
| out_dtype, | ||
| ) | ||
| cache_key = _mm_fp4_cache_key(sf_vec_size, tactic, enable_pdl, out_dtype) |
dhiraj113
reviewed
Jul 20, 2026
Collaborator
|
[FAILED] Pipeline #58812151: 10/20 passed |
dhiraj113
approved these changes
Jul 21, 2026
5 tasks
5 tasks
aleozlx
added a commit
that referenced
this pull request
Aug 4, 2026
Mirrors the v0.6.16 revert (#4246, commit 0a27ba3) so that 0.6.17 does not ship the API-breaking ring-buffer cache contract that 0.6.16 left out. Reverts: afd4754 mamba checkpointing SSU: two-kernel split + ring-buffer cache for checkpointing SSU (#3975) f90e9c4 docs(mamba): document checkpointing varlen arguments (#4129) Unlike release-v0.6.16, this branch also carries #4129, so both are reverted here; 0.6.16 only needed #3975. Verified no collateral damage: the SM107 changes from #4280 in tests/mamba/conftest.py and the #4029 changes in flashinfer/utils.py are preserved. NOTE: like the 0.6.16 revert, this also reverts the is_cvt_rs_supported correctness fix that rode along in #3975 (back to `major in (10, 11)`). See the release notes discussion -- that hunk is a candidate to keep. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
bkryu
added a commit
that referenced
this pull request
Aug 8, 2026
## 📌 Description Rolls the on-disk CuTe-DSL kernel cache out to the three sm12x NVFP4 fused-MoE kernels (micro, static, dynamic), addressing #4317. The infrastructure landed in #3874 / #4029; @bkryu noted in #4317 that it still had to reach individual kernels and that the core team lacked bandwidth, so this applies it to b12x MoE. Follows the rollout note in `docs/design_docs/cute_dsl_kernel_cache.md`: each `cute.compile` becomes a closure passed to `build_and_load_cute_dsl_kernel`. The three cache-key tuples become named functions so both cache levels derive from one source of truth, and the artifact name appends a digest of that tuple — the keys hold floats and `None` (`swiglu_*`), and `1.5` / `-1.5` both sanitize to `1_5`, so a formatted name would not be injective. As in the existing adopters, the kernels now compile against `make_fake_stream(use_tvm_ffi_env_stream=True)`, so TVM-FFI supplies the caller's current stream and the two launch sites no longer pass one (compiled signatures: 24 / 24 / 32 parameters). **Not covered:** the direct-micro kernel (`compile_direct_micro_kernel`), which this module started dispatching to recently. It compiles without `--enable-tvm-ffi` and launches through the DSL rather than a TVM-FFI callable, so caching it is a separate change — #4317 is only partly closed by this PR. ### Measured — GB10 (sm_121), DSL 4.6.0, five kernel shapes per process | process | total compile + load | |---|---:| | before, two runs | 17.92 s / 17.95 s | | after, cold first run | 18.91 s | | after, two later runs | **0.130 s / 0.134 s** (~135×) | Warm, per kernel: dynamic 7.7 s → 1 ms; static 3.2 s → 1 ms; micro 2.3–3.2 s → 1 ms (the first warm kernel pays 0.13 s of one-off module setup). The cold run costs a few percent for the export. Artifacts are 161–262 KB each. ## 🔍 Related Issues #4317 (partly — see the scope note above). Builds on the cache infrastructure from #3874 and #4029. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] `pre-commit` installed. - [x] Hooks installed. - [x] `pre-commit run` on the two changed files: all hooks pass, no files modified. (I ran it file-scoped rather than `--all-files`, since the remaining hooks are repo-wide.) ## 🧪 Tests - [x] Tests added. - [x] The new tests pass; see the caveat below for what does not run on my hardware. - New `tests/moe/test_b12x_moe_kernel_cache.py`: 61 naming-contract tests — signature coverage, per-argument perturbation, symbol safety, cross-family collision — replicating `tests/jit/test_cute_dsl_cache.py` as the design doc asks of new adopters (happy to fold them into that file instead if you would rather keep all adopters together). - Also checked on the same host: a corrupt artifact and a read-only cache directory each fall back to compiling with a warning, and editing a key source invalidates the module and recompiles it once. **Not verified — please check on sm120 hardware.** `tests/moe/test_b12x_fused_moe.py`'s numerical tests do not run on my GB10: the nvcc reference ops fail to build (`CUDA compiler and CUDA toolkit headers are incompatible`). The suite gives an identical `142 failed, 15 passed` — the same 142 test ids — on unmodified `main` and on this branch, so I have no accuracy signal either way. All timings above are compile/load time; steady-state kernel performance should be unchanged, since it is the same binary. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added shared on-disk caching for static, micro, and dynamic MoE kernels. * Improved cache invalidation when source dependencies change. * Kernel launches now automatically use the caller’s active CUDA stream. * **Tests** * Added comprehensive coverage for cache-key completeness, naming stability, symbol safety, and uniqueness across kernel types. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Han-Yin Chang <nick20350@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Brian K. Ryu <bryu@nvidia.com>
bkryu
added a commit
that referenced
this pull request
Aug 13, 2026
…r mm_bf16_fp4 (#4038) ## 📌 Description #3597 added `mm_bf16_fp4` (bf16 activations x nvfp4 weights) so vLLM can replace Marlin, its default W4A16 backend, with a FlashInfer kernel; DGX Spark is the lead target. On decode shapes the kernel trailed Marlin for two reasons: small-n grids underfill the GPU, and at single-token batches (m=1) too few resident warps per SM hide DRAM latency. This PR addresses both: m=1 decode beats or matches Marlin on every part we measured, and end-to-end serving of Qwen3.6-27B-NVFP4 flips from behind to ahead at batch 1 with no regression elsewhere. **What it does** - Adds split-K tactics (2/4/8 splits) to the autotuner space, offered only when splitting shortens the grid's last wave by at least 25%. Splits write fp32 partials, and a PDL-chained reduce kernel sums them in fixed order, so results are deterministic run to run. - Adds occupancy tactics: 2 or 3 co-resident CTAs per SM trade pipeline depth for latency hiding on weight-bound grids, plus occupancy 2 combined with split-K for narrow-n shapes. - Adds a streaming GEMV kernel (`gemv_bf16_fp4_sm12x.py`) for the bandwidth-bound m=1 case: no shared memory, no tensor cores, weights stream from global memory to registers with latency hidden by warp count. It reads the same packed operands as the MMA kernel, so the autotuner picks between the two per shape. - Sizes GEMV split-K from the device: alongside the power-of-2 splits, the menu carries a split targeting ~20 warps/SM (the measured saturation point). Tactic indices become device-scoped, so the autotuner cache key now carries the SM count. - Routes the no-autotune m=1 fallback onto the GEMV. Serving stacks do not tune every shape (vLLM's warmup never captures the logits GEMM), so the lm_head always takes this path. - Fixes kernel launch to pass no cluster dimensions: the boilerplate `cluster=[1,1,1]` routed launches through the cluster work distributor, whose co-residency cap silently defeated the occupancy tactics on SM12x. No public API changes. The one observable behavior change: untuned m=1 calls on SM12x now run the GEMV, whose output is bitwise different from the MMA heuristic's but equally accurate and still deterministic. ### Performance #### Split-K on the #3597 decode shapes (RTX PRO 6000 / DGX Spark / RTX 5080) Single-token decode GEMMs (m=1). The first table covers serving-class shapes (Llama-8B projection layers plus the #3597 example shape); the second covers #3597's own benchmark grid. Median GPU time over CUDA-graph replays with a cold L2 cache, as in serving. Baseline is vLLM's Marlin on the same GPU; FlashInfer runs with autotuning. Speedup = Marlin time / FlashInfer time. Each cell is a RTX PRO 6000 / DGX Spark / RTX 5080 triple. | n x k | vs Marlin, before this PR | vs Marlin, with this PR | |--:|:--:|:--:| | 2048x7168 | 0.57 / 1.00 / 0.47 | **1.36** / **1.02** / **0.78** | | 4096x4096 | 0.75 / 0.91 / 0.90 | **1.06** / **1.00** / 0.90 | | 4096x14336 | 0.60 / 0.91 / 0.78 | **0.98** / **1.00** / 0.78 | | 14336x4096 | 0.97 / 1.00 / 0.86 | 0.97 / 1.00 / 0.86 | | 10304x2688 | 0.78 / 0.98 / 0.96 | 0.78 / 0.98 / 0.96 | The same comparison over #3597's benchmark grid (4096x4096 appears in the table above): | n x k | vs Marlin, before this PR | vs Marlin, with this PR | |--:|:--:|:--:| | 512x2048 | 2.09 / 0.95 / 0.95 | **4.48** / **1.43** / **2.30** | | 512x4096 | 1.19 / 0.70 / 0.56 | **3.61** / **1.26** / **1.88** | | 1024x2048 | 1.26 / 0.98 / 0.70 | **2.39** / **1.03** / **1.32** | | 1024x4096 | 0.99 / 0.86 / 0.49 | **2.72** / **1.04** / **1.21** | | 2048x512 | 1.86 / 1.25 / 1.19 | 1.86 / 1.25 / 1.19 | | 2048x1024 | 1.36 / 1.16 / 0.90 | **1.40** / 1.16 / **1.02** | | 2048x2048 | 1.17 / 1.11 / 0.70 | **1.72** / **1.02**\* (1.11) / **0.97** | | 2048x4096 | 0.73 / 1.08 / 0.54 | **1.48** / **1.04**\* (1.08) / **0.84** | | 4096x512 | 1.40 / 1.01 / 1.38 | 1.40 / **0.95**\* (1.01) / 1.38 | | 4096x1024 | 1.41 / 0.93 / 1.18 | 1.41 / **0.95** / 1.18 | | 4096x2048 | 0.95 / 0.95 / 1.06 | 0.95 / **1.00** / 1.06 | | 131072x2048 | 0.98 / 0.97 / 0.90 | 0.98 / 0.97 / 0.90 | | 248320x2048 | 0.98 / 1.00 / 0.93 | 0.98 / 1.00 / 0.93 | - Bold marks the cells this PR changes (the tuner picks a new split-K tactic); unbolded picks perform as before. - \* These three cells are tuner mis-picks, not kernel regressions: an accurate pick would keep #3597's pre-existing non-split config, and the value in parentheses is what that config achieves. The Reviewer Notes explain the cause. - On the serving shapes, the RTX 5080 column stays below 1.0 even where this PR helps. Profiling of the larger losses points to activation re-reads through L2, a separate problem from grid fill and out of scope here. #### Qwen3.6-27B-NVFP4 decode GEMMs at m=1 (RTX 5080, DGX Spark) Same methodology as above; speedup = Marlin time / FlashInfer time. | GEMM (n x k) | RTX 5080 | DGX Spark | |--|:--:|:--:| | gate_up 34816x5120 | **1.03** | 1.00 | | down 5120x17408 | **1.03** | 1.00 | | lm_head 248320x5120 | **1.04**\* | **1.04** | \* Marlin's lm_head repack does not fit on the 16 GB RTX 5080, so this cell compares against the best in-tree MMA tactic. Spark ties at its bandwidth floor on the first two shapes. #### End-to-end vLLM serving of Qwen3.6-27B-NVFP4 (RTX PRO 6000) Full serving A/B, FlashInfer leg vs Marlin leg under identical settings, aiperf output-token throughput, ratio = FlashInfer / Marlin: batch-1 decode 1.011x, low-concurrency speculative decode 1.08 to 1.11x ahead, and every other cell at parity within run-to-run noise (0.978 to 1.014x) with no regression beyond it. ## 🔍 Related Issues Follow-up to #3597. ## 🚀 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.). New tests: - Every enumerated tactic (MMA and GEMV) is checked against a reference and for bit-exact run-to-run determinism. - Unit tests pin the fallback selectors' picks; one test drives tactic=-1 through the GEMV fallback end to end. - Full test file passes on RTX 5080, RTX PRO 6000, and GB10. ## Reviewer Notes - Most gains require autotuning, which serving frameworks run at startup. The no-autotune fallback picks match the tuner's choices on every part we measured. - The autotuner times candidates with a warm L2 while decode serving runs cold, so it can over-rank split tactics; the 25% last-wave guard compensates but does not fully close it (the three Spark cells in the grid table). This measurement gap is general and deserves its own issue. - The fallback picks add JIT-compiled kernel variants per decode shape class, cached in-process only; that cost amortizes to once per machine when this module adopts the #3874 CuTe-DSL disk cache, as #4029 did for the sibling `mm_fp4` path. The GEMV's device-derived splits widen this surface, so the follow-up is worth prioritizing. - Other FlashInfer cute-dsl kernels also pass `cluster=[1,1,1]` at launch and inherit the same co-residency cap; they are worth a separate audit. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added split-K support for bf16 × fp4 matrix multiplication to improve performance across varying workloads. - Added a dedicated SM12x GEMV path for efficient single-row operations. - Added automatic tuning for split counts, occupancy, and device-specific execution strategies. - Added support for FP16 GEMV outputs and deterministic partial-result reduction. - **Bug Fixes** - Improved handling of GEMV and split-K fallback selection across supported shapes and GPU configurations. - **Tests** - Added coverage for accuracy, determinism, GEMV correctness, and split-K selection. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Brian K. Ryu <bryu@nvidia.com>
11 tasks
jefby
pushed a commit
to jefby/flashinfer
that referenced
this pull request
Aug 19, 2026
…r mm_bf16_fp4 (flashinfer-ai#4038) ## 📌 Description flashinfer-ai#3597 added `mm_bf16_fp4` (bf16 activations x nvfp4 weights) so vLLM can replace Marlin, its default W4A16 backend, with a FlashInfer kernel; DGX Spark is the lead target. On decode shapes the kernel trailed Marlin for two reasons: small-n grids underfill the GPU, and at single-token batches (m=1) too few resident warps per SM hide DRAM latency. This PR addresses both: m=1 decode beats or matches Marlin on every part we measured, and end-to-end serving of Qwen3.6-27B-NVFP4 flips from behind to ahead at batch 1 with no regression elsewhere. **What it does** - Adds split-K tactics (2/4/8 splits) to the autotuner space, offered only when splitting shortens the grid's last wave by at least 25%. Splits write fp32 partials, and a PDL-chained reduce kernel sums them in fixed order, so results are deterministic run to run. - Adds occupancy tactics: 2 or 3 co-resident CTAs per SM trade pipeline depth for latency hiding on weight-bound grids, plus occupancy 2 combined with split-K for narrow-n shapes. - Adds a streaming GEMV kernel (`gemv_bf16_fp4_sm12x.py`) for the bandwidth-bound m=1 case: no shared memory, no tensor cores, weights stream from global memory to registers with latency hidden by warp count. It reads the same packed operands as the MMA kernel, so the autotuner picks between the two per shape. - Sizes GEMV split-K from the device: alongside the power-of-2 splits, the menu carries a split targeting ~20 warps/SM (the measured saturation point). Tactic indices become device-scoped, so the autotuner cache key now carries the SM count. - Routes the no-autotune m=1 fallback onto the GEMV. Serving stacks do not tune every shape (vLLM's warmup never captures the logits GEMM), so the lm_head always takes this path. - Fixes kernel launch to pass no cluster dimensions: the boilerplate `cluster=[1,1,1]` routed launches through the cluster work distributor, whose co-residency cap silently defeated the occupancy tactics on SM12x. No public API changes. The one observable behavior change: untuned m=1 calls on SM12x now run the GEMV, whose output is bitwise different from the MMA heuristic's but equally accurate and still deterministic. ### Performance #### Split-K on the flashinfer-ai#3597 decode shapes (RTX PRO 6000 / DGX Spark / RTX 5080) Single-token decode GEMMs (m=1). The first table covers serving-class shapes (Llama-8B projection layers plus the flashinfer-ai#3597 example shape); the second covers flashinfer-ai#3597's own benchmark grid. Median GPU time over CUDA-graph replays with a cold L2 cache, as in serving. Baseline is vLLM's Marlin on the same GPU; FlashInfer runs with autotuning. Speedup = Marlin time / FlashInfer time. Each cell is a RTX PRO 6000 / DGX Spark / RTX 5080 triple. | n x k | vs Marlin, before this PR | vs Marlin, with this PR | |--:|:--:|:--:| | 2048x7168 | 0.57 / 1.00 / 0.47 | **1.36** / **1.02** / **0.78** | | 4096x4096 | 0.75 / 0.91 / 0.90 | **1.06** / **1.00** / 0.90 | | 4096x14336 | 0.60 / 0.91 / 0.78 | **0.98** / **1.00** / 0.78 | | 14336x4096 | 0.97 / 1.00 / 0.86 | 0.97 / 1.00 / 0.86 | | 10304x2688 | 0.78 / 0.98 / 0.96 | 0.78 / 0.98 / 0.96 | The same comparison over flashinfer-ai#3597's benchmark grid (4096x4096 appears in the table above): | n x k | vs Marlin, before this PR | vs Marlin, with this PR | |--:|:--:|:--:| | 512x2048 | 2.09 / 0.95 / 0.95 | **4.48** / **1.43** / **2.30** | | 512x4096 | 1.19 / 0.70 / 0.56 | **3.61** / **1.26** / **1.88** | | 1024x2048 | 1.26 / 0.98 / 0.70 | **2.39** / **1.03** / **1.32** | | 1024x4096 | 0.99 / 0.86 / 0.49 | **2.72** / **1.04** / **1.21** | | 2048x512 | 1.86 / 1.25 / 1.19 | 1.86 / 1.25 / 1.19 | | 2048x1024 | 1.36 / 1.16 / 0.90 | **1.40** / 1.16 / **1.02** | | 2048x2048 | 1.17 / 1.11 / 0.70 | **1.72** / **1.02**\* (1.11) / **0.97** | | 2048x4096 | 0.73 / 1.08 / 0.54 | **1.48** / **1.04**\* (1.08) / **0.84** | | 4096x512 | 1.40 / 1.01 / 1.38 | 1.40 / **0.95**\* (1.01) / 1.38 | | 4096x1024 | 1.41 / 0.93 / 1.18 | 1.41 / **0.95** / 1.18 | | 4096x2048 | 0.95 / 0.95 / 1.06 | 0.95 / **1.00** / 1.06 | | 131072x2048 | 0.98 / 0.97 / 0.90 | 0.98 / 0.97 / 0.90 | | 248320x2048 | 0.98 / 1.00 / 0.93 | 0.98 / 1.00 / 0.93 | - Bold marks the cells this PR changes (the tuner picks a new split-K tactic); unbolded picks perform as before. - \* These three cells are tuner mis-picks, not kernel regressions: an accurate pick would keep flashinfer-ai#3597's pre-existing non-split config, and the value in parentheses is what that config achieves. The Reviewer Notes explain the cause. - On the serving shapes, the RTX 5080 column stays below 1.0 even where this PR helps. Profiling of the larger losses points to activation re-reads through L2, a separate problem from grid fill and out of scope here. #### Qwen3.6-27B-NVFP4 decode GEMMs at m=1 (RTX 5080, DGX Spark) Same methodology as above; speedup = Marlin time / FlashInfer time. | GEMM (n x k) | RTX 5080 | DGX Spark | |--|:--:|:--:| | gate_up 34816x5120 | **1.03** | 1.00 | | down 5120x17408 | **1.03** | 1.00 | | lm_head 248320x5120 | **1.04**\* | **1.04** | \* Marlin's lm_head repack does not fit on the 16 GB RTX 5080, so this cell compares against the best in-tree MMA tactic. Spark ties at its bandwidth floor on the first two shapes. #### End-to-end vLLM serving of Qwen3.6-27B-NVFP4 (RTX PRO 6000) Full serving A/B, FlashInfer leg vs Marlin leg under identical settings, aiperf output-token throughput, ratio = FlashInfer / Marlin: batch-1 decode 1.011x, low-concurrency speculative decode 1.08 to 1.11x ahead, and every other cell at parity within run-to-run noise (0.978 to 1.014x) with no regression beyond it. ## 🔍 Related Issues Follow-up to flashinfer-ai#3597. ## 🚀 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.). New tests: - Every enumerated tactic (MMA and GEMV) is checked against a reference and for bit-exact run-to-run determinism. - Unit tests pin the fallback selectors' picks; one test drives tactic=-1 through the GEMV fallback end to end. - Full test file passes on RTX 5080, RTX PRO 6000, and GB10. ## Reviewer Notes - Most gains require autotuning, which serving frameworks run at startup. The no-autotune fallback picks match the tuner's choices on every part we measured. - The autotuner times candidates with a warm L2 while decode serving runs cold, so it can over-rank split tactics; the 25% last-wave guard compensates but does not fully close it (the three Spark cells in the grid table). This measurement gap is general and deserves its own issue. - The fallback picks add JIT-compiled kernel variants per decode shape class, cached in-process only; that cost amortizes to once per machine when this module adopts the flashinfer-ai#3874 CuTe-DSL disk cache, as flashinfer-ai#4029 did for the sibling `mm_fp4` path. The GEMV's device-derived splits widen this surface, so the follow-up is worth prioritizing. - Other FlashInfer cute-dsl kernels also pass `cluster=[1,1,1]` at launch and inherit the same co-residency cap; they are worth a separate audit. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added split-K support for bf16 × fp4 matrix multiplication to improve performance across varying workloads. - Added a dedicated SM12x GEMV path for efficient single-row operations. - Added automatic tuning for split counts, occupancy, and device-specific execution strategies. - Added support for FP16 GEMV outputs and deterministic partial-result reduction. - **Bug Fixes** - Improved handling of GEMV and split-K fallback selection across supported shapes and GPU configurations. - **Tests** - Added coverage for accuracy, determinism, GEMV correctness, and split-K selection. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Brian K. Ryu <bryu@nvidia.com>
kahyunnam
added a commit
that referenced
this pull request
Aug 19, 2026
Mirrors the v0.6.16 revert (#4246, commit 0a27ba3) and its v0.6.17 counterpart (09fd5fc) so that 0.6.18 does not ship the API-breaking ring-buffer cache contract that neither 0.6.16 nor 0.6.17 shipped. Reverts: afd4754 mamba checkpointing SSU: two-kernel split + ring-buffer cache for checkpointing SSU (#3975) f90e9c4 docs(mamba): document checkpointing varlen arguments (#4129) Like release-v0.6.17, this branch carries #4129 as well, so both are reverted; 0.6.16 only needed #3975. main still carries both. Verified no collateral damage: the SM107 change from #4280 in tests/mamba/conftest.py and the #4029/#4078 changes in flashinfer/utils.py are preserved, and the five core reverted files now byte-match release-v0.6.17. NOTE: like both earlier reverts, this also reverts the cvt_rs fix that rode along in #3975 -- is_cvt_rs_supported goes back to `major in (10, 11)` (wrong for SM110a) and the CUDA guard back to SM100_ALL only (B300/sm_103a falls to software emulation). That matches what 0.6.17 shipped, but the hunk remains a candidate to keep.
5 tasks
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
4 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.
📌 Description
Follow-up to #3948 (top-N tactic ranking) and built on #3874 (CuTe-DSL disk cache). Autotuned
mm_fp4(backend='cute-dsl')still:This PR:
JITLinkcached artifacts from disk is measured to be done in ~10 ms per kernel.Autotune time improvements
Evaluated with
python3 benchmarks/flashinfer_benchmark.py --routine mm_fp4 --m 256 --n 1024 --k 7168 --out_dtype bfloat16 --backends cute-dsl --use_128x4_sf_layout --use_nvfp4 --refcheck --autotuneEnd-to-end process wall time on B200, 3 runs each:
Autotuned results are unchanged: the same tactics are profiled with bitwise identical kernel and output; just changes in the compilation infra.
🔍 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit