Implementation:Sgl project Sglang Elementwise Ops

Knowledge Sources	Sgl_project_Sglang
Domains	GPU Kernels, Normalization, Activation Functions, Rotary Embeddings
Last Updated	2026-02-10 00:00 GMT

Overview

Python interface for fused elementwise GPU kernels including RMS normalization, activation functions, rotary position embeddings, MLA concatenation, and data movement operations.

Description

elementwise.py provides performance-critical fused elementwise kernels that avoid multiple GPU memory round-trips during the forward pass of transformer models. The module wraps C++ CUDA kernels via torch.ops.sgl_kernel.* and draws extensively from the FlashInfer project.

The module is organized into several functional groups:

Normalization Kernels:

• rmsnorm -- Standard Root Mean Square normalization: out[i] = (input[i] / RMS(input)) * weight[i]. Supports optional output tensor for in-place updates and PDL (Programmatic Dependent Launch) for Hopper overlap.
• fused_add_rmsnorm -- Fused residual add + RMSNorm that performs residual += input then normalizes, avoiding a separate addition kernel.
• gemma_rmsnorm -- Gemma-style variant: out[i] = (input[i] / RMS(input)) * (weight[i] + 1), where the +1 bias is applied to the weight.
• gemma_fused_add_rmsnorm -- Combined Gemma residual add + RMSNorm.

Activation Functions:

• silu_and_mul -- SiLU gating: out = silu(input[..., :d]) * input[..., d:] where the last dimension is split in half. Requires 16-byte aligned pointers.
• gelu_tanh_and_mul -- GELU-tanh gating variant with the same split pattern.
• gelu_and_mul -- Standard GELU gating variant.
• gelu_quick -- ROCm-only quick GELU: y = x * sigmoid(1.702 * x).

Rotary Position Embeddings:

• apply_rope_with_cos_sin_cache_inplace -- Applies rotary embeddings using precomputed cos/sin cache. Supports both NeoX-style (split half dimensions) and interleaved-style (even/odd dimensions). Accepts an optional FusedSetKVBufferArg dataclass to fuse the KV buffer write into the same kernel, eliminating a separate memory pass.
• rotary_embedding -- Simpler rotary embedding wrapper compatible with vLLM-style interfaces.

MLA Operations:

• concat_mla_k -- Concatenates MLA key components (k_nope and k_rope) into a single tensor.
• concat_mla_absorb_q -- Concatenates two tensors along the last dimension for MLA absorbed query computation.

Data Movement:

• downcast_fp8 -- Downcasts key and value tensors to FP8 with per-tensor scaling.
• copy_to_gpu_no_ce -- Copies data to GPU without using the Copy Engine, useful for overlapping compute and data transfer.

The FusedSetKVBufferArg dataclass bundles parameters for fusing KV cache writes into the RoPE kernel: value tensor, k/v buffers, optional k/v scales, and cache location indices.

Usage

Use these functions as drop-in replacements for naive PyTorch implementations of normalization, activation, and rotary embedding operations. They are particularly beneficial for inference where kernel launch overhead matters and memory bandwidth is the bottleneck.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/elementwise.py
• Lines: 1-406

Signature

def rmsnorm(
    input: torch.Tensor,
    weight: torch.Tensor,
    eps: float = 1e-6,
    out: Optional[torch.Tensor] = None,
    enable_pdl: Optional[bool] = None,
) -> torch.Tensor:

def fused_add_rmsnorm(
    input: torch.Tensor,
    residual: torch.Tensor,
    weight: torch.Tensor,
    eps: float = 1e-6,
    enable_pdl: Optional[bool] = None,
) -> None:

def gemma_rmsnorm(
    input: torch.Tensor,
    weight: torch.Tensor,
    eps: float = 1e-6,
    out: Optional[torch.Tensor] = None,
    enable_pdl: Optional[bool] = None,
) -> torch.Tensor:

def gemma_fused_add_rmsnorm(
    input: torch.Tensor,
    residual: torch.Tensor,
    weight: torch.Tensor,
    eps: float = 1e-6,
    enable_pdl: Optional[bool] = None,
) -> None:

def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:

def gelu_tanh_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:

def gelu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:

def apply_rope_with_cos_sin_cache_inplace(
    positions: torch.Tensor,
    query: torch.Tensor,
    key: torch.Tensor,
    head_size: int,
    cos_sin_cache: torch.Tensor,
    is_neox: bool = True,
    fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
    enable_pdl: Optional[bool] = None,
) -> None:

def rotary_embedding(
    positions: torch.Tensor,
    query: torch.Tensor,
    key: torch.Tensor,
    head_size: int,
    cos_sin_cache: torch.Tensor,
    is_neox: bool = True,
):

def downcast_fp8(k, v, k_out, v_out, k_scale, v_scale, loc, mult=1, offset=0) -> None:

def copy_to_gpu_no_ce(input: torch.Tensor, output: torch.Tensor):

def concat_mla_k(k, k_nope, k_rope):

def concat_mla_absorb_q(a, b) -> torch.Tensor:

Import

from sgl_kernel import rmsnorm, fused_add_rmsnorm, silu_and_mul
from sgl_kernel import apply_rope_with_cos_sin_cache_inplace
from sgl_kernel.elementwise import FusedSetKVBufferArg

I/O Contract

Inputs

Name	Type	Required	Description
input	torch.Tensor	Yes	Input tensor, typically shape (batch_size, hidden_size)
weight	torch.Tensor	Yes	Weight tensor for normalization, shape (hidden_size,)
eps	float	No	Epsilon for numerical stability, default 1e-6
out	torch.Tensor	No	Pre-allocated output tensor for in-place operation
enable_pdl	Optional[bool]	No	Enable Programmatic Dependent Launch (auto-detected on Hopper)
positions	torch.Tensor	Yes (RoPE)	Position indices, shape (nnz,)
cos_sin_cache	torch.Tensor (float32)	Yes (RoPE)	Precomputed cos/sin cache, shape (max_seq_len, rotary_dim)
fused_set_kv_buffer_arg	FusedSetKVBufferArg	No	Optional fused KV buffer write parameters

Outputs

Name	Type	Description
output	torch.Tensor	Normalized/activated tensor, same shape as input (for rmsnorm, activations)
None	void	In-place operations (fused_add_rmsnorm, apply_rope) modify inputs directly

Usage Examples

import torch
from sgl_kernel import rmsnorm, silu_and_mul, fused_add_rmsnorm
from sgl_kernel.elementwise import apply_rope_with_cos_sin_cache_inplace

# RMS normalization
hidden = torch.randn(32, 4096, device="cuda", dtype=torch.bfloat16)
weight = torch.ones(4096, device="cuda", dtype=torch.bfloat16)
output = rmsnorm(hidden, weight, eps=1e-6)

# Fused residual add + RMSNorm (in-place)
fused_add_rmsnorm(hidden, residual, weight, eps=1e-6)
# After: residual = residual + hidden_orig, hidden = RMSNorm(residual) * weight

# SiLU gating activation
gate_up = torch.randn(32, 8192, device="cuda", dtype=torch.bfloat16)
activated = silu_and_mul(gate_up)  # shape: (32, 4096)

# Rotary embedding with precomputed cache
apply_rope_with_cos_sin_cache_inplace(
    positions=pos_ids,     # (nnz,)
    query=query,           # (nnz, num_q_heads * head_size)
    key=key,               # (nnz, num_k_heads * head_size)
    head_size=128,
    cos_sin_cache=cache,   # (max_seq_len, rotary_dim), float32
    is_neox=True,
)

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime