[Kernel] add Flashinfer cutedsl w4a16 linear - #53014
Merged
vadiklyutiy merged 4 commits intoSep 1, 2026
Merged
Conversation
IwakuraRein
requested review from
AndreasKaratzas,
WoosukKwon,
mgoin,
pavanimajety,
robertgshaw2-redhat,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
August 19, 2026 22:38
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
IwakuraRein
requested review from
ProExpertProg,
houseroad and
youkaichao
as code owners
August 19, 2026 22:47
Contributor
|
Documentation preview: https://vllm--53014.org.readthedocs.build/en/53014/ |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
IwakuraRein
force-pushed
the
flashinfer-cutedsl-w4a16
branch
from
August 31, 2026 16:46
3999934 to
d91d18f
Compare
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
Member
|
/ci run |
vadiklyutiy
approved these changes
Aug 31, 2026
|
✅ Triggered Buildkite CI #86433 for commit |
vadiklyutiy
enabled auto-merge (squash)
August 31, 2026 22:17
github-actions
Bot
added
the
ready
vadiklyutiy
added
ready
|
✅ @IwakuraRein, CI is now available for this PR.
|
Contributor
Author
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #86433. |
am-cohere
pushed a commit
to am-cohere/vllm
that referenced
this pull request
Sep 1, 2026
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
puririshi98
added a commit
to puririshi98/vllm
that referenced
this pull request
Sep 1, 2026
…op (dequant + 16-bit GEMM, opt-in)
At prefill-sized M the weight-only Marlin kernel runs far below dense
16-bit rate (Nemotron-3.5-Lightning-30B-A3B-NVFP4 cache-miss profile:
137 ms GPU per 6576-token chunk, 96.5% GEMM, marlin-dominated at ~300
TFLOPS effective). Weight-only quantization only has to win while the
GEMM is memory-bound; at chunked-prefill M it is compute-bound and the
dequant tax is paid once per K x N instead of M times through the slow
kernel. When M >= threshold, apply_fp8_marlin_linear /
apply_fp4_marlin_linear dequantize the original-format weight into a
shared 16-bit workspace reserved at weight-loading time and run the
full-rate GEMM (torch.mm / F.linear); below the threshold every call
stays on Marlin, so decode and spec-verify launches (M <= 8, including
inside captured decode graphs) keep the memory-bound winner.
The M-vs-threshold branch lives behind a torch custom op
(vllm::marlin_large_m_gemm) because a Python
`reshaped_x.size(0) >= threshold` in the traced apply is decided ONCE
by dynamo (a BACKED symint guard vLLM drops with evaluate_guards=False)
and baked into every compiled artifact; inductor compile-range
endpoints cannot fix that (measured falsification: decode collapsed
1.3568 -> 3.1145 ms/tok with the dequant copy_mul kernels x23/x21 in
the decode window under every range config). As an opaque op the branch
is re-decided eagerly per call from the real x.size(0): register_fake
returns the single meta valid for BOTH branches (a fresh contiguous
[M, N] tensor in the activation dtype, enforced in the impl), both the
Marlin lock workspace and the shared dequant workspace are declared in
mutates_args (hidden mutation would let the compiler cache or reorder
around the op), and the threshold is an op argument wired to
VLLM_MARLIN_LARGE_M_BF16 (0 = off, default; 1 = 512; >= 16 sets it
directly, so 512 and 4096 are both reachable for the sweep; values in
[2, 16) are rejected to protect the decode guard). The branch is a pure
function of M, so a cudagraph captured at some padded M legally freezes
exactly the branch eager execution would take at that M -- FULL capture
at decode shapes stays legal -- and the small-M arm inside the op is
the same factored-out 16-bit Marlin call sequence
(_marlin_16bit_gemm) the non-dispatched apply path runs, by
construction.
Per-decode-call added overhead (the poison axis): with the flag OFF
(default) the only addition is one getattr per apply call,
constant-folded away at trace time -- the compiled artifact is
unchanged. With the flag ON, decode steps replayed from captured
cudagraphs pay nothing (the op's Python body runs at capture time
only); non-captured regions pay one dispatcher round-trip plus the
branch compare per dispatched dense layer call (order 1-2 us of host
work, x23 dense layers, i.e. sub-0.01 ms/tok class at AL 3.00) -- which
is exactly why the serve A/B's decode acceptance criterion is every-rep
ITL within +-0.01 ms/tok backed by a decode-window kernel diff, not an
assumption.
Admission and memory semantics are unchanged from the eager dispatch:
fail-closed at prepare time to dense FP8 non-blockwise with 16-bit
activations and dense NVFP4 g16 with a scalar global scale only; a
256 MiB per-layer cap excludes lm_head-class layers; MoE stays excluded
by measurement (honest per-call expert dequant loses 0.77x/0.29x at
miss/hit chunks). Enabling retains the original quantized tensor and
dequant-ready 16-bit scales of each dispatched layer, plus one shared
per-(device, dtype) workspace sized by the largest dispatched K x N,
grown only during weight loading through a stable holder object (layer
contexts keep a reference across boot-time growth) and never allocated
in the forward; if the reserved workspace cannot host a layer in the
activation dtype the op fails closed to Marlin. Numerics reproduce the
Marlin kernel's effective dequantization (NVFP4 processed-scale
clamp-to-zero, FP8 pre-exponent-fusion scales); accumulation order
differs, so outputs match at kernel-test tolerance, not bitwise.
Falsifiable predictions for the serve A/B: on the DSpark recipe (32K shared
prefix, C=1, AL exactly 3.00) with the flag on, steady cache-miss TTFT
lands in the 634 ms class measured in the first serve A/B (-260 ms vs
its same-allocation 894.55 ms baseline), rehit -44, hit -11.5; every-rep decode ITL
inside +-0.01 ms/tok of the same-allocation bracket with copy_mul_slice_view
x23/x21 ABSENT from the decode-window kernel diff and small-M marlin
launches PRESENT, on sm103 and sm100.
Tested (tests/kernels/quantization/test_marlin_gemm.py): both dispatch
arms against the dequantized reference and each other (FP8 both weight
layouts x bias x fp16/bf16; NVFP4 both dtypes), NVFP4 clamp semantics,
a marlin_gemm spy asserting M <= 8 always selects Marlin while
dispatched M never calls it, threshold floor and workspace-cap
fail-closed gates, custom-threshold boundary, 512/4096 sweep-point
boundary dispatch, capture-vs-eager bitwise equality at M in {1, 2, 4}
(fp8 + nvfp4, atomic-add reduce forced off), opcheck schema/fake on
both branches, and a compiled toy region (fullgraph, dynamic=True)
proving a single artifact serves both branches with error_on_recompile
armed -- the traced-branch failure mode reproduces as a hard test
failure. The dynamo mechanism check also passes CPU-side through the
real apply path with the kernel stubbed; the GPU suite runs in-container
on the measurement hardware.
Measured A/B (on main @ 85c1365; identical serve recipe and marlin
kernel build in both arms, max_num_batched_tokens 8770, AL 3.00
synthetic dspark, in-allocation brackets, one node per run): steady cache-miss TTFT 648.69 vs 911.21 ms
(-262.5) at threshold 512 and 651.73 (-259.5) at 4096 on GB300 sm103;
685.69 vs 917.34 (-231.7) at 4096 on GB200 sm100. Every-rep decode ITL within the +-0.01 ms/tok band at 4096
(worst +0.0014); at 512 every rep measured FASTER than the same-run
baseline (1.3976/1.3967 vs min 1.4101) -- 512 is the shipped default
when armed, 4096 the documented hit-neutral alternative. Decode-window
kernel diff: copy_mul_dequant 0, nvjet_splitK 0, dense marlin 7700 on
both thresholds and both arches; miss-window engagement
vllm::marlin_large_m_gemm x545 + fused triton call-site kernels x870 in
the traced 5-chunk miss window (plan [8768, 8768, 8768, 6576, 1952];
545 = 5 x 109 dispatch-eligible linears). Both engagement counts are
call-site counts and therefore threshold-invariant: the M-vs-threshold
branch lives inside the opaque op, which is invoked for every eligible
linear regardless of the branch taken. The per-threshold branch receipt
is the dense-marlin count in the same window -- 635 at threshold 512 vs
844 at 4096, every other kernel label identical: the plan's 1952-token
tail chunk is below 4096, stays on Marlin there (+209 marlin kernels),
and dispatches at 512, consistent with the 2192/1952 hit chunks
dispatching at 512 only. Cost: KV headroom
-1.71 GiB absolute (177.40 -> 175.69); hit TTFT at 512 +7.6 ms
standalone (dequant on the 2192/1952 hit chunks) -- composed with the
final-split elision the merged 4144-token hit chunk inverts this into
a win (combined candidate: hit 100.27 vs 136.86 in-allocation).
Engaged accuracy adjudication (run on the pre-rebase
base tree -- this patch's runtime (vllm/) hunks are patch-identical on
both bases (vllm/-restricted patch-id 1bb0141f); its test/bench files
later took a disclosed 9-line lint-only amend; gsm8k 5-shot
n=250 temp 0, real rejection sampling, APC off, every prompt through
the dispatched path): flex -0.0040 (+-0.0612), strict -0.0120 (+-0.0633), 109/250
byte-identical, flips symmetric, real acceptance length 3.5183 vs
3.5159.
Interaction with vllm-project#53014 (FlashInfer cute-dsl W4A16 NVFP4 default):
upstream now prefers FlashInferCuteDslNvFp4W4A16LinearKernel over
Marlin for W4A16_NVFP4 under --linear-backend=auto on SM100/103, so on
those parts the NVFP4 half of this patch no longer engages by default.
The measured recipe pins --linear-backend=marlin, whose backend-map
entry (MarlinNvFp4LinearKernel -> apply_fp4_marlin_linear, where the
large-M context attaches during Marlin weight processing) is untouched
by vllm-project#53014; under that pin -- and everywhere Marlin remains the
selected kernel -- the dispatch engages exactly as measured. The FP8
channelwise half (marlin_utils_fp8), which carries most of the win, is
unaffected by vllm-project#53014 entirely.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Rishi Puri <riship@nvidia.com>
puririshi98
added a commit
to puririshi98/vllm
that referenced
this pull request
Sep 1, 2026
…op (dequant + 16-bit GEMM, opt-in)
At prefill-sized M the weight-only Marlin kernel runs far below dense
16-bit rate (Nemotron-3.5-Lightning-30B-A3B-NVFP4 cache-miss profile:
137 ms GPU per 6576-token chunk, 96.5% GEMM, marlin-dominated at ~300
TFLOPS effective). Weight-only quantization only has to win while the
GEMM is memory-bound; at chunked-prefill M it is compute-bound and the
dequant tax is paid once per K x N instead of M times through the slow
kernel. When M >= threshold, apply_fp8_marlin_linear /
apply_fp4_marlin_linear dequantize the original-format weight into a
shared 16-bit workspace reserved at weight-loading time and run the
full-rate GEMM (torch.mm / F.linear); below the threshold every call
stays on Marlin, so decode and spec-verify launches (M <= 8, including
inside captured decode graphs) keep the memory-bound winner.
The M-vs-threshold branch lives behind a torch custom op
(vllm::marlin_large_m_gemm) because a Python
`reshaped_x.size(0) >= threshold` in the traced apply is decided ONCE
by dynamo (a BACKED symint guard vLLM drops with evaluate_guards=False)
and baked into every compiled artifact; inductor compile-range
endpoints cannot fix that (measured falsification: decode collapsed
1.3568 -> 3.1145 ms/tok with the dequant copy_mul kernels x23/x21 in
the decode window under every range config). As an opaque op the branch
is re-decided eagerly per call from the real x.size(0): register_fake
returns the single meta valid for BOTH branches (a fresh contiguous
[M, N] tensor in the activation dtype, enforced in the impl), both the
Marlin lock workspace and the shared dequant workspace are declared in
mutates_args (hidden mutation would let the compiler cache or reorder
around the op), and the threshold is an op argument wired to
VLLM_MARLIN_LARGE_M_BF16 (0 = off, default; 1 = 512; >= 16 sets it
directly, so 512 and 4096 are both reachable for the sweep; values in
[2, 16) are rejected to protect the decode guard). The branch is a pure
function of M, so a cudagraph captured at some padded M legally freezes
exactly the branch eager execution would take at that M -- FULL capture
at decode shapes stays legal -- and the small-M arm inside the op is
the same factored-out 16-bit Marlin call sequence
(_marlin_16bit_gemm) the non-dispatched apply path runs, by
construction.
Per-decode-call added overhead (the poison axis): with the flag OFF
(default) the only addition is one getattr per apply call,
constant-folded away at trace time -- the compiled artifact is
unchanged. With the flag ON, decode steps replayed from captured
cudagraphs pay nothing (the op's Python body runs at capture time
only); non-captured regions pay one dispatcher round-trip plus the
branch compare per dispatched dense layer call (order 1-2 us of host
work, x23 dense layers, i.e. sub-0.01 ms/tok class at AL 3.00) -- which
is exactly why the serve A/B's decode acceptance criterion is every-rep
ITL within +-0.01 ms/tok backed by a decode-window kernel diff, not an
assumption.
Admission and memory semantics are unchanged from the eager dispatch:
fail-closed at prepare time to dense FP8 non-blockwise with 16-bit
activations and dense NVFP4 g16 with a scalar global scale only; a
256 MiB per-layer cap excludes lm_head-class layers; MoE stays excluded
by measurement (honest per-call expert dequant loses 0.77x/0.29x at
miss/hit chunks). Enabling retains the original quantized tensor and
dequant-ready 16-bit scales of each dispatched layer, plus one shared
per-(device, dtype) workspace sized by the largest dispatched K x N,
grown only during weight loading through a stable holder object (layer
contexts keep a reference across boot-time growth) and never allocated
in the forward; if the reserved workspace cannot host a layer in the
activation dtype the op fails closed to Marlin. Numerics reproduce the
Marlin kernel's effective dequantization (NVFP4 processed-scale
clamp-to-zero, FP8 pre-exponent-fusion scales); accumulation order
differs, so outputs match at kernel-test tolerance, not bitwise.
Falsifiable predictions for the serve A/B: on the DSpark recipe (32K shared
prefix, C=1, AL exactly 3.00) with the flag on, steady cache-miss TTFT
lands in the 634 ms class measured in the first serve A/B (-260 ms vs
its same-allocation 894.55 ms baseline), rehit -44, hit -11.5; every-rep decode ITL
inside +-0.01 ms/tok of the same-allocation bracket with copy_mul_slice_view
x23/x21 ABSENT from the decode-window kernel diff and small-M marlin
launches PRESENT, on sm103 and sm100.
Tested (tests/kernels/quantization/test_marlin_gemm.py): both dispatch
arms against the dequantized reference and each other (FP8 both weight
layouts x bias x fp16/bf16; NVFP4 both dtypes), NVFP4 clamp semantics,
a marlin_gemm spy asserting M <= 8 always selects Marlin while
dispatched M never calls it, threshold floor and workspace-cap
fail-closed gates, custom-threshold boundary, 512/4096 sweep-point
boundary dispatch, capture-vs-eager bitwise equality at M in {1, 2, 4}
(fp8 + nvfp4, atomic-add reduce forced off), opcheck schema/fake on
both branches, and a compiled toy region (fullgraph, dynamic=True)
proving a single artifact serves both branches with error_on_recompile
armed -- the traced-branch failure mode reproduces as a hard test
failure. The dynamo mechanism check also passes CPU-side through the
real apply path with the kernel stubbed; the GPU suite runs in-container
on the measurement hardware.
Measured A/B (on main @ 85c1365; identical serve recipe and marlin
kernel build in both arms, max_num_batched_tokens 8770, AL 3.00
synthetic dspark, in-allocation brackets, one node per run): steady cache-miss TTFT 648.69 vs 911.21 ms
(-262.5) at threshold 512 and 651.73 (-259.5) at 4096 on GB300 sm103;
685.69 vs 917.34 (-231.7) at 4096 on GB200 sm100. Every-rep decode ITL within the +-0.01 ms/tok band at 4096
(worst +0.0014); at 512 every rep measured FASTER than the same-run
baseline (1.3976/1.3967 vs min 1.4101) -- 512 is the shipped default
when armed, 4096 the documented hit-neutral alternative. Decode-window
kernel diff: copy_mul_dequant 0, nvjet_splitK 0, dense marlin 7700 on
both thresholds and both arches; miss-window engagement
vllm::marlin_large_m_gemm x545 + fused triton call-site kernels x870 in
the traced 5-chunk miss window (plan [8768, 8768, 8768, 6576, 1952];
545 = 5 x 109 dispatch-eligible linears). Both engagement counts are
call-site counts and therefore threshold-invariant: the M-vs-threshold
branch lives inside the opaque op, which is invoked for every eligible
linear regardless of the branch taken. The per-threshold branch receipt
is the dense-marlin count in the same window -- 635 at threshold 512 vs
844 at 4096, every other kernel label identical: the plan's 1952-token
tail chunk is below 4096, stays on Marlin there (+209 marlin kernels),
and dispatches at 512, consistent with the 2192/1952 hit chunks
dispatching at 512 only. Cost: KV headroom
-1.71 GiB absolute (177.40 -> 175.69); hit TTFT at 512 +7.6 ms
standalone (dequant on the 2192/1952 hit chunks) -- composed with the
final-split elision the merged 4144-token hit chunk inverts this into
a win (combined candidate: hit 100.27 vs 136.86 in-allocation).
Engaged accuracy adjudication (run on the pre-rebase
base tree -- this patch's runtime (vllm/) hunks are patch-identical on
both bases (vllm/-restricted patch-id 1bb0141f); its test/bench files
later took a disclosed 9-line lint-only amend; gsm8k 5-shot
n=250 temp 0, real rejection sampling, APC off, every prompt through
the dispatched path): flex -0.0040 (+-0.0612), strict -0.0120 (+-0.0633), 109/250
byte-identical, flips symmetric, real acceptance length 3.5183 vs
3.5159.
Interaction with vllm-project#53014 (FlashInfer cute-dsl W4A16 NVFP4 default):
upstream now prefers FlashInferCuteDslNvFp4W4A16LinearKernel over
Marlin for W4A16_NVFP4 under --linear-backend=auto on SM100/103, so on
those parts the NVFP4 half of this patch no longer engages by default.
The measured recipe pins --linear-backend=marlin, whose backend-map
entry (MarlinNvFp4LinearKernel -> apply_fp4_marlin_linear, where the
large-M context attaches during Marlin weight processing) is untouched
by vllm-project#53014; under that pin -- and everywhere Marlin remains the
selected kernel -- the dispatch engages exactly as measured. The FP8
channelwise half (marlin_utils_fp8), which carries most of the win, is
unaffected by vllm-project#53014 entirely.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Rishi Puri <riship@nvidia.com>
puririshi98
added a commit
to puririshi98/vllm
that referenced
this pull request
Sep 1, 2026
…op (dequant + 16-bit GEMM, opt-in)
At prefill-sized M the weight-only Marlin kernel runs far below dense
16-bit rate (Nemotron-3.5-Lightning-30B-A3B-NVFP4 cache-miss profile:
137 ms GPU per 6576-token chunk, 96.5% GEMM, marlin-dominated at ~300
TFLOPS effective). Weight-only quantization only has to win while the
GEMM is memory-bound; at chunked-prefill M it is compute-bound and the
dequant tax is paid once per K x N instead of M times through the slow
kernel. When M >= threshold, apply_fp8_marlin_linear /
apply_fp4_marlin_linear dequantize the original-format weight into a
shared 16-bit workspace reserved at weight-loading time and run the
full-rate GEMM (torch.mm / F.linear); below the threshold every call
stays on Marlin, so decode and spec-verify launches (M <= 8, including
inside captured decode graphs) keep the memory-bound winner.
The M-vs-threshold branch lives behind a torch custom op
(vllm::marlin_large_m_gemm) because a Python
`reshaped_x.size(0) >= threshold` in the traced apply is decided ONCE
by dynamo (a BACKED symint guard vLLM drops with evaluate_guards=False)
and baked into every compiled artifact; inductor compile-range
endpoints cannot fix that (measured falsification: decode collapsed
1.3568 -> 3.1145 ms/tok with the dequant copy_mul kernels x23/x21 in
the decode window under every range config). As an opaque op the branch
is re-decided eagerly per call from the real x.size(0): register_fake
returns the single meta valid for BOTH branches (a fresh contiguous
[M, N] tensor in the activation dtype, enforced in the impl), both the
Marlin lock workspace and the shared dequant workspace are declared in
mutates_args (hidden mutation would let the compiler cache or reorder
around the op), and the threshold is an op argument wired to
VLLM_MARLIN_LARGE_M_BF16 (0 = off, default; 1 = 512; >= 16 sets it
directly, so 512 and 4096 are both reachable for the sweep; values in
[2, 16) are rejected to protect the decode guard). The branch is a pure
function of M, so a cudagraph captured at some padded M legally freezes
exactly the branch eager execution would take at that M -- FULL capture
at decode shapes stays legal -- and the small-M arm inside the op is
the same factored-out 16-bit Marlin call sequence
(_marlin_16bit_gemm) the non-dispatched apply path runs, by
construction.
Per-decode-call added overhead (the poison axis): with the flag OFF
(default) the only addition is one getattr per apply call,
constant-folded away at trace time -- the compiled artifact is
unchanged. With the flag ON, decode steps replayed from captured
cudagraphs pay nothing (the op's Python body runs at capture time
only); non-captured regions pay one dispatcher round-trip plus the
branch compare per dispatched dense layer call (order 1-2 us of host
work, x23 dense layers, i.e. sub-0.01 ms/tok class at AL 3.00) -- which
is exactly why the serve A/B's decode acceptance criterion is every-rep
ITL within +-0.01 ms/tok backed by a decode-window kernel diff, not an
assumption.
Admission and memory semantics are unchanged from the eager dispatch:
fail-closed at prepare time to dense FP8 non-blockwise with 16-bit
activations and dense NVFP4 g16 with a scalar global scale only; a
256 MiB per-layer cap excludes lm_head-class layers; MoE stays excluded
by measurement (honest per-call expert dequant loses 0.77x/0.29x at
miss/hit chunks). Enabling retains the original quantized tensor and
dequant-ready 16-bit scales of each dispatched layer, plus one shared
per-(device, dtype) workspace sized by the largest dispatched K x N,
grown only during weight loading through a stable holder object (layer
contexts keep a reference across boot-time growth) and never allocated
in the forward; if the reserved workspace cannot host a layer in the
activation dtype the op fails closed to Marlin. Numerics reproduce the
Marlin kernel's effective dequantization (NVFP4 processed-scale
clamp-to-zero, FP8 pre-exponent-fusion scales); accumulation order
differs, so outputs match at kernel-test tolerance, not bitwise.
Falsifiable predictions for the serve A/B: on the DSpark recipe (32K shared
prefix, C=1, AL exactly 3.00) with the flag on, steady cache-miss TTFT
lands in the 634 ms class measured in the first serve A/B (-260 ms vs
its same-allocation 894.55 ms baseline), rehit -44, hit -11.5; every-rep decode ITL
inside +-0.01 ms/tok of the same-allocation bracket with copy_mul_slice_view
x23/x21 ABSENT from the decode-window kernel diff and small-M marlin
launches PRESENT, on sm103 and sm100.
Tested (tests/kernels/quantization/test_marlin_gemm.py): both dispatch
arms against the dequantized reference and each other (FP8 both weight
layouts x bias x fp16/bf16; NVFP4 both dtypes), NVFP4 clamp semantics,
a marlin_gemm spy asserting M <= 8 always selects Marlin while
dispatched M never calls it, threshold floor and workspace-cap
fail-closed gates, custom-threshold boundary, 512/4096 sweep-point
boundary dispatch, capture-vs-eager bitwise equality at M in {1, 2, 4}
(fp8 + nvfp4, atomic-add reduce forced off), opcheck schema/fake on
both branches, and a compiled toy region (fullgraph, dynamic=True)
proving a single artifact serves both branches with error_on_recompile
armed -- the traced-branch failure mode reproduces as a hard test
failure. The dynamo mechanism check also passes CPU-side through the
real apply path with the kernel stubbed; the GPU suite runs in-container
on the measurement hardware.
Measured A/B (on main @ 85c1365; identical serve recipe and marlin
kernel build in both arms, max_num_batched_tokens 8770, AL 3.00
synthetic dspark, in-allocation brackets, one node per run): steady cache-miss TTFT 648.69 vs 911.21 ms
(-262.5) at threshold 512 and 651.73 (-259.5) at 4096 on GB300 sm103;
685.69 vs 917.34 (-231.7) at 4096 on GB200 sm100. Every-rep decode ITL within the +-0.01 ms/tok band at 4096
(worst +0.0014); at 512 every rep measured FASTER than the same-run
baseline (1.3976/1.3967 vs min 1.4101) -- 512 is the shipped default
when armed, 4096 the documented hit-neutral alternative. Decode-window
kernel diff: copy_mul_dequant 0, nvjet_splitK 0, dense marlin 7700 on
both thresholds and both arches; miss-window engagement
vllm::marlin_large_m_gemm x545 + fused triton call-site kernels x870 in
the traced 5-chunk miss window (plan [8768, 8768, 8768, 6576, 1952];
545 = 5 x 109 dispatch-eligible linears). Both engagement counts are
call-site counts and therefore threshold-invariant: the M-vs-threshold
branch lives inside the opaque op, which is invoked for every eligible
linear regardless of the branch taken. The per-threshold branch receipt
is the dense-marlin count in the same window -- 635 at threshold 512 vs
844 at 4096, every other kernel label identical: the plan's 1952-token
tail chunk is below 4096, stays on Marlin there (+209 marlin kernels),
and dispatches at 512, consistent with the 2192/1952 hit chunks
dispatching at 512 only. Cost: KV headroom
-1.71 GiB absolute (177.40 -> 175.69); hit TTFT at 512 +7.6 ms
standalone (dequant on the 2192/1952 hit chunks) -- composed with the
final-split elision the merged 4144-token hit chunk inverts this into
a win (combined candidate: hit 100.27 vs 136.86 in-allocation).
Engaged accuracy adjudication (run on the pre-rebase
base tree -- this patch's runtime (vllm/) hunks are patch-identical on
both bases (vllm/-restricted patch-id 1bb0141f); its test/bench files
later took a disclosed 9-line lint-only amend; gsm8k 5-shot
n=250 temp 0, real rejection sampling, APC off, every prompt through
the dispatched path): flex -0.0040 (+-0.0612), strict -0.0120 (+-0.0633), 109/250
byte-identical, flips symmetric, real acceptance length 3.5183 vs
3.5159.
Interaction with vllm-project#53014 (FlashInfer cute-dsl W4A16 NVFP4 default):
upstream now prefers FlashInferCuteDslNvFp4W4A16LinearKernel over
Marlin for W4A16_NVFP4 under --linear-backend=auto on SM100/103, so on
those parts the NVFP4 half of this patch no longer engages by default.
The measured recipe pins --linear-backend=marlin, whose backend-map
entry (MarlinNvFp4LinearKernel -> apply_fp4_marlin_linear, where the
large-M context attaches during Marlin weight processing) is untouched
by vllm-project#53014; under that pin -- and everywhere Marlin remains the
selected kernel -- the dispatch engages exactly as measured. The FP8
channelwise half (marlin_utils_fp8), which carries most of the win, is
unaffected by vllm-project#53014 entirely.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Rishi Puri <riship@nvidia.com>
puririshi98
added a commit
to puririshi98/vllm
that referenced
this pull request
Sep 2, 2026
…op (dequant + 16-bit GEMM, opt-in)
At prefill-sized M the weight-only Marlin kernel runs far below dense
16-bit rate (Nemotron-3.5-Lightning-30B-A3B-NVFP4 cache-miss profile:
137 ms GPU per 6576-token chunk, 96.5% GEMM, marlin-dominated at ~300
TFLOPS effective). Weight-only quantization only has to win while the
GEMM is memory-bound; at chunked-prefill M it is compute-bound and the
dequant tax is paid once per K x N instead of M times through the slow
kernel. When M >= threshold, apply_fp8_marlin_linear /
apply_fp4_marlin_linear dequantize the original-format weight into a
shared 16-bit workspace reserved at weight-loading time and run the
full-rate GEMM (torch.mm / F.linear); below the threshold every call
stays on Marlin, so decode and spec-verify launches (M <= 8, including
inside captured decode graphs) keep the memory-bound winner.
The M-vs-threshold branch lives behind a torch custom op
(vllm::marlin_large_m_gemm) because a Python
`reshaped_x.size(0) >= threshold` in the traced apply is decided ONCE
by dynamo (a BACKED symint guard vLLM drops with evaluate_guards=False)
and baked into every compiled artifact; inductor compile-range
endpoints cannot fix that (measured falsification: decode collapsed
1.3568 -> 3.1145 ms/tok with the dequant copy_mul kernels x23/x21 in
the decode window under every range config). As an opaque op the branch
is re-decided eagerly per call from the real x.size(0): register_fake
returns the single meta valid for BOTH branches (a fresh contiguous
[M, N] tensor in the activation dtype, enforced in the impl), both the
Marlin lock workspace and the shared dequant workspace are declared in
mutates_args (hidden mutation would let the compiler cache or reorder
around the op), and the threshold is an op argument wired to
VLLM_MARLIN_LARGE_M_BF16 (0 = off, default; 1 = 512; >= 16 sets it
directly, so 512 and 4096 are both reachable for the sweep; values in
[2, 16) are rejected to protect the decode guard). The branch is a pure
function of M, so a cudagraph captured at some padded M legally freezes
exactly the branch eager execution would take at that M -- FULL capture
at decode shapes stays legal -- and the small-M arm inside the op is
the same factored-out 16-bit Marlin call sequence
(_marlin_16bit_gemm) the non-dispatched apply path runs, by
construction.
Per-decode-call added overhead (the poison axis): with the flag OFF
(default) the only addition is one getattr per apply call,
constant-folded away at trace time -- the compiled artifact is
unchanged. With the flag ON, decode steps replayed from captured
cudagraphs pay nothing (the op's Python body runs at capture time
only); non-captured regions pay one dispatcher round-trip plus the
branch compare per dispatched dense layer call (order 1-2 us of host
work, x23 dense layers, i.e. sub-0.01 ms/tok class at AL 3.00) -- which
is exactly why the serve A/B's decode acceptance criterion is every-rep
ITL within +-0.01 ms/tok backed by a decode-window kernel diff, not an
assumption.
Admission and memory semantics are unchanged from the eager dispatch:
fail-closed at prepare time to dense FP8 non-blockwise with 16-bit
activations and dense NVFP4 g16 with a scalar global scale only; a
256 MiB per-layer cap excludes lm_head-class layers; MoE stays excluded
by measurement (honest per-call expert dequant loses 0.77x/0.29x at
miss/hit chunks). Enabling retains the original quantized tensor and
dequant-ready 16-bit scales of each dispatched layer, plus one shared
per-(device, dtype) workspace sized by the largest dispatched K x N,
grown only during weight loading through a stable holder object (layer
contexts keep a reference across boot-time growth) and never allocated
in the forward; if the reserved workspace cannot host a layer in the
activation dtype the op fails closed to Marlin. Numerics reproduce the
Marlin kernel's effective dequantization (NVFP4 processed-scale
clamp-to-zero, FP8 pre-exponent-fusion scales); accumulation order
differs, so outputs match at kernel-test tolerance, not bitwise.
Falsifiable predictions for the serve A/B: on the DSpark recipe (32K shared
prefix, C=1, AL exactly 3.00) with the flag on, steady cache-miss TTFT
lands in the 634 ms class measured in the first serve A/B (-260 ms vs
its same-allocation 894.55 ms baseline), rehit -44, hit -11.5; every-rep decode ITL
inside +-0.01 ms/tok of the same-allocation bracket with copy_mul_slice_view
x23/x21 ABSENT from the decode-window kernel diff and small-M marlin
launches PRESENT, on sm103 and sm100.
Tested (tests/kernels/quantization/test_marlin_gemm.py): both dispatch
arms against the dequantized reference and each other (FP8 both weight
layouts x bias x fp16/bf16; NVFP4 both dtypes), NVFP4 clamp semantics,
a marlin_gemm spy asserting M <= 8 always selects Marlin while
dispatched M never calls it, threshold floor and workspace-cap
fail-closed gates, custom-threshold boundary, 512/4096 sweep-point
boundary dispatch, capture-vs-eager bitwise equality at M in {1, 2, 4}
(fp8 + nvfp4, atomic-add reduce forced off), opcheck schema/fake on
both branches, and a compiled toy region (fullgraph, dynamic=True)
proving a single artifact serves both branches with error_on_recompile
armed -- the traced-branch failure mode reproduces as a hard test
failure. The dynamo mechanism check also passes CPU-side through the
real apply path with the kernel stubbed; the GPU suite runs in-container
on the measurement hardware.
Measured A/B (on main @ 85c1365; identical serve recipe and marlin
kernel build in both arms, max_num_batched_tokens 8770, AL 3.00
synthetic dspark, in-allocation brackets, one node per run): steady cache-miss TTFT 648.69 vs 911.21 ms
(-262.5) at threshold 512 and 651.73 (-259.5) at 4096 on GB300 sm103;
685.69 vs 917.34 (-231.7) at 4096 on GB200 sm100. Every-rep decode ITL within the +-0.01 ms/tok band at 4096
(worst +0.0014); at 512 every rep measured FASTER than the same-run
baseline (1.3976/1.3967 vs min 1.4101) -- 512 is the shipped default
when armed, 4096 the documented hit-neutral alternative. Decode-window
kernel diff: copy_mul_dequant 0, nvjet_splitK 0, dense marlin 7700 on
both thresholds and both arches; miss-window engagement
vllm::marlin_large_m_gemm x545 + fused triton call-site kernels x870 in
the traced 5-chunk miss window (plan [8768, 8768, 8768, 6576, 1952];
545 = 5 x 109 dispatch-eligible linears). Both engagement counts are
call-site counts and therefore threshold-invariant: the M-vs-threshold
branch lives inside the opaque op, which is invoked for every eligible
linear regardless of the branch taken. The per-threshold branch receipt
is the dense-marlin count in the same window -- 635 at threshold 512 vs
844 at 4096, every other kernel label identical: the plan's 1952-token
tail chunk is below 4096, stays on Marlin there (+209 marlin kernels),
and dispatches at 512, consistent with the 2192/1952 hit chunks
dispatching at 512 only. Cost: KV headroom
-1.71 GiB absolute (177.40 -> 175.69); hit TTFT at 512 +7.6 ms
standalone (dequant on the 2192/1952 hit chunks) -- composed with the
final-split elision the merged 4144-token hit chunk inverts this into
a win (combined candidate: hit 100.27 vs 136.86 in-allocation).
Engaged accuracy adjudication (run on the pre-rebase
base tree -- this patch's runtime (vllm/) hunks are patch-identical on
both bases (vllm/-restricted patch-id 1bb0141f); its test/bench files
later took a disclosed 9-line lint-only amend; gsm8k 5-shot
n=250 temp 0, real rejection sampling, APC off, every prompt through
the dispatched path): flex -0.0040 (+-0.0612), strict -0.0120 (+-0.0633), 109/250
byte-identical, flips symmetric, real acceptance length 3.5183 vs
3.5159.
Interaction with vllm-project#53014 (FlashInfer cute-dsl W4A16 NVFP4 default):
upstream now prefers FlashInferCuteDslNvFp4W4A16LinearKernel over
Marlin for W4A16_NVFP4 under --linear-backend=auto on SM100/103, so on
those parts the NVFP4 half of this patch no longer engages by default.
The measured recipe pins --linear-backend=marlin, whose backend-map
entry (MarlinNvFp4LinearKernel -> apply_fp4_marlin_linear, where the
large-M context attaches during Marlin weight processing) is untouched
by vllm-project#53014; under that pin -- and everywhere Marlin remains the
selected kernel -- the dispatch engages exactly as measured. The FP8
channelwise half (marlin_utils_fp8), which carries most of the win, is
unaffected by vllm-project#53014 entirely.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Rishi Puri <riship@nvidia.com>
mylibrar
pushed a commit
to tanyuqian/vllm
that referenced
this pull request
Sep 3, 2026
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
4 tasks
sheralskumar
pushed a commit
to sheralskumar/vllm
that referenced
this pull request
Sep 8, 2026
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com>
D-G-Dimitrov
pushed a commit
to D-G-Dimitrov/vllm
that referenced
this pull request
Sep 8, 2026
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com> (cherry picked from commit 882ca8d)
puririshi98
added a commit
to puririshi98/vllm
that referenced
this pull request
Sep 9, 2026
…op (dequant + 16-bit GEMM, opt-in)
At prefill-sized M the weight-only Marlin kernel runs far below dense
16-bit rate (Nemotron-3.5-Lightning-30B-A3B-NVFP4 cache-miss profile:
137 ms GPU per 6576-token chunk, 96.5% GEMM, marlin-dominated at ~300
TFLOPS effective). Weight-only quantization only has to win while the
GEMM is memory-bound; at chunked-prefill M it is compute-bound and the
dequant tax is paid once per K x N instead of M times through the slow
kernel. When M >= threshold, apply_fp8_marlin_linear /
apply_fp4_marlin_linear dequantize the original-format weight into a
shared 16-bit workspace reserved at weight-loading time and run the
full-rate GEMM (torch.mm / F.linear); below the threshold every call
stays on Marlin, so decode and spec-verify launches (M <= 8, including
inside captured decode graphs) keep the memory-bound winner.
The M-vs-threshold branch lives behind a torch custom op
(vllm::marlin_large_m_gemm) because a Python
`reshaped_x.size(0) >= threshold` in the traced apply is decided ONCE
by dynamo (a BACKED symint guard vLLM drops with evaluate_guards=False)
and baked into every compiled artifact; inductor compile-range
endpoints cannot fix that (measured falsification: decode collapsed
1.3568 -> 3.1145 ms/tok with the dequant copy_mul kernels x23/x21 in
the decode window under every range config). As an opaque op the branch
is re-decided eagerly per call from the real x.size(0): register_fake
returns the single meta valid for BOTH branches (a fresh contiguous
[M, N] tensor in the activation dtype, enforced in the impl), both the
Marlin lock workspace and the shared dequant workspace are declared in
mutates_args (hidden mutation would let the compiler cache or reorder
around the op), and the threshold is an op argument wired to
VLLM_MARLIN_LARGE_M_BF16 (0 = off, default; 1 = 512; >= 16 sets it
directly, so 512 and 4096 are both reachable for the sweep; values in
[2, 16) are rejected to protect the decode guard). The branch is a pure
function of M, so a cudagraph captured at some padded M legally freezes
exactly the branch eager execution would take at that M -- FULL capture
at decode shapes stays legal -- and the small-M arm inside the op is
the same factored-out 16-bit Marlin call sequence
(_marlin_16bit_gemm) the non-dispatched apply path runs, by
construction.
Per-decode-call added overhead (the poison axis): with the flag OFF
(default) the only addition is one getattr per apply call,
constant-folded away at trace time -- the compiled artifact is
unchanged. With the flag ON, decode steps replayed from captured
cudagraphs pay nothing (the op's Python body runs at capture time
only); non-captured regions pay one dispatcher round-trip plus the
branch compare per dispatched dense layer call (order 1-2 us of host
work, x23 dense layers, i.e. sub-0.01 ms/tok class at AL 3.00) -- which
is exactly why the serve A/B's decode acceptance criterion is every-rep
ITL within +-0.01 ms/tok backed by a decode-window kernel diff, not an
assumption.
Admission and memory semantics are unchanged from the eager dispatch:
fail-closed at prepare time to dense FP8 non-blockwise with 16-bit
activations and dense NVFP4 g16 with a scalar global scale only; a
256 MiB per-layer cap excludes lm_head-class layers; MoE stays excluded
by measurement (honest per-call expert dequant loses 0.77x/0.29x at
miss/hit chunks). Enabling retains the original quantized tensor and
dequant-ready 16-bit scales of each dispatched layer, plus one shared
per-(device, dtype) workspace sized by the largest dispatched K x N,
grown only during weight loading through a stable holder object (layer
contexts keep a reference across boot-time growth) and never allocated
in the forward; if the reserved workspace cannot host a layer in the
activation dtype the op fails closed to Marlin. Numerics reproduce the
Marlin kernel's effective dequantization (NVFP4 processed-scale
clamp-to-zero, FP8 pre-exponent-fusion scales); accumulation order
differs, so outputs match at kernel-test tolerance, not bitwise.
Falsifiable predictions for the serve A/B: on the DSpark recipe (32K shared
prefix, C=1, AL exactly 3.00) with the flag on, steady cache-miss TTFT
lands in the 634 ms class measured in the first serve A/B (-260 ms vs
its same-allocation 894.55 ms baseline), rehit -44, hit -11.5; every-rep decode ITL
inside +-0.01 ms/tok of the same-allocation bracket with copy_mul_slice_view
x23/x21 ABSENT from the decode-window kernel diff and small-M marlin
launches PRESENT, on sm103 and sm100.
Tested (tests/kernels/quantization/test_marlin_gemm.py): both dispatch
arms against the dequantized reference and each other (FP8 both weight
layouts x bias x fp16/bf16; NVFP4 both dtypes), NVFP4 clamp semantics,
a marlin_gemm spy asserting M <= 8 always selects Marlin while
dispatched M never calls it, threshold floor and workspace-cap
fail-closed gates, custom-threshold boundary, 512/4096 sweep-point
boundary dispatch, capture-vs-eager bitwise equality at M in {1, 2, 4}
(fp8 + nvfp4, atomic-add reduce forced off), opcheck schema/fake on
both branches, and a compiled toy region (fullgraph, dynamic=True)
proving a single artifact serves both branches with error_on_recompile
armed -- the traced-branch failure mode reproduces as a hard test
failure. The dynamo mechanism check also passes CPU-side through the
real apply path with the kernel stubbed; the GPU suite runs in-container
on the measurement hardware.
Measured A/B (on main @ 85c1365; identical serve recipe and marlin
kernel build in both arms, max_num_batched_tokens 8770, AL 3.00
synthetic dspark, in-allocation brackets, one node per run): steady cache-miss TTFT 648.69 vs 911.21 ms
(-262.5) at threshold 512 and 651.73 (-259.5) at 4096 on GB300 sm103;
685.69 vs 917.34 (-231.7) at 4096 on GB200 sm100. Every-rep decode ITL within the +-0.01 ms/tok band at 4096
(worst +0.0014); at 512 every rep measured FASTER than the same-run
baseline (1.3976/1.3967 vs min 1.4101) -- 512 is the shipped default
when armed, 4096 the documented hit-neutral alternative. Decode-window
kernel diff: copy_mul_dequant 0, nvjet_splitK 0, dense marlin 7700 on
both thresholds and both arches; miss-window engagement
vllm::marlin_large_m_gemm x545 + fused triton call-site kernels x870 in
the traced 5-chunk miss window (plan [8768, 8768, 8768, 6576, 1952];
545 = 5 x 109 dispatch-eligible linears). Both engagement counts are
call-site counts and therefore threshold-invariant: the M-vs-threshold
branch lives inside the opaque op, which is invoked for every eligible
linear regardless of the branch taken. The per-threshold branch receipt
is the dense-marlin count in the same window -- 635 at threshold 512 vs
844 at 4096, every other kernel label identical: the plan's 1952-token
tail chunk is below 4096, stays on Marlin there (+209 marlin kernels),
and dispatches at 512, consistent with the 2192/1952 hit chunks
dispatching at 512 only. Cost: KV headroom
-1.71 GiB absolute (177.40 -> 175.69); hit TTFT at 512 +7.6 ms
standalone (dequant on the 2192/1952 hit chunks) -- composed with the
final-split elision the merged 4144-token hit chunk inverts this into
a win (combined candidate: hit 100.27 vs 136.86 in-allocation).
Engaged accuracy adjudication (run on the pre-rebase
base tree, where this patch's runtime (vllm/) hunks had vllm/-restricted
patch-id 1bb0141f; its test/bench files later took a disclosed 9-line
lint-only amend. On this base the vllm/ patch-id is 78146196: the only
interdiff is dropping the `g_idx=None, perm=None` kwargs at the
`ops.marlin_gemm` call sites in `_marlin_16bit_gemm`, which vllm-project#54809
removed from the op signature -- they selected the non-act-order path,
the only one vllm-project#54809 left, so the patch's Python is behaviorally
unchanged. vllm-project#54809 also rewrote the Marlin kernel itself; the gsm8k
receipt has not been re-run on it. gsm8k 5-shot
n=250 temp 0, real rejection sampling, APC off, every prompt through
the dispatched path): flex -0.0040 (+-0.0612), strict -0.0120 (+-0.0633), 109/250
byte-identical, flips symmetric, real acceptance length 3.5183 vs
3.5159.
Interaction with vllm-project#53014 (FlashInfer cute-dsl W4A16 NVFP4 default):
upstream now prefers FlashInferCuteDslNvFp4W4A16LinearKernel over
Marlin for W4A16_NVFP4 under --linear-backend=auto on SM100/103, so on
those parts the NVFP4 half of this patch no longer engages by default.
The measured recipe pins --linear-backend=marlin, whose backend-map
entry (MarlinNvFp4LinearKernel -> apply_fp4_marlin_linear, where the
large-M context attaches during Marlin weight processing) is untouched
by vllm-project#53014; under that pin -- and everywhere Marlin remains the
selected kernel -- the dispatch engages exactly as measured. The FP8
channelwise half (marlin_utils_fp8), which carries most of the win, is
unaffected by vllm-project#53014 entirely.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Rishi Puri <riship@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
Purpose
This PR adds
FlashInferCuteDslNvFp4W4A16LinearKernel. Flashinfer already includes a sm12x cute-dsl w4a16 linear kernel.flashinfer-ai/flashinfer#4466 has added a sm100 cute-dsl w4a16 linear kernel and is included in 0.6.18. vLLM will default to flashinfer cute-dsl instead of Marlin when using sm100/103.
Test Plan
Test the accuracy and performance of qwen 3.6 27B nvfp4 checkpoint. It is quantized by ModelOpt thus contains w4a16 nvfp4 layer.
Test Result
Tested on a single B200:
Accuracy Tests
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)