unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to archive

[Performance] Add Triton kernel for Gemma3n sparse GELU

Language: Python+555 -21performance
The Issue & Context

The existing implementation of the sparse GELU activation in vLLM used a multi-kernel native or Inductor approach, which caused inefficiencies in CUDA execution and limited performance on supported hardware.

Root Cause Analysis

The native multi-kernel path did not fuse key operations like per-row statistics, Gaussian thresholding, tanh GELU, and gated multiplication, resulting in overhead and suboptimal GPU utilization.

Solution Strategy & Architectural Pattern

Developed a dedicated Triton kernel that fuses these operations into a single CUDA kernel, registered the activation with the vLLM IR, prioritized Triton over native on CUDA for supported configurations, and added extensive testing, benchmarking, and AOT compilation support to ensure correctness and performance.

Key Takeaway for Contributors

Fusing complex activation operations into a single custom Triton kernel can significantly improve GPU efficiency and throughput, especially for sparse activations, and integrating this with IR dispatch and tooling ensures maintainability and broad hardware support.