wehub-resource-sync
341 lines
14 KiB
Python
341 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Gemma4 MTP (Multi-Token Prediction) proposer for speculative decoding.
|
|
|
|
The Gemma4 assistant model runs all decoder layers per draft step
|
|
(producing one token), and all its attention layers share KV cache
|
|
with the target model via cross-model KV sharing.
|
|
"""
|
|
|
|
from collections import defaultdict
|
|
from copy import copy
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from vllm.config import VllmConfig, get_layers_from_vllm_config, replace
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
|
from vllm.v1.attention.backend import CommonAttentionMetadata
|
|
from vllm.v1.kv_cache_interface import (
|
|
KVCacheConfig,
|
|
KVCacheSpec,
|
|
UniformTypeKVCacheSpecs,
|
|
)
|
|
from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
|
|
from vllm.v1.worker.utils import AttentionGroup
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class Gemma4Proposer(SpecDecodeBaseProposer):
|
|
def __init__(
|
|
self,
|
|
vllm_config: VllmConfig,
|
|
device: torch.device,
|
|
runner=None,
|
|
):
|
|
super().__init__(
|
|
vllm_config,
|
|
device,
|
|
pass_hidden_states_to_model=True,
|
|
runner=runner,
|
|
)
|
|
# All draft steps predict from the same position (the last
|
|
# target-model position), so positions and seq_lens must not
|
|
# advance between steps.
|
|
self.constant_draft_positions = True
|
|
|
|
# Per-group block tables for multi-group KV cache models.
|
|
# Populated by gpu_model_runner during _prepare_inputs.
|
|
self._per_group_block_tables: dict[int, torch.Tensor] = {}
|
|
|
|
# Centroids CUDA graphs — populated in load_model if centroids
|
|
# masking is active. _centroids_sizes is pre-sorted for fast
|
|
# lookup in _greedy_sample.
|
|
self._centroids_sizes: list[int] = []
|
|
self._centroids_graphs: dict[int, torch.cuda.CUDAGraph] = {}
|
|
self._centroids_inputs: dict[int, torch.Tensor] = {}
|
|
self._centroids_outputs: dict[int, torch.Tensor] = {}
|
|
|
|
def set_per_group_block_table(self, gid: int, block_table: torch.Tensor) -> None:
|
|
self._per_group_block_tables[gid] = block_table
|
|
|
|
def model_returns_tuple(self) -> bool:
|
|
# forward() returns (draft_hidden_states, backbone_hidden_states).
|
|
# The proposer uses draft_hidden_states for compute_logits and
|
|
# backbone_hidden_states for the hidden-state feedback buffer.
|
|
return True
|
|
|
|
def build_per_group_and_layer_attn_metadata(
|
|
self,
|
|
common_attn_metadata: CommonAttentionMetadata,
|
|
draft_index: int = 0,
|
|
) -> tuple[list[object], dict[str, object]]:
|
|
"""Build attention metadata using the correct block table per group.
|
|
|
|
Gemma4 has multiple KV cache groups (sliding vs full attention)
|
|
with different block tables. The base class receives a single
|
|
common_attn_metadata whose block_table belongs to one group.
|
|
We swap in the correct block table for each draft attention group.
|
|
"""
|
|
per_group_attn_metadata: list[object] = []
|
|
per_layer_attn_metadata: dict[str, object] = {}
|
|
batch_size = common_attn_metadata.batch_size()
|
|
for attn_group in self.draft_attn_groups:
|
|
gid = attn_group.kv_cache_group_id
|
|
if gid in self._per_group_block_tables:
|
|
cm = copy(common_attn_metadata)
|
|
# Slice to actual batch size to match cu_seqlens_q dimension.
|
|
# The stored block tables may be padded (num_reqs_padded) from
|
|
# the target forward pass, but the drafter operates on the
|
|
# unpadded batch.
|
|
cm.block_table_tensor = self._per_group_block_tables[gid][:batch_size]
|
|
else:
|
|
cm = common_attn_metadata
|
|
attn_metadata = attn_group.get_metadata_builder().build_for_drafting(
|
|
common_attn_metadata=cm, draft_index=draft_index
|
|
)
|
|
per_group_attn_metadata.append(attn_metadata)
|
|
for layer_name in attn_group.layer_names:
|
|
per_layer_attn_metadata[layer_name] = attn_metadata
|
|
return per_group_attn_metadata, per_layer_attn_metadata
|
|
|
|
def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
if self._centroids_sizes:
|
|
T = hidden_states.shape[0]
|
|
for size in self._centroids_sizes:
|
|
if size >= T:
|
|
self._centroids_inputs[size][:T].copy_(hidden_states)
|
|
self._centroids_graphs[size].replay()
|
|
return self._centroids_outputs[size][:T].clone()
|
|
return self.model.get_top_tokens(hidden_states)
|
|
return super()._greedy_sample(hidden_states)
|
|
|
|
def _setup_centroids_cuda_graphs(self) -> None:
|
|
"""Capture CUDA graphs for centroids get_top_tokens at key sizes."""
|
|
masked_emb = self.model.masked_embedding
|
|
lm_head_weight = self.model._get_full_lm_head_weight()
|
|
|
|
for size in [1, 2, 4, 8, 16, 32, 64]:
|
|
static_input = torch.zeros(
|
|
size,
|
|
masked_emb.hidden_size,
|
|
dtype=self.dtype,
|
|
device=self.device,
|
|
)
|
|
for _ in range(3):
|
|
masked_emb.get_top_tokens(static_input, lm_head_weight)
|
|
torch.accelerator.synchronize()
|
|
|
|
g = torch.cuda.CUDAGraph()
|
|
with torch.cuda.graph(g):
|
|
static_output = masked_emb.get_top_tokens(
|
|
static_input,
|
|
lm_head_weight,
|
|
)
|
|
self._centroids_graphs[size] = g
|
|
self._centroids_inputs[size] = static_input
|
|
self._centroids_outputs[size] = static_output
|
|
|
|
self._centroids_sizes = sorted(self._centroids_graphs)
|
|
logger.info(
|
|
"Gemma4 MTP: captured centroids CUDA graphs for sizes %s.",
|
|
self._centroids_sizes,
|
|
)
|
|
|
|
def _create_draft_vllm_config(self) -> VllmConfig:
|
|
"""Preserve the target's forced TRITON_ATTN backend for draft layers.
|
|
|
|
Gemma4 forces TRITON_ATTN due to heterogeneous head dimensions
|
|
(head_dim=256 sliding, global_head_dim=512 full). The base class
|
|
resets attention_config.backend to None for draft models, causing
|
|
sliding layers to fall back to FLASH_ATTN which cannot handle
|
|
KV-shared cache. Override to carry the target's backend through.
|
|
"""
|
|
base = super()._create_draft_vllm_config()
|
|
target_backend = self.vllm_config.attention_config.backend
|
|
if target_backend is not None:
|
|
base = replace(
|
|
base,
|
|
attention_config=replace(
|
|
base.attention_config,
|
|
backend=target_backend,
|
|
),
|
|
)
|
|
return base
|
|
|
|
def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None:
|
|
"""Gemma4 MTP always keeps its own draft-dim lm_head.
|
|
|
|
The draft model's lm_head operates in draft hidden_size (e.g. 256),
|
|
which differs from the target's backbone hidden_size (e.g. 1536).
|
|
Sharing would break compute_logits (and centroids masking when
|
|
use_ordered_embeddings is enabled).
|
|
"""
|
|
logger.info(
|
|
"Gemma4 MTP: keeping draft model's own lm_head (draft_dim != backbone_dim)."
|
|
)
|
|
|
|
def load_model(self, target_model: nn.Module) -> None:
|
|
target_attn_layer_names = set(
|
|
get_layers_from_vllm_config(
|
|
self.vllm_config,
|
|
AttentionLayerBase, # type: ignore[type-abstract]
|
|
).keys()
|
|
)
|
|
|
|
super().load_model(target_model)
|
|
|
|
self._setup_gemma4_kv_sharing(target_attn_layer_names)
|
|
|
|
if getattr(self.model, "masked_embedding", None) is not None:
|
|
self._setup_centroids_cuda_graphs()
|
|
|
|
def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
|
|
"""Draft layers span multiple KV cache groups (sliding + full
|
|
attention with different head dimensions), so skip the base
|
|
class single-group assertion."""
|
|
|
|
def initialize_attn_backend(
|
|
self,
|
|
kv_cache_config: KVCacheConfig,
|
|
kernel_block_sizes: list[int] | None = None,
|
|
) -> None:
|
|
"""Create separate AttentionGroup objects per KV cache spec
|
|
so that each head-dim variant gets its own metadata builder."""
|
|
all_attn_layers = get_layers_from_vllm_config(
|
|
self.vllm_config,
|
|
AttentionLayerBase, # type: ignore[type-abstract]
|
|
)
|
|
|
|
layer_to_gid: dict[str, int] = {}
|
|
layer_to_spec: dict[str, KVCacheSpec] = {}
|
|
for gid, group in enumerate(kv_cache_config.kv_cache_groups):
|
|
group_spec = group.kv_cache_spec
|
|
for ln in group.layer_names:
|
|
layer_to_gid[ln] = gid
|
|
if isinstance(group_spec, UniformTypeKVCacheSpecs):
|
|
if ln in group_spec.kv_cache_specs:
|
|
layer_to_spec[ln] = group_spec.kv_cache_specs[ln]
|
|
else:
|
|
tgt = getattr(
|
|
all_attn_layers.get(ln),
|
|
"kv_sharing_target_layer_name",
|
|
None,
|
|
)
|
|
if tgt and tgt in group_spec.kv_cache_specs:
|
|
layer_to_spec[ln] = group_spec.kv_cache_specs[tgt]
|
|
else:
|
|
layer_to_spec[ln] = group_spec
|
|
else:
|
|
layer_to_spec[ln] = group_spec
|
|
|
|
attention_groups: dict[tuple[tuple[str, str], KVCacheSpec], AttentionGroup] = {}
|
|
for layer_name in self._draft_attn_layer_names:
|
|
if layer_name not in layer_to_spec:
|
|
continue
|
|
attn_layer = all_attn_layers[layer_name]
|
|
attn_backend = attn_layer.get_attn_backend()
|
|
spec = layer_to_spec[layer_name]
|
|
gid = layer_to_gid[layer_name]
|
|
group_key = (attn_backend.full_cls_name(), spec)
|
|
|
|
if group_key not in attention_groups:
|
|
kernel_block_size = (
|
|
kernel_block_sizes[gid]
|
|
if kernel_block_sizes is not None and gid < len(kernel_block_sizes)
|
|
else None
|
|
)
|
|
attn_group = AttentionGroup(
|
|
backend=attn_backend,
|
|
layer_names=[layer_name],
|
|
kv_cache_spec=spec,
|
|
kv_cache_group_id=gid,
|
|
)
|
|
attn_group.create_metadata_builders(
|
|
self.vllm_config,
|
|
self.device,
|
|
kernel_block_size=kernel_block_size,
|
|
)
|
|
attention_groups[group_key] = attn_group
|
|
else:
|
|
attention_groups[group_key].layer_names.append(layer_name)
|
|
|
|
self.draft_attn_groups = list(attention_groups.values())
|
|
if self.draft_attn_groups:
|
|
self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id
|
|
self.block_size = (
|
|
self.draft_attn_groups[0]
|
|
.get_metadata_builder()
|
|
.kv_cache_spec.block_size
|
|
)
|
|
else:
|
|
self.kv_cache_gid = 0
|
|
self.block_size = kv_cache_config.kv_cache_groups[
|
|
0
|
|
].kv_cache_spec.block_size
|
|
logger.debug("Using block size %d for drafting layers", self.block_size)
|
|
|
|
def _setup_gemma4_kv_sharing(
|
|
self,
|
|
target_attn_layer_names: set[str],
|
|
) -> None:
|
|
"""Wire draft layers to share KV with the target model.
|
|
|
|
Each draft decoder layer is mapped to the last non-KV-shared
|
|
target layer of the same attention type (sliding or full).
|
|
"""
|
|
draft_config = self.speculative_config.draft_model_config.hf_config
|
|
draft_text_config = draft_config.get_text_config()
|
|
target_config = self.vllm_config.model_config.hf_config
|
|
target_text_config = target_config.get_text_config()
|
|
target_layer_types = getattr(target_text_config, "layer_types", [])
|
|
|
|
if not (hasattr(self.model, "model") and hasattr(self.model.model, "layers")):
|
|
return
|
|
|
|
target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0)
|
|
num_non_shared = len(target_layer_types) - target_num_kv_shared
|
|
type_to_target_indices: dict[str, list[int]] = defaultdict(list)
|
|
for idx, lt in enumerate(target_layer_types[:num_non_shared]):
|
|
type_to_target_indices[lt].append(idx)
|
|
|
|
target_prefix = "model.layers"
|
|
for name in target_attn_layer_names:
|
|
if ".layers." in name:
|
|
target_prefix = name.split(".layers.")[0] + ".layers"
|
|
break
|
|
|
|
draft_layer_types = getattr(draft_text_config, "layer_types", [])
|
|
for draft_idx, layer in enumerate(self.model.model.layers):
|
|
if not hasattr(layer, "self_attn"):
|
|
continue
|
|
attn = getattr(layer.self_attn, "attn", None)
|
|
if attn is None:
|
|
continue
|
|
|
|
draft_layer_type = (
|
|
draft_layer_types[draft_idx]
|
|
if draft_idx < len(draft_layer_types)
|
|
else "full_attention"
|
|
)
|
|
candidates = type_to_target_indices.get(draft_layer_type, [])
|
|
if not candidates:
|
|
logger.warning(
|
|
"No target layer of type '%s' for draft layer %d",
|
|
draft_layer_type,
|
|
draft_idx,
|
|
)
|
|
continue
|
|
|
|
target_idx = candidates[-1]
|
|
target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"
|
|
attn.kv_sharing_target_layer_name = target_layer_name
|
|
logger.info(
|
|
"Gemma4 MTP: draft layer %d (%s) -> %s",
|
|
draft_idx,
|
|
draft_layer_type,
|
|
target_layer_name,
|
|
)
|