Repository navigation
[Kernel][Fusion] Fuse GELU-tanh+Mul with dynamic per-token FP8 quantization - #64212
Open
subhashpolisetti wants to merge 1 commit into
Open
subhashpolisetti wants to merge 1 commit into
subhashpolisetti wants to merge 1 commit into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
subhashpolisetti
requested review from
AndreasKaratzas,
BoyuanFeng,
Harry-Chen,
LucasWilkinson,
ProExpertProg,
WoosukKwon,
khluu,
mgoin,
pavanimajety,
robertgshaw2-redhat,
tlrmchlsmth,
yewentao256,
youkaichao,
zou3519 and
zyongye
as code owners
October 28, 2026 06:09
Contributor
|
Documentation preview: https://vllm--64212.org.readthedocs.build/en/64212/ |
mergify
Bot
added
documentation
1 task done
t239573445
added a commit
to t239573445/vllm
that referenced
this pull request
Oct 31, 2026
GELU-family models (Gemma 1/2/3, Gemma3n, Rnj1) run `gelu_tanh_and_mul` followed by `scaled_fp8_quant` when the checkpoint carries static FP8 per-tensor quantization. That path writes a bf16 activation of `intermediate_size` elements per token and then reads it back for the quantization, so the same data crosses DRAM twice. This adds `torch.ops._C.gelu_tanh_and_mul_quant`. The kernel reuses the existing `act_and_mul_quant_kernel` template with a tanh-GELU functor, so only the activation differs from `silu_and_mul_quant`. The op is registered in the manual act+quant fusion registry for `(GeluAndMul, kFp8StaticTensorSym)`, gated on `approximate == "tanh"` so `GeluAndMul(approximate="none")` keeps the unfused path. Numerics: bit-identical to `F.gelu(x, approximate="tanh") * up -> scaled_fp8_quant` over 33.5M elements (verified up to 8192x4096). An algebraically equivalent erfc-based formulation was measured to differ in 8 of 12.5M elements, so the tanh formulation already used by `gelu_tanh_kernel` is kept, which also keeps fused and unfused paths identical. Performance (RTX 5090 D v2, SM120, bf16 in / fp8 out, median of 200 iters after 50 warmup iters): | shape | unfused | fused | speedup | |-------------|---------|---------|---------| | 1x4096 | 6.0 us | 2.7 us | 2.21x | | 64x4096 | 5.9 us | 2.7 us | 2.15x | | 512x4096 | 8.2 us | 4.1 us | 2.00x | | 2048x4096 | 16.3 us | 12.3 us | 1.33x | | 2048x14336 | 165.6 us| 105.3 us| 1.57x | geomean 1.82x. 2048x14336 matches the down_proj input width of a Gemma-4-31B style MLP. Tests: op-level bit-exactness test, a `maybe_fused_act_quant` registry-level test, and a fallback test for `approximate="none"`. `tests/kernels/test_fused_quant_activation.py`: 204 passed. Also updates `docs/design/fusions.md` and adds `benchmarks/kernels/benchmark_gelu_tanh_fp8_quant.py` (--check mode asserts bit-exactness against the unfused path). Adds the `fix_functionalization` branch for the new op, mirroring `silu_and_mul_quant`; without it the op falls into the kwargs path, which the existing comment there documents as returning the wrong answer. Wires `maybe_fused_act_quant` into `Gemma2MLP` and `Gemma4MLP` (one line each, the same pattern `llama.py` uses); without it the registry row is inert on real models. `gemma3.py` is intentionally left to vllm-project#64212, which already touches that file for the dynamic-per-token variant, to avoid duplicate wiring. Signed-off-by: Alan <239573445@qq.com>
t239573445
added a commit
to t239573445/vllm
that referenced
this pull request
Oct 31, 2026
GELU-family models (Gemma 1/2/3, Gemma3n, Rnj1) run `gelu_tanh_and_mul` followed by `scaled_fp8_quant` when the checkpoint carries static FP8 per-tensor quantization. That path writes a bf16 activation of `intermediate_size` elements per token and then reads it back for the quantization, so the same data crosses DRAM twice. This adds `torch.ops._C.gelu_tanh_and_mul_quant`. The kernel reuses the existing `act_and_mul_quant_kernel` template with a tanh-GELU functor, so only the activation differs from `silu_and_mul_quant`. The op is registered in the manual act+quant fusion registry for `(GeluAndMul, kFp8StaticTensorSym)`, gated on `approximate == "tanh"` so `GeluAndMul(approximate="none")` keeps the unfused path. Numerics: bit-identical to `F.gelu(x, approximate="tanh") * up -> scaled_fp8_quant` over 33.5M elements (verified up to 8192x4096). An algebraically equivalent erfc-based formulation was measured to differ in 8 of 12.5M elements, so the tanh formulation already used by `gelu_tanh_kernel` is kept, which also keeps fused and unfused paths identical. Performance (RTX 5090 D v2, SM120, bf16 in / fp8 out, median of 200 iters after 50 warmup iters): | shape | unfused | fused | speedup | |-------------|---------|---------|---------| | 1x4096 | 6.0 us | 2.7 us | 2.21x | | 64x4096 | 5.9 us | 2.7 us | 2.15x | | 512x4096 | 8.2 us | 4.1 us | 2.00x | | 2048x4096 | 16.3 us | 12.3 us | 1.33x | | 2048x14336 | 165.6 us| 105.3 us| 1.57x | geomean 1.82x. 2048x14336 matches the down_proj input width of a Gemma-4-31B style MLP. Tests: op-level bit-exactness test, a `maybe_fused_act_quant` registry-level test, and a fallback test for `approximate="none"`. `tests/kernels/test_fused_quant_activation.py`: 204 passed. Also updates `docs/design/fusions.md` and adds `benchmarks/kernels/benchmark_gelu_tanh_fp8_quant.py` (--check mode asserts bit-exactness against the unfused path). Adds the `fix_functionalization` branch for the new op, mirroring `silu_and_mul_quant`; without it the op falls into the kwargs path, which the existing comment there documents as returning the wrong answer. Wires `maybe_fused_act_quant` into `Gemma2MLP` and `Gemma4MLP` (one line each, the same pattern `llama.py` uses); without it the registry row is inert on real models. `gemma3.py` is intentionally left to vllm-project#64212, which already touches that file for the dynamic-per-token variant, to avoid duplicate wiring. Write the tanh-GELU inner term as `KAPPA * (f * f * f)` with an explicit `x_cube`, matching `gelu_tanh_kernel` in `csrc/activation_kernels.cu` term-for-term. The previous `KAPPA * f * f * f` parses as `((KAPPA*f)*f)*f`, which is a different floating-point association; it measured identical after bf16 rounding and fp8 quantization, but there is no reason to rely on that. Signed-off-by: Alan <239573445@qq.com>
t239573445
added a commit
to t239573445/vllm
that referenced
this pull request
Oct 31, 2026
GELU-family models (Gemma 1/2/3, Gemma3n, Rnj1) run `gelu_tanh_and_mul` followed by `scaled_fp8_quant` when the checkpoint carries static FP8 per-tensor quantization. That path writes a bf16 activation of `intermediate_size` elements per token and then reads it back for the quantization, so the same data crosses DRAM twice. This adds `torch.ops._C.gelu_tanh_and_mul_quant`. The kernel reuses the existing `act_and_mul_quant_kernel` template with a tanh-GELU functor, so only the activation differs from `silu_and_mul_quant`. The op is registered in the manual act+quant fusion registry for `(GeluAndMul, kFp8StaticTensorSym)`, gated on `approximate == "tanh"` so `GeluAndMul(approximate="none")` keeps the unfused path. Numerics: bit-identical to `F.gelu(x, approximate="tanh") * up -> scaled_fp8_quant` over 33.5M elements (verified up to 8192x4096). An algebraically equivalent erfc-based formulation was measured to differ in 8 of 12.5M elements, so the tanh formulation already used by `gelu_tanh_kernel` is kept, which also keeps fused and unfused paths identical. Performance (RTX 5090 D v2, SM120, bf16 in / fp8 out, median of 200 iters after 50 warmup iters): | shape | unfused | fused | speedup | |-------------|---------|---------|---------| | 1x4096 | 6.0 us | 2.7 us | 2.21x | | 64x4096 | 5.9 us | 2.7 us | 2.15x | | 512x4096 | 8.2 us | 4.1 us | 2.00x | | 2048x4096 | 16.3 us | 12.3 us | 1.33x | | 2048x14336 | 165.6 us| 105.3 us| 1.57x | geomean 1.82x. 2048x14336 matches the down_proj input width of a Gemma-4-31B style MLP. Tests: op-level bit-exactness test, a `maybe_fused_act_quant` registry-level test, and a fallback test for `approximate="none"`. `tests/kernels/test_fused_quant_activation.py`: 204 passed. Also updates `docs/design/fusions.md` and adds `benchmarks/kernels/benchmark_gelu_tanh_fp8_quant.py` (--check mode asserts bit-exactness against the unfused path). Adds the `fix_functionalization` branch for the new op, mirroring `silu_and_mul_quant`; without it the op falls into the kwargs path, which the existing comment there documents as returning the wrong answer. Wires `maybe_fused_act_quant` into `Gemma2MLP` and `Gemma4MLP` (one line each, the same pattern `llama.py` uses); without it the registry row is inert on real models. `gemma3.py` is intentionally left to vllm-project#64212, which already touches that file for the dynamic-per-token variant, to avoid duplicate wiring. Write the tanh-GELU inner term as `KAPPA * (f * f * f)` with an explicit `x_cube`, matching `gelu_tanh_kernel` in `csrc/activation_kernels.cu` term-for-term. The previous `KAPPA * f * f * f` parses as `((KAPPA*f)*f)*f`, which is a different floating-point association; it measured identical after bf16 rounding and fp8 quantization, but there is no reason to rely on that. Adds tests that pin the fused op against the actual unfused path (`gelu_tanh_and_mul` + `scaled_fp8_quant`), plus odd/sub-vector-width hidden sizes and non-2D inputs. Signed-off-by: Alan <239573445@qq.com>
t239573445
added a commit
to t239573445/vllm
that referenced
this pull request
Oct 31, 2026
GELU-family models (Gemma 1/2/3, Gemma3n, Rnj1) run `gelu_tanh_and_mul` followed by `scaled_fp8_quant` when the checkpoint carries static FP8 per-tensor quantization. That path writes a bf16 activation of `intermediate_size` elements per token and then reads it back for the quantization, so the same data crosses DRAM twice. This adds `torch.ops._C.gelu_tanh_and_mul_quant`. The kernel reuses the existing `act_and_mul_quant_kernel` template with a tanh-GELU functor, so only the activation differs from `silu_and_mul_quant`. The op is registered in the manual act+quant fusion registry for `(GeluAndMul, kFp8StaticTensorSym)`, gated on `approximate == "tanh"` so `GeluAndMul(approximate="none")` keeps the unfused path. Numerics: bit-identical to `F.gelu(x, approximate="tanh") * up -> scaled_fp8_quant` over 33.5M elements (verified up to 8192x4096). An algebraically equivalent erfc-based formulation was measured to differ in 8 of 12.5M elements, so the tanh formulation already used by `gelu_tanh_kernel` is kept, which also keeps fused and unfused paths identical. Performance (RTX 5090 D v2, SM120, bf16 in / fp8 out, median of 200 iters after 50 warmup iters): | shape | unfused | fused | speedup | |-------------|---------|---------|---------| | 1x4096 | 6.0 us | 2.7 us | 2.21x | | 64x4096 | 5.9 us | 2.7 us | 2.15x | | 512x4096 | 8.2 us | 4.1 us | 2.00x | | 2048x4096 | 16.3 us | 12.3 us | 1.33x | | 2048x14336 | 165.6 us| 105.3 us| 1.57x | geomean 1.82x. 2048x14336 matches the down_proj input width of a Gemma-4-31B style MLP. Tests: op-level bit-exactness test, a `maybe_fused_act_quant` registry-level test, and a fallback test for `approximate="none"`. `tests/kernels/test_fused_quant_activation.py`: 204 passed. Also updates `docs/design/fusions.md` and adds `benchmarks/kernels/benchmark_gelu_tanh_fp8_quant.py` (--check mode asserts bit-exactness against the unfused path). Adds the `fix_functionalization` branch for the new op, mirroring `silu_and_mul_quant`; without it the op falls into the kwargs path, which the existing comment there documents as returning the wrong answer. Wires `maybe_fused_act_quant` into `Gemma2MLP` and `Gemma4MLP` (one line each, the same pattern `llama.py` uses); without it the registry row is inert on real models. `gemma3.py` is intentionally left to vllm-project#64212, which already touches that file for the dynamic-per-token variant, to avoid duplicate wiring. Write the tanh-GELU inner term as `KAPPA * (f * f * f)` with an explicit `x_cube`, matching `gelu_tanh_kernel` in `csrc/activation_kernels.cu` term-for-term. The previous `KAPPA * f * f * f` parses as `((KAPPA*f)*f)*f`, which is a different floating-point association; it measured identical after bf16 rounding and fp8 quantization, but there is no reason to rely on that. Adds tests that pin the fused op against the actual unfused path (`gelu_tanh_and_mul` + `scaled_fp8_quant`), plus odd/sub-vector-width hidden sizes and non-2D inputs. Reworks the docs entry: the GELU row no longer links at the SiLU section and no longer claims to be driven by the `fuse_act_quant` PassConfig flag. It is a manual fusion applied in the model forward, so it is active at every optimization level and gets its own section. Signed-off-by: Alan <239573445@qq.com>
t239573445
added a commit
to t239573445/vllm
that referenced
this pull request
Oct 31, 2026
GELU-family models (Gemma 1/2/3, Gemma3n, Rnj1) run `gelu_tanh_and_mul` followed by `scaled_fp8_quant` when the checkpoint carries static FP8 per-tensor quantization. That path writes a bf16 activation of `intermediate_size` elements per token and then reads it back for the quantization, so the same data crosses DRAM twice. This adds `torch.ops._C.gelu_tanh_and_mul_quant`. The kernel reuses the existing `act_and_mul_quant_kernel` template with a tanh-GELU functor, so only the activation differs from `silu_and_mul_quant`. The op is registered in the manual act+quant fusion registry for `(GeluAndMul, kFp8StaticTensorSym)`, gated on `approximate == "tanh"` so `GeluAndMul(approximate="none")` keeps the unfused path. Numerics: bit-identical to `F.gelu(x, approximate="tanh") * up -> scaled_fp8_quant` over 33.5M elements (verified up to 8192x4096). An algebraically equivalent erfc-based formulation was measured to differ in 8 of 12.5M elements, so the tanh formulation already used by `gelu_tanh_kernel` is kept, which also keeps fused and unfused paths identical. Performance (RTX 5090 D v2, SM120, bf16 in / fp8 out, median of 200 iters after 50 warmup iters): | shape | unfused | fused | speedup | |-------------|---------|---------|---------| | 1x4096 | 6.0 us | 2.7 us | 2.21x | | 64x4096 | 5.9 us | 2.7 us | 2.15x | | 512x4096 | 8.2 us | 4.1 us | 2.00x | | 2048x4096 | 16.3 us | 12.3 us | 1.33x | | 2048x14336 | 165.6 us| 105.3 us| 1.57x | geomean 1.82x. 2048x14336 matches the down_proj input width of a Gemma-4-31B style MLP. Tests: op-level bit-exactness test, a `maybe_fused_act_quant` registry-level test, and a fallback test for `approximate="none"`. `tests/kernels/test_fused_quant_activation.py`: 204 passed. Also updates `docs/design/fusions.md` and adds `benchmarks/kernels/benchmark_gelu_tanh_fp8_quant.py` (--check mode asserts bit-exactness against the unfused path). Adds the `fix_functionalization` branch for the new op, mirroring `silu_and_mul_quant`; without it the op falls into the kwargs path, which the existing comment there documents as returning the wrong answer. Wires `maybe_fused_act_quant` into `Gemma2MLP` and `Gemma4MLP` (one line each, the same pattern `llama.py` uses); without it the registry row is inert on real models. `gemma3.py` is intentionally left to vllm-project#64212, which already touches that file for the dynamic-per-token variant, to avoid duplicate wiring. Write the tanh-GELU inner term as `KAPPA * (f * f * f)` with an explicit `x_cube`, matching `gelu_tanh_kernel` in `csrc/activation_kernels.cu` term-for-term. The previous `KAPPA * f * f * f` parses as `((KAPPA*f)*f)*f`, which is a different floating-point association; it measured identical after bf16 rounding and fp8 quantization, but there is no reason to rely on that. Adds tests that pin the fused op against the actual unfused path (`gelu_tanh_and_mul` + `scaled_fp8_quant`), plus odd/sub-vector-width hidden sizes and non-2D inputs. Reworks the docs entry: the GELU row no longer links at the SiLU section and no longer claims to be driven by the `fuse_act_quant` PassConfig flag. It is a manual fusion applied in the model forward, so it is active at every optimization level and gets its own section. Use the anchor form that mkdocs actually renders (single hyphen): its slugifier collapses runs of separators, so the `--` used by the other rows in that table does not resolve. Signed-off-by: Alan <239573445@qq.com>
t239573445
added a commit
to t239573445/vllm
that referenced
this pull request
Nov 1, 2026
GELU-family models (Gemma 1/2/3, Gemma3n, Rnj1) run `gelu_tanh_and_mul` followed by `scaled_fp8_quant` when the checkpoint carries static FP8 per-tensor quantization. That path writes a bf16 activation of `intermediate_size` elements per token and then reads it back for the quantization, so the same data crosses DRAM twice. This adds `torch.ops._C.gelu_tanh_and_mul_quant`. The kernel reuses the existing `act_and_mul_quant_kernel` template with a tanh-GELU functor, so only the activation differs from `silu_and_mul_quant`. The op is registered in the manual act+quant fusion registry for `(GeluAndMul, kFp8StaticTensorSym)`, gated on `approximate == "tanh"` so `GeluAndMul(approximate="none")` keeps the unfused path. Numerics: bit-identical to `F.gelu(x, approximate="tanh") * up -> scaled_fp8_quant` over 33.5M elements (verified up to 8192x4096). An algebraically equivalent erfc-based formulation was measured to differ in 8 of 12.5M elements, so the tanh formulation already used by `gelu_tanh_kernel` is kept, which also keeps fused and unfused paths identical. Performance (RTX 5090 D v2, SM120, bf16 in / fp8 out, median of 200 iters after 50 warmup iters): | shape | unfused | fused | speedup | |-------------|---------|---------|---------| | 1x4096 | 6.0 us | 2.7 us | 2.21x | | 64x4096 | 5.9 us | 2.7 us | 2.15x | | 512x4096 | 8.2 us | 4.1 us | 2.00x | | 2048x4096 | 16.3 us | 12.3 us | 1.33x | | 2048x14336 | 165.6 us| 105.3 us| 1.57x | geomean 1.82x. 2048x14336 matches the down_proj input width of a Gemma-4-31B style MLP. Tests: op-level bit-exactness test, a `maybe_fused_act_quant` registry-level test, and a fallback test for `approximate="none"`. `tests/kernels/test_fused_quant_activation.py`: 204 passed. Also updates `docs/design/fusions.md` and adds `benchmarks/kernels/benchmark_gelu_tanh_fp8_quant.py` (--check mode asserts bit-exactness against the unfused path). Adds the `fix_functionalization` branch for the new op, mirroring `silu_and_mul_quant`; without it the op falls into the kwargs path, which the existing comment there documents as returning the wrong answer. Wires `maybe_fused_act_quant` into `Gemma2MLP` and `Gemma4MLP` (one line each, the same pattern `llama.py` uses); without it the registry row is inert on real models. `gemma3.py` is intentionally left to vllm-project#64212, which already touches that file for the dynamic-per-token variant, to avoid duplicate wiring. Write the tanh-GELU inner term as `KAPPA * (f * f * f)` with an explicit `x_cube`, matching `gelu_tanh_kernel` in `csrc/activation_kernels.cu` term-for-term. The previous `KAPPA * f * f * f` parses as `((KAPPA*f)*f)*f`, which is a different floating-point association; it measured identical after bf16 rounding and fp8 quantization, but there is no reason to rely on that. Adds tests that pin the fused op against the actual unfused path (`gelu_tanh_and_mul` + `scaled_fp8_quant`), plus odd/sub-vector-width hidden sizes and non-2D inputs. Reworks the docs entry: the GELU row no longer links at the SiLU section and no longer claims to be driven by the `fuse_act_quant` PassConfig flag. It is a manual fusion applied in the model forward, so it is active at every optimization level and gets its own section. Use the anchor form that mkdocs actually renders (single hyphen): its slugifier collapses runs of separators, so the `--` used by the other rows in that table does not resolve. Add a non-contiguous-input test: like `silu_and_mul_quant`, this op does not guard contiguity, so pin that the fused path agrees with the unfused path on such a view instead of drifting silently. Signed-off-by: Alan <239573445@qq.com>
subhashpolisetti
force-pushed
the
feat-gelu-fp8-fusion-v3
branch
2 times, most recently
from
November 1, 2026 09:53
413cc3d to
d44e150
Compare
…zation Add gelu_tanh_and_mul_per_token_quant, a CUDA kernel that computes gelu_tanh(gate) * up and quantizes it to FP8 with one scale per token in a single pass over the input, bit-identical to gelu_tanh_and_mul followed by dynamic_per_token_scaled_fp8_quant. Register it as the (GeluAndMul, per-token FP8) row of the manual activation+quant fusion, used when GELU runs as the CUDA op, and wire it into Gemma3 and Gemma 4. The CUTLASS FP8 linear now accepts the per-token input quantization key. Signed-off-by: subhash polisetti <subhashr161347@gmail.com>
subhashpolisetti
force-pushed
the
feat-gelu-fp8-fusion-v3
branch
from
November 1, 2026 18:01
d44e150 to
e12b1ed
Compare
This branch has not been deployed
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
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.
Overview
Adds a CUDA kernel,
gelu_tanh_and_mul_per_token_quant, and the(GeluAndMul, per-token dynamic FP8)row in the manual activation+quant fusion registry, as part of the manual fusion port (#43501, #43224). When GELU runs as the CUDAgelu_and_mulop, Gemma3 and Gemma 4 FP8-dynamic MLPs produce the down projection's input in one kernel instead ofgelu_tanh_and_mulfollowed bydynamic_per_token_scaled_fp8_quant.Claims
torch.compiledefault path is unchanged. The row only replaces the CUDAgelu_and_mulop.Validation
CUDA 13.0, torch 2.14, bf16,
--kernel-config.linear_backend=cutlass,RedHatAI/gemma-*-it-FP8-dynamic.H100 ran these on this commit's base,
5d55ef1590. Generation also passes withVLLM_GPU_SYNC_CHECK=error, with CUDA custom ops and with--enforce-eager, on Gemma3-1B and Gemma-4-31B.Exactness. During a real generation (128 gsm8k prompts × 128 tokens), every call also ran the two original kernels on the same input, and FP8 values and scales were compared byte for byte:
On Gemma-4-31B, greedy token IDs also match with and without the row on 5 prompts (H100, CUDA ops).
End to end, CUDA ops, main → this PR, mean of 2 runs in main-PR-PR-main order:
Throughput varies 2.6-3% between runs of one setup on H100 and up to 7% on B200, so I don't count those H100 and B200 cells as gains. The Gemma-4-31B run hit its time limit before the compiled arm and the last main throughput run.
GPU time,
torch.profiler, 32 prompts × 1024 in / 128 out, ms:Kernel benchmark and gsm8k
L4, µs, median of
triton.testing.do_bench_cudagraph, bf16, the two kernels → this kernel. All 20 shapes are 1.12-3.2× faster, and fp16 is within 2%.Standalone comparison at the same widths and 1-8192 tokens, L2-warm / flushed: H100 1.22-1.72× / 1.28-1.72×, B200 1.08-1.54× / 1.11-1.50×, faster on 24/24 shapes on both. Random inputs are bit-exact on 72/72 cases on both (bf16 and fp16, with and without
scale_ub).gsm8k, greedy, 1319 questions, CUDA ops, L4. End-to-end runs are not deterministic on main, and these stay within its spread:
Benchmarks and exactness ran on older trees with this same kernel file:
e96b490b12(H100, Gemma3) ande0d00e4b8e(B200, and H100 Gemma-4-31B).Details
scaled_fp8_conversion, from shared memory when the row fits. The GELU expression and rounding are copied fromgelu_tanh_kernelandact_and_mul_kernel, which is what keeps the output identical.torch.compile, Inductor fuses GELU with the quantization using its own FP8 rounding, so the row skips that path and output stays unchanged.Gemma3MLPandGemma4MLP. Gemma 4 MoE layers fuse the activation in their expert kernels.Gemma4Unifiedis not checked.gemma3.pyline, and [Kernel][Fusion] Fuse GELU-tanh+Mul with NVFP4 quantization (Gemma 2/3/4) #59467 and [Kernel][Fusion] Fuse tanh-GELU gated activation with static FP8 quant #64854 add the samegemma4.pyline. I can merge this kernel with [Kernels][Fusion] Fused silu mul per token quant #58664's into one activation-templated kernel as a follow-up.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.