| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Gemma4 MTP (Multi-Token Prediction) proposer for speculative decoding. |
| 4 | |
| 5 | The Gemma4 assistant model runs all decoder layers per draft step |
| 6 | (producing one token), and all its attention layers share KV cache |
| 7 | with the target model via cross-model KV sharing. |
| 8 | """ |
| 9 | |
| 10 | from collections import defaultdict |
| 11 | from copy import copy |
| 12 | |
| 13 | import torch |
| 14 | import torch.nn as nn |
| 15 | |
| 16 | from vllm.config import VllmConfig, get_layers_from_vllm_config, replace |
| 17 | from vllm.distributed.parallel_state import get_pp_group |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 20 | from vllm.utils.torch_utils import current_stream |
| 21 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 22 | from vllm.v1.kv_cache_interface import ( |
| 23 | KVCacheConfig, |
| 24 | KVCacheSpec, |
| 25 | UniformTypeKVCacheSpecs, |
| 26 | ) |
| 27 | from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer |
| 28 | from vllm.v1.worker.utils import AttentionGroup |
| 29 | |
| 30 | logger = init_logger(__name__) |
| 31 | |
| 32 | |
| 33 | class Gemma4Proposer(SpecDecodeBaseProposer): |
| 34 | def __init__( |
| 35 | self, |
| 36 | vllm_config: VllmConfig, |
| 37 | device: torch.device, |
| 38 | runner=None, |
| 39 | ): |
| 40 | super().__init__( |
| 41 | vllm_config, |
| 42 | device, |
| 43 | pass_hidden_states_to_model=True, |
| 44 | runner=runner, |
| 45 | ) |
| 46 | # All draft steps predict from the same position (the last |
| 47 | # target-model position), so positions and seq_lens must not |
| 48 | # advance between steps. |
| 49 | self.constant_draft_positions = True |
| 50 | |
| 51 | # Per-group block tables for multi-group KV cache models. |
| 52 | # Populated by gpu_model_runner during _prepare_inputs. |
| 53 | self._per_group_block_tables: dict[int, torch.Tensor] = {} |
| 54 | |
| 55 | # Centroids CUDA graphs — populated in load_model if centroids |
| 56 | # masking is active. _centroids_sizes is pre-sorted for fast |
| 57 | # lookup in _greedy_sample. |
| 58 | self._centroids_sizes: list[int] = [] |
| 59 | self._centroids_graphs: dict[int, torch.cuda.CUDAGraph] = {} |
| 60 | self._centroids_inputs: dict[int, torch.Tensor] = {} |
| 61 | self._centroids_outputs: dict[int, torch.Tensor] = {} |
| 62 | |
| 63 | def set_per_group_block_table(self, gid: int, block_table: torch.Tensor) -> None: |
| 64 | self._per_group_block_tables[gid] = block_table |
| 65 | |
| 66 | def model_returns_tuple(self) -> bool: |
| 67 | # forward() returns (draft_hidden_states, backbone_hidden_states). |
| 68 | # The proposer uses draft_hidden_states for compute_logits and |
| 69 | # backbone_hidden_states for the hidden-state feedback buffer. |
| 70 | return True |
| 71 | |
| 72 | def build_per_group_and_layer_attn_metadata( |
| 73 | self, |
| 74 | common_attn_metadata: CommonAttentionMetadata, |
| 75 | draft_index: int = 0, |
| 76 | ) -> tuple[list[object], dict[str, object]]: |
| 77 | """Build attention metadata using the correct block table per group. |
| 78 | |
| 79 | Gemma4 has multiple KV cache groups (sliding vs full attention) |
| 80 | with different block tables. The base class receives a single |
| 81 | common_attn_metadata whose block_table belongs to one group. |
| 82 | We swap in the correct block table for each draft attention group. |
| 83 | """ |
| 84 | per_group_attn_metadata: list[object] = [] |
| 85 | per_layer_attn_metadata: dict[str, object] = {} |
| 86 | batch_size = common_attn_metadata.batch_size() |
| 87 | for attn_group in self.draft_attn_groups: |
| 88 | gid = attn_group.kv_cache_group_id |
| 89 | if gid in self._per_group_block_tables: |
| 90 | cm = copy(common_attn_metadata) |
| 91 | # Slice to actual batch size to match cu_seqlens_q dimension. |
| 92 | # The stored block tables may be padded (num_reqs_padded) from |
| 93 | # the target forward pass, but the drafter operates on the |
| 94 | # unpadded batch. |
| 95 | cm.block_table_tensor = self._per_group_block_tables[gid][:batch_size] |
| 96 | else: |
| 97 | cm = common_attn_metadata |
| 98 | attn_metadata = attn_group.get_metadata_builder().build_for_drafting( |
| 99 | common_attn_metadata=cm, draft_index=draft_index |
| 100 | ) |
| 101 | per_group_attn_metadata.append(attn_metadata) |
| 102 | for layer_name in attn_group.layer_names: |
| 103 | per_layer_attn_metadata[layer_name] = attn_metadata |
| 104 | return per_group_attn_metadata, per_layer_attn_metadata |
| 105 | |
| 106 | def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 107 | if self._centroids_sizes: |
| 108 | T = hidden_states.shape[0] |
| 109 | for size in self._centroids_sizes: |
| 110 | if size >= T: |
| 111 | self._centroids_inputs[size][:T].copy_(hidden_states) |
| 112 | self._centroids_graphs[size].replay() |
| 113 | return self._centroids_outputs[size][:T].clone() |
| 114 | return self.model.get_top_tokens(hidden_states) |
| 115 | return super()._greedy_sample(hidden_states) |
| 116 | |
| 117 | def _setup_centroids_cuda_graphs(self) -> None: |
| 118 | """Capture CUDA graphs for centroids get_top_tokens at key sizes.""" |
| 119 | masked_emb = self.model.masked_embedding |
| 120 | lm_head_weight = self.model._get_full_lm_head_weight() |
| 121 | capture_stream = current_stream() |
| 122 | |
| 123 | for size in [1, 2, 4, 8, 16, 32, 64]: |
| 124 | static_input = torch.zeros( |
| 125 | size, |
| 126 | masked_emb.hidden_size, |
| 127 | dtype=self.dtype, |
| 128 | device=self.device, |
| 129 | ) |
| 130 | for _ in range(3): |
| 131 | masked_emb.get_top_tokens(static_input, lm_head_weight) |
| 132 | torch.accelerator.synchronize() |
| 133 | |
| 134 | g = torch.cuda.CUDAGraph() |
| 135 | with torch.cuda.graph(g, stream=capture_stream): |
| 136 | static_output = masked_emb.get_top_tokens( |
| 137 | static_input, |
| 138 | lm_head_weight, |
| 139 | ) |
| 140 | self._centroids_graphs[size] = g |
| 141 | self._centroids_inputs[size] = static_input |
| 142 | self._centroids_outputs[size] = static_output |
| 143 | |
| 144 | self._centroids_sizes = sorted(self._centroids_graphs) |
| 145 | logger.info( |
| 146 | "Gemma4 MTP: captured centroids CUDA graphs for sizes %s.", |
| 147 | self._centroids_sizes, |
| 148 | ) |
| 149 | |
| 150 | def _create_draft_vllm_config(self) -> VllmConfig: |
| 151 | """Preserve the target's forced TRITON_ATTN backend for draft layers. |
| 152 | |
| 153 | Gemma4 forces TRITON_ATTN due to heterogeneous head dimensions |
| 154 | (head_dim=256 sliding, global_head_dim=512 full). The base class |
| 155 | resets attention_config.backend to None for draft models, causing |
| 156 | sliding layers to fall back to FLASH_ATTN which cannot handle |
| 157 | KV-shared cache. Override to carry the target's backend through. |
| 158 | """ |
| 159 | base = super()._create_draft_vllm_config() |
| 160 | target_backend = self.vllm_config.attention_config.backend |
| 161 | if target_backend is not None: |
| 162 | base = replace( |
| 163 | base, |
| 164 | attention_config=replace( |
| 165 | base.attention_config, |
| 166 | backend=target_backend, |
| 167 | ), |
| 168 | ) |
| 169 | return base |
| 170 | |
| 171 | def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None: |
| 172 | """Gemma4 MTP requires dim-mismatched embedding sharing. |
| 173 | |
| 174 | The draft checkpoint's embed_tokens is a draft-dim placeholder |
| 175 | (tied to lm_head so load_weights populates both); the model |
| 176 | expects it to be replaced by the target's backbone-dim embedding, |
| 177 | so the base class's embedding-dim equality guard must not apply. |
| 178 | """ |
| 179 | if get_pp_group().world_size != 1: |
| 180 | return |
| 181 | inner_model = getattr(target_language_model, "model", None) |
| 182 | target_embed_tokens = getattr(inner_model, "embed_tokens", None) |
| 183 | if target_embed_tokens is None: |
| 184 | raise AttributeError( |
| 185 | "Target model does not have an 'embed_tokens' attribute" |
| 186 | ) |
| 187 | del self.model.model.embed_tokens |
| 188 | self.model.model.embed_tokens = target_embed_tokens |
| 189 | logger.info( |
| 190 | "Gemma4 MTP: sharing target model's backbone-dim embed_tokens " |
| 191 | "with the draft model." |
| 192 | ) |
| 193 | |
| 194 | def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None: |
| 195 | """Gemma4 MTP always keeps its own draft-dim lm_head. |
| 196 | |
| 197 | The draft model's lm_head operates in draft hidden_size (e.g. 256), |
| 198 | which differs from the target's backbone hidden_size (e.g. 1536). |
| 199 | Sharing would break compute_logits (and centroids masking when |
| 200 | use_ordered_embeddings is enabled). |
| 201 | """ |
| 202 | logger.info( |
| 203 | "Gemma4 MTP: keeping draft model's own lm_head (draft_dim != backbone_dim)." |
| 204 | ) |
| 205 | |
| 206 | def load_model(self, target_model: nn.Module) -> None: |
| 207 | target_attn_layer_names = set( |
| 208 | get_layers_from_vllm_config( |
| 209 | self.vllm_config, |
| 210 | AttentionLayerBase, # type: ignore[type-abstract] |
| 211 | ).keys() |
| 212 | ) |
| 213 | |
| 214 | super().load_model(target_model) |
| 215 | |
| 216 | self._setup_gemma4_kv_sharing(target_attn_layer_names) |
| 217 | |
| 218 | if getattr(self.model, "masked_embedding", None) is not None: |
| 219 | self._setup_centroids_cuda_graphs() |
| 220 | |
| 221 | def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None: |
| 222 | """Draft layers span multiple KV cache groups (sliding + full |
| 223 | attention with different head dimensions), so skip the base |
| 224 | class single-group assertion.""" |
| 225 | |
| 226 | def initialize_attn_backend( |
| 227 | self, |
| 228 | kv_cache_config: KVCacheConfig, |
| 229 | kernel_block_sizes: list[int] | None = None, |
| 230 | ) -> None: |
| 231 | """Create separate AttentionGroup objects per KV cache spec |
| 232 | so that each head-dim variant gets its own metadata builder.""" |
| 233 | all_attn_layers = get_layers_from_vllm_config( |
| 234 | self.vllm_config, |
| 235 | AttentionLayerBase, # type: ignore[type-abstract] |
| 236 | ) |
| 237 | |
| 238 | layer_to_gid: dict[str, int] = {} |
| 239 | layer_to_spec: dict[str, KVCacheSpec] = {} |
| 240 | for gid, group in enumerate(kv_cache_config.kv_cache_groups): |
| 241 | group_spec = group.kv_cache_spec |
| 242 | for ln in group.layer_names: |
| 243 | layer_to_gid[ln] = gid |
| 244 | if isinstance(group_spec, UniformTypeKVCacheSpecs): |
| 245 | if ln in group_spec.kv_cache_specs: |
| 246 | layer_to_spec[ln] = group_spec.kv_cache_specs[ln] |
| 247 | else: |
| 248 | tgt = getattr( |
| 249 | all_attn_layers.get(ln), |
| 250 | "kv_sharing_target_layer_name", |
| 251 | None, |
| 252 | ) |
| 253 | if tgt and tgt in group_spec.kv_cache_specs: |
| 254 | layer_to_spec[ln] = group_spec.kv_cache_specs[tgt] |
| 255 | else: |
| 256 | layer_to_spec[ln] = group_spec |
| 257 | else: |
| 258 | layer_to_spec[ln] = group_spec |
| 259 | |
| 260 | attention_groups: dict[tuple[tuple[str, str], KVCacheSpec], AttentionGroup] = {} |
| 261 | for layer_name in self._draft_attn_layer_names: |
| 262 | if layer_name not in layer_to_spec: |
| 263 | continue |
| 264 | attn_layer = all_attn_layers[layer_name] |
| 265 | attn_backend = attn_layer.get_attn_backend() |
| 266 | spec = layer_to_spec[layer_name] |
| 267 | gid = layer_to_gid[layer_name] |
| 268 | group_key = (attn_backend.full_cls_name(), spec) |
| 269 | |
| 270 | if group_key not in attention_groups: |
| 271 | kernel_block_size = ( |
| 272 | kernel_block_sizes[gid] |
| 273 | if kernel_block_sizes is not None and gid < len(kernel_block_sizes) |
| 274 | else None |
| 275 | ) |
| 276 | attn_group = AttentionGroup( |
| 277 | backend=attn_backend, |
| 278 | layer_names=[layer_name], |
| 279 | kv_cache_spec=spec, |
| 280 | kv_cache_group_id=gid, |
| 281 | ) |
| 282 | attn_group.create_metadata_builders( |
| 283 | self.vllm_config, |
| 284 | self.device, |
| 285 | kernel_block_size=kernel_block_size, |
| 286 | ) |
| 287 | attention_groups[group_key] = attn_group |
| 288 | else: |
| 289 | attention_groups[group_key].layer_names.append(layer_name) |
| 290 | |
| 291 | self.draft_attn_groups = list(attention_groups.values()) |
| 292 | if self.draft_attn_groups: |
| 293 | self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id |
| 294 | self.block_size = ( |
| 295 | self.draft_attn_groups[0] |
| 296 | .get_metadata_builder() |
| 297 | .kv_cache_spec.block_size |
| 298 | ) |
| 299 | else: |
| 300 | self.kv_cache_gid = 0 |
| 301 | self.block_size = kv_cache_config.kv_cache_groups[ |
| 302 | 0 |
| 303 | ].kv_cache_spec.block_size |
| 304 | logger.debug("Using block size %d for drafting layers", self.block_size) |
| 305 | |
| 306 | def _setup_gemma4_kv_sharing( |
| 307 | self, |
| 308 | target_attn_layer_names: set[str], |
| 309 | ) -> None: |
| 310 | """Wire draft layers to share KV with the target model. |
| 311 | |
| 312 | Each draft decoder layer is mapped to the last non-KV-shared |
| 313 | target layer of the same attention type (sliding or full). |
| 314 | """ |
| 315 | draft_config = self.speculative_config.draft_model_config.hf_config |
| 316 | draft_text_config = draft_config.get_text_config() |
| 317 | target_config = self.vllm_config.model_config.hf_config |
| 318 | target_text_config = target_config.get_text_config() |
| 319 | target_layer_types = getattr(target_text_config, "layer_types", []) |
| 320 | |
| 321 | if not (hasattr(self.model, "model") and hasattr(self.model.model, "layers")): |
| 322 | return |
| 323 | |
| 324 | target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0) |
| 325 | num_non_shared = len(target_layer_types) - target_num_kv_shared |
| 326 | type_to_target_indices: dict[str, list[int]] = defaultdict(list) |
| 327 | for idx, lt in enumerate(target_layer_types[:num_non_shared]): |
| 328 | type_to_target_indices[lt].append(idx) |
| 329 | |
| 330 | target_prefix = "model.layers" |
| 331 | for name in target_attn_layer_names: |
| 332 | if ".layers." in name: |
| 333 | target_prefix = name.split(".layers.")[0] + ".layers" |
| 334 | break |
| 335 | |
| 336 | draft_layer_types = getattr(draft_text_config, "layer_types", []) |
| 337 | for draft_idx, layer in enumerate(self.model.model.layers): |
| 338 | if not hasattr(layer, "self_attn"): |
| 339 | continue |
| 340 | attn = getattr(layer.self_attn, "attn", None) |
| 341 | if attn is None: |
| 342 | continue |
| 343 | |
| 344 | draft_layer_type = ( |
| 345 | draft_layer_types[draft_idx] |
| 346 | if draft_idx < len(draft_layer_types) |
| 347 | else "full_attention" |
| 348 | ) |
| 349 | candidates = type_to_target_indices.get(draft_layer_type, []) |
| 350 | if not candidates: |
| 351 | logger.warning( |
| 352 | "No target layer of type '%s' for draft layer %d", |
| 353 | draft_layer_type, |
| 354 | draft_idx, |
| 355 | ) |
| 356 | continue |
| 357 | |
| 358 | target_idx = candidates[-1] |
| 359 | target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn" |
| 360 | attn.kv_sharing_target_layer_name = target_layer_name |
| 361 | logger.info( |
| 362 | "Gemma4 MTP: draft layer %d (%s) -> %s", |
| 363 | draft_idx, |
| 364 | draft_layer_type, |
| 365 | target_layer_name, |
| 366 | ) |