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

optimize: optimize EmbeddingGemma prefill performance (#32383)

Browse Source
This commit is contained in:
Mick
2026-07-27 17:34:29 +08:00
committed by GitHub
parent 9a0bd24bed
commit 08af5aea57
7 changed files with 285 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashattention_backend.py
+12 -5
View File
@@ -1192,7 +1192,18 @@ class FlashAttentionBackend(AttentionBackend):
is_swa_layer = (
layer.sliding_window_size is not None and layer.sliding_window_size > -1
)
window_size = (layer.sliding_window_size, 0) if is_swa_layer else (-1, -1)
causal = not (
layer.is_cross_attention
or layer.attn_type
in (AttentionType.ENCODER_ONLY, AttentionType.DECODER_BIDIRECTIONAL)
)
# FlashAttention's sliding-window tuple is (left, right). Bidirectional
# encoder layers must see the same local context on both sides.
window_size = (
(layer.sliding_window_size, 0 if causal else layer.sliding_window_size)
if is_swa_layer
else (-1, -1)
)
fa_k_descale, fa_v_descale = None, None
# only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
# has corresponding quantization method so that layer.k_scale is not None,
@@ -1211,10 +1222,6 @@ class FlashAttentionBackend(AttentionBackend):
q = q.to(self.kv_cache_dtype)
q_rope = q_rope.to(self.kv_cache_dtype) if q_rope is not None else None
k_rope = k_rope.to(self.kv_cache_dtype) if k_rope is not None else None
causal = True
if layer.is_cross_attention or layer.attn_type == AttentionType.ENCODER_ONLY:
causal = False
# Check if we should use local attention
use_local_attn = (
self.has_local_attention
python/sglang/srt/layers/layernorm.py
+21 -5
View File
@@ -1058,22 +1058,38 @@ class Gemma3RMSNorm(MultiPlatformOp):
def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward_native(self, x):
def forward_native(self, x, residual: Optional[torch.Tensor] = None):
if residual is not None:
residual = x + residual
x = residual
output = self._norm(x.float())
# Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16)
# See https://github.com/huggingface/transformers/pull/29402
output = output * (1.0 + self.weight.float())
return output.type_as(x)
output = output.type_as(x)
return output if residual is None else (output, residual)
def forward_cpu(self, x):
def forward_cpu(self, x, residual: Optional[torch.Tensor] = None):
if residual is not None:
return self.forward_native(x, residual)
if _is_cpu_amx_available and x.stride(-1) == 1:
return torch.ops.sgl_kernel.gemma3_rmsnorm_cpu(x, self.weight, self.eps)
return self.forward_native(x)
def forward_cuda(self, x):
def forward_cuda(self, x, residual: Optional[torch.Tensor] = None):
if residual is not None:
# The decoder residual is token-major and contiguous. The fused
# kernel updates both tensors in place: x becomes the normalized
# output and residual becomes x + residual for the next layer.
gemma_fused_add_rmsnorm(x, residual, self.weight.data, self.eps)
return x, residual
if x.dim() == 2:
return gemma_rmsnorm(x, self.weight.data, self.eps)
return self.forward_native(x)
def forward_npu(self, x):
def forward_npu(self, x, residual: Optional[torch.Tensor] = None):
if residual is not None:
return self.forward_native(x, residual)
output, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.eps)
return output
python/sglang/srt/managers/tokenizer_manager.py
+17
View File
@@ -868,6 +868,23 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
encoded.get("token_type_ids") if is_cross_encoder else None
)
# vLLM's OpenAI embeddings endpoint includes special tokens for
# encoder models. EmbeddingGemma's restored Gemma tokenizer adds BOS
# but, by its checkpoint default, omits EOS. Add EOS explicitly here
# rather than mutating tokenizer-global post-processing state.
if (
self.model_config.is_embedding_gemma
and self.tokenizer.eos_token_id is not None
):
input_ids = [
(
ids
if ids and ids[-1] == self.tokenizer.eos_token_id
else [*ids, self.tokenizer.eos_token_id]
)
for ids in input_ids
]
# Step 4: Extract results based on input format
return self._extract_tokenizer_results(
input_ids, token_type_ids, input_format, original_batch_size
python/sglang/srt/models/gemma3_causal.py
+99 -16
View File
@@ -14,10 +14,12 @@
# limitations under the License.
# ==============================================================================
import copy
import json
from typing import Iterable, List, Optional, Set, Tuple
import einops
import torch
import torch.nn.functional as F
from torch import nn
from transformers import (
ROPE_INIT_FUNCTIONS,
@@ -46,6 +48,7 @@ from sglang.srt.model_loader.weight_utils import (
)
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import add_prefix, cpu_has_amx_support, is_cpu, make_layers
from sglang.srt.utils.hf_transformers.common import _resolve_local_or_cached_file
_is_cpu = is_cpu()
_is_cpu_amx_available = cpu_has_amx_support()
@@ -374,12 +377,19 @@ class Gemma3DecoderLayer(nn.Module):
position_embeddings_global: torch.Tensor,
position_embeddings_local: torch.Tensor,
forward_batch: ForwardBatch,
residual: Optional[torch.Tensor] = None,
**kwargs,
) -> tuple[
torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]
]:
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
# Keep the residual live across layers so the add preceding the next
# RMSNorm is fused by Gemma3RMSNorm. This matches the upstream Gemma3
# residual layout and is safe to capture in a breakable CUDA graph.
if residual is None:
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
else:
hidden_states, residual = self.input_layernorm(hidden_states, residual)
# apply global RoPE to non-sliding layer only
if self.self_attn.is_sliding:
@@ -395,15 +405,13 @@ class Gemma3DecoderLayer(nn.Module):
**kwargs,
)
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = residual + hidden_states
residual = hidden_states
hidden_states = self.pre_feedforward_layernorm(hidden_states)
hidden_states, residual = self.pre_feedforward_layernorm(
hidden_states, residual
)
hidden_states = self.mlp(hidden_states)
hidden_states = self.post_feedforward_layernorm(hidden_states)
hidden_states = residual + hidden_states
outputs = (hidden_states,)
outputs = (hidden_states, residual)
return outputs
@@ -626,21 +634,22 @@ class Gemma3TextModel(PreTrainedModel):
hidden_states = input_embeds
aux_hidden_states = []
residual = None
num_layers = len(self.layers)
if _is_cpu and _is_cpu_amx_available:
for i, layer in enumerate(self.layers):
if i in self.layers_to_capture:
aux_hidden_states.append(hidden_states)
layer_outputs = layer(
hidden_states, residual = layer(
positions=positions,
position_embeddings_global=None,
position_embeddings_local=None,
hidden_states=hidden_states,
forward_batch=forward_batch,
residual=residual,
**kwargs,
)
hidden_states = layer_outputs[0]
else:
if positions.dim() == 1:
positions = einops.rearrange(positions, "s -> 1 s")
@@ -650,15 +659,15 @@ class Gemma3TextModel(PreTrainedModel):
for i, layer in enumerate(self.layers):
if i in self.layers_to_capture:
aux_hidden_states.append(hidden_states)
layer_outputs = layer(
hidden_states, residual = layer(
positions=positions,
position_embeddings_global=position_embeddings_global,
position_embeddings_local=position_embeddings_local,
hidden_states=hidden_states,
forward_batch=forward_batch,
residual=residual,
**kwargs,
)
hidden_states = layer_outputs[0]
# Capture the output of the last layer if requested.
# layers_to_capture uses +1 offset (captures input of layer i = output of i-1),
@@ -666,7 +675,7 @@ class Gemma3TextModel(PreTrainedModel):
if num_layers in self.layers_to_capture:
aux_hidden_states.append(hidden_states)
hidden_states = self.norm(hidden_states)
hidden_states, _ = self.norm(hidden_states, residual)
if len(aux_hidden_states) == 0:
return hidden_states
@@ -960,9 +969,50 @@ class EmbeddingGemmaModel(Gemma3ForCausalLM):
self.model = Gemma3TextModel(
config, quant_config, prefix=add_prefix("model", prefix)
)
self.pooler = Pooler(pooling_type=PoolingType.MEAN, normalize=True)
# SentenceTransformers applies mean pooling, then its optional Dense
# projector modules, then L2 normalization. Keep normalization outside
# Pooler so this ordering is preserved.
self.pooler = Pooler(pooling_type=PoolingType.MEAN, normalize=False)
self.projector = self._build_sentence_transformer_projector(config)
self.capture_aux_hidden_states = False
@staticmethod
def _build_sentence_transformer_projector(config: Gemma3TextConfig):
"""Create the checkpoint's SentenceTransformers Dense tail, if present."""
model_path = getattr(config, "_name_or_path", "")
try:
modules_path = _resolve_local_or_cached_file(model_path, "modules.json")
with open(modules_path) as f:
module_specs = json.load(f)
layers = []
for spec in module_specs:
if spec.get("type") != "sentence_transformers.models.Dense":
continue
dense_config_path = _resolve_local_or_cached_file(
model_path, f"{spec['path']}/config.json"
)
with open(dense_config_path) as f:
dense_config = json.load(f)
if dense_config.get("activation_function") not in (
None,
"torch.nn.modules.linear.Identity",
):
raise ValueError(
"EmbeddingGemma only supports identity SentenceTransformers "
"Dense activations"
)
layers.append(
nn.Linear(
dense_config["in_features"],
dense_config["out_features"],
bias=dense_config.get("bias", True),
)
)
return nn.Sequential(*layers) if layers else None
except (FileNotFoundError, OSError, ValueError, KeyError, json.JSONDecodeError):
return None
@torch.no_grad()
def forward(
self,
@@ -977,7 +1027,10 @@ class EmbeddingGemmaModel(Gemma3ForCausalLM):
hidden_states = self.model(
input_ids, positions, forward_batch, input_embeds, **kwargs
)
return self.pooler(hidden_states, forward_batch)
pooled = self.pooler(hidden_states, forward_batch).embeddings
if self.projector is not None:
pooled = self.projector(pooled)
return EmbeddingPoolerOutput(embeddings=F.normalize(pooled, p=2, dim=-1))
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
"""Load both native Gemma3 and Sentence Transformers checkpoints.
@@ -998,7 +1051,37 @@ class EmbeddingGemmaModel(Gemma3ForCausalLM):
for name, weight in weights
if name.startswith("model.") or name.startswith(backbone_prefixes)
)
return super().load_weights(remapped_weights)
loaded_params = super().load_weights(remapped_weights)
if self.projector is not None:
model_path = getattr(self.config, "_name_or_path", "")
modules_path = _resolve_local_or_cached_file(model_path, "modules.json")
with open(modules_path) as f:
module_specs = json.load(f)
dense_specs = [
spec
for spec in module_specs
if spec.get("type") == "sentence_transformers.models.Dense"
]
from safetensors.torch import load_file
for layer, spec in zip(self.projector, dense_specs):
weights_path = _resolve_local_or_cached_file(
model_path, f"{spec['path']}/model.safetensors"
)
weights = load_file(weights_path, device="cpu")
weight_key = next(
key
for key in ("weight", "linear.weight", "dense.weight")
if key in weights
)
layer.weight.data.copy_(weights[weight_key].to(layer.weight.device))
if layer.bias is not None:
layer.bias.data.copy_(
weights[weight_key.replace("weight", "bias")].to(
layer.bias.device
)
)
return loaded_params
EntryClass = [Gemma3ForCausalLM, EmbeddingGemmaModel]
python/sglang/srt/server_args.py
+48 -11
View File
@@ -3644,23 +3644,60 @@ class ServerArgs:
# Breakable CUDA Graph captures one complete prefill and is the graph
# mode validated for this encoder-style attention.
if getattr(model_config, "is_embedding_gemma", False):
# This is an encoder-only model even though its HF architecture is
# named Gemma3TextModel. Marking it as embedding mode enables the
# FlashAttention raw-K/V fast path, which does not write or read
# the paged KV cache during its single prefill forward.
self.is_embedding = True
self.disable_radix_cache = True
self.chunked_prefill_size = -1
# Submit a list-valued embeddings request atomically so BCG can
# replay its full prefill batch instead of starting item zero
# while the remaining texts are still being tokenized.
self.enable_tokenizer_batch_encode = True
requested_prefill_backend = (
self.prefill_attention_backend or self.attention_backend
)
if (
is_cuda()
and (is_sm90_supported() or is_sm100_supported())
and requested_prefill_backend in (None, "fa3", "fa4")
):
# Hopper/Blackwell's default FA backend can consume raw K/V
# tensors for a single embedding prefill. Enable its no-KV
# pool path before memory-pool sizing; an explicit non-FA
# backend retains the existing paged-KV behavior.
self.prefill_only_disable_kv_cache = True
self._validate_prefill_only_disable_kv_cache_args()
self.cuda_graph_config.decode.backend = Backend.DISABLED
if is_cuda() and self.cuda_graph_config.prefill.backend != Backend.DISABLED:
self.cuda_graph_config.prefill.backend = Backend.BREAKABLE
# CUDA-graph sizing has already run by this point. With
# chunked prefill disabled its generic default is -1, which
# otherwise leaves BCG with no shapes to capture. Use the
# model's maximum request length as the safe default; callers
# can still raise it for larger aggregate prefill batches.
if (self.cuda_graph_config.prefill.max_bs or 0) <= 0:
self.cuda_graph_config.prefill.max_bs = model_config.context_len
self.cuda_graph_config.prefill.bs = (
self._generate_prefill_cuda_graph_batch_sizes(
model_config.context_len
)
# CUDA-graph sizing has already run by this point and derives
# its generic maximum from the 8K chunked-prefill default.
# On the Hopper/Blackwell FA raw-K/V path, raise the unlocked
# default to a full eight-way 2K embedding batch; callers can
# still override this for larger aggregate prefills.
prefill_config = self.cuda_graph_config.prefill
# Unit-level capability tests may invoke this hook without
# running the full CUDA-graph configuration parser, which is
# where this internal lock set is normally initialized.
# Treat that minimal construction as having no user-locked
# graph settings.
cuda_graph_config_locked = getattr(
self, "_cuda_graph_config_locked", set()
)
if (Phase.PREFILL, "max_bs") not in cuda_graph_config_locked:
prefill_config.max_bs = max(
prefill_config.max_bs or 0,
model_config.context_len,
16384,
)
if (Phase.PREFILL, "bs") not in cuda_graph_config_locked:
prefill_config.bs = (
self._generate_prefill_cuda_graph_batch_sizes(
prefill_config.max_bs
)
)
elif not is_cuda():
# BCG is CUDA-only. Other graph backends do not support this
# encoder-style prefill, so retain the eager Triton path.