# Copyright 2026 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== from __future__ import annotations import copy import logging from typing import Dict, Iterable, Optional, Tuple import torch from torch import nn from transformers import PretrainedConfig, PreTrainedModel from sglang.srt.distributed import get_pp_group from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.logits_processor import ( LogitsMetadata, LogitsProcessor, LogitsProcessorOutput, ) from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.mem_cache.memory_pool import KVCache from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.models.gemma4_causal import Gemma4ForCausalLM, Gemma4TextModel from sglang.srt.speculative.frozen_kv_mtp_info import FrozenKVMTPContext from sglang.srt.utils import add_prefix logger = logging.getLogger(__name__) def _get_text_config(model_or_config) -> PretrainedConfig: """Normalize either a model or a (possibly wrapped) config to ``Gemma4TextConfig``.""" cfg = getattr(model_or_config, "config", model_or_config) return getattr(cfg, "text_config", cfg) def _resolve_target_text_model(target_model): for attr in ("language_model", "model"): candidate = getattr(target_model, attr, None) if candidate is not None and hasattr(candidate, "layers"): return candidate raise AttributeError( f"Frozen-KV MTP cannot locate the target trunk on " f"{type(target_model).__name__}; expected ``.language_model`` " "(multimodal) or ``.model`` (text-only) with a ``.layers`` attribute." ) class Gemma4AssistantForCausalLM(Gemma4ForCausalLM): """Gemma 4 MTP assistant: target embed + recurrent hidden through pre/post projection; own ``lm_head``.""" base_model_prefix = "model" def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: text_config = copy.deepcopy(_get_text_config(config)) text_config.num_kv_shared_layers = 0 PreTrainedModel.__init__(self, config=text_config) self.assistant_config = config self.config = text_config self.quant_config = quant_config self.pp_group = get_pp_group() self.vocab_size = text_config.vocab_size self.hidden_size = text_config.hidden_size self.backbone_hidden_size = config.backbone_hidden_size self.target_embed_scale = self.backbone_hidden_size**0.5 self.use_ordered_embeddings = bool( getattr(config, "use_ordered_embeddings", False) ) self.centroid_intermediate_top_k = int( getattr(config, "centroid_intermediate_top_k", 32) ) self.target_embed_weight: Optional[torch.Tensor] = None self.pre_projection = ReplicatedLinear( 2 * self.backbone_hidden_size, self.hidden_size, bias=False, quant_config=None, prefix=add_prefix("pre_projection", prefix), ) self.model = Gemma4TextModel( config=text_config, quant_config=quant_config, prefix=add_prefix("model", prefix), ) self.post_projection = ReplicatedLinear( self.hidden_size, self.backbone_hidden_size, bias=False, quant_config=None, prefix=add_prefix("post_projection", prefix), ) if text_config.tie_word_embeddings: self.lm_head = self.model.embed_tokens else: self.lm_head = nn.Linear(self.hidden_size, self.vocab_size, bias=False) self.logits_processor = LogitsProcessor(text_config, skip_all_gather=True) if self.use_ordered_embeddings: self.num_centroids = int(config.num_centroids) self.vocab_size_per_centroid, rem = divmod( self.vocab_size, self.num_centroids ) if rem: raise ValueError( "Frozen-KV MTP centroid head requires vocab_size to be a " f"multiple of num_centroids (vocab={self.vocab_size}, " f"num_centroids={self.num_centroids})." ) self.centroids = nn.Linear(self.hidden_size, self.num_centroids, bias=False) self.register_buffer( "token_ordering", torch.zeros(self.vocab_size, dtype=torch.long), persistent=True, ) else: self.num_centroids = self.vocab_size_per_centroid = self.centroids = None self.register_buffer("token_ordering", None, persistent=False) self.kv_context: Optional[FrozenKVMTPContext] = None self.post_init() def bind_frozen_kv_context(self, ctx: FrozenKVMTPContext) -> None: """Bind assistant attention to target-owned KV and suppress assistant KV writes.""" for assistant_logical, layer in enumerate(self.model.layers): target_phys = ctx.get_physical_layer_id(assistant_logical) layer.self_attn.is_kv_shared_layer = True layer.self_attn.kv_shared_layer_index = target_phys layer.self_attn.attn.layer_id = target_phys layer.self_attn.layer_id = assistant_logical self.kv_context = ctx def build_frozen_kv_mtp_context( self, target_model, target_token_to_kv_pool: KVCache, ) -> FrozenKVMTPContext: """Map each assistant layer to the target physical layer that owns its K/V. HF Gemma 4 ties each typed (sliding/full) assistant layer to the target's last layer of the same type; that layer is itself KV-shared with an earlier non-shared layer (via ``kv_shared_layer_index``). We collapse those two hops once so attention can hand a direct ``layer_id`` to ``RadixAttention`` at bind time. """ target_text = _get_text_config(target_model) assistant_text = _get_text_config(self) layers = _resolve_target_text_model(target_model).layers def kv_owner(idx: int) -> int: attn = layers[idx].self_attn owner = ( getattr(attn, "kv_shared_layer_index", None) if getattr(attn, "is_kv_shared_layer", False) else idx ) if owner is None or getattr( layers[owner].self_attn, "is_kv_shared_layer", False ): raise RuntimeError( f"Frozen-KV MTP: target layer {idx} resolved to physical " f"{owner!r}, which is missing or itself KV-shared " "(HF invariant changed?)." ) return owner L = target_text.num_hidden_layers by_type = {target_text.layer_types[i]: kv_owner(i) for i in (L - 2, L - 1)} physical: Dict[int, int] = {} for i, t in enumerate(assistant_text.layer_types): if t not in by_type: raise ValueError( f"Frozen-KV MTP assistant layer {i} has type {t!r}, " f"expected one of {sorted(by_type)}." ) physical[i] = by_type[t] return FrozenKVMTPContext( target_token_to_kv_pool=target_token_to_kv_pool, physical_layer_ids=physical, ) def get_embed_and_head(self) -> Tuple[torch.Tensor, torch.Tensor]: if self.target_embed_weight is None: raise RuntimeError( "Gemma4AssistantForCausalLM target embedding is not bound yet." ) return self.target_embed_weight, self.lm_head.weight def set_embed_and_head(self, embed: torch.Tensor, head: torch.Tensor) -> None: """Rebind target embedding; ``head`` ignored (assistant keeps ``lm_head``).""" del head self.target_embed_weight = embed if torch.cuda.is_available(): torch.cuda.empty_cache() def get_attention_sliding_window_size(self) -> int: # Gemma 4 config treats the bound as inclusive; SGLang attention metadata # uses an exclusive window size, matching the target Gemma 4 models. return self.config.sliding_window - 1 @torch.no_grad() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: Optional[torch.Tensor] = None, **kwargs, ) -> LogitsProcessorOutput: if input_embeds is None: if self.target_embed_weight is None: raise RuntimeError( "Gemma4AssistantForCausalLM requires set_embed_and_head() " "before token-id forward." ) token_embed = ( torch.nn.functional.embedding(input_ids, self.target_embed_weight) * self.target_embed_scale ) else: token_embed = input_embeds if forward_batch.spec_info is None or not hasattr( forward_batch.spec_info, "hidden_states" ): raise RuntimeError( "Frozen-KV MTP forward requires forward_batch.spec_info." "hidden_states to carry the recurrent state. The worker's " "_frozen_kv_target_view context manager must be exited " "before model forward, leaving spec_info populated." ) prev_hidden = forward_batch.spec_info.hidden_states if token_embed.shape != prev_hidden.shape: raise ValueError( "Frozen-KV MTP forward: token_embed and prev_hidden must have " f"the same shape (got {token_embed.shape} vs {prev_hidden.shape})." ) z, _ = self.pre_projection(torch.cat([token_embed, prev_hidden], dim=-1)) hidden_states = self.model( input_ids=None, positions=positions, forward_batch=forward_batch, input_embeds=z, per_layer_inputs=None, **kwargs, ) projected_states, _ = self.post_projection(hidden_states) if self.use_ordered_embeddings: return self._centroid_logits_processor( input_ids, hidden_states, projected_states, forward_batch ) return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch, hidden_states_before_norm=projected_states, ) def _apply_centroid_masking(self, hidden_states: torch.Tensor) -> torch.Tensor: """Centroid-masked logits for E2B/E4B assistant heads.""" if self.centroids is None or self.token_ordering is None: raise RuntimeError( "Frozen-KV MTP centroid head invoked but centroid weights " "are not initialized." ) prefix_shape = hidden_states.shape[:-1] flat_hidden = hidden_states.reshape(-1, hidden_states.shape[-1]) num_tokens = flat_hidden.shape[0] _, top_k_indices = torch.topk( self.centroids(flat_hidden), k=self.centroid_intermediate_top_k, dim=-1, ) # Contiguous gather: [C, vpc, H] indexed by centroid IDs. num_selected = self.centroid_intermediate_top_k * self.vocab_size_per_centroid selected_embeddings = self.lm_head.weight.view( self.num_centroids, self.vocab_size_per_centroid, self.hidden_size, )[top_k_indices].reshape(num_tokens, num_selected, self.hidden_size) selected_logits = torch.bmm( flat_hidden.unsqueeze(1), selected_embeddings.transpose(1, 2), ).squeeze(1) # Scatter to real vocab positions via token_ordering. centroid_vocab_indices = ( self.token_ordering.long() .view(self.num_centroids, self.vocab_size_per_centroid)[top_k_indices] .view(num_tokens, -1) ) mask_value = torch.finfo(selected_logits.dtype).min / 2 output = torch.full( (num_tokens, self.vocab_size), mask_value, dtype=selected_logits.dtype, device=selected_logits.device, ) output.scatter_(dim=-1, index=centroid_vocab_indices, src=selected_logits) return output.view(*prefix_shape, self.vocab_size) def _centroid_logits_processor( self, input_ids: torch.Tensor, hidden_states: torch.Tensor, projected_states: torch.Tensor, forward_batch: ForwardBatch, ) -> LogitsProcessorOutput: logits_metadata = LogitsMetadata.from_forward_batch(forward_batch) if logits_metadata.extend_return_logprob: raise NotImplementedError( "Frozen-KV MTP centroid head does not support input logprobs yet." ) ( pruned_states, pruned_states_before_norm, aux_pruned_states, sample_indices, *_, ) = self.logits_processor._get_pruned_states( hidden_states, projected_states, None, logits_metadata ) hidden_states_to_store = self.logits_processor._get_hidden_states_to_store( hidden_states, projected_states, None, pruned_states, pruned_states_before_norm, aux_pruned_states, sample_indices, logits_metadata, ) del input_ids, hidden_states, projected_states logits = self._apply_centroid_masking(pruned_states) sampled_logits = ( logits[sample_indices] if sample_indices is not None else logits ) return LogitsProcessorOutput( next_token_logits=sampled_logits, hidden_states=hidden_states_to_store, mm_input_embeds=logits_metadata.mm_input_embeds, ) def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): def remap_assistant_weights(): for name, weight in weights: if name.startswith("masked_embedding."): name = name.removeprefix("masked_embedding.") yield name, weight result = super().load_weights(remap_assistant_weights()) if self.use_ordered_embeddings: self._reorder_embedding_to_centroid_order() return result @torch.no_grad() def _reorder_embedding_to_centroid_order(self) -> None: """Reorder lm_head.weight from natural vocab order to centroid order.""" if self.token_ordering is None: return ordering = self.token_ordering.long() lm_head_w = self.lm_head.weight reordered = lm_head_w.data[ordering] lm_head_w.data.copy_(reordered) logger.info( "Reordered lm_head/embed_tokens (%s) to centroid order " "for contiguous centroid masking.", list(lm_head_w.shape), ) class Gemma4UnifiedAssistantForCausalLM(Gemma4AssistantForCausalLM): """Gemma 4 unified MTP assistant; text path identical to the gemma4 assistant.""" EntryClass = [Gemma4AssistantForCausalLM, Gemma4UnifiedAssistantForCausalLM]