perf(gemm): split-K, occupancy, and m=1 stream-GEMV decode tactics for mm_bf16_fp4 by yichengj0 · Pull Request #4038 · flashinfer-ai/flashinfer

yichengj0 · 2026-07-17T22:38:05Z

📌 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 feat: Add BF16_FP4 GEMM with cuDNN and CuTe-DSL backends for SM120/121 for W4A16 workloads #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

I have installed pre-commit by running pip install pre-commit (or used your preferred method).
I have installed the hooks with pre-commit install.
I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

Tests have been added or updated as needed.
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 feat(jit): JitSpec ABC + disk cache for JIT-compiled CuTe-DSL kernels #3874 CuTe-DSL disk cache, as perf(gemm): Improve mm_fp4 cute-dsl autotune time via disk-cache and parallel compilation #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.

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.

gemini-code-assist · 2026-07-17T22:38:08Z

Caution

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

coderabbitai · 2026-07-17T22:38:13Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

The BF16×FP4 Cute-DSL backend adds split-K GEMM and an SM12x m=1 GEMV kernel. Tactic selection, compilation, partial reduction, device-specific caching, and CUDA validation now cover these paths.

BF16-FP4 Cute-DSL execution

Layer / File(s)	Summary
Tactic selection and dispatch `flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py`	Adds split, occupancy, kernel-kind, device-specific cache, GEMV dispatch, fallback selection, and partial-buffer handling.
Dense GEMM split-K execution `flashinfer/gemm/kernels/cute_dsl/dense_gemm_bf16_fp4_blackwell.py`	Partitions K work, writes FP32 partials, and reduces them into the output.
SM12x GEMV kernel `flashinfer/gemm/kernels/cute_dsl/gemv_bf16_fp4_sm12x.py`	Adds streaming FP4 GEMV execution with direct output or deterministic split reduction.
Tactic and fallback validation `tests/gemm/test_mm_bf16_fp4.py`	Tests tactic accuracy, determinism, GEMV output types, fallback selection, and split heuristics.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant GEMMForward
  participant TacticSelector
  participant CuteDSLCompiler
  participant Kernel
  participant PartialReducer
  GEMMForward->>TacticSelector: select GEMM or GEMV tactic
  TacticSelector->>CuteDSLCompiler: compile split and occupancy configuration
  CuteDSLCompiler->>Kernel: launch with partial buffer
  Kernel->>PartialReducer: write FP32 partials when split
  PartialReducer->>GEMMForward: write reduced output

Possibly related PRs

flashinfer-ai/flashinfer#4266: Extends the same Blackwell CuTeDSL BF16 split-K GEMM backend.

Suggested reviewers: aleozlx, dhiraj113, yzh119

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 61.54% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main changes: split-K, occupancy, and stream-GEMV tactics for bf16×fp4 GEMM performance.
Description check	✅ Passed	The description includes the required sections, explains the changes and rationale, links the related issue, documents tests, and completes the checklist.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

Create stacked PR
Commit on current branch

🧪 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-21T23:53:45Z

Caution

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

Direct-load GEMV for bf16 x nvfp4 at tiny M: weights stream GMEM->registers (no SMEM staging, no tensor cores), one output row per warp, warp-shuffle reduction, schedule modeled on MsaProxyScoreDecodeStreamSm12x. Consumes the canonical (N, K/2) packed layout + linear E4M3 scales, so no weight repack. Dev-5080 dram-cold (box capped ~300 GB/s): m=1 parity with the autotuned MMA kernel on all three Qwen3.6-27B shapes (298/292/273 GB/s eff); decisive test is prod 5080 where the MMA kernel sits at 77-85% of the 828 GB/s roof. m=2/m=4 collapse (0.5x/0.22x) because A reloads scale with M on the serial path -- v1 is an m=1 specialist; A-in-SMEM staging is the known v2 lever. Not wired into dispatch yet; standalone cute.compile prototype. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Replace bare literals with constants derived from I32_PER_LANE / CODES_PER_I32 / SF_VEC_SIZE (the 64/2 in the scale-block indexing, the 4/8 fragment sizes) and comment the log2(warp) butterfly count. No codegen change: dev-5080 timing and rel_l2 identical. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

… layout Replace the canonical-layout prototype with a warp-cooperative walk of the backend's prepared operands (tile-packed int32 weight + S0E5M3 scales), so the GEMV and the MMA kernel share one weight format and dispatch per call. Each warp owns a 64-wide N tile and streams its 512B packed K-tiles fully coalesced; the pack's thread mapping fixes per lane which (k, n) each byte holds, leaving four f32 partials per lane plus one xor-1 butterfly per dot. fp4 and scale decode use the hw cvt.f16x2.e2m1x2 and f16 bit-place ops, dropping the prototype's power-of-two unfolds and its K % 1024 restriction. splits shards K across grid.y for grid fill, reusing the MMA kernel's fp32 workspace + fixed-order reduce for determinism; PDL chains match too. Dev-5080 dram-cold m=1 matches the canonical-layout prototype (294/294/271 GB/s eff on the Qwen shapes) that won the m=1 slot on prod 5080. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Tactic configs gain a kind field ('mma' | 'gemv'); gemv entries carry only a splits knob. get_valid_tactics gates them to the m=1 bucket (buckets round up, so no m > 1 call can land there) and to a 12-96 warps-per-SM residency window. forward dispatches gemv tactics to the new compile cache entry and raises on m != 1 as a guard; the fallback path is unchanged. The every-tactic test now skips gemv configs at m > 1 and covers all gemv splits (correctness + determinism) at the m=1 cell. Autotuner verified end to end on dev 5080: picks gemv-s32 on down and gemv-s1 on lm_head unaided (gate_up ties under warm-L2 at this box's DRAM cap; prod dram-cold data says gemv there too). AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- Reuse the shared f16x2_to_f32x2 from fp4_common instead of a local copy; drop the now-unused Tuple import. - Gate the gemv tactic to SM12x (compute capability major 12), the only arch it is validated on -- the autotuner ranks by time, so an untested arch must not be offered a kernel whose correctness is unchecked there. - Build the tactic-config list once per forward() call instead of two or three times in the gemv branch. - Pin the M == 1 invariant with a comment (M is symbolic in the traced __call__, so the runner's host-side m != 1 guard is the real check). - Cover fp16 output through the gemv path (test_cute_dsl_gemv_fp16_out); the every-tactic test only drove bf16. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The power-of-2 split menu quantizes wave fill (gate_up lands at 0.54 or 1.08 waves on 84 SMs, nothing in between); derive one extra split from the device SM count so tiles x splits lands at ~0.9-0.95 waves, leaving slack for stragglers and the PDL-overlapped tail of the prior kernel. The reduce loop becomes dynamic since fill splits reach ~100+. Measured on the DRAM-capped dev 5080: fill splits sit exactly on the pow2 baselines (s7 342.3us vs s4 341.0, s47 170.5 vs s32 170.7) while 1.08 waves costs +3%; the upside case needs a production-bandwidth part. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Serving stacks do not tune every shape (vLLM's warmup autotune never captures compute_logits, so lm_head always takes the fallback), and the tactic=-1 path only knew the MMA heuristic. Pick the gemv instead where it applies, at the smallest split reaching ~20 warps/SM: measured perf saturates there and deeper splits only add partial traffic (P6K lm_head s2 costs 2% over s1). The fill-split formula moves into a shared helper. Recovers +3.1% per-op on P6K lm_head under vLLM as-shipped. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The P6K probe showed throughput saturates at ~20 warps/SM and deeper splits pay partial traffic, so the knee split (down s47 on 188 SMs, 35.97us vs pow2-best s32 36.50us dram-cold) is the shape the menu was missing. Shared helper feeds both the menu and the m=1 fallback; the fallback docstring drops its duplicate rationale. Comment polish per dev-writing-style (no bench numbers in docstrings). AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- pin the tactic-index space to the device: the autotuner cache key now carries the SM count, since the config list holds device-derived gemv splits and a tuned index must not replay on a different part - keep do_preparation on the MMA heuristic: the autotuner issues it outside its per-tactic exception guard, so a gemv failure there would abort tuning for the whole op - cache the tactic-config builder (all-int args; ~7us -> 0.1us per forward call) and make sm_count required so no caller can silently enumerate a misaligned index space - share the 12-96 warps/SM residency window between the tactic filter and the fallback selector, and enforce the ceiling in the fallback - unroll the partial reduce for the common small split counts again; dynamic loop only for the ~100+ device-derived splits - tests: derive the fallback skip guard from the test's own shape, inline the one-use device helper AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Measured on three SKUs (dev 5080, RTX PRO 6000, GB10), the max-fill split never beats the knee split and loses 1-3% where bandwidth allows any spread, so its menu entry was tuner-rejected dead weight and one compile key per shape. The wave cap folds into the knee helper; the menu is now pow2 + knee. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

coderabbitai · 2026-08-05T00:59:32Z

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

coderabbitai

Actionable comments posted: 1

🤖 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 `@tests/gemm/test_mm_bf16_fp4.py`:
- Around line 305-309: The GEMV tests must skip tactics on devices without
compute capability 12. In tests/gemm/test_mm_bf16_fp4.py lines 305-309, update
the condition around _bf16_fp4_cute_dsl_tactic_configs to skip GEMV when m is
not 1 or get_compute_capability(device)[0] is not 12; at lines 350-357, add
pytest.skip before the tactic loop when the device major capability is not 12.

🪄 Autofix

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 Plus

Run ID: b5318e40-37a3-45af-bbb0-4d246155cb80

📥 Commits

Reviewing files that changed from the base of the PR and between 0961af9 and d074a6b.

📒 Files selected for processing (4)

flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py
flashinfer/gemm/kernels/cute_dsl/dense_gemm_bf16_fp4_blackwell.py
flashinfer/gemm/kernels/cute_dsl/gemv_bf16_fp4_sm12x.py
tests/gemm/test_mm_bf16_fp4.py

The tests enumerate the tactic list directly, bypassing get_valid_tactics' arch gate, so on other supported compute capabilities they would run the gemv where the production path never selects it. Addresses the CodeRabbit review finding. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Cut restated claims, cross-file narration, and rationale stated in more than one place (determinism, the occupancy trade, the m=1-only rule now each have one owner); comment-only. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

bkryu · 2026-08-12T23:30:01Z

@flashinfer-bot run

bkryu · 2026-08-12T23:30:12Z

/bot run tests/gemm

flashinfer-bot · 2026-08-12T23:31:45Z

GitLab MR !1119 has been updated with latest changes, and the CI pipeline #62425362 is currently running. I'll report back once the pipeline job completes.

…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.  ## 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.  --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Brian K. Ryu <bryu@nvidia.com>

flashinfer-bot added the op: gemm label Jul 17, 2026

yichengj0 force-pushed the w4a16-persistent-decode branch from 03e1f42 to 8d37c5c Compare July 20, 2026 17:58

yichengj0 marked this pull request as ready for review July 21, 2026 23:53

yichengj0 requested review from aleozlx, bkryu, cyx-6, dhiraj113, kahyunnam, saltyminty, sricketts, yongwww, yyihuang and yzh119 as code owners July 21, 2026 23:53

yichengj0 requested review from jimmyzho, kahyunnam, nv-yunzheq, qsang-nv, saltyminty, samuellees and yyihuang as code owners July 21, 2026 23:53

yichengj0 and others added 10 commits August 5, 2026 00:58

yichengj0 force-pushed the w4a16-persistent-decode branch from 5d22eb4 to d074a6b Compare August 5, 2026 00:59

yichengj0 changed the title ~~perf: split-K decode tactics for the mm_bf16_fp4 cute-dsl backend~~ perf(gemm): split-K, occupancy, and m=1 stream-GEMV decode tactics for mm_bf16_fp4 Aug 5, 2026

coderabbitai Bot reviewed Aug 5, 2026

View reviewed changes

Comment thread tests/gemm/test_mm_bf16_fp4.py

yichengj0 and others added 4 commits August 5, 2026 02:44

style: apply pinned ruff-format

3dc4a94

AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Merge branch 'main' into w4a16-persistent-decode

a8e9944

bkryu requested a review from StudyingShao as a code owner August 12, 2026 23:29

flashinfer-bot added run-ci and removed run-ci labels Aug 12, 2026

bkryu approved these changes Aug 13, 2026

View reviewed changes

bkryu enabled auto-merge (squash) August 13, 2026 03:53

bkryu merged commit 53a1c3b into flashinfer-ai:main Aug 13, 2026
27 of 28 checks passed

coderabbitai Bot mentioned this pull request Aug 15, 2026

feat: sm100 cute_dsl w4a16 gemm #4466

Merged

5 tasks

kahyunnam added the model: qwen3.5 / 3.6 / 3.8 label Sep 2, 2026

Navigation Menu

Uh oh!

perf(gemm): split-K, occupancy, and m=1 stream-GEMV decode tactics for mm_bf16_fp4 - #4038

perf(gemm): split-K, occupancy, and m=1 stream-GEMV decode tactics for mm_bf16_fp4#4038
bkryu merged 29 commits into
flashinfer-ai:mainfrom
yichengj0:w4a16-persistent-decode

yichengj0 commented Jul 17, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jul 17, 2026

Uh oh!

coderabbitai Bot commented Jul 17, 2026 •

edited

Loading

Reviews paused

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

gemini-code-assist Bot commented Jul 21, 2026

Uh oh!

coderabbitai Bot commented Aug 5, 2026

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

bkryu commented Aug 12, 2026

Uh oh!

bkryu commented Aug 12, 2026

Uh oh!

flashinfer-bot commented Aug 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

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

Uh oh!

📌 Description

Performance

Split-K on the #3597 decode shapes (RTX PRO 6000 / DGX Spark / RTX 5080)

Qwen3.6-27B-NVFP4 decode GEMMs at m=1 (RTX 5080, DGX Spark)

End-to-end vLLM serving of Qwen3.6-27B-NVFP4 (RTX PRO 6000)

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Summary by CodeRabbit

Uh oh!

gemini-code-assist Bot commented Jul 17, 2026

Uh oh!

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

Uh oh!

Reviews paused

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

gemini-code-assist Bot commented Jul 21, 2026

Uh oh!

coderabbitai Bot commented Aug 5, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bkryu commented Aug 12, 2026

Uh oh!

bkryu commented Aug 12, 2026

Uh oh!

flashinfer-bot commented Aug 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yichengj0 commented Jul 17, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 17, 2026 •

edited

Loading