| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Copyright 2025 The vLLM team. |
| 4 | # Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved. |
| 5 | # |
| 6 | # |
| 7 | |
| 8 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 9 | # you may not use this file except in compliance with the License. |
| 10 | # You may obtain a copy of the License at |
| 11 | # |
| 12 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 13 | # |
| 14 | # Unless required by applicable law or agreed to in writing, software |
| 15 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 16 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 17 | # See the License for the specific language governing permissions and |
| 18 | # limitations under the License. |
| 19 | """Gemma 4 model implementation for vLLM.""" |
| 20 | |
| 21 | from collections.abc import Iterable |
| 22 | from dataclasses import replace |
| 23 | from itertools import islice |
| 24 | |
| 25 | import torch |
| 26 | from torch import nn |
| 27 | |
| 28 | from vllm.compilation.decorators import support_torch_compile |
| 29 | from vllm.config import CacheConfig, VllmConfig |
| 30 | from vllm.config.utils import getattr_iter |
| 31 | from vllm.distributed import ( |
| 32 | get_pp_group, |
| 33 | get_tensor_model_parallel_rank, |
| 34 | get_tensor_model_parallel_world_size, |
| 35 | ) |
| 36 | from vllm.forward_context import get_forward_context |
| 37 | from vllm.logger import init_logger |
| 38 | from vllm.model_executor.layers.activation import get_act_and_mul_fn |
| 39 | from vllm.model_executor.layers.attention import Attention |
| 40 | from vllm.model_executor.layers.fused_moe import ( |
| 41 | FusedMoEFactory, |
| 42 | GateLinear, |
| 43 | MoERunner, |
| 44 | ) |
| 45 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 46 | from vllm.model_executor.layers.linear import ( |
| 47 | ColumnParallelLinear, |
| 48 | MergedColumnParallelLinear, |
| 49 | QKVParallelLinear, |
| 50 | ReplicatedLinear, |
| 51 | RowParallelLinear, |
| 52 | ) |
| 53 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 54 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 55 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 56 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 57 | ParallelLMHead, |
| 58 | VocabParallelEmbedding, |
| 59 | ) |
| 60 | from vllm.platforms import current_platform |
| 61 | from vllm.sequence import IntermediateTensors |
| 62 | from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config |
| 63 | from vllm.triton_utils import tl, triton |
| 64 | from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata |
| 65 | |
| 66 | from .interfaces import ( |
| 67 | EagleModelMixin, |
| 68 | MixtureOfExperts, |
| 69 | SupportsEagle3, |
| 70 | SupportsLoRA, |
| 71 | SupportsPP, |
| 72 | SupportsQuant, |
| 73 | ) |
| 74 | from .utils import ( |
| 75 | AutoWeightsLoader, |
| 76 | PPMissingLayer, |
| 77 | ShardIds, |
| 78 | WeightsMapper, |
| 79 | extract_layer_index, |
| 80 | make_layers, |
| 81 | maybe_prefix, |
| 82 | ) |
| 83 | |
| 84 | logger = init_logger(__name__) |
| 85 | |
| 86 | |
| 87 | def _gemma4_layer_weights_mapper(config) -> WeightsMapper: |
| 88 | """Stack each decoder layer's projections the way its modules are built. |
| 89 | |
| 90 | gate/up pack into `mlp.gate_up_proj`. Layers with a qkv_proj pack q/k/v |
| 91 | into it; `attention_k_eq_v` full-attention layers ship k_proj only, so K |
| 92 | is loaded into both the K and V shards and no v_proj is expected. KV-shared |
| 93 | layers only have q_proj; the K/V tensors original checkpoints still ship |
| 94 | for them are dropped. |
| 95 | """ |
| 96 | num_layers = config.num_hidden_layers |
| 97 | first_kv_shared = num_layers - getattr(config, "num_kv_shared_layers", 0) |
| 98 | k_eq_v_layers = ( |
| 99 | { |
| 100 | i |
| 101 | for i, layer_type in enumerate(config.layer_types[:first_kv_shared]) |
| 102 | if layer_type == "full_attention" |
| 103 | } |
| 104 | if getattr(config, "attention_k_eq_v", False) |
| 105 | else set() |
| 106 | ) |
| 107 | qkv_shards: tuple[tuple[str, ShardIds], ...] = (("q", "q"), ("k", "k"), ("v", "v")) |
| 108 | k_eq_v_shards: tuple[tuple[str, ShardIds], ...] = (("q", "q"), ("k", ["k", "v"])) |
| 109 | stacked: dict[str, tuple[str, ShardIds]] = { |
| 110 | ".mlp.gate_proj.": (".mlp.gate_up_proj.", 0), |
| 111 | ".mlp.up_proj.": (".mlp.gate_up_proj.", 1), |
| 112 | } |
| 113 | stacked |= { |
| 114 | f"layers.{i}.self_attn.{shard}_proj.": ( |
| 115 | f"layers.{i}.self_attn.qkv_proj.", |
| 116 | shard_id, |
| 117 | ) |
| 118 | for i in range(first_kv_shared) |
| 119 | for shard, shard_id in (k_eq_v_shards if i in k_eq_v_layers else qkv_shards) |
| 120 | } |
| 121 | return WeightsMapper( |
| 122 | orig_to_new_substr={ |
| 123 | f"layers.{i}.self_attn.{name}.": None |
| 124 | for i in range(first_kv_shared, num_layers) |
| 125 | for name in ("k_proj", "v_proj", "k_norm") |
| 126 | }, |
| 127 | orig_to_new_stacked=stacked, |
| 128 | ) |
| 129 | |
| 130 | |
| 131 | @triton.jit |
| 132 | def _gemma4_routing_kernel( |
| 133 | gating_ptr, |
| 134 | per_expert_scale_ptr, |
| 135 | topk_weights_ptr, |
| 136 | topk_ids_ptr, |
| 137 | E: tl.constexpr, |
| 138 | K: tl.constexpr, |
| 139 | BLOCK_E: tl.constexpr, |
| 140 | ): |
| 141 | pid = tl.program_id(0) |
| 142 | offs_e = tl.arange(0, BLOCK_E) |
| 143 | valid = offs_e < E |
| 144 | |
| 145 | logits = tl.load( |
| 146 | gating_ptr + pid * E + offs_e, |
| 147 | mask=valid, |
| 148 | other=-float("inf"), |
| 149 | ).to(tl.float32) |
| 150 | |
| 151 | max_l = tl.max(logits, axis=0) |
| 152 | |
| 153 | # Float32 → ascending-sortable bijection |
| 154 | MIN32 = -2147483648 |
| 155 | logit_bits = logits.to(tl.int32, bitcast=True) |
| 156 | sign_b = logit_bits >> 31 |
| 157 | key = tl.where(sign_b == 0, logit_bits ^ -1, logit_bits ^ MIN32) |
| 158 | key = tl.where(valid, key, 0x7FFFFFFF) |
| 159 | sk64 = key.to(tl.int64) & 0x00000000FFFFFFFF |
| 160 | packed = (sk64 << 32) | offs_e.to(tl.int64) |
| 161 | sorted_p = tl.sort(packed, descending=False) |
| 162 | |
| 163 | # Vectorized extraction of ALL sorted elements — no K-loop, no cross-lane reductions |
| 164 | all_keys = ((sorted_p >> 32) & 0x00000000FFFFFFFF).to(tl.int32) |
| 165 | all_ids = (sorted_p & 0x00000000FFFFFFFF).to(tl.int32) |
| 166 | |
| 167 | # Inverse bijection: recover original logit bits |
| 168 | sign_k = all_keys >> 31 |
| 169 | all_bits = tl.where(sign_k < 0, all_keys ^ -1, all_keys ^ MIN32) |
| 170 | all_logits = all_bits.to(tl.float32, bitcast=True) |
| 171 | |
| 172 | # Compute raw_exp for ALL BLOCK_E elements — vectorized, ~2 VALU clocks |
| 173 | all_raw_exp = tl.math.exp2((all_logits - max_l) * 1.4426950408889634) |
| 174 | |
| 175 | # Sum only top-K for renorm — ONE masked reduction |
| 176 | top_mask = offs_e < K |
| 177 | renorm_raw = tl.sum(tl.where(top_mask, all_raw_exp, 0.0), axis=0) |
| 178 | renorm_raw = tl.where(renorm_raw > 0.0, renorm_raw, 1.0) |
| 179 | inv_renorm = 1.0 / renorm_raw |
| 180 | |
| 181 | # Load scales for top-K only (masked gather; scale array is tiny → L1 cached) |
| 182 | all_scales = tl.load( |
| 183 | per_expert_scale_ptr + all_ids.to(tl.int64), |
| 184 | mask=top_mask, |
| 185 | other=1.0, |
| 186 | ).to(tl.float32) |
| 187 | |
| 188 | # Final weights: vectorized multiply (only top-K will be stored) |
| 189 | all_weights = (all_raw_exp * inv_renorm * all_scales).to(tl.float32) |
| 190 | |
| 191 | # Write results with TWO masked stores — replaces K × 2 serial scalar stores |
| 192 | base_off = pid * K + offs_e |
| 193 | tl.store(topk_ids_ptr + base_off, all_ids, mask=top_mask) |
| 194 | tl.store(topk_weights_ptr + base_off, all_weights, mask=top_mask) |
| 195 | |
| 196 | |
| 197 | def gemma4_fused_routing_kernel_triton( |
| 198 | gating_output: torch.Tensor, |
| 199 | topk: int, |
| 200 | per_expert_scale: torch.Tensor, |
| 201 | num_warps: int = 1, |
| 202 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 203 | gating_output = gating_output.contiguous() |
| 204 | per_expert_scale = per_expert_scale.contiguous() |
| 205 | T, E = gating_output.shape |
| 206 | weights = torch.empty(T, topk, dtype=torch.float32, device=gating_output.device) |
| 207 | ids = torch.empty(T, topk, dtype=torch.int32, device=gating_output.device) |
| 208 | BLOCK_E = triton.next_power_of_2(E) |
| 209 | _gemma4_routing_kernel[(T,)]( |
| 210 | gating_output, |
| 211 | per_expert_scale, |
| 212 | weights, |
| 213 | ids, |
| 214 | E, |
| 215 | topk, |
| 216 | BLOCK_E, |
| 217 | num_warps=num_warps, |
| 218 | ) |
| 219 | return weights, ids |
| 220 | |
| 221 | |
| 222 | def gemma4_routing_function_torch( |
| 223 | gating_output: torch.Tensor, |
| 224 | topk: int, |
| 225 | per_expert_scale: torch.Tensor, |
| 226 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 227 | _, topk_ids = torch.topk(gating_output, k=topk, dim=-1) |
| 228 | router_probabilities = torch.nn.functional.softmax(gating_output, dim=-1) |
| 229 | indicator = torch.nn.functional.one_hot( |
| 230 | topk_ids, num_classes=gating_output.size(-1) |
| 231 | ).sum(dim=-2) |
| 232 | gate_weights = indicator * router_probabilities |
| 233 | renorm_factor = torch.sum(gate_weights, dim=-1, keepdim=True) |
| 234 | renorm_factor = torch.where(renorm_factor > 0.0, renorm_factor, 1.0) |
| 235 | dispatch_weights = gate_weights / renorm_factor |
| 236 | |
| 237 | topk_weights = dispatch_weights.gather(1, topk_ids) |
| 238 | |
| 239 | # Fold per_expert_scale into routing weights |
| 240 | expert_scales = per_expert_scale[topk_ids].to(topk_weights.dtype) |
| 241 | topk_weights = topk_weights * expert_scales |
| 242 | return topk_weights.to(torch.float32), topk_ids.to(torch.int32) |
| 243 | |
| 244 | |
| 245 | def _get_text_config(config): |
| 246 | """Dereference text_config if config is a nested Gemma4Config. |
| 247 | |
| 248 | Gemma4 checkpoints use architectures=["Gemma4ForConditionalGeneration"] |
| 249 | which yields a Gemma4Config with nested text_config. This function |
| 250 | transparently returns the text config regardless of nesting. |
| 251 | """ |
| 252 | if hasattr(config, "text_config"): |
| 253 | return config.text_config |
| 254 | return config |
| 255 | |
| 256 | |
| 257 | class Gemma4MLP(nn.Module): |
| 258 | def __init__( |
| 259 | self, |
| 260 | hidden_size: int, |
| 261 | intermediate_size: int, |
| 262 | hidden_activation: str, |
| 263 | quant_config: QuantizationConfig | None = None, |
| 264 | prefix: str = "", |
| 265 | ) -> None: |
| 266 | super().__init__() |
| 267 | self.gate_up_proj = MergedColumnParallelLinear( |
| 268 | hidden_size, |
| 269 | [intermediate_size] * 2, |
| 270 | bias=False, |
| 271 | quant_config=quant_config, |
| 272 | prefix=f"{prefix}.gate_up_proj", |
| 273 | ) |
| 274 | self.down_proj = RowParallelLinear( |
| 275 | intermediate_size, |
| 276 | hidden_size, |
| 277 | bias=False, |
| 278 | quant_config=quant_config, |
| 279 | prefix=f"{prefix}.down_proj", |
| 280 | ) |
| 281 | self.act_fn = get_act_and_mul_fn(hidden_activation) |
| 282 | |
| 283 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 284 | gate_up, _ = self.gate_up_proj(x) |
| 285 | x = self.act_fn(gate_up) |
| 286 | x, _ = self.down_proj(x) |
| 287 | return x |
| 288 | |
| 289 | |
| 290 | class Gemma4Router(nn.Module): |
| 291 | """Router for Gemma4 MoE that preprocesses input before projection. |
| 292 | |
| 293 | Applies RMSNorm (no learned weight), root_size scaling |
| 294 | (hidden_size^{-0.5}), then a learned per-dimension scale before |
| 295 | projecting to expert logits. |
| 296 | |
| 297 | This preprocessing is applied ONLY to the router's input, not to |
| 298 | the expert MLPs' input. |
| 299 | """ |
| 300 | |
| 301 | def __init__( |
| 302 | self, |
| 303 | config, |
| 304 | quant_config: QuantizationConfig | None = None, |
| 305 | prefix: str = "", |
| 306 | ) -> None: |
| 307 | super().__init__() |
| 308 | self.hidden_size = config.hidden_size |
| 309 | |
| 310 | # RMSNorm without learned weight — pure normalization only |
| 311 | self.norm = RMSNorm(self.hidden_size, eps=config.rms_norm_eps, has_weight=False) |
| 312 | # Per-dimension learned scale, applied after norm + root_size |
| 313 | self.scale = nn.Parameter(torch.ones(self.hidden_size)) |
| 314 | # Per-expert output scale folded into routing weights so that |
| 315 | # MoERunner's fused kernel computes: Σ_e (expert_e * w_e * scale_e) |
| 316 | self.per_expert_scale = nn.Parameter(torch.ones(config.num_experts)) |
| 317 | # Constant 1/sqrt(hidden_size) scaling factor |
| 318 | self.register_buffer( |
| 319 | "root_size", |
| 320 | torch.tensor(self.hidden_size**-0.5), |
| 321 | persistent=False, |
| 322 | ) |
| 323 | # Project to expert logits; replicated across TP for consistent routing |
| 324 | # GateLinear supports bf16 W/A → fp32 output, which is important |
| 325 | # because the topk kernel often needs fp32 for stable routing. |
| 326 | self.proj = GateLinear( |
| 327 | self.hidden_size, |
| 328 | config.num_experts, |
| 329 | bias=False, |
| 330 | out_dtype=torch.float32, |
| 331 | prefix=f"{prefix}.proj", |
| 332 | ) |
| 333 | |
| 334 | def routing_function( |
| 335 | self, |
| 336 | hidden_states: torch.Tensor, |
| 337 | gating_output: torch.Tensor, |
| 338 | topk: int, |
| 339 | renormalize: bool, |
| 340 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 341 | """Softmax over ALL experts → top-k → renormalize. |
| 342 | |
| 343 | MoERunner's built-in fused_topk scopes softmax differently, so Gemma4 |
| 344 | needs its own routing for numerical correctness. |
| 345 | |
| 346 | NOTE: self.per_expert_scale is read at call time (not captured into a |
| 347 | local) so that torch.func.functional_call parameter substitution |
| 348 | reaches the routing function correctly. |
| 349 | """ |
| 350 | if current_platform.is_cuda_alike() or current_platform.is_xpu(): |
| 351 | return gemma4_fused_routing_kernel_triton( |
| 352 | gating_output, topk, self.per_expert_scale |
| 353 | ) |
| 354 | |
| 355 | return gemma4_routing_function_torch(gating_output, topk, self.per_expert_scale) |
| 356 | |
| 357 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 358 | """Returns raw router logits [T, E].""" |
| 359 | x = self.norm(x) |
| 360 | x = x * self.root_size.to(x.dtype) |
| 361 | x = x * self.scale.to(x.dtype) |
| 362 | router_logits, _ = self.proj(x) |
| 363 | return router_logits |
| 364 | |
| 365 | |
| 366 | class Gemma4Attention(nn.Module): |
| 367 | def __init__( |
| 368 | self, |
| 369 | config, |
| 370 | hidden_size: int, |
| 371 | num_heads: int, |
| 372 | num_kv_heads: int, |
| 373 | head_dim: int, |
| 374 | max_position_embeddings: int, |
| 375 | cache_config: CacheConfig | None = None, |
| 376 | quant_config: QuantizationConfig | None = None, |
| 377 | attn_logits_soft_cap: float | None = None, |
| 378 | prefix: str = "", |
| 379 | ) -> None: |
| 380 | super().__init__() |
| 381 | self.config = config |
| 382 | self.hidden_size = hidden_size |
| 383 | |
| 384 | tp_size = get_tensor_model_parallel_world_size() |
| 385 | self.tp_rank = get_tensor_model_parallel_rank() |
| 386 | self.total_num_heads = num_heads |
| 387 | assert self.total_num_heads % tp_size == 0 |
| 388 | self.num_heads = self.total_num_heads // tp_size |
| 389 | self.total_num_kv_heads = num_kv_heads |
| 390 | if self.total_num_kv_heads >= tp_size: |
| 391 | assert self.total_num_kv_heads % tp_size == 0 |
| 392 | else: |
| 393 | assert tp_size % self.total_num_kv_heads == 0 |
| 394 | self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) |
| 395 | self.head_dim = head_dim |
| 396 | self.q_size = self.num_heads * self.head_dim |
| 397 | self.kv_size = self.num_kv_heads * self.head_dim |
| 398 | # Gemma4 uses scaling=1.0. |
| 399 | # Unlike Gemma2/3, query_pre_attn_scalar is NOT used here; |
| 400 | # Q/K norms with learnable weights handle scaling implicitly. |
| 401 | self.scaling = 1.0 |
| 402 | |
| 403 | layer_idx = extract_layer_index(prefix) |
| 404 | num_kv_shared_layers = getattr(config, "num_kv_shared_layers", 0) |
| 405 | first_kv_shared_layer_idx = config.num_hidden_layers - num_kv_shared_layers |
| 406 | self.is_kv_shared_layer = ( |
| 407 | num_kv_shared_layers > 0 and layer_idx >= first_kv_shared_layer_idx |
| 408 | ) |
| 409 | |
| 410 | if self.is_kv_shared_layer: |
| 411 | # K/V come from the target layer's cache, so like HF this layer |
| 412 | # only has q_proj (and no k_norm/v_norm). |
| 413 | self.q_proj = ColumnParallelLinear( |
| 414 | hidden_size, |
| 415 | self.total_num_heads * self.head_dim, |
| 416 | bias=config.attention_bias, |
| 417 | quant_config=quant_config, |
| 418 | prefix=f"{prefix}.q_proj", |
| 419 | ) |
| 420 | else: |
| 421 | # QKVParallelLinear handles GQA correctly for all layer types. |
| 422 | # k_eq_v layers load K weights into both K and V slots via |
| 423 | # hf_to_vllm_mapper — no structural difference needed. |
| 424 | self.qkv_proj = QKVParallelLinear( |
| 425 | hidden_size, |
| 426 | self.head_dim, |
| 427 | self.total_num_heads, |
| 428 | self.total_num_kv_heads, |
| 429 | bias=config.attention_bias, |
| 430 | quant_config=quant_config, |
| 431 | prefix=f"{prefix}.qkv_proj", |
| 432 | ) |
| 433 | self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) |
| 434 | # V norm: no learnable scale (pure normalization only) |
| 435 | self.v_norm = RMSNorm( |
| 436 | self.head_dim, eps=config.rms_norm_eps, has_weight=False |
| 437 | ) |
| 438 | self.o_proj = RowParallelLinear( |
| 439 | self.total_num_heads * self.head_dim, |
| 440 | hidden_size, |
| 441 | bias=config.attention_bias, |
| 442 | quant_config=quant_config, |
| 443 | prefix=f"{prefix}.o_proj", |
| 444 | ) |
| 445 | |
| 446 | # Q norm: output = norm(x) * weight (learnable per-head scale) |
| 447 | self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) |
| 448 | |
| 449 | # Determine layer type and sliding window |
| 450 | layer_type = config.layer_types[layer_idx] |
| 451 | self.is_sliding = layer_type == "sliding_attention" |
| 452 | sliding_window = config.sliding_window if self.is_sliding else None |
| 453 | |
| 454 | # Initialize RoPE based on layer type. |
| 455 | # Gemma4 uses different RoPE parameters for sliding vs full attention. |
| 456 | if layer_type in config.rope_parameters: |
| 457 | # Per-layer-type rope config (dict format). |
| 458 | # rope_parameters already contains the correct |
| 459 | # partial_rotary_factor per layer type (1.0 for full |
| 460 | # attention, 1.0 for sliding). Do NOT override with |
| 461 | # global_partial_rotary_factor — that config key is |
| 462 | # not needed for Gemma4 — config uses per-layer rope_parameters. |
| 463 | rope_parameters = dict(config.rope_parameters[layer_type]) |
| 464 | else: |
| 465 | # Legacy config format fallback. |
| 466 | rope_parameters = dict(config.rope_parameters.copy()) |
| 467 | if self.is_sliding: |
| 468 | rope_parameters["rope_theta"] = getattr( |
| 469 | config, "rope_local_base_freq", 10000.0 |
| 470 | ) |
| 471 | |
| 472 | # KV sharing: layers in the last `num_kv_shared_layers` share KV |
| 473 | # cache with earlier layers of the same type. |
| 474 | kv_sharing_target_layer_name = None |
| 475 | if self.is_kv_shared_layer: |
| 476 | # Find the last non-shared layer of the same attention type |
| 477 | prev_layers = config.layer_types[:first_kv_shared_layer_idx] |
| 478 | current_layer_type = config.layer_types[layer_idx] |
| 479 | kv_shared_layer_index = ( |
| 480 | len(prev_layers) - 1 - prev_layers[::-1].index(current_layer_type) |
| 481 | ) |
| 482 | if kv_shared_layer_index >= 0: |
| 483 | if ".layers." in prefix: |
| 484 | param_name_before_layers = prefix.split(".layers.")[0] |
| 485 | else: |
| 486 | raise ValueError( |
| 487 | "Unexpected prefix format for Gemma4Attention: " |
| 488 | f"'{prefix}'. Expected to contain '.layers.'." |
| 489 | ) |
| 490 | kv_sharing_target_layer_name = ( |
| 491 | f"{param_name_before_layers}.layers." |
| 492 | f"{kv_shared_layer_index}.self_attn.attn" |
| 493 | ) |
| 494 | |
| 495 | self.rotary_emb = get_rope( |
| 496 | self.head_dim, |
| 497 | max_position=max_position_embeddings, |
| 498 | rope_parameters=rope_parameters, |
| 499 | is_neox_style=True, |
| 500 | ) |
| 501 | |
| 502 | self.attn = Attention( |
| 503 | self.num_heads, |
| 504 | self.head_dim, |
| 505 | self.scaling, |
| 506 | num_kv_heads=self.num_kv_heads, |
| 507 | cache_config=cache_config, |
| 508 | quant_config=quant_config, |
| 509 | logits_soft_cap=attn_logits_soft_cap, |
| 510 | per_layer_sliding_window=sliding_window, |
| 511 | kv_sharing_target_layer_name=kv_sharing_target_layer_name, |
| 512 | # Gemma4 vision bidi: on sliding layers the bidirectional image |
| 513 | # block must stay within the sliding window, matching HF's |
| 514 | # (causal OR blockwise) AND sliding_window. Without this the image |
| 515 | # span (~1100 soft tokens at max_soft_tokens=1120) exceeds the 1024 |
| 516 | # window; the runner keeps the full range and the kernel bounds it |
| 517 | # per-query here. |
| 518 | mm_prefix_clamp_sliding_window=self.is_sliding, |
| 519 | prefix=f"{prefix}.attn", |
| 520 | ) |
| 521 | |
| 522 | def forward( |
| 523 | self, |
| 524 | positions: torch.Tensor, |
| 525 | hidden_states: torch.Tensor, |
| 526 | **kwargs, |
| 527 | ) -> torch.Tensor: |
| 528 | if self.is_kv_shared_layer: |
| 529 | # Shared KV: only Q is projected; K/V come from the target layer. |
| 530 | q, _ = self.q_proj(hidden_states) |
| 531 | q = q.unflatten(-1, (self.num_heads, self.head_dim)) |
| 532 | q = self.q_norm(q) |
| 533 | q = q.flatten(-2, -1) |
| 534 | q, _ = self.rotary_emb(positions, q, None) |
| 535 | attn_output = self.attn(q, None, None) |
| 536 | else: |
| 537 | # For k_eq_v, K weights are loaded into both K and V slots of |
| 538 | # qkv_proj, so V == K automatically. |
| 539 | qkv, _ = self.qkv_proj(hidden_states) |
| 540 | q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 541 | |
| 542 | q = q.unflatten(-1, (self.num_heads, self.head_dim)) |
| 543 | q = self.q_norm(q) |
| 544 | q = q.flatten(-2, -1) |
| 545 | |
| 546 | k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)) |
| 547 | k = self.k_norm(k) |
| 548 | k = k.flatten(-2, -1) |
| 549 | q, k = self.rotary_emb(positions, q, k) |
| 550 | |
| 551 | v = v.unflatten(-1, (self.num_kv_heads, self.head_dim)) |
| 552 | v = self.v_norm(v) |
| 553 | v = v.flatten(-2, -1) |
| 554 | attn_output = self.attn(q, k, v) |
| 555 | |
| 556 | output, _ = self.o_proj(attn_output) |
| 557 | |
| 558 | return output |
| 559 | |
| 560 | |
| 561 | class Gemma4DecoderLayer(nn.Module): |
| 562 | def __init__( |
| 563 | self, |
| 564 | config, |
| 565 | cache_config: CacheConfig | None = None, |
| 566 | quant_config: QuantizationConfig | None = None, |
| 567 | prefix: str = "", |
| 568 | ) -> None: |
| 569 | super().__init__() |
| 570 | self.hidden_size = config.hidden_size |
| 571 | self.hidden_size_per_layer_input = getattr( |
| 572 | config, "hidden_size_per_layer_input", 0 |
| 573 | ) |
| 574 | |
| 575 | layer_idx = extract_layer_index(prefix) |
| 576 | self.layer_idx = layer_idx |
| 577 | |
| 578 | # Gemma4 uses different head dimensions for sliding vs full attention |
| 579 | layer_config = gemma4_layer_config(config, layer_idx) |
| 580 | head_dim = layer_config.head_dim |
| 581 | num_kv_heads = layer_config.num_key_value_heads |
| 582 | |
| 583 | self.self_attn = Gemma4Attention( |
| 584 | config=config, |
| 585 | hidden_size=self.hidden_size, |
| 586 | num_heads=config.num_attention_heads, |
| 587 | num_kv_heads=num_kv_heads, |
| 588 | head_dim=head_dim, |
| 589 | max_position_embeddings=config.max_position_embeddings, |
| 590 | cache_config=cache_config, |
| 591 | quant_config=quant_config, |
| 592 | attn_logits_soft_cap=getattr(config, "attn_logit_softcapping", None), |
| 593 | prefix=f"{prefix}.self_attn", |
| 594 | ) |
| 595 | |
| 596 | # Compute per-layer intermediate_size from config. |
| 597 | # When use_double_wide_mlp is set, intermediate_size doubles for |
| 598 | # KV-shared layers (layers >= first_kv_shared_layer_idx). |
| 599 | first_kv_shared_layer_idx = config.num_hidden_layers - getattr( |
| 600 | config, "num_kv_shared_layers", 0 |
| 601 | ) |
| 602 | is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0 |
| 603 | use_double_wide_mlp = ( |
| 604 | getattr(config, "use_double_wide_mlp", False) and is_kv_shared_layer |
| 605 | ) |
| 606 | layer_intermediate_size = config.intermediate_size * ( |
| 607 | 2 if use_double_wide_mlp else 1 |
| 608 | ) |
| 609 | |
| 610 | self.mlp = Gemma4MLP( |
| 611 | hidden_size=self.hidden_size, |
| 612 | intermediate_size=layer_intermediate_size, |
| 613 | hidden_activation=config.hidden_activation, |
| 614 | quant_config=quant_config, |
| 615 | prefix=f"{prefix}.mlp", |
| 616 | ) |
| 617 | |
| 618 | # Layer norms: output = norm(x) * weight |
| 619 | self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 620 | self.post_attention_layernorm = RMSNorm( |
| 621 | config.hidden_size, eps=config.rms_norm_eps |
| 622 | ) |
| 623 | self.pre_feedforward_layernorm = RMSNorm( |
| 624 | config.hidden_size, eps=config.rms_norm_eps |
| 625 | ) |
| 626 | self.post_feedforward_layernorm = RMSNorm( |
| 627 | config.hidden_size, eps=config.rms_norm_eps |
| 628 | ) |
| 629 | |
| 630 | # MoE (Mixture of Experts) — router + expert block parallel to MLP |
| 631 | self.enable_moe_block = getattr(config, "enable_moe_block", False) or getattr( |
| 632 | config, "use_second_mlp_block", False |
| 633 | ) |
| 634 | self.router: Gemma4Router | None |
| 635 | self.experts: MoERunner | None |
| 636 | self.post_feedforward_layernorm_1: RMSNorm | None |
| 637 | self.post_feedforward_layernorm_2: RMSNorm | None |
| 638 | self.pre_feedforward_layernorm_2: RMSNorm | None |
| 639 | if self.enable_moe_block: |
| 640 | self.router = Gemma4Router( |
| 641 | config, |
| 642 | quant_config=quant_config, |
| 643 | prefix=f"{prefix}.router", |
| 644 | ) |
| 645 | names = ("moe_intermediate_size", "expert_intermediate_size") |
| 646 | intermediate_size = getattr_iter(config, names) |
| 647 | if intermediate_size is None: |
| 648 | raise ValueError("Gemma4 MoE requires an expert intermediate size") |
| 649 | self.experts = FusedMoEFactory( |
| 650 | num_experts=config.num_experts, |
| 651 | top_k=config.top_k_experts, |
| 652 | hidden_size=config.hidden_size, |
| 653 | intermediate_size=intermediate_size, |
| 654 | renormalize=True, |
| 655 | quant_config=quant_config, |
| 656 | prefix=f"{prefix}.experts", |
| 657 | custom_routing_function=self.router.routing_function, |
| 658 | activation="gelu_tanh", |
| 659 | ) |
| 660 | self.post_feedforward_layernorm_1 = RMSNorm( |
| 661 | config.hidden_size, eps=config.rms_norm_eps |
| 662 | ) |
| 663 | self.post_feedforward_layernorm_2 = RMSNorm( |
| 664 | config.hidden_size, eps=config.rms_norm_eps |
| 665 | ) |
| 666 | self.pre_feedforward_layernorm_2 = RMSNorm( |
| 667 | config.hidden_size, eps=config.rms_norm_eps |
| 668 | ) |
| 669 | else: |
| 670 | self.router = None |
| 671 | self.experts = None |
| 672 | self.post_feedforward_layernorm_1 = None |
| 673 | self.post_feedforward_layernorm_2 = None |
| 674 | self.pre_feedforward_layernorm_2 = None |
| 675 | |
| 676 | # Per-Layer Embedding (PLE) components — present in each decoder layer |
| 677 | self.per_layer_input_gate: ReplicatedLinear | None |
| 678 | self.per_layer_projection: ReplicatedLinear | None |
| 679 | self.post_per_layer_input_norm: RMSNorm | None |
| 680 | if ( |
| 681 | self.hidden_size_per_layer_input is not None |
| 682 | and self.hidden_size_per_layer_input > 0 |
| 683 | ): |
| 684 | # Gate: projects hidden_states → per-layer dim for gating |
| 685 | self.per_layer_input_gate = ReplicatedLinear( |
| 686 | self.hidden_size, |
| 687 | self.hidden_size_per_layer_input, |
| 688 | bias=False, |
| 689 | quant_config=quant_config, |
| 690 | prefix=f"{prefix}.per_layer_input_gate", |
| 691 | return_bias=False, |
| 692 | ) |
| 693 | # Projection: projects gated per-layer input back → hidden size |
| 694 | self.per_layer_projection = ReplicatedLinear( |
| 695 | self.hidden_size_per_layer_input, |
| 696 | self.hidden_size, |
| 697 | bias=False, |
| 698 | quant_config=quant_config, |
| 699 | prefix=f"{prefix}.per_layer_projection", |
| 700 | return_bias=False, |
| 701 | ) |
| 702 | # Post-PLE norm: output = norm(x) * weight |
| 703 | self.post_per_layer_input_norm = RMSNorm( |
| 704 | config.hidden_size, eps=config.rms_norm_eps |
| 705 | ) |
| 706 | else: |
| 707 | self.per_layer_input_gate = None |
| 708 | self.per_layer_projection = None |
| 709 | self.post_per_layer_input_norm = None |
| 710 | |
| 711 | # Layer scalar (loaded from checkpoint) — applies to ALL text layers |
| 712 | self.register_buffer("layer_scalar", torch.ones(1)) |
| 713 | |
| 714 | def forward( |
| 715 | self, |
| 716 | positions: torch.Tensor, |
| 717 | hidden_states: torch.Tensor, |
| 718 | residual: torch.Tensor | None, |
| 719 | per_layer_input: torch.Tensor | None = None, |
| 720 | **kwargs, |
| 721 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 722 | # Gemma4 residual pattern: |
| 723 | # 1. input_norm(x) → attn → post_attn_norm → ADD residual |
| 724 | # 2. pre_ff_norm → mlp → post_ff_norm → ADD residual |
| 725 | residual = hidden_states |
| 726 | |
| 727 | hidden_states = self.input_layernorm(residual) |
| 728 | |
| 729 | hidden_states = self.self_attn( |
| 730 | positions=positions, |
| 731 | hidden_states=hidden_states, |
| 732 | **kwargs, |
| 733 | ) |
| 734 | |
| 735 | hidden_states = self.post_attention_layernorm(hidden_states) |
| 736 | hidden_states = hidden_states + residual |
| 737 | residual = hidden_states |
| 738 | |
| 739 | # MLP runs unconditionally (same inputs for MoE and non-MoE) |
| 740 | hidden_states = self.pre_feedforward_layernorm(hidden_states) |
| 741 | hidden_states = self.mlp(hidden_states) |
| 742 | |
| 743 | if self.enable_moe_block: |
| 744 | assert self.post_feedforward_layernorm_1 is not None |
| 745 | assert self.pre_feedforward_layernorm_2 is not None |
| 746 | assert self.router is not None |
| 747 | assert self.experts is not None |
| 748 | assert self.post_feedforward_layernorm_2 is not None |
| 749 | hidden_states_1 = self.post_feedforward_layernorm_1(hidden_states) |
| 750 | |
| 751 | hidden_states_2 = self.pre_feedforward_layernorm_2(residual) |
| 752 | router_logits = self.router(residual) |
| 753 | hidden_states_2 = self.experts(hidden_states_2, router_logits) |
| 754 | hidden_states_2 = self.post_feedforward_layernorm_2(hidden_states_2) |
| 755 | |
| 756 | # Combine MLP and MoE outputs |
| 757 | hidden_states = hidden_states_1 + hidden_states_2 |
| 758 | |
| 759 | hidden_states = self.post_feedforward_layernorm(hidden_states) |
| 760 | hidden_states = hidden_states + residual |
| 761 | |
| 762 | # Apply PLE (Per-Layer Embedding) if configured |
| 763 | if per_layer_input is not None and self.per_layer_input_gate is not None: |
| 764 | assert self.per_layer_projection is not None |
| 765 | assert self.post_per_layer_input_norm is not None |
| 766 | gate = self.per_layer_input_gate(hidden_states) |
| 767 | gate = torch.nn.functional.gelu(gate, approximate="tanh") |
| 768 | gated_per_layer = gate * per_layer_input |
| 769 | per_layer_contribution = self.per_layer_projection(gated_per_layer) |
| 770 | per_layer_contribution = self.post_per_layer_input_norm( |
| 771 | per_layer_contribution |
| 772 | ) |
| 773 | hidden_states = hidden_states + per_layer_contribution |
| 774 | |
| 775 | # Apply layer scalar for full-attention layers |
| 776 | # Apply per-layer scalar (all text layers) |
| 777 | hidden_states = hidden_states * self.layer_scalar |
| 778 | |
| 779 | return hidden_states, None |
| 780 | |
| 781 | |
| 782 | def _run_decoder_layers( |
| 783 | decoder_layers: list[Gemma4DecoderLayer], |
| 784 | layer_idx_start: int, |
| 785 | positions: torch.Tensor, |
| 786 | hidden_states: torch.Tensor, |
| 787 | per_layer_inputs: torch.Tensor | None = None, |
| 788 | **kwargs, |
| 789 | ) -> torch.Tensor: |
| 790 | """Run a slice of decoder layers with PLE extraction.""" |
| 791 | residual = None |
| 792 | for idx, layer in enumerate(decoder_layers): |
| 793 | layer_idx = idx + layer_idx_start |
| 794 | layer_per_input = ( |
| 795 | per_layer_inputs[:, layer_idx, :] if per_layer_inputs is not None else None |
| 796 | ) |
| 797 | hidden_states, residual = layer( |
| 798 | positions, |
| 799 | hidden_states, |
| 800 | residual, |
| 801 | per_layer_input=layer_per_input, |
| 802 | **kwargs, |
| 803 | ) |
| 804 | return hidden_states |
| 805 | |
| 806 | |
| 807 | @support_torch_compile( |
| 808 | enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill |
| 809 | ) |
| 810 | class Gemma4SelfDecoderLayers(nn.Module): |
| 811 | """Compiled wrapper: embedding + non-KV-shared layers (YOCO first half). |
| 812 | |
| 813 | Owns the embedding and PLE modules so they are inside the compiled |
| 814 | graph. Gemma4Model delegates embedding methods here. |
| 815 | """ |
| 816 | |
| 817 | def __init__( |
| 818 | self, |
| 819 | *, |
| 820 | vllm_config: VllmConfig, |
| 821 | prefix: str = "", |
| 822 | decoder_layers: list[Gemma4DecoderLayer], |
| 823 | layer_idx_start: int, |
| 824 | embed_tokens: VocabParallelEmbedding, |
| 825 | normalizer: torch.Tensor, |
| 826 | embed_tokens_per_layer: VocabParallelEmbedding | None, |
| 827 | embed_scale_per_layer: torch.Tensor | None, |
| 828 | per_layer_model_projection: ColumnParallelLinear | None, |
| 829 | per_layer_projection_norm: RMSNorm | None, |
| 830 | per_layer_input_scale: torch.Tensor | None, |
| 831 | per_layer_projection_scale: torch.Tensor | None, |
| 832 | ): |
| 833 | super().__init__() |
| 834 | self.decoder_layers = decoder_layers |
| 835 | self.layer_idx_start = layer_idx_start |
| 836 | |
| 837 | config = _get_text_config(vllm_config.model_config.hf_config) |
| 838 | self.config = config |
| 839 | self.hidden_size_per_layer_input = getattr( |
| 840 | config, "hidden_size_per_layer_input", 0 |
| 841 | ) |
| 842 | self.vocab_size_per_layer_input = getattr( |
| 843 | config, "vocab_size_per_layer_input", config.vocab_size |
| 844 | ) |
| 845 | |
| 846 | # Shared references to modules owned by Gemma4Model — must be |
| 847 | # inside this nn.Module so torch.compile captures them. |
| 848 | self.embed_tokens = embed_tokens |
| 849 | self.embed_tokens_per_layer = embed_tokens_per_layer |
| 850 | self.per_layer_model_projection = per_layer_model_projection |
| 851 | self.per_layer_projection_norm = per_layer_projection_norm |
| 852 | |
| 853 | # Register the shared scalars as buffers so they follow the model on |
| 854 | # .to(device); plain attributes stay on CPU. |
| 855 | self.register_buffer("normalizer", normalizer, persistent=False) |
| 856 | self.register_buffer( |
| 857 | "embed_scale_per_layer", embed_scale_per_layer, persistent=False |
| 858 | ) |
| 859 | self.register_buffer( |
| 860 | "per_layer_input_scale", per_layer_input_scale, persistent=False |
| 861 | ) |
| 862 | self.register_buffer( |
| 863 | "per_layer_projection_scale", per_layer_projection_scale, persistent=False |
| 864 | ) |
| 865 | |
| 866 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 867 | return self.embed_tokens(input_ids) * self.normalizer |
| 868 | |
| 869 | def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None: |
| 870 | """Get per-layer embeddings from embed_tokens_per_layer. |
| 871 | |
| 872 | Returns: |
| 873 | Per-layer embeddings (num_tokens, num_layers, |
| 874 | hidden_size_per_layer_input) |
| 875 | |
| 876 | """ |
| 877 | if self.embed_tokens_per_layer is None: |
| 878 | return None |
| 879 | per_layer_inputs_mask = torch.logical_and( |
| 880 | input_ids >= 0, |
| 881 | input_ids < self.vocab_size_per_layer_input, |
| 882 | ) |
| 883 | per_layer_inputs_tokens = torch.where( |
| 884 | per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids) |
| 885 | ) |
| 886 | per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens) |
| 887 | per_layer_embeds = per_layer_embeds * self.embed_scale_per_layer |
| 888 | return per_layer_embeds.reshape( |
| 889 | *input_ids.shape, |
| 890 | self.config.num_hidden_layers, |
| 891 | self.hidden_size_per_layer_input, |
| 892 | ) |
| 893 | |
| 894 | def project_per_layer_inputs( |
| 895 | self, |
| 896 | inputs_embeds: torch.Tensor, |
| 897 | per_layer_inputs: torch.Tensor | None, |
| 898 | ) -> torch.Tensor | None: |
| 899 | """Project inputs_embeds and combine with per_layer_inputs. |
| 900 | |
| 901 | Steps: |
| 902 | 1. Project inputs_embeds: hidden_size → total_ple_dim |
| 903 | 2. Scale by hidden_size^{-0.5} |
| 904 | 3. Reshape to (num_tokens, num_layers, per_layer_dim) |
| 905 | 4. Normalize with per_layer_projection_norm |
| 906 | 5. Combine: (projection + per_layer_inputs) * 1/sqrt(2) |
| 907 | """ |
| 908 | if self.per_layer_model_projection is None: |
| 909 | return None |
| 910 | assert self.per_layer_projection_norm is not None |
| 911 | per_layer_projection = self.per_layer_model_projection(inputs_embeds) |
| 912 | per_layer_projection = per_layer_projection * self.per_layer_projection_scale |
| 913 | per_layer_projection = per_layer_projection.reshape( |
| 914 | *inputs_embeds.shape[:-1], |
| 915 | self.config.num_hidden_layers, |
| 916 | self.hidden_size_per_layer_input, |
| 917 | ) |
| 918 | per_layer_projection = self.per_layer_projection_norm(per_layer_projection) |
| 919 | if per_layer_inputs is None: |
| 920 | return per_layer_projection |
| 921 | return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale |
| 922 | |
| 923 | def forward( |
| 924 | self, |
| 925 | input_ids: torch.Tensor | None, |
| 926 | positions: torch.Tensor, |
| 927 | inputs_embeds: torch.Tensor | None = None, |
| 928 | per_layer_inputs: torch.Tensor | None = None, |
| 929 | **kwargs, |
| 930 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 931 | if inputs_embeds is not None: |
| 932 | hidden_states = inputs_embeds |
| 933 | per_layer_inputs = self.project_per_layer_inputs( |
| 934 | hidden_states, per_layer_inputs |
| 935 | ) |
| 936 | else: |
| 937 | hidden_states = self.embed_input_ids(input_ids) |
| 938 | per_layer_embeds = self.get_per_layer_inputs(input_ids) |
| 939 | per_layer_inputs = self.project_per_layer_inputs( |
| 940 | hidden_states, per_layer_embeds |
| 941 | ) |
| 942 | |
| 943 | hidden_states = _run_decoder_layers( |
| 944 | self.decoder_layers, |
| 945 | self.layer_idx_start, |
| 946 | positions, |
| 947 | hidden_states, |
| 948 | per_layer_inputs, |
| 949 | **kwargs, |
| 950 | ) |
| 951 | return hidden_states, per_layer_inputs |
| 952 | |
| 953 | |
| 954 | @support_torch_compile( |
| 955 | enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill |
| 956 | ) |
| 957 | class Gemma4CrossDecoderLayers(nn.Module): |
| 958 | """Cross-decoder layers (YOCO second half, KV-shared).""" |
| 959 | |
| 960 | def __init__( |
| 961 | self, |
| 962 | *, |
| 963 | vllm_config: VllmConfig, |
| 964 | prefix: str = "", |
| 965 | decoder_layers: list[Gemma4DecoderLayer], |
| 966 | layer_idx_start: int, |
| 967 | ): |
| 968 | super().__init__() |
| 969 | self.decoder_layers = decoder_layers |
| 970 | self.layer_idx_start = layer_idx_start |
| 971 | |
| 972 | def forward( |
| 973 | self, |
| 974 | positions: torch.Tensor, |
| 975 | hidden_states: torch.Tensor, |
| 976 | per_layer_inputs: torch.Tensor | None = None, |
| 977 | **kwargs, |
| 978 | ) -> torch.Tensor: |
| 979 | return _run_decoder_layers( |
| 980 | self.decoder_layers, |
| 981 | self.layer_idx_start, |
| 982 | positions, |
| 983 | hidden_states, |
| 984 | per_layer_inputs, |
| 985 | **kwargs, |
| 986 | ) |
| 987 | |
| 988 | |
| 989 | @support_torch_compile( |
| 990 | enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill |
| 991 | ) |
| 992 | class Gemma4Model(nn.Module, EagleModelMixin, SupportsQuant): |
| 993 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 994 | super().__init__() |
| 995 | config = _get_text_config(vllm_config.model_config.hf_config) |
| 996 | cache_config = vllm_config.cache_config |
| 997 | quant_config = vllm_config.quant_config |
| 998 | self.config = config |
| 999 | self.quant_config = quant_config |
| 1000 | self.hf_to_vllm_mapper = _gemma4_layer_weights_mapper(config) |
| 1001 | |
| 1002 | # PLE config values (default to 0 if not present — disables PLE) |
| 1003 | self.hidden_size_per_layer_input = getattr( |
| 1004 | config, "hidden_size_per_layer_input", 0 |
| 1005 | ) |
| 1006 | self.vocab_size_per_layer_input = getattr( |
| 1007 | config, "vocab_size_per_layer_input", config.vocab_size |
| 1008 | ) |
| 1009 | |
| 1010 | self.embed_tokens = VocabParallelEmbedding( |
| 1011 | config.vocab_size, |
| 1012 | config.hidden_size, |
| 1013 | quant_config=quant_config, |
| 1014 | prefix=f"{prefix}.embed_tokens", |
| 1015 | ) |
| 1016 | |
| 1017 | # Per-Layer Embedding (PLE) components |
| 1018 | self.embed_tokens_per_layer: VocabParallelEmbedding | None |
| 1019 | self.per_layer_model_projection: ColumnParallelLinear | None |
| 1020 | self.per_layer_projection_norm: RMSNorm | None |
| 1021 | if ( |
| 1022 | self.hidden_size_per_layer_input is not None |
| 1023 | and self.hidden_size_per_layer_input > 0 |
| 1024 | ): |
| 1025 | total_ple_dim = self.hidden_size_per_layer_input * config.num_hidden_layers |
| 1026 | self.embed_tokens_per_layer = VocabParallelEmbedding( |
| 1027 | self.vocab_size_per_layer_input, |
| 1028 | total_ple_dim, |
| 1029 | quant_config=quant_config, |
| 1030 | prefix=f"{prefix}.embed_tokens_per_layer", |
| 1031 | ) |
| 1032 | # Scaled embedding factor (from config, not hardcoded) |
| 1033 | # Register as buffer so it moves to GPU with the model |
| 1034 | # and interacts correctly with torch.compile AOT caching. |
| 1035 | self.register_buffer( |
| 1036 | "embed_scale_per_layer", |
| 1037 | torch.tensor(self.hidden_size_per_layer_input**0.5), |
| 1038 | persistent=False, |
| 1039 | ) |
| 1040 | # Projection: hidden_size → total_ple_dim |
| 1041 | # ColumnParallelLinear with gather_output=True |
| 1042 | self.per_layer_model_projection = ColumnParallelLinear( |
| 1043 | config.hidden_size, |
| 1044 | total_ple_dim, |
| 1045 | bias=False, |
| 1046 | gather_output=True, |
| 1047 | return_bias=False, |
| 1048 | quant_config=quant_config, |
| 1049 | prefix=f"{prefix}.per_layer_model_projection", |
| 1050 | ) |
| 1051 | # PLE projection norm: output = norm(x) * weight |
| 1052 | self.per_layer_projection_norm = RMSNorm( |
| 1053 | self.hidden_size_per_layer_input, |
| 1054 | eps=config.rms_norm_eps, |
| 1055 | ) |
| 1056 | # Scale factor for combining projection + per_layer_inputs |
| 1057 | # Register as buffer so it moves to GPU with the model |
| 1058 | # and interacts correctly with torch.compile AOT caching. |
| 1059 | self.register_buffer( |
| 1060 | "per_layer_input_scale", |
| 1061 | torch.rsqrt(torch.tensor(2.0)), |
| 1062 | persistent=False, |
| 1063 | ) |
| 1064 | # Scaled projection: multiply output by hidden_size**-0.5. |
| 1065 | # Register as buffer for GPU placement and torch.compile. |
| 1066 | self.register_buffer( |
| 1067 | "per_layer_projection_scale", |
| 1068 | torch.tensor(config.hidden_size**-0.5), |
| 1069 | persistent=False, |
| 1070 | ) |
| 1071 | else: |
| 1072 | self.embed_tokens_per_layer = None |
| 1073 | self.embed_scale_per_layer = None |
| 1074 | self.per_layer_model_projection = None |
| 1075 | self.per_layer_projection_norm = None |
| 1076 | self.per_layer_input_scale = None |
| 1077 | self.per_layer_projection_scale = None |
| 1078 | |
| 1079 | self.start_layer, self.end_layer, self.layers = make_layers( |
| 1080 | config.num_hidden_layers, |
| 1081 | lambda prefix: Gemma4DecoderLayer( |
| 1082 | config, |
| 1083 | cache_config=cache_config, |
| 1084 | quant_config=quant_config, |
| 1085 | prefix=prefix, |
| 1086 | ), |
| 1087 | prefix=f"{prefix}.layers", |
| 1088 | ) |
| 1089 | # Final norm: output = norm(x) * weight |
| 1090 | self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 1091 | |
| 1092 | # Embedding scale = sqrt(hidden_size), cast to model dtype to avoid |
| 1093 | # mixed-precision drift from bf16 * fp32 across deep stacks. |
| 1094 | self.register_buffer( |
| 1095 | "normalizer", |
| 1096 | torch.tensor( |
| 1097 | config.hidden_size**0.5, |
| 1098 | dtype=vllm_config.model_config.dtype, |
| 1099 | ), |
| 1100 | persistent=False, |
| 1101 | ) |
| 1102 | |
| 1103 | # --- You Only Cache Once (YOCO) split for fast prefill --- |
| 1104 | first_kv_shared_layer_idx = config.num_hidden_layers - getattr( |
| 1105 | config, "num_kv_shared_layers", 0 |
| 1106 | ) |
| 1107 | |
| 1108 | from vllm.compilation.backends import set_model_tag |
| 1109 | |
| 1110 | # Layers 0..(K-1) are self-decoder layers in YOCO |
| 1111 | with set_model_tag("self_decoder"): |
| 1112 | self.self_decoder = Gemma4SelfDecoderLayers( |
| 1113 | vllm_config=vllm_config, |
| 1114 | prefix=f"{prefix}.self_decoder", |
| 1115 | decoder_layers=self.layers[:first_kv_shared_layer_idx], |
| 1116 | layer_idx_start=0, |
| 1117 | embed_tokens=self.embed_tokens, |
| 1118 | normalizer=self.normalizer, |
| 1119 | embed_tokens_per_layer=getattr(self, "embed_tokens_per_layer", None), |
| 1120 | embed_scale_per_layer=getattr(self, "embed_scale_per_layer", None), |
| 1121 | per_layer_model_projection=getattr( |
| 1122 | self, "per_layer_model_projection", None |
| 1123 | ), |
| 1124 | per_layer_projection_norm=getattr( |
| 1125 | self, "per_layer_projection_norm", None |
| 1126 | ), |
| 1127 | per_layer_input_scale=getattr(self, "per_layer_input_scale", None), |
| 1128 | per_layer_projection_scale=getattr( |
| 1129 | self, "per_layer_projection_scale", None |
| 1130 | ), |
| 1131 | ) |
| 1132 | # Layers K..(N-1) are cross-decoder layers in YOCO |
| 1133 | with set_model_tag("cross_decoder"): |
| 1134 | self.cross_decoder = Gemma4CrossDecoderLayers( |
| 1135 | vllm_config=vllm_config, |
| 1136 | prefix=f"{prefix}.cross_decoder", |
| 1137 | decoder_layers=self.layers[first_kv_shared_layer_idx:], |
| 1138 | layer_idx_start=first_kv_shared_layer_idx, |
| 1139 | ) |
| 1140 | |
| 1141 | self.fast_prefill_enabled = cache_config.kv_sharing_fast_prefill |
| 1142 | |
| 1143 | if self.fast_prefill_enabled: |
| 1144 | # Allocate static buffers for CUDAGraph |
| 1145 | max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 1146 | device = next(self.parameters()).device |
| 1147 | self.positions = torch.zeros( |
| 1148 | max_num_tokens, dtype=torch.int64, device=device |
| 1149 | ) |
| 1150 | self.hidden_states = torch.zeros( |
| 1151 | (max_num_tokens, config.hidden_size), |
| 1152 | dtype=vllm_config.model_config.dtype, |
| 1153 | device=device, |
| 1154 | ) |
| 1155 | if ( |
| 1156 | self.hidden_size_per_layer_input |
| 1157 | and self.hidden_size_per_layer_input > 0 |
| 1158 | ): |
| 1159 | self.per_layer_inputs = torch.zeros( |
| 1160 | ( |
| 1161 | max_num_tokens, |
| 1162 | config.num_hidden_layers, |
| 1163 | self.hidden_size_per_layer_input, |
| 1164 | ), |
| 1165 | dtype=vllm_config.model_config.dtype, |
| 1166 | device=device, |
| 1167 | ) |
| 1168 | else: |
| 1169 | self.per_layer_inputs = None |
| 1170 | |
| 1171 | # Custom factory that includes per_layer_inputs for PLE-enabled PP. |
| 1172 | # per_layer_inputs has shape (batch, num_layers, per_layer_dim), |
| 1173 | # which differs from the standard (batch, hidden_size) shape, |
| 1174 | # so we can't use the default factory. |
| 1175 | ple_dim = self.hidden_size_per_layer_input |
| 1176 | num_layers = config.num_hidden_layers |
| 1177 | hidden_size = config.hidden_size |
| 1178 | |
| 1179 | def _make_empty_intermediate_tensors( |
| 1180 | batch_size: int, |
| 1181 | dtype: torch.dtype, |
| 1182 | device: torch.device, |
| 1183 | ) -> IntermediateTensors: |
| 1184 | tensors: dict[str, torch.Tensor] = { |
| 1185 | "hidden_states": torch.zeros( |
| 1186 | (batch_size, hidden_size), |
| 1187 | dtype=dtype, |
| 1188 | device=device, |
| 1189 | ), |
| 1190 | } |
| 1191 | if ple_dim and ple_dim > 0: |
| 1192 | tensors["per_layer_inputs"] = torch.zeros( |
| 1193 | (batch_size, num_layers, ple_dim), |
| 1194 | dtype=dtype, |
| 1195 | device=device, |
| 1196 | ) |
| 1197 | return IntermediateTensors(tensors) |
| 1198 | |
| 1199 | self.make_empty_intermediate_tensors = _make_empty_intermediate_tensors |
| 1200 | |
| 1201 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 1202 | return self.self_decoder.embed_input_ids(input_ids) |
| 1203 | |
| 1204 | def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None: |
| 1205 | """Get per-layer embeddings from embed_tokens_per_layer. |
| 1206 | |
| 1207 | Returns: |
| 1208 | Per-layer embeddings (num_tokens, num_layers, |
| 1209 | hidden_size_per_layer_input) |
| 1210 | |
| 1211 | """ |
| 1212 | return self.self_decoder.get_per_layer_inputs(input_ids) |
| 1213 | |
| 1214 | def project_per_layer_inputs( |
| 1215 | self, |
| 1216 | inputs_embeds: torch.Tensor, |
| 1217 | per_layer_inputs: torch.Tensor | None, |
| 1218 | ) -> torch.Tensor | None: |
| 1219 | """Project inputs_embeds and combine with per_layer_inputs. |
| 1220 | |
| 1221 | Steps: |
| 1222 | 1. Project inputs_embeds: hidden_size → total_ple_dim |
| 1223 | 2. Scale by hidden_size^{-0.5} |
| 1224 | 3. Reshape to (num_tokens, num_layers, per_layer_dim) |
| 1225 | 4. Normalize with per_layer_projection_norm |
| 1226 | 5. Combine: (projection + per_layer_inputs) * 1/sqrt(2) |
| 1227 | """ |
| 1228 | return self.self_decoder.project_per_layer_inputs( |
| 1229 | inputs_embeds, per_layer_inputs |
| 1230 | ) |
| 1231 | |
| 1232 | def fast_prefill_forward( |
| 1233 | self, |
| 1234 | input_ids: torch.Tensor | None, |
| 1235 | positions: torch.Tensor, |
| 1236 | inputs_embeds: torch.Tensor | None = None, |
| 1237 | per_layer_inputs: torch.Tensor | None = None, |
| 1238 | **kwargs, |
| 1239 | ) -> torch.Tensor: |
| 1240 | logits_indices_padded, num_logits_indices = None, None |
| 1241 | attn_metadata = get_forward_context().attn_metadata |
| 1242 | |
| 1243 | if attn_metadata is not None: |
| 1244 | assert isinstance(attn_metadata, dict) |
| 1245 | layer_attn_metadata = attn_metadata[ |
| 1246 | self.layers[-1].self_attn.attn.layer_name |
| 1247 | ] |
| 1248 | if isinstance(layer_attn_metadata, KVSharingFastPrefillMetadata): |
| 1249 | logits_indices_padded = layer_attn_metadata.logits_indices_padded |
| 1250 | num_logits_indices = layer_attn_metadata.num_logits_indices |
| 1251 | |
| 1252 | batch_size = positions.size(0) |
| 1253 | self.positions[:batch_size].copy_(positions) |
| 1254 | self_decoder_hidden_states, per_layer_inputs = self.self_decoder( |
| 1255 | input_ids=input_ids, |
| 1256 | positions=self.positions[:batch_size], |
| 1257 | inputs_embeds=inputs_embeds, |
| 1258 | per_layer_inputs=per_layer_inputs, |
| 1259 | **kwargs, |
| 1260 | ) |
| 1261 | |
| 1262 | if logits_indices_padded is None: |
| 1263 | logits_indices_padded = torch.arange( |
| 1264 | batch_size, |
| 1265 | dtype=positions.dtype, |
| 1266 | device=positions.device, |
| 1267 | ) |
| 1268 | |
| 1269 | # NOTE: Keep .clone() until fix in |
| 1270 | # https://github.com/vllm-project/vllm/pull/22282 |
| 1271 | hidden_states = self_decoder_hidden_states.clone() |
| 1272 | |
| 1273 | num_padded = logits_indices_padded.size(0) |
| 1274 | self.positions[:num_padded].copy_(positions[logits_indices_padded]) |
| 1275 | self.hidden_states[:num_padded].copy_( |
| 1276 | self_decoder_hidden_states[logits_indices_padded] |
| 1277 | ) |
| 1278 | if self.per_layer_inputs is not None and per_layer_inputs is not None: |
| 1279 | self.per_layer_inputs[:num_padded].copy_( |
| 1280 | per_layer_inputs[logits_indices_padded] |
| 1281 | ) |
| 1282 | |
| 1283 | # Update batch_descriptor so the cross-decoder's piecewise |
| 1284 | # CUDAGraphWrapper dispatches to the correct (reduced) batch size. |
| 1285 | forward_context = get_forward_context() |
| 1286 | orig_batch_desc = forward_context.batch_descriptor |
| 1287 | if orig_batch_desc is not None: |
| 1288 | forward_context.batch_descriptor = replace( |
| 1289 | orig_batch_desc, num_tokens=num_padded |
| 1290 | ) |
| 1291 | |
| 1292 | cross_per_layer = ( |
| 1293 | self.per_layer_inputs[:num_padded] |
| 1294 | if self.per_layer_inputs is not None |
| 1295 | else None |
| 1296 | ) |
| 1297 | cross_hidden_states = self.cross_decoder( |
| 1298 | self.positions[:num_padded], |
| 1299 | self.hidden_states[:num_padded], |
| 1300 | cross_per_layer, |
| 1301 | **kwargs, |
| 1302 | ) |
| 1303 | |
| 1304 | # Restore the original batch_descriptor |
| 1305 | forward_context.batch_descriptor = orig_batch_desc |
| 1306 | |
| 1307 | if num_logits_indices is not None: |
| 1308 | assert num_logits_indices > 0 |
| 1309 | hidden_states[logits_indices_padded[:num_logits_indices]] = ( |
| 1310 | cross_hidden_states[:num_logits_indices] |
| 1311 | ) |
| 1312 | else: |
| 1313 | hidden_states = cross_hidden_states |
| 1314 | |
| 1315 | return hidden_states |
| 1316 | |
| 1317 | def forward( |
| 1318 | self, |
| 1319 | input_ids: torch.Tensor | None, |
| 1320 | positions: torch.Tensor, |
| 1321 | intermediate_tensors: IntermediateTensors | None, |
| 1322 | inputs_embeds: torch.Tensor | None = None, |
| 1323 | per_layer_inputs: torch.Tensor | None = None, |
| 1324 | **kwargs, |
| 1325 | ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: |
| 1326 | if self.fast_prefill_enabled: |
| 1327 | hidden_states = self.fast_prefill_forward( |
| 1328 | input_ids, |
| 1329 | positions, |
| 1330 | inputs_embeds, |
| 1331 | per_layer_inputs, |
| 1332 | **kwargs, |
| 1333 | ) |
| 1334 | hidden_states = self.norm(hidden_states) |
| 1335 | return hidden_states |
| 1336 | |
| 1337 | # Normal (non-fast-prefill) path with PP support |
| 1338 | if get_pp_group().is_first_rank: |
| 1339 | if inputs_embeds is not None: |
| 1340 | hidden_states = inputs_embeds |
| 1341 | # When called from the multimodal wrapper, raw PLE |
| 1342 | # embeddings are pre-computed and passed explicitly. |
| 1343 | # Project them through per_layer_model_projection. |
| 1344 | per_layer_inputs = self.project_per_layer_inputs( |
| 1345 | hidden_states, per_layer_inputs |
| 1346 | ) |
| 1347 | else: |
| 1348 | hidden_states = self.embed_input_ids(input_ids) |
| 1349 | # Compute per-layer inputs for PLE |
| 1350 | per_layer_embeds = self.get_per_layer_inputs(input_ids) |
| 1351 | per_layer_inputs = self.project_per_layer_inputs( |
| 1352 | hidden_states, per_layer_embeds |
| 1353 | ) |
| 1354 | else: |
| 1355 | assert intermediate_tensors is not None |
| 1356 | hidden_states = intermediate_tensors["hidden_states"] |
| 1357 | if per_layer_inputs is not None: |
| 1358 | per_layer_inputs = intermediate_tensors["per_layer_inputs"] |
| 1359 | residual = None |
| 1360 | aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) |
| 1361 | for layer_idx, layer in enumerate( |
| 1362 | islice(self.layers, self.start_layer, self.end_layer) |
| 1363 | ): |
| 1364 | # Extract the per-layer embedding for this specific layer |
| 1365 | if per_layer_inputs is not None: |
| 1366 | actual_layer_idx = self.start_layer + layer_idx |
| 1367 | layer_per_input = per_layer_inputs[ |
| 1368 | :, actual_layer_idx, : |
| 1369 | ] # (num_tokens, per_layer_dim) |
| 1370 | else: |
| 1371 | layer_per_input = None |
| 1372 | hidden_states, residual = layer( |
| 1373 | positions, |
| 1374 | hidden_states, |
| 1375 | residual, |
| 1376 | per_layer_input=layer_per_input, |
| 1377 | **kwargs, |
| 1378 | ) |
| 1379 | self._maybe_add_hidden_state( |
| 1380 | aux_hidden_states, layer_idx + 1, hidden_states, residual |
| 1381 | ) |
| 1382 | if not get_pp_group().is_last_rank: |
| 1383 | tensors: dict[str, torch.Tensor] = { |
| 1384 | "hidden_states": hidden_states, |
| 1385 | } |
| 1386 | if per_layer_inputs is not None: |
| 1387 | tensors["per_layer_inputs"] = per_layer_inputs |
| 1388 | return IntermediateTensors(tensors) |
| 1389 | # Gemma4 incorporates residual into hidden_states directly |
| 1390 | # Apply norm without residual fusion when possible. |
| 1391 | if residual is None: |
| 1392 | hidden_states = self.norm(hidden_states) |
| 1393 | else: |
| 1394 | hidden_states, _ = self.norm(hidden_states, residual) |
| 1395 | |
| 1396 | if len(aux_hidden_states) > 0: |
| 1397 | return hidden_states, aux_hidden_states |
| 1398 | return hidden_states |
| 1399 | |
| 1400 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1401 | loader = AutoWeightsLoader(self) |
| 1402 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |
| 1403 | |
| 1404 | |
| 1405 | class Gemma4ForCausalLM( |
| 1406 | nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3 |
| 1407 | ): |
| 1408 | hf_to_vllm_mapper = WeightsMapper( |
| 1409 | orig_to_new_substr={ |
| 1410 | # Multimodal weights are handled by the multimodal wrapper. |
| 1411 | "audio_tower.": None, |
| 1412 | "vision_tower.": None, |
| 1413 | "embed_audio.": None, |
| 1414 | "embed_vision.": None, |
| 1415 | }, |
| 1416 | orig_to_new_prefix={ |
| 1417 | # Gemma4ForConditionalGeneration already loads the text stack |
| 1418 | # from `model.language_model.*`. We reuse that same checkpoint |
| 1419 | # and adapter naming for the text-only Gemma4ForCausalLM path, |
| 1420 | # so LoRA keys from the conditional wrapper map onto `model.*`. |
| 1421 | "model.language_model.": "model.", |
| 1422 | }, |
| 1423 | ) |
| 1424 | # KV-shared layers only have q_proj, so qkv_proj packing applies to |
| 1425 | # the other layers. |
| 1426 | packed_modules_mapping = { |
| 1427 | "qkv_proj": [ |
| 1428 | "q_proj", |
| 1429 | "k_proj", |
| 1430 | "v_proj", |
| 1431 | ], |
| 1432 | "gate_up_proj": [ |
| 1433 | "gate_proj", |
| 1434 | "up_proj", |
| 1435 | ], |
| 1436 | } |
| 1437 | |
| 1438 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 1439 | config = _get_text_config(vllm_config.model_config.hf_config) |
| 1440 | quant_config = vllm_config.quant_config |
| 1441 | |
| 1442 | super().__init__() |
| 1443 | self.config = config |
| 1444 | self.quant_config = quant_config |
| 1445 | self.model = Gemma4Model( |
| 1446 | vllm_config=vllm_config, |
| 1447 | prefix=maybe_prefix(prefix, "model"), |
| 1448 | ) |
| 1449 | |
| 1450 | self.lm_head = ParallelLMHead( |
| 1451 | config.vocab_size, |
| 1452 | config.hidden_size, |
| 1453 | quant_config=quant_config, |
| 1454 | prefix=maybe_prefix(prefix, "lm_head"), |
| 1455 | ) |
| 1456 | if config.tie_word_embeddings: |
| 1457 | self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) |
| 1458 | |
| 1459 | self.logits_processor = LogitsProcessor( |
| 1460 | config.vocab_size, |
| 1461 | soft_cap=getattr(config, "final_logit_softcapping", None), |
| 1462 | ) |
| 1463 | self.make_empty_intermediate_tensors = ( |
| 1464 | self.model.make_empty_intermediate_tensors |
| 1465 | ) |
| 1466 | |
| 1467 | # --- MixtureOfExperts protocol --- |
| 1468 | self.moe_layers: list[nn.Module] = [ |
| 1469 | layer.experts |
| 1470 | for layer in self.model.layers |
| 1471 | if not isinstance(layer, PPMissingLayer) and layer.experts is not None |
| 1472 | ] |
| 1473 | self.num_moe_layers = len(self.moe_layers) |
| 1474 | |
| 1475 | num_experts = config.num_experts if self.num_moe_layers else 0 |
| 1476 | self.num_logical_experts = num_experts |
| 1477 | self.num_physical_experts = num_experts |
| 1478 | self.num_local_physical_experts = num_experts |
| 1479 | self.num_routed_experts = num_experts |
| 1480 | |
| 1481 | self.num_expert_groups = 1 |
| 1482 | self.num_shared_experts = 0 |
| 1483 | self.num_redundant_experts = 0 |
| 1484 | |
| 1485 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 1486 | return self.model.embed_input_ids(input_ids) |
| 1487 | |
| 1488 | def forward( |
| 1489 | self, |
| 1490 | input_ids: torch.Tensor, |
| 1491 | positions: torch.Tensor, |
| 1492 | intermediate_tensors: IntermediateTensors | None = None, |
| 1493 | inputs_embeds: torch.Tensor | None = None, |
| 1494 | **kwargs, |
| 1495 | ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: |
| 1496 | hidden_states = self.model( |
| 1497 | input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs |
| 1498 | ) |
| 1499 | return hidden_states |
| 1500 | |
| 1501 | def compute_logits( |
| 1502 | self, |
| 1503 | hidden_states: torch.Tensor, |
| 1504 | ) -> torch.Tensor | None: |
| 1505 | return self.logits_processor(self.lm_head, hidden_states) |
| 1506 | |
| 1507 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 1508 | loader = AutoWeightsLoader(self) |
| 1509 | return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) |