feat: Add BF16_FP4 GEMM with cuDNN and CuTe-DSL backends for SM120/121 for W4A16 workloads by bkryu · Pull Request #3597 · flashinfer-ai/flashinfer

bkryu · 2026-06-11T23:17:51Z

📌 Description

Summary

This PR adds mm_bf16_fp4 -- A bf16xfp4 mixed-precision GEMM for W4A16 inference workloads. Supports two backends: cuDNN and CuTe-DSL kernel. Supported on SM100/103/110/120/121. cuDNN requires backend version 9.23.1 or later.

The computation is out = (a @ dequant(b, b_sf).T) * alpha, where

a is a bfloat16 tensor,
b is an nvfp4 tensor,
b_sf is the per-16-element FP8-E4M3 block scales, and
alpha is an optional fp32 global scale.

Kernel is optimized for SM121 (DGX Spark). SM120 requires further optimization.

CuTe-DSL kernel adopted from the CUTLASS bf16 dense GEMM example. Traces were auto-generated with AI

API

Weight-only quantized GEMMs need a backend-specific weight layout, so the API is a two-step flow:

prepare_bf16_fp4_weights --> For preprocessing nvfp4 weights (B operand)
- The preparation step is a performance optimization to accelerate the dequantization at GEMM runtime.
mm_bf16_fp4 → execute the bf16xfp4 GEMM.

import torch
import flashinfer

## Prepare inputs and quantize the B operand (weights)
M, N, K = 128, 4096, 4096
a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)  # activations
w = torch.randn(N, K, device="cuda", dtype=torch.bfloat16)  # weight, row-major (N, K)
g_w = (448 * 6) / w.float().abs().max()
b_fp4, b_sf = flashinfer.nvfp4_quantize( 
    w, g_w, sfLayout=flashinfer.SfLayout.layout_128x4, do_shuffle=False, backend="cute-dsl"
)
alpha = torch.tensor([1.0 / g_w.item()], device="cuda", dtype=torch.float32)

backend = "cute-dsl"
## Prepare weights at model load: repack canonical nvfp4 weights for a backend.
b_p, sf_p, alpha_p = flashinfer.prepare_bf16_fp4_weights(b_fp4, b_sf, alpha, backend=backend)

## At inference time, run with a different A operand for each batch of tokens.
out = flashinfer.mm_bf16_fp4(a, b_p, sf_p, alpha_p, backend=backend)

## Refcheck against the unquantized weight (error dominated by FP4 weight quantization):
ref = (a.float() @ w.float().T).to(torch.bfloat16)
rel_l2 = (out.float() - ref.float()).norm() / ref.float().norm()
cos = torch.nn.functional.cosine_similarity(
    out.float().flatten(), ref.float().flatten(), dim=0
)
assert rel_l2 < 0.15 and cos > 0.99

Testing

Adds a test_mm_bf16_fp4.py unit test script. Example command and output collected on DGX Spark

$ pytest tests/gemm/test_mm_bf16_fp4.py
================================================================= test session starts =================================================================  
platform linux -- Python 3.12.3, pytest-9.0.3, pluggy-1.6.0
rootdir: ...
configfile: pytest.ini
plugins: anyio-4.13.0
collected 234 items

tests/gemm/test_mm_bf16_fp4.py ............................................................................................................... [ 47%]
..................................................................................................................s........                     [100%]

===================================================== 233 passed, 1 skipped in 132.32s (0:02:12) ======================================================

The skipped test is:

tests/gemm/test_mm_bf16_fp4.py::test_backend_out_dtype_override[cudnn] PASSED                                                                  [ 96%]
tests/gemm/test_mm_bf16_fp4.py::test_backend_out_dtype_override[cute-dsl] SKIPPED (cute-dsl requires out_dtype == a.dtype)                     [ 96%]
where cuDNN allows different a and out dtypes (bf16 and fp16) while cute-dsl does not.

Benchmarks

Adds a mm_bf16_fp4 microbenchmark routine . Example command:

python benchmarks/flashinfer_benchmark.py --routine mm_bf16_fp4 --backends cute-dsl cudnn --m 1 --k 2688 --n 10304 --refcheck --no_cuda_graph --autotune

Benchmark results comparing backend='cute-dsl' with Marlin kernels called via vLLM on DGX Spark (SM121).

Benchmarks runnable via side branch with Marlin enabled in microbenchmarks.

Click to show reproducer microbenchmark commands

Note: The following commands used in the mm_fp4_w4a16_cudnn_and_dsl_with_marlin branch has a different routine name as it comes from a previous iteration.

python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 2048 --n 512 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 512 --n 2048 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 2048 --n 1024 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 1024 --n 2048 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 2048 --n 2048 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 4096 --n 512 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 512 --n 4096 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 4096 --n 1024 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 1024 --n 4096 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 4096 --n 2048 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 2048 --n 4096 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 4096 --n 4096 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 2048 --n 131072 --refcheck --no_cuda_graph
python benchmarks/flashinfer_benchmark.py --routine mm_w4a16_fp4 --backends cute-dsl marlin --m 1 --k 2048 --n 248320 --refcheck --no_cuda_graph

M	N	K	cute-dsl (ms)	cute-dsl TB/s	Marlin (ms)	Marlin TB/s	Speedup
1	2048	512	0.008	0.072	0.010	0.058	1.25×
1	4096	512	0.014	0.085	0.016	0.073	1.14×
1	2048	1024	0.013	0.088	0.017	0.072	1.31×
1	4096	1024	0.024	0.097	0.025	0.093	1.04×
1	512	2048	0.010	0.061	0.010	0.059	1.00×
1	1024	2048	0.013	0.090	0.016	0.076	1.23×
1	2048	2048	0.023	0.104	0.027	0.088	1.17×
1	4096	2048	0.045	0.105	0.048	0.099	1.07×
1	131072	2048	0.717	0.211	0.710	0.213	0.99×
1	248320	2048	1.199	0.239	1.227	0.234	1.02×
1	512	4096	0.017	0.070	0.016	0.075	0.94×
1	1024	4096	0.023	0.104	0.025	0.094	1.09×
1	2048	4096	0.041	0.115	0.046	0.102	1.12×
1	4096	4096	0.080	0.118	0.085	0.111	1.06×

cute-dsl matches or beats Marlin on 12 of 14 shapes (geomean ~1.10×, range 0.94×–1.31×); the two large-N cases and 512×4096 are within noise.

🔍 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

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.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

Tests have been added or updated as needed.
All tests are passing (unittest, etc.).

Reviewer Notes

Summary by CodeRabbit

New Features
- Added W4A16 (FP4-weight, 16-bit activation) support to the mm_fp4 GEMM API with backend dispatch and public weight-prep/export.
Performance / Kernels
- Introduced a new Blackwell-optimized dense GEMM kernel and enhanced Cute-DSL backend preparation/launch with caching and autotuning.
Benchmarks
- Added an mm_fp4_w4a16 benchmark routine; benchmark CSV output now fills all expected result columns with CLI-backed defaults.
Trace / Templates
- Added W4A16 FP4 trace templates and generated trace fixtures for both backends.
Tests
- Added cross-backend API/trace reference correctness coverage and reference-validation tests.

Drop the Marlin nvfp4 W4A16 adapter and its hooks in the mm_w4a16_fp4 benchmark routine before opening the PR. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

bkryu · 2026-06-16T16:50:48Z

/bot run

flashinfer-bot · 2026-06-16T16:51:44Z

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

dhiraj113 · 2026-06-16T19:47:09Z

+# backend gets its own template and ``mm_fp4``'s ``trace=`` is a dispatch
+# callable (same pattern as the trtllm MoE routing templates).
+
+# E2M1 nibble value table (low 3 bits magnitude, bit 3 sign).


Please remove AI slop.

Discussed this separately. Can change wording

dhiraj113 · 2026-06-16T21:38:50Z

+_CUTE_DSL_PACK_INTS_PER_TILE: int = 128  # int32s per (16K x 64N) repack block
+
+
+def _cute_dsl_pack_fp4_weight(b: torch.Tensor) -> torch.Tensor:


If possible, please split this cute_dsl and cudnn specific code into their own files/modules, so that only the gemm api, runners remain in this file, and the implementation detail of preparing cudnn graph, preparing arguments for cute dsl kernel ete are cleanly moved to their own modules/files.

bkryu · 2026-06-16T22:37:26Z

/bot run

flashinfer-bot · 2026-06-16T22:38:16Z

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

flashinfer-bot · 2026-06-17T20:52:18Z

[FAILED] Pipeline #55051715: 10/20 passed

…_fp4_weights (#3710) ## Summary - and were introduced in #3597 (BF16×FP4 GEMM with cuDNN and CuTe-DSL backends for SM120/121, W4A16 workloads) - Both are decorated with but were absent from , causing API-coverage alerts - Adds a new **BF16×FP4 GEMM (W4A16)** section in listing both functions under ## Changes - : new section inserted after *FP4 GEMM* with and ## Test plan - [ ] Verify picks up both symbols without error - [ ] Confirm API-coverage alert clears on next daily run  ## Summary by CodeRabbit * **Documentation** * Added a new API documentation subsection for **BF16×FP4 GEMM (W4A16)**, including generated reference pages for the weight preparation and matrix multiplication operations.  --------- Co-authored-by: cindyz <cindyz@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

…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.  ## 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>

…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>

…ht-only linear layers Wire FlashInfer's mm_bf16_fp4 (flashinfer-ai/flashinfer#3597, in FlashInfer v0.6.14, vLLM's current pin) into the NVFP4 linear kernel registry as FlashInferW4A16NvFp4LinearKernel. Like Marlin, the weight is repacked once at load time via prepare_bf16_fp4_weights; at run time the kernel dequantizes the FP4 weight and runs the matrix multiply in bf16, so activations are never quantized. The kernel is the default on SM121 (DGX Spark), where it is tuned, and opt-in via --linear-backend flashinfer_cutedsl on SM100/SM110/SM120. Marlin remains the default elsewhere and for fp16-activation models. The ModelOpt W4A16 method routes through init_nvfp4_linear_kernel(use_a16=True) so both it and the compressed-tensors W4A16 scheme share the selection. The GEMM is exposed as torch custom op vllm::flashinfer_mm_bf16_fp4 so it works under torch.compile and CUDA graph capture, and the existing FlashInfer autotune warmup tunes it during engine startup. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

bkryu and others added 18 commits May 27, 2026 14:53

First commit

1c1e7b2

some cleanup and more scaffolding

8082252

add_cudnn_w4a16

0ee01fc

Merge branch 'main' into mm_fp4_w4a16_cudnn_and_dsl

c124a7d

Add cute-dsl kernel

bec10ef

Improve cute dsl kernel

57f6b9c

Relax N % 128

8d90190

cute-dsl autotuning

b7afa24

Checkpoint

3079dc5

Finalize kernel. Now need cleanup

fc33c4e

Add missing file

2afdccd

Temporarily add Marlin benchmark

ee22fbf

Cleanup torch backend

8d00aa9

Merge branch 'main' into mm_fp4_w4a16_cudnn_and_dsl

5d8ff5e

Cleanup1

8e0c566

Cleanup2

d2964c2

Merge branch 'main' into mm_fp4_w4a16_cudnn_and_dsl

823fb46

Remove temporary vLLM Marlin comparison benchmark

95e1f0e

Drop the Marlin nvfp4 W4A16 adapter and its hooks in the mm_w4a16_fp4 benchmark routine before opening the PR. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

bkryu requested review from aleozlx, cyx-6, dhiraj113, jimmyzho, kahyunnam, nv-yunzheq, saltyminty, samuellees, sricketts, yongwww, yyihuang and yzh119 as code owners June 11, 2026 23:17

bkryu added 3 commits June 16, 2026 16:10

Address review comments. Refactor into mm_bf16_fp4

a1070f5

Cleanup

6dc1e98

Rename W4A16 to bf16_fp4

3e02472

bkryu changed the title ~~feat: Add W4A16 FP4 GEMM with cuDNN and CuTe-DSL backends for SM120/121~~ feat: Add BF16_FP4 GEMM with cuDNN and CuTe-DSL backends for SM120/121 for W4A16 workloads Jun 16, 2026