From a98b456c706febe45ed6e8b80b69ee17dce7a807 Mon Sep 17 00:00:00 2001 From: blzheng Date: Thu, 19 Mar 2026 14:02:26 +0800 Subject: [PATCH] [CPU] Add frontend support for Gemma (#12590) --- python/sglang/lang/chat_template.py | 2 + python/sglang/srt/models/gemma3_causal.py | 108 ++++++++++++++++++---- python/sglang/srt/models/siglip.py | 2 +- 3 files changed, 92 insertions(+), 20 deletions(-) diff --git a/python/sglang/lang/chat_template.py b/python/sglang/lang/chat_template.py index 212d07e0b..850270308 100644 --- a/python/sglang/lang/chat_template.py +++ b/python/sglang/lang/chat_template.py @@ -398,6 +398,8 @@ register_chat_template( "user": ("user\n", "\n"), "assistant": ("model\n", "\n"), }, + image_token="", + audio_token="", style=ChatTemplateStyle.PLAIN, ) ) diff --git a/python/sglang/srt/models/gemma3_causal.py b/python/sglang/srt/models/gemma3_causal.py index acd6def10..0481cae0e 100644 --- a/python/sglang/srt/models/gemma3_causal.py +++ b/python/sglang/srt/models/gemma3_causal.py @@ -35,14 +35,17 @@ from sglang.srt.layers.linear import ( from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import AttentionType, RadixAttention -from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb +from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb, get_rope from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) -from sglang.srt.utils import add_prefix, make_layers +from sglang.srt.utils import add_prefix, cpu_has_amx_support, is_cpu, make_layers + +_is_cpu = is_cpu() +_is_cpu_amx_available = cpu_has_amx_support() # Aligned with HF's implementation, using sliding window inclusive with the last token @@ -141,7 +144,8 @@ class Gemma3Attention(nn.Module): config, "head_dim", hidden_size // config.num_attention_heads ) self.head_dim = head_dim - + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) + self.rotary_dim = int(partial_rotary_factor * self.head_dim) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim @@ -197,7 +201,14 @@ class Gemma3Attention(nn.Module): ) self.rope_scaling = {"rope_type": "default"} self.sliding_window = None - + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=max_position_embeddings, + base=self.rope_theta, + rope_scaling=self.rope_scaling, + is_neox_style=getattr(config, "rope_is_neox_style", True), + ) self.attn = RadixAttention( self.num_heads, self.head_dim, @@ -217,8 +228,39 @@ class Gemma3Attention(nn.Module): self.q_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps) self.k_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps) - def forward( + def forward_cpu( self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + forward_batch: ForwardBatch, + **kwargs, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + # [s, h * head_dim] + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + # [s, h, head_dim] + q = q.unflatten(-1, (self.num_heads, self.head_dim)).unsqueeze(0) + q = self.q_norm(q) + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)).unsqueeze(0) + k = self.k_norm(k) + q, k = self.rotary_emb(positions, q, k) + + attn_output = self.attn(q, k, v, forward_batch=forward_batch) + + # Compatible with triton backend which returns [1, s, h, head_dim] + if attn_output.dim() == 4 and attn_output.shape[0] == 1: + attn_output = attn_output.squeeze(0) + attn_output = attn_output.flatten(-2, -1) + # [s, h * head_dim] + + output, _ = self.o_proj(attn_output) + return output + + def forward_native( + self, + positions: torch.Tensor, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], forward_batch: ForwardBatch, @@ -257,6 +299,22 @@ class Gemma3Attention(nn.Module): output, _ = self.o_proj(attn_output) return output + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + forward_batch: ForwardBatch, + **kwargs, + ) -> torch.Tensor: + if _is_cpu and _is_cpu_amx_available: + return self.forward_cpu( + positions, hidden_states, position_embeddings, forward_batch, **kwargs + ) + return self.forward_native( + positions, hidden_states, position_embeddings, forward_batch, **kwargs + ) + class Gemma3DecoderLayer(nn.Module): def __init__( @@ -555,21 +613,33 @@ class Gemma3TextModel(PreTrainedModel): else: hidden_states = input_embeds - if positions.dim() == 1: - positions = einops.rearrange(positions, "s -> 1 s") + if _is_cpu and _is_cpu_amx_available: + for layer in self.layers: + layer_outputs = layer( + positions=positions, + position_embeddings_global=None, + position_embeddings_local=None, + hidden_states=hidden_states, + forward_batch=forward_batch, + **kwargs, + ) + hidden_states = layer_outputs[0] + else: + if positions.dim() == 1: + positions = einops.rearrange(positions, "s -> 1 s") - position_embeddings_global = self.rotary_emb(hidden_states, positions) - position_embeddings_local = self.rotary_emb_local(hidden_states, positions) - for layer in self.layers: - layer_outputs = layer( - positions=positions, - position_embeddings_global=position_embeddings_global, - position_embeddings_local=position_embeddings_local, - hidden_states=hidden_states, - forward_batch=forward_batch, - **kwargs, - ) - hidden_states = layer_outputs[0] + position_embeddings_global = self.rotary_emb(hidden_states, positions) + position_embeddings_local = self.rotary_emb_local(hidden_states, positions) + for layer in self.layers: + layer_outputs = layer( + positions=positions, + position_embeddings_global=position_embeddings_global, + position_embeddings_local=position_embeddings_local, + hidden_states=hidden_states, + forward_batch=forward_batch, + **kwargs, + ) + hidden_states = layer_outputs[0] hidden_states = self.norm(hidden_states) diff --git a/python/sglang/srt/models/siglip.py b/python/sglang/srt/models/siglip.py index 989fcd9fc..bd57dc581 100644 --- a/python/sglang/srt/models/siglip.py +++ b/python/sglang/srt/models/siglip.py @@ -51,7 +51,7 @@ class SiglipVisionEmbeddings(nn.Module): patch_embeds = self.patch_embedding( pixel_values.to(dtype=target_dtype) ) # shape = [*, width, grid, grid] - embeddings = patch_embeds.flatten(2).transpose(1, 2) + embeddings = patch_embeds.flatten(2).transpose(1, 2).contiguous() # interpolate_pos_encoding is never used in sglang embeddings = embeddings + self.position_embedding(self.position_ids)