[Performance] Add Triton kernel for Gemma3n sparse GELU

[BWAAEEEK]

Purpose

This PR adds a dedicated Triton provider for Gemma3n's gelu_and_mul_sparse activation. The kernel fuses per-row mean and variance, Gaussian thresholding, tanh GELU, and gated multiplication, avoiding the multi-kernel native/Inductor path.

It also:

• registers gelu_and_mul_sparse with vLLM IR and routes GeluAndMulSparse through IR dispatch;
• selects triton before native on CUDA when the dtype, layout, approximation mode, and intermediate size are supported;
• exposes the kernel through torch.library.triton_op and wrap_triton, allowing torch.compile to trace the launch while preserving standalone AOT dependencies;
• adds correctness, fallback, opcheck, IR lowering, CUDA Graph, and standalone AOT regression coverage;
• adds a reusable kernel benchmark covering token counts 1-2048 and intermediate sizes 2048-32768.

Related to #32676 and #41903.

Why this is not a duplicate of #41903

#41903 is a broad IR migration for four activation families and provides only the native implementation for gelu_and_mul_sparse. This PR overlaps with its IR registration for that one op, but is materially performance-focused: it adds a dedicated Triton provider, CUDA priority/configuration, an AOT dependency fix, cross-architecture benchmarks, and model-level validation. I recorded the overlap in #41903 and requested that it retain the other activation migrations while excluding gelu_and_mul_sparse.

AI assistance was used and I reviewed every changed line and ran the tests and evaluations reported below.

Test Plan

Tested on:

• NVIDIA B200 (SM100), BF16, PyTorch 2.11.0+cu130, CUDA 13.0
• NVIDIA A100-SXM4 40GB (SM80), BF16, PyTorch 2.11.0+cu129

Commands:

.venv/bin/python -m pytest \
  tests/kernels/ir/test_activation.py \
  tests/compile/passes/ir/test_lowering.py::test_lowering_gelu_and_mul_sparse \
  tests/compile/test_aot_compile.py::test_sparse_activation_standalone_aot_preserves_dependencies \
  -q

.venv/bin/python -m pytest \
  tests/kernels/core/test_activation.py \
  tests/kernels/ir/test_ir_ops.py \
  tests/model_executor/test_enabled_custom_ops.py \
  -q

.venv/bin/python -m pytest \
  'tests/models/test_initialization.py::test_can_initialize_small_subset[Gemma3nForCausalLM]' \
  'tests/models/test_initialization.py::test_can_initialize_large_subset[Gemma3nForConditionalGeneration]' \
  -q

.venv/bin/python benchmarks/kernels/benchmark_gelu_and_mul_sparse.py \
  --dtype bfloat16

.venv/bin/pre-commit run --files \
  benchmarks/kernels/benchmark_gelu_and_mul_sparse.py \
  tests/compile/passes/ir/test_lowering.py \
  tests/compile/test_aot_compile.py \
  tests/kernels/ir/test_activation.py \
  vllm/config/kernel.py \
  vllm/ir/ops/__init__.py \
  vllm/ir/ops/activation.py \
  vllm/kernels/__init__.py \
  vllm/kernels/triton/__init__.py \
  vllm/kernels/triton/activation.py \
  vllm/model_executor/layers/activation.py \
  vllm/platforms/cuda.py

.venv/bin/pre-commit run mypy-3.12 --all-files --hook-stage manual

Model validation used Gemma3n E2B with greedy decoding, prefix caching disabled, and CUDA Graph enabled. GSM8K used 1,319 questions, 5-shot prompting, max_tokens=256, temperature 0, and seed 42.

Test Result

Correctness and integration

• B200 sparse activation, IR lowering, and standalone AOT: 76 passed (74 activation + 2 compiler)
• B200 adjacent activation/IR/custom-op regression suites: 274 passed
• Official gated Gemma3n initialization tests: 2 passed
• A100 sparse activation, opcheck, CUDA Graph, IR lowering, and standalone AOT: 76 passed (74 activation + 2 compiler)
• Changed-file pre-commit hooks: passed
• Python 3.12 mypy manual hook: passed
• git diff --check: passed

The standalone AOT regression specifically guards against dropping the upstream projection or Triton write before the downstream projection.

Cross-architecture warp and dtype validation

Following review, the fixed eight-warp launch was replaced with a reduction-width heuristic (1/2/4/8 warps, capped at eight). Production Gemma dimensions (d >= 2048) retain eight warps.

• On A100, the selected configurations improved the geometric mean over fixed eight warps by 1.022x across the tested small dimensions.
• On B200, the selected configurations improved the geometric mean over fixed eight warps by 1.134x across 18 small-dimension/token-count combinations.
• On A100, the current kernel was faster than compiled native in 18/18 tested production shapes, with speedups from 1.143x to 4.209x.
• At d=32768, A100 reports a 48-byte spill, but eight warps remains the fastest tested configuration. No architecture restriction is needed based on the measured SM80 and SM100 results.

Correctness was rechecked separately for each supported dtype:

Dtype	atol	rtol	Maximum absolute error	Zero-mask mismatches
FP32	1e-5	1.3e-6	1.43e-6	0
FP16	1e-3	1e-3	0	0
BF16	1e-3	1.6e-2	0	0

The input-dtype roundtrips are required to reproduce native PyTorch rounding at operation boundaries. Removing them on A100 caused 14/24 FP16 and 21/24 BF16 cases to exceed tolerance, with 281 and 2,700 sparse-mask mismatches respectively.

BF16 kernel performance

Each of the 42 shapes was checked against the native implementation before timing.

GPU	Triton vs compiled native	Geometric mean	Median	Range
A100-SXM4 40GB	42/42 faster	2.59x	2.49x	1.15x-4.06x
B200	42/42 faster	-	-	1.11x-4.06x

On A100, Triton was also 4.85x faster than eager native by geometric mean.

A100 Gemma3n E2B end-to-end

Workload: 4 prompts, 32 output tokens each, 128 output tokens total, warmup followed by 30 measurements.

Provider	Mean	Median	p95	Median throughput
Triton	302.95 ms	307.41 ms	318.91 ms	416.38 tok/s
Native	304.43 ms	305.34 ms	315.71 ms	419.20 tok/s

• Triton was 0.48% faster by mean and 0.68% slower by median.
• Mean difference 95% CI: -7.53 ms to +4.59 ms.
• Welch test: p=0.627.

The result is performance parity with no statistically significant A100 end-to-end regression. Three of four prompts matched token-for-token in the first comparison. The remaining prompt diverged at a newline but retained the same meaning and final answer; both providers also showed within-provider token variation on repeated runs.

B200 Gemma3n E2B end-to-end and evaluation

• 32 short prompts x 16 output tokens: median latency was effectively unchanged (172.79 ms native vs 172.92 ms Triton).
• 8 prompts x 2,048 input tokens x 1 output token: median latency improved from 150.14 ms to 147.81 ms (1.58%); all greedy tokens and selected logprobs matched exactly.
• GSM8K accuracy: 0.5262 native vs 0.5307 Triton.
• GSM8K invalid rate: 0.00227 native vs 0.00152 Triton.
• GSM8K latency: 13.24 s native vs 12.75 s Triton.
• GSM8K output throughput: 11,322.2 tok/s native vs 11,757.0 tok/s Triton.

The model evaluation shows no output-quality regression.

---

Essential Elements of an Effective PR Description Checklist

• The purpose and relationship to existing work are documented.
• Exact test commands are provided.
• Correctness, kernel performance, end-to-end performance, and model evaluation results are provided.
• No documentation update is required; this optimizes an existing supported model and operation.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[tjtanaa]

@BWAAEEEK how does it fair with various concurrencies e.g. 1, 64, 128? Is this always faster even when the number of tokens are large?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @BWAAEEEK.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[BWAAEEEK]

Thanks @tjtanaa. I rebased the PR onto the latest main and changed the wrapper from an opaque torch.library.custom_op to compiler-visible torch.library.triton_op + wrap_triton.

The short answer is: the end-to-end path is not always faster on every GPU/concurrency, but the updated implementation no longer shows a statistically significant regression in the tested cases. It is performance-neutral on A100 and improves the higher-concurrency B200 cases.

Gemma3n E2B, BF16, CUDA Graph enabled, prefix caching disabled, 14 input / 128 output tokens, 20 warmups + 30 measured runs:

GPU	Concurrency	Native mean	Triton mean	Result	p-value
A100	1	917.230 ms	916.698 ms	0.058% faster	0.405
A100	64	1284.185 ms	1285.307 ms	0.087% slower	0.695
A100	128	1661.086 ms	1658.934 ms	0.130% faster	0.323
B200	1	790.452 ms	791.535 ms	0.137% slower	0.855
B200	64	986.389 ms	955.814 ms	3.100% faster	~1.1e-10
B200	128	1191.494 ms	1166.609 ms	2.089% faster	0.00357

The A100 differences are all statistically neutral. On B200, C1 is neutral while C64 and C128 improve significantly.

I also rechecked the shorter 14-input / 16-output C64 case:

• A100: 1.453% higher mean latency, p=0.296, 95% CI [-2.667, +8.598] ms
• B200: 1.271% higher mean latency, p=0.658, 95% CI [-6.97, +10.95] ms

Neither short-case difference is statistically significant. The A100 short-run output tokens matched exactly.

The reason for the wrapper change was visible in the A100 Nsight comparison. Relative Triton-vs-native deltas changed from:

• cudaEventSynchronize: +28.488 ms to -25.791 ms
• GPU idle gap: +20.043 ms to -26.317 ms
• NVTX wall time: +19.731 ms to -24.379 ms

Both providers still have the same 52,875 total kernels, 48,330 graph kernels, and 45 graph launches, so the compiler-visible wrapper does not add graph partitions or launches.

Post-rebase validation:

• sparse activation correctness/opcheck/CUDA Graph: 68 passed
• IR lowering + standalone AOT: 2 passed
• applicable pre-commit hooks: passed

Long greedy generations were not bitwise deterministic even across repeated runs of the same provider, so I am not treating cross-provider long-output differences as a kernel regression. Numerical correctness/opcheck passed, and the short 16-token outputs matched exactly.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @BWAAEEEK.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tjtanaa]

@BadrBasowid PTAL as well

[BadrBasowid]

looks good, I added a few comments

[BadrBasowid]

Is fixed num_warps=8 optimal across the supported CUDA architectures and all intermediate sizes? The kernel covers d=1 through 32768, so eight warps feels a bit too high, at least to cover all CUDA archs. If this was tuned for one specific GPU, should the optimized implementation be restricted to that hardware until broader results are available?

[BadrBasowid]

I also say this because i am worried about spills on archs that haven't been tested

[BWAAEEEK]

Thanks for raising this. I replaced the fixed num_warps=8 launch with a reduction-width heuristic:

block_size = triton.next_power_of_2(d)
num_warps = min(max(block_size // 256, 1), 8)

This keeps eight warps for the production Gemma dimensions (d >= 2048) while avoiding over-provisioning for small reductions.

I validated the selected configurations on both A100 (SM80) and B200 (SM100). Relative to fixed eight warps, the small-dimension geometric mean improved by 1.022x on A100 and 1.134x on B200. On A100, the current kernel was faster than compiled native in all 18 tested production shapes (1.143x-4.209x).

At d=32768, A100 reports a 48-byte spill, but eight warps is still the fastest tested configuration. The B200 resource sweep leads to the same selection. Based on the SM80 and SM100 correctness, resource, and performance results, I do not think the provider needs to be restricted to one architecture.

[BadrBasowid]

I see, i think delaying the up load might help resolve the spills in some cases, but i don't want to block the PR for that reason, especially since you are seeing perf gain despite the spills. Perhaps you can just add a TODO to let others, or yourself, know about the spills and to try and take a stab at it at a later time.

[BadrBasowid]

Do note that delaying the up loads might reduce spills, but it also might hurt perf in some cases, i just mentioned it as one of many possible solutions. So if you do attempt it in this PR, make sure you do not introduce perf regression for smaller D sizes. But its also fine if you want to work on it later, or let others attempt to do it! @tjtanaa you can also let us know what you think.

[BWAAEEEK]

@BadrBasowid Thanks for the suggestion. I explored several up load placements across the full d=1..32768 range on B200.

Loading up immediately after the mean reduction preserved bitwise correctness and introduced no regressions above 2%, but it did not meaningfully reduce register/spill pressure. Delaying it until after the variance or near the final multiplication substantially reduced the SM80 stack frame for d=32768 (1312B to 592B in offline SM80 compilation), but caused approximately 2–3% regressions in some B200 shapes.

I also tried delaying only the FP32 conversion, but Triton generated effectively the same resource usage and performance as the current implementation.

Since there is no universally better placement and the current kernel remains faster despite the A100 spill, I don’t think adding an architecture- or shape-specific heuristic here would be justified. I’ll keep the current implementation for now, leave the spill optimization as a TODO, and continue exploring it separately as a follow-up.

[BadrBasowid]

Thanks for verifying, yes just add a TODO

[BadrBasowid]

Are the dtype roundtrips necessary?

[BWAAEEEK]

Yes. The roundtrips reproduce the input-dtype rounding performed at native PyTorch operation boundaries; keeping the intermediates in FP32 changes both the values and the sparsity decision.

On A100, removing the roundtrips caused 14/24 FP16 and 21/24 BF16 cases to exceed tolerance, with 281 and 2,700 zero-mask mismatches respectively. The same ablation on B200 produced 25 failing cases and 3,292 zero-mask mismatches. I added a short comment to make this intent explicit in the kernel.

[BadrBasowid]

can you report tolerances separately for FP32, FP16, and BF16.

[BWAAEEEK]

The per-dtype tolerances and observed maximum absolute errors on both A100 and B200 are:

dtype	atol	rtol	max abs error
FP32	1e-5	1.3e-6	1.43e-6
FP16	1e-3	1e-3	0
BF16	1e-3	1.6e-2	0

All tested cases passed their dtype-specific tolerance, and there were no sparse zero-mask mismatches.

[BadrBasowid]

what about fp16 and fp32?

[BWAAEEEK]

Added. The small-size test is now parameterized over FP16, BF16, and FP32 for d = 1, 31, 128, and it also requires exact agreement of the sparse zero mask.

The updated activation suite passes 74/74 tests on both A100 and B200; IR lowering and standalone AOT add another 2/2 passing tests on each GPU.

[BWAAEEEK]

@BadrBasowid Thanks for the review. I have addressed the four inline comments and updated the PR description with the additional A100 and B200 validation results.

The fixed eight-warp launch now uses a reduction-width heuristic, the dtype roundtrips are supported by ablation results on both architectures, the per-dtype tolerances are reported separately, and the small-size tests now cover FP16, BF16, and FP32 with exact sparse-mask checks.

When you have a chance, could you take another look at the updates? Thanks!

[DarkLight1337]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87881 for commit d151c957d69d.

[DarkLight1337]

/ci retry

[github-actions]

✅ Triggered Buildkite CI #88040 for commit 758b3b6de509, running 2 failed step(s) from Buildkite CI #87881.

[DarkLight1337]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88043 for commit 758b3b6de509.

[BWAAEEEK]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #88043.

[tjtanaa]

LGTM now.

[tjtanaa]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88303 for commit e51eb64947f6.