minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 15 Packages Projects Releases Wiki Activity

[2/3] [EAGLE] perf: Fuse TP vocab-parallel embedding (#30948)

Browse Source
This commit is contained in:
Kaixi
2026-07-16 15:57:51 -07:00
committed by GitHub
parent d539bf2cda
commit dc0b3eb68f
7 changed files with 509 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/__init__.py
+1
View File
@@ -19,6 +19,7 @@ _GROUPS = (
"attention",
"communication",
"diffusion",
"embeddings",
"gemm",
"grammar",
"kvcache",
python/sglang/kernels/ops/embeddings/__init__.py
+17
View File
@@ -0,0 +1,17 @@
"""Embedding kernels."""
from sglang.kernels.registry import register_kernel
from sglang.kernels.spec import KernelBackend, KernelSpec
register_kernel(
KernelSpec(
op="embeddings.vocab_parallel_embedding",
backend=KernelBackend.TRITON,
target=(
"sglang.kernels.ops.embeddings.vocab_parallel_embedding:"
"vocab_parallel_embedding"
),
)
)
__all__ = []
python/sglang/kernels/ops/embeddings/vocab_parallel_embedding.py
+96
View File
@@ -0,0 +1,96 @@
"""Fused Triton vocabulary-parallel embedding lookup."""
import torch
import triton
import triton.language as tl
@triton.jit
def _vocab_parallel_embedding_kernel(
input_ptr,
weight_ptr,
out_ptr,
# The scalar params are tl.constexpr on purpose: it lets the compiler fold
# the vocab-window comparisons into a single range check (measured ~5%
# faster at large token counts), and each embedding layer has one fixed
# (hidden_dim, stride, shard-window) tuple, so the specialization costs one
# compile per layer at warmup.
hidden_dim: tl.constexpr,
weight_stride0: tl.constexpr,
org_vocab_start_index: tl.constexpr,
org_vocab_end_index: tl.constexpr,
num_org_vocab_padding: tl.constexpr,
added_vocab_start_index: tl.constexpr,
added_vocab_end_index: tl.constexpr,
BLOCK_H: tl.constexpr,
):
row = tl.program_id(0).to(tl.int64)
col_block = tl.program_id(1)
cols = col_block * BLOCK_H + tl.arange(0, BLOCK_H)
col_mask = cols < hidden_dim
token = tl.load(input_ptr + row).to(tl.int64)
org_vocab_mask = (token >= org_vocab_start_index) & (token < org_vocab_end_index)
added_vocab_mask = (token >= added_vocab_start_index) & (
token < added_vocab_end_index
)
valid = org_vocab_mask | added_vocab_mask
added_offset = (
added_vocab_start_index
- (org_vocab_end_index - org_vocab_start_index)
- num_org_vocab_padding
)
local_id = tl.where(
org_vocab_mask, token - org_vocab_start_index, token - added_offset
)
local_id = tl.where(valid, local_id, 0)
vals = tl.load(
weight_ptr + local_id * weight_stride0 + cols,
mask=col_mask & valid,
other=0.0,
)
tl.store(out_ptr + row * hidden_dim + cols, vals, mask=col_mask)
def vocab_parallel_embedding(
input_: torch.Tensor,
weight: torch.Tensor,
org_vocab_start_index: int,
org_vocab_end_index: int,
num_org_vocab_padding: int,
added_vocab_start_index: int,
added_vocab_end_index: int,
) -> torch.Tensor:
assert input_.is_cuda
assert input_.is_contiguous()
assert input_.dtype in (torch.int32, torch.int64)
assert weight.is_cuda
assert weight.ndim == 2
assert weight.stride(1) == 1
hidden_dim = weight.shape[1]
output = torch.empty(
(*input_.shape, hidden_dim), dtype=weight.dtype, device=weight.device
)
n_tokens = input_.numel()
if n_tokens == 0:
return output
block_h = min(1024, triton.next_power_of_2(hidden_dim))
grid = (n_tokens, triton.cdiv(hidden_dim, block_h))
_vocab_parallel_embedding_kernel[grid](
input_,
weight,
output,
hidden_dim,
weight.stride(0),
org_vocab_start_index,
org_vocab_end_index,
num_org_vocab_padding,
added_vocab_start_index,
added_vocab_end_index,
BLOCK_H=block_h,
num_warps=8,
)
return output
python/sglang/srt/layers/vocab_parallel_embedding.py
+66 -28
View File
@@ -9,6 +9,9 @@ from typing import List, Optional, Sequence, Tuple
import torch
from torch.nn.parameter import Parameter, UninitializedParameter
from sglang.kernels.ops.embeddings.vocab_parallel_embedding import (
vocab_parallel_embedding as fused_vocab_parallel_embedding,
)
from sglang.srt.distributed import (
divide,
get_tp_group,
@@ -498,40 +501,75 @@ class VocabParallelEmbedding(torch.nn.Module):
param[: loaded_weight.shape[0]].data.copy_(loaded_weight)
param[loaded_weight.shape[0] :].data.fill_(0)
def _use_triton_embedding(self, input_: torch.Tensor) -> bool:
"""Whether the fused Triton kernel can replace the mask+gather+fill unit."""
if self.tp_size == 1:
return False
if not isinstance(self.quant_method, UnquantizedEmbeddingMethod):
return False
if not input_.is_cuda or not input_.is_contiguous():
return False
if input_.dtype not in (torch.int32, torch.int64):
return False
return (
self.weight.is_cuda
and self.weight.ndim == 2
and self.weight.stride(1) == 1
and self.weight.dtype in (torch.float16, torch.bfloat16, torch.float32)
)
def _embed_local_shard(self, input_: torch.Tensor) -> torch.Tensor:
"""Embed against the local vocab shard; out-of-shard rows are zero
(identity when tp_size == 1).
The output must be allocated inside the symmetric-memory context so
the caller's all-reduce can use it; the mask temporaries and the
in-place fill deliberately stay outside the pool.
"""
symm_alloc = use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric()
)
if self.tp_size == 1:
with symm_alloc:
return self.quant_method.embedding(self, input_.long())
if self._use_triton_embedding(input_):
with symm_alloc:
return fused_vocab_parallel_embedding(
input_,
self.weight,
self.shard_indices.org_vocab_start_index,
self.shard_indices.org_vocab_end_index,
self.shard_indices.num_org_vocab_padding,
self.shard_indices.added_vocab_start_index,
self.shard_indices.added_vocab_end_index,
)
# Map out-of-shard ids to index 0, gather, then zero those rows.
masked_input, input_mask = get_masked_input_and_mask(
input_,
self.shard_indices.org_vocab_start_index,
self.shard_indices.org_vocab_end_index,
self.shard_indices.num_org_vocab_padding,
self.shard_indices.added_vocab_start_index,
self.shard_indices.added_vocab_end_index,
)
with symm_alloc:
output_parallel = self.quant_method.embedding(self, masked_input.long())
output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
return output_parallel
def forward(self, input_):
# Surface a bad token id (>= vocab_size, or a negative / unmasked sentinel) as a
# located async assert instead of a silent OOB embedding gather (tp=1 does not mask).
maybe_detect_oob(
input_, 0, self.num_embeddings, "VocabParallelEmbedding input id"
)
if self.tp_size > 1:
# Build the mask.
masked_input, input_mask = get_masked_input_and_mask(
input_,
self.shard_indices.org_vocab_start_index,
self.shard_indices.org_vocab_end_index,
self.shard_indices.num_org_vocab_padding,
self.shard_indices.added_vocab_start_index,
self.shard_indices.added_vocab_end_index,
)
else:
masked_input = input_
# Get the embeddings.
with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric()
):
output_parallel = self.quant_method.embedding(self, masked_input.long())
if self.tp_size > 1:
# Mask the output embedding.
output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
if not get_attn_tp_context().input_scattered:
if self.use_attn_tp_group:
output_parallel = attn_tp_all_reduce(output_parallel)
else:
# Reduce across all the model parallel GPUs.
output_parallel = tensor_model_parallel_all_reduce(output_parallel)
output_parallel = self._embed_local_shard(input_)
if self.tp_size > 1 and not get_attn_tp_context().input_scattered:
if self.use_attn_tp_group:
output_parallel = attn_tp_all_reduce(output_parallel)
else:
# Reduce across all the model parallel GPUs.
output_parallel = tensor_model_parallel_all_reduce(output_parallel)
return output_parallel
def extra_repr(self) -> str: