| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inference-only Gemma4 MTP (Multi-Token Prediction) model. |
| 4 | |
| 5 | The Gemma4 assistant model is a lightweight decoder that shares KV cache |
| 6 | with the target (backbone) model. All assistant decoder layers are |
| 7 | KV-shared: they only have Q projections (no K/V projections or norms), |
| 8 | and read K/V from the target model's cache at runtime. |
| 9 | |
| 10 | Checkpoint layout (``gemma4_assistant``):: |
| 11 | |
| 12 | model.embed_tokens.* -- token embeddings |
| 13 | model.layers.{i}.* -- decoder layers (Q-only attention + MLP) |
| 14 | model.norm.* -- final RMSNorm |
| 15 | pre_projection.* -- Linear(2 * backbone_hidden_size, hidden_size) |
| 16 | post_projection.* -- Linear(hidden_size, backbone_hidden_size) |
| 17 | lm_head.* -- language model head (tied to embed_tokens) |
| 18 | masked_embedding.centroids.* -- centroid projection (when use_ordered_embeddings) |
| 19 | masked_embedding.token_ordering -- token-to-centroid mapping buffer |
| 20 | """ |
| 21 | |
| 22 | from collections.abc import Iterable |
| 23 | |
| 24 | import torch |
| 25 | from torch import nn |
| 26 | |
| 27 | from vllm.compilation.decorators import support_torch_compile |
| 28 | from vllm.config import CacheConfig, VllmConfig |
| 29 | from vllm.distributed import ( |
| 30 | get_tensor_model_parallel_world_size, |
| 31 | tensor_model_parallel_all_gather, |
| 32 | ) |
| 33 | from vllm.logger import init_logger |
| 34 | from vllm.model_executor.layers.attention import Attention |
| 35 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 36 | from vllm.model_executor.layers.linear import ( |
| 37 | ColumnParallelLinear, |
| 38 | RowParallelLinear, |
| 39 | ) |
| 40 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 41 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 42 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 43 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 44 | ParallelLMHead, |
| 45 | VocabParallelEmbedding, |
| 46 | ) |
| 47 | from vllm.sequence import IntermediateTensors |
| 48 | from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config |
| 49 | |
| 50 | from .gemma4 import Gemma4MLP, _get_text_config |
| 51 | from .utils import ( |
| 52 | AutoWeightsLoader, |
| 53 | WeightsMapper, |
| 54 | extract_layer_index, |
| 55 | get_draft_quant_config, |
| 56 | maybe_prefix, |
| 57 | ) |
| 58 | |
| 59 | logger = init_logger(__name__) |
| 60 | |
| 61 | |
| 62 | class Gemma4MTPMaskedEmbedder(nn.Module): |
| 63 | """Sparse logit computation via centroid-based vocabulary masking. |
| 64 | |
| 65 | Instead of computing logits against the full vocabulary, projects |
| 66 | hidden states to centroid scores, selects top-K centroids, and |
| 67 | computes logits only for the ~top_k * (vocab_size / num_centroids) |
| 68 | tokens belonging to those centroids. |
| 69 | """ |
| 70 | |
| 71 | token_ordering: torch.Tensor |
| 72 | |
| 73 | def __init__( |
| 74 | self, |
| 75 | hidden_size: int, |
| 76 | vocab_size: int, |
| 77 | num_centroids: int, |
| 78 | centroid_intermediate_top_k: int, |
| 79 | ) -> None: |
| 80 | super().__init__() |
| 81 | self.hidden_size = hidden_size |
| 82 | self.vocab_size = vocab_size |
| 83 | self.num_centroids = num_centroids |
| 84 | self.centroid_intermediate_top_k = centroid_intermediate_top_k |
| 85 | self.vocab_size_per_centroid = vocab_size // num_centroids |
| 86 | self.num_selected = centroid_intermediate_top_k * self.vocab_size_per_centroid |
| 87 | |
| 88 | self.centroids = nn.Linear(hidden_size, num_centroids, bias=False) |
| 89 | self.register_buffer( |
| 90 | "token_ordering", |
| 91 | torch.empty(vocab_size, dtype=torch.long), |
| 92 | ) |
| 93 | |
| 94 | def _select_and_score( |
| 95 | self, |
| 96 | hidden_states: torch.Tensor, |
| 97 | lm_head_weight: torch.Tensor, |
| 98 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 99 | """Centroid selection + sparse dot product. |
| 100 | |
| 101 | Returns: |
| 102 | logits: (num_tokens, num_selected) sparse logits. |
| 103 | indices: (num_tokens, num_selected) corresponding vocab indices. |
| 104 | |
| 105 | """ |
| 106 | num_tokens = hidden_states.shape[0] |
| 107 | _, top_k_indices = torch.topk( |
| 108 | self.centroids(hidden_states), |
| 109 | k=self.centroid_intermediate_top_k, |
| 110 | dim=-1, |
| 111 | ) |
| 112 | clusters = self.token_ordering.view( |
| 113 | self.num_centroids, |
| 114 | self.vocab_size_per_centroid, |
| 115 | ) |
| 116 | selected = clusters[top_k_indices] |
| 117 | embeddings = lm_head_weight[selected.reshape(-1)].view( |
| 118 | num_tokens, |
| 119 | self.num_selected, |
| 120 | self.hidden_size, |
| 121 | ) |
| 122 | logits = torch.einsum("td,tsd->ts", hidden_states, embeddings) |
| 123 | return logits, selected.view(num_tokens, -1) |
| 124 | |
| 125 | def forward( |
| 126 | self, |
| 127 | hidden_states: torch.Tensor, |
| 128 | lm_head_weight: torch.Tensor, |
| 129 | ) -> torch.Tensor: |
| 130 | """Full-vocab logits with non-selected positions masked to -inf.""" |
| 131 | logits, indices = self._select_and_score(hidden_states, lm_head_weight) |
| 132 | output = torch.full( |
| 133 | (hidden_states.shape[0], self.vocab_size), |
| 134 | fill_value=torch.finfo(hidden_states.dtype).min, |
| 135 | dtype=hidden_states.dtype, |
| 136 | device=hidden_states.device, |
| 137 | ) |
| 138 | return output.scatter_(-1, indices, logits) |
| 139 | |
| 140 | def get_top_tokens( |
| 141 | self, |
| 142 | hidden_states: torch.Tensor, |
| 143 | lm_head_weight: torch.Tensor, |
| 144 | ) -> torch.Tensor: |
| 145 | """Sparse argmax — returns vocab token IDs without full-vocab tensor.""" |
| 146 | logits, indices = self._select_and_score(hidden_states, lm_head_weight) |
| 147 | return indices.gather(-1, logits.argmax(-1, keepdim=True)).squeeze(-1) |
| 148 | |
| 149 | |
| 150 | class Gemma4MTPAttention(nn.Module): |
| 151 | """Q-only attention for Gemma4 MTP layers. |
| 152 | |
| 153 | K/V come from the target model's KV cache via |
| 154 | ``kv_sharing_target_layer_name`` (set by the proposer after |
| 155 | model construction). |
| 156 | """ |
| 157 | |
| 158 | def __init__( |
| 159 | self, |
| 160 | config, |
| 161 | hidden_size: int, |
| 162 | num_heads: int, |
| 163 | num_kv_heads: int, |
| 164 | head_dim: int, |
| 165 | max_position_embeddings: int, |
| 166 | cache_config: CacheConfig | None = None, |
| 167 | quant_config: QuantizationConfig | None = None, |
| 168 | attn_logits_soft_cap: float | None = None, |
| 169 | prefix: str = "", |
| 170 | ) -> None: |
| 171 | super().__init__() |
| 172 | self.config = config |
| 173 | self.hidden_size = hidden_size |
| 174 | |
| 175 | tp_size = get_tensor_model_parallel_world_size() |
| 176 | self.total_num_heads = num_heads |
| 177 | self.num_heads = self.total_num_heads // tp_size |
| 178 | self.total_num_kv_heads = num_kv_heads |
| 179 | self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) |
| 180 | self.head_dim = head_dim |
| 181 | self.q_size = self.num_heads * self.head_dim |
| 182 | self.scaling = 1.0 |
| 183 | |
| 184 | self.q_proj = ColumnParallelLinear( |
| 185 | hidden_size, |
| 186 | self.total_num_heads * self.head_dim, |
| 187 | bias=config.attention_bias, |
| 188 | quant_config=quant_config, |
| 189 | prefix=f"{prefix}.q_proj", |
| 190 | ) |
| 191 | self.o_proj = RowParallelLinear( |
| 192 | self.total_num_heads * self.head_dim, |
| 193 | hidden_size, |
| 194 | bias=config.attention_bias, |
| 195 | quant_config=quant_config, |
| 196 | prefix=f"{prefix}.o_proj", |
| 197 | ) |
| 198 | self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) |
| 199 | |
| 200 | layer_idx = extract_layer_index(prefix) |
| 201 | layer_type = config.layer_types[layer_idx] |
| 202 | self.is_sliding = layer_type == "sliding_attention" |
| 203 | sliding_window = config.sliding_window if self.is_sliding else None |
| 204 | |
| 205 | if layer_type in config.rope_parameters: |
| 206 | rope_parameters = dict(config.rope_parameters[layer_type]) |
| 207 | else: |
| 208 | rope_parameters = dict(config.rope_parameters.copy()) |
| 209 | if self.is_sliding: |
| 210 | rope_parameters["rope_theta"] = getattr( |
| 211 | config, "rope_local_base_freq", 10000.0 |
| 212 | ) |
| 213 | |
| 214 | self.rotary_emb = get_rope( |
| 215 | self.head_dim, |
| 216 | max_position=max_position_embeddings, |
| 217 | rope_parameters=rope_parameters, |
| 218 | is_neox_style=True, |
| 219 | ) |
| 220 | |
| 221 | # kv_sharing_target_layer_name is set after model construction |
| 222 | # by Gemma4Proposer._setup_gemma4_kv_sharing(). |
| 223 | self.is_kv_shared_layer = True |
| 224 | self.attn = Attention( |
| 225 | self.num_heads, |
| 226 | self.head_dim, |
| 227 | self.scaling, |
| 228 | num_kv_heads=self.num_kv_heads, |
| 229 | cache_config=cache_config, |
| 230 | quant_config=quant_config, |
| 231 | logits_soft_cap=attn_logits_soft_cap, |
| 232 | per_layer_sliding_window=sliding_window, |
| 233 | prefix=f"{prefix}.attn", |
| 234 | ) |
| 235 | |
| 236 | def forward( |
| 237 | self, |
| 238 | positions: torch.Tensor, |
| 239 | hidden_states: torch.Tensor, |
| 240 | **kwargs, |
| 241 | ) -> torch.Tensor: |
| 242 | q, _ = self.q_proj(hidden_states) |
| 243 | |
| 244 | q = q.unflatten(-1, (self.num_heads, self.head_dim)) |
| 245 | q = self.q_norm(q) |
| 246 | q = q.flatten(-2, -1) |
| 247 | |
| 248 | q, _ = self.rotary_emb(positions, q, None) |
| 249 | |
| 250 | # Attention reads K/V from the target's cache via KV sharing. |
| 251 | attn_output = self.attn(q, None, None) |
| 252 | output, _ = self.o_proj(attn_output) |
| 253 | return output |
| 254 | |
| 255 | |
| 256 | class Gemma4MTPDecoderLayer(nn.Module): |
| 257 | def __init__( |
| 258 | self, |
| 259 | config, |
| 260 | cache_config: CacheConfig | None = None, |
| 261 | quant_config: QuantizationConfig | None = None, |
| 262 | prefix: str = "", |
| 263 | ) -> None: |
| 264 | super().__init__() |
| 265 | self.hidden_size = config.hidden_size |
| 266 | |
| 267 | layer_idx = extract_layer_index(prefix) |
| 268 | layer_config = gemma4_layer_config(config, layer_idx) |
| 269 | head_dim = layer_config.head_dim |
| 270 | num_kv_heads = layer_config.num_key_value_heads |
| 271 | |
| 272 | self.self_attn = Gemma4MTPAttention( |
| 273 | config=config, |
| 274 | hidden_size=self.hidden_size, |
| 275 | num_heads=config.num_attention_heads, |
| 276 | num_kv_heads=num_kv_heads, |
| 277 | head_dim=head_dim, |
| 278 | max_position_embeddings=config.max_position_embeddings, |
| 279 | cache_config=cache_config, |
| 280 | quant_config=quant_config, |
| 281 | attn_logits_soft_cap=getattr(config, "attn_logit_softcapping", None), |
| 282 | prefix=f"{prefix}.self_attn", |
| 283 | ) |
| 284 | |
| 285 | text_config = _get_text_config(config) |
| 286 | self.mlp = Gemma4MLP( |
| 287 | hidden_size=self.hidden_size, |
| 288 | intermediate_size=text_config.intermediate_size, |
| 289 | hidden_activation=text_config.hidden_activation, |
| 290 | quant_config=quant_config, |
| 291 | prefix=f"{prefix}.mlp", |
| 292 | ) |
| 293 | |
| 294 | self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 295 | self.post_attention_layernorm = RMSNorm( |
| 296 | config.hidden_size, eps=config.rms_norm_eps |
| 297 | ) |
| 298 | self.pre_feedforward_layernorm = RMSNorm( |
| 299 | config.hidden_size, eps=config.rms_norm_eps |
| 300 | ) |
| 301 | self.post_feedforward_layernorm = RMSNorm( |
| 302 | config.hidden_size, eps=config.rms_norm_eps |
| 303 | ) |
| 304 | |
| 305 | self.register_buffer("layer_scalar", torch.ones(1)) |
| 306 | |
| 307 | def forward( |
| 308 | self, |
| 309 | positions: torch.Tensor, |
| 310 | hidden_states: torch.Tensor, |
| 311 | residual: torch.Tensor | None, |
| 312 | **kwargs, |
| 313 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 314 | residual = hidden_states |
| 315 | hidden_states = self.input_layernorm(residual) |
| 316 | |
| 317 | hidden_states = self.self_attn( |
| 318 | positions=positions, |
| 319 | hidden_states=hidden_states, |
| 320 | **kwargs, |
| 321 | ) |
| 322 | |
| 323 | hidden_states = self.post_attention_layernorm(hidden_states) |
| 324 | hidden_states = hidden_states + residual |
| 325 | residual = hidden_states |
| 326 | |
| 327 | hidden_states = self.pre_feedforward_layernorm(hidden_states) |
| 328 | hidden_states = self.mlp(hidden_states) |
| 329 | |
| 330 | hidden_states = self.post_feedforward_layernorm(hidden_states) |
| 331 | hidden_states = hidden_states + residual |
| 332 | |
| 333 | hidden_states = hidden_states * self.layer_scalar |
| 334 | return hidden_states, None |
| 335 | |
| 336 | |
| 337 | class Gemma4MultiTokenPredictor(nn.Module): |
| 338 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 339 | super().__init__() |
| 340 | |
| 341 | speculative_config = vllm_config.speculative_config |
| 342 | assert speculative_config is not None |
| 343 | config = speculative_config.draft_model_config.hf_config |
| 344 | text_config = _get_text_config(config) |
| 345 | quant_config = get_draft_quant_config(vllm_config) |
| 346 | self.config = text_config |
| 347 | self.quant_config = quant_config |
| 348 | |
| 349 | self.hidden_size = text_config.hidden_size |
| 350 | self.backbone_hidden_size = getattr( |
| 351 | config, "backbone_hidden_size", self.hidden_size |
| 352 | ) |
| 353 | self.vocab_size = text_config.vocab_size |
| 354 | self.num_mtp_layers = text_config.num_hidden_layers |
| 355 | |
| 356 | self.embed_tokens = VocabParallelEmbedding( |
| 357 | self.vocab_size, |
| 358 | self.hidden_size, |
| 359 | quant_config=quant_config, |
| 360 | prefix=f"{prefix}.embed_tokens", |
| 361 | ) |
| 362 | |
| 363 | self.pre_projection = ColumnParallelLinear( |
| 364 | 2 * self.backbone_hidden_size, |
| 365 | self.hidden_size, |
| 366 | bias=False, |
| 367 | gather_output=True, |
| 368 | quant_config=quant_config, |
| 369 | prefix=f"{prefix}.pre_projection", |
| 370 | ) |
| 371 | |
| 372 | self.post_projection = RowParallelLinear( |
| 373 | self.hidden_size, |
| 374 | self.backbone_hidden_size, |
| 375 | bias=False, |
| 376 | input_is_parallel=False, |
| 377 | quant_config=quant_config, |
| 378 | prefix=f"{prefix}.post_projection", |
| 379 | ) |
| 380 | |
| 381 | self.layers = nn.ModuleList( |
| 382 | Gemma4MTPDecoderLayer( |
| 383 | text_config, |
| 384 | cache_config=vllm_config.cache_config, |
| 385 | quant_config=quant_config, |
| 386 | prefix=f"{prefix}.layers.{idx}", |
| 387 | ) |
| 388 | for idx in range(self.num_mtp_layers) |
| 389 | ) |
| 390 | |
| 391 | self.norm = RMSNorm(self.hidden_size, eps=text_config.rms_norm_eps) |
| 392 | |
| 393 | # After embedding sharing, embed_tokens is replaced with the |
| 394 | # target model's backbone-dim embedding. Scale by |
| 395 | # sqrt(backbone_hidden_size) to match the target's convention. |
| 396 | self.register_buffer( |
| 397 | "normalizer", |
| 398 | torch.tensor(self.backbone_hidden_size**0.5), |
| 399 | persistent=False, |
| 400 | ) |
| 401 | |
| 402 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 403 | return self.embed_tokens(input_ids) * self.normalizer |
| 404 | |
| 405 | def forward( |
| 406 | self, |
| 407 | input_ids: torch.Tensor, |
| 408 | positions: torch.Tensor, |
| 409 | hidden_states: torch.Tensor, |
| 410 | intermediate_tensors: IntermediateTensors | None = None, |
| 411 | inputs_embeds: torch.Tensor | None = None, |
| 412 | spec_step_idx: int = 0, |
| 413 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 414 | """Returns (draft_hidden_states, backbone_hidden_states). |
| 415 | |
| 416 | draft_hidden_states: draft-dim, used by compute_logits via lm_head. |
| 417 | backbone_hidden_states: backbone-dim, stored in the proposer's |
| 418 | hidden-state buffer and fed back as input to the next step. |
| 419 | """ |
| 420 | if inputs_embeds is None: |
| 421 | inputs_embeds = self.embed_input_ids(input_ids) |
| 422 | |
| 423 | combined = torch.cat([inputs_embeds, hidden_states], dim=-1) |
| 424 | hidden_states, _ = self.pre_projection(combined) |
| 425 | |
| 426 | residual = None |
| 427 | for layer in self.layers: |
| 428 | hidden_states, residual = layer( |
| 429 | positions=positions, |
| 430 | hidden_states=hidden_states, |
| 431 | residual=residual, |
| 432 | ) |
| 433 | |
| 434 | draft_hidden_states = self.norm(hidden_states) |
| 435 | |
| 436 | backbone_hidden_states, _ = self.post_projection(draft_hidden_states) |
| 437 | return draft_hidden_states, backbone_hidden_states |
| 438 | |
| 439 | |
| 440 | @support_torch_compile |
| 441 | class Gemma4MTP(nn.Module): |
| 442 | """Gemma4 Multi-Token Prediction model for speculative decoding. |
| 443 | |
| 444 | forward() returns (draft_hidden_states, backbone_hidden_states). |
| 445 | The proposer uses draft_hidden_states for compute_logits (via |
| 446 | the draft-dim lm_head) and backbone_hidden_states for the |
| 447 | hidden-state feedback buffer. |
| 448 | """ |
| 449 | |
| 450 | has_own_lm_head = True |
| 451 | |
| 452 | hf_to_vllm_mapper = WeightsMapper( |
| 453 | orig_to_new_prefix={ |
| 454 | "pre_projection.": "model.pre_projection.", |
| 455 | "post_projection.": "model.post_projection.", |
| 456 | }, |
| 457 | orig_to_new_stacked={ |
| 458 | ".gate_proj": (".gate_up_proj", 0), |
| 459 | ".up_proj": (".gate_up_proj", 1), |
| 460 | }, |
| 461 | ) |
| 462 | |
| 463 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 464 | super().__init__() |
| 465 | speculative_config = vllm_config.speculative_config |
| 466 | assert speculative_config is not None |
| 467 | config = speculative_config.draft_model_config.hf_config |
| 468 | text_config = _get_text_config(config) |
| 469 | self.quant_config = get_draft_quant_config(vllm_config) |
| 470 | self.config = config |
| 471 | self._stable_full_lm_head_weight: torch.Tensor | None = None |
| 472 | |
| 473 | self.model = Gemma4MultiTokenPredictor( |
| 474 | vllm_config=vllm_config, |
| 475 | prefix=maybe_prefix(prefix, "draft_model"), |
| 476 | ) |
| 477 | |
| 478 | # lm_head operates in draft-dim. Tied to embed_tokens at init |
| 479 | # so load_weights populates both from a single checkpoint entry. |
| 480 | # After embedding sharing, lm_head.weight still references the |
| 481 | # original draft-dim tensor. |
| 482 | self.lm_head = ParallelLMHead( |
| 483 | text_config.vocab_size, |
| 484 | text_config.hidden_size, |
| 485 | quant_config=self.quant_config, |
| 486 | prefix=maybe_prefix(prefix, "lm_head"), |
| 487 | ) |
| 488 | if getattr(config, "tie_word_embeddings", True): |
| 489 | self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) |
| 490 | |
| 491 | self.logits_processor = LogitsProcessor( |
| 492 | text_config.vocab_size, |
| 493 | soft_cap=getattr(text_config, "final_logit_softcapping", None), |
| 494 | ) |
| 495 | |
| 496 | self.masked_embedding: Gemma4MTPMaskedEmbedder | None |
| 497 | if getattr(config, "use_ordered_embeddings", False): |
| 498 | num_centroids = getattr(config, "num_centroids", 2048) |
| 499 | top_k = getattr(config, "centroid_intermediate_top_k", 32) |
| 500 | self.masked_embedding = Gemma4MTPMaskedEmbedder( |
| 501 | hidden_size=text_config.hidden_size, |
| 502 | vocab_size=text_config.vocab_size, |
| 503 | num_centroids=num_centroids, |
| 504 | centroid_intermediate_top_k=top_k, |
| 505 | ) |
| 506 | logger.info( |
| 507 | "Gemma4 MTP: centroids masking enabled " |
| 508 | "(num_centroids=%d, top_k=%d, active_tokens=%d/%d).", |
| 509 | num_centroids, |
| 510 | top_k, |
| 511 | top_k * (text_config.vocab_size // num_centroids), |
| 512 | text_config.vocab_size, |
| 513 | ) |
| 514 | else: |
| 515 | self.masked_embedding = None |
| 516 | |
| 517 | draft_cfg = speculative_config.draft_model_config |
| 518 | gen_cfg = draft_cfg.try_get_generation_config() |
| 519 | self._suppress_token_ids = gen_cfg.get("suppress_tokens") if gen_cfg else None |
| 520 | # Materialized on-device in load_weights: compute_logits runs under CUDA |
| 521 | # graph capture in the V2 speculator, where indexing with a Python list |
| 522 | # would issue an unpinned H2D copy (illegal during capture). |
| 523 | self._suppress_idx: torch.Tensor | None = None |
| 524 | |
| 525 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 526 | return self.model.embed_input_ids(input_ids) |
| 527 | |
| 528 | def forward( |
| 529 | self, |
| 530 | input_ids: torch.Tensor | None, |
| 531 | positions: torch.Tensor, |
| 532 | hidden_states: torch.Tensor, |
| 533 | intermediate_tensors: IntermediateTensors | None = None, |
| 534 | inputs_embeds: torch.Tensor | None = None, |
| 535 | spec_step_idx: int = 0, |
| 536 | **kwargs: object, |
| 537 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 538 | return self.model( |
| 539 | input_ids, |
| 540 | positions, |
| 541 | hidden_states, |
| 542 | intermediate_tensors, |
| 543 | inputs_embeds, |
| 544 | spec_step_idx, |
| 545 | ) |
| 546 | |
| 547 | def _get_full_lm_head_weight(self) -> torch.Tensor: |
| 548 | if self._stable_full_lm_head_weight is not None: |
| 549 | return self._stable_full_lm_head_weight |
| 550 | assert self.masked_embedding is not None |
| 551 | lm_head_weight = self.lm_head.weight |
| 552 | tp_size = get_tensor_model_parallel_world_size() |
| 553 | if tp_size > 1: |
| 554 | lm_head_weight = tensor_model_parallel_all_gather( |
| 555 | lm_head_weight, |
| 556 | dim=0, |
| 557 | ) |
| 558 | lm_head_weight = lm_head_weight[: self.masked_embedding.vocab_size] |
| 559 | if tp_size > 1: |
| 560 | lm_head_weight = lm_head_weight.contiguous() |
| 561 | self._stable_full_lm_head_weight = lm_head_weight |
| 562 | return lm_head_weight |
| 563 | |
| 564 | def compute_logits( |
| 565 | self, |
| 566 | hidden_states: torch.Tensor, |
| 567 | spec_step_idx: int = 0, |
| 568 | ) -> torch.Tensor | None: |
| 569 | if self.masked_embedding is not None: |
| 570 | logits = self.masked_embedding( |
| 571 | hidden_states, |
| 572 | self._get_full_lm_head_weight(), |
| 573 | ) |
| 574 | else: |
| 575 | logits = self.logits_processor(self.lm_head, hidden_states) |
| 576 | if logits is not None and self._suppress_idx is not None: |
| 577 | logits.index_fill_(1, self._suppress_idx, -float("inf")) |
| 578 | return logits |
| 579 | |
| 580 | def get_top_tokens( |
| 581 | self, |
| 582 | hidden_states: torch.Tensor, |
| 583 | ) -> torch.Tensor: |
| 584 | """Sparse argmax via centroids masking. Returns token IDs directly.""" |
| 585 | assert self.masked_embedding is not None |
| 586 | return self.masked_embedding.get_top_tokens( |
| 587 | hidden_states, |
| 588 | self._get_full_lm_head_weight(), |
| 589 | ) |
| 590 | |
| 591 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 592 | self._stable_full_lm_head_weight = None |
| 593 | loader = AutoWeightsLoader(self) |
| 594 | loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 595 | if self._suppress_token_ids: |
| 596 | self._suppress_idx = torch.tensor( |
| 597 | self._suppress_token_ids, |
| 598 | dtype=torch.long, |
| 599 | device=next(self.parameters()).device, |
| 600 | ) |
| 601 | return loaded |