| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Qwen4Exp weight-free QSA indexer.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from typing import cast |
| 8 | |
| 9 | import torch |
| 10 | from torch import nn |
| 11 | from transformers import Qwen4ExpTextConfig |
| 12 | |
| 13 | from vllm.config import VllmConfig |
| 14 | from vllm.forward_context import get_forward_context |
| 15 | from vllm.model_executor.layers.layernorm import GemmaRMSNorm |
| 16 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 17 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 18 | from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope |
| 19 | |
| 20 | from ..common.qsa_cache import ( |
| 21 | QSACompressedKeyCache, |
| 22 | QSAForwardMetadata, |
| 23 | QSAKeyStateCache, |
| 24 | canonical_qsa_rope_positions, |
| 25 | ) |
| 26 | |
| 27 | |
| 28 | def apply_qsa_rope( |
| 29 | rotary_emb: nn.Module, |
| 30 | positions: torch.Tensor, |
| 31 | tensor: torch.Tensor, |
| 32 | ) -> torch.Tensor: |
| 33 | """Apply the main attention's exact 1D/MRoPE composition to QSA heads.""" |
| 34 | num_tokens, _, head_dim = tensor.shape |
| 35 | rotary_dim = rotary_emb.rotary_dim |
| 36 | cache = rotary_emb._match_cos_sin_cache_dtype(tensor) # noqa: SLF001 |
| 37 | cos_sin = cache[positions] |
| 38 | cos, sin = cos_sin.chunk(2, dim=-1) |
| 39 | if positions.ndim == 2: |
| 40 | shape = tensor.shape |
| 41 | tensor, _ = triton_mrope( |
| 42 | tensor.reshape(num_tokens, -1), |
| 43 | tensor.new_empty((num_tokens, head_dim)), |
| 44 | cos, |
| 45 | sin, |
| 46 | rotary_emb.mrope_section, |
| 47 | head_dim, |
| 48 | rotary_dim, |
| 49 | rotary_emb.mrope_interleaved, |
| 50 | rotary_emb.is_neox_style, |
| 51 | ) |
| 52 | return tensor.reshape(shape) |
| 53 | |
| 54 | rotated = rotary_emb.apply_rotary_emb( |
| 55 | tensor[..., :rotary_dim], |
| 56 | cos, |
| 57 | sin, |
| 58 | ) |
| 59 | return torch.cat((rotated, tensor[..., rotary_dim:]), dim=-1) |
| 60 | |
| 61 | |
| 62 | def apply_qsa_rmsnorm( |
| 63 | norm: GemmaRMSNorm, |
| 64 | tensor: torch.Tensor, |
| 65 | ) -> torch.Tensor: |
| 66 | """Use vLLM's portable RMSNorm implementation on ROCm.""" |
| 67 | return cast(torch.Tensor, norm(tensor)) |
| 68 | |
| 69 | |
| 70 | class QSAIndexer(nn.Module): |
| 71 | """Replicated Q/K projection plus paged, weight-free QSA selection. |
| 72 | |
| 73 | ``prefix`` must be the checkpoint's indexer prefix, normally |
| 74 | ``model.layers.N.self_attn.indexer``. Consequently the trainable names are |
| 75 | ``index_qk_proj``, ``q_layernorm`` and ``k_layernorm`` under that prefix. |
| 76 | """ |
| 77 | |
| 78 | def __init__( |
| 79 | self, |
| 80 | *, |
| 81 | vllm_config: VllmConfig, |
| 82 | config: Qwen4ExpTextConfig, |
| 83 | layer_id: int, |
| 84 | rotary_emb: nn.Module, |
| 85 | quant_config: QuantizationConfig | None = None, |
| 86 | prefix: str = "", |
| 87 | ) -> None: |
| 88 | super().__init__() |
| 89 | if vllm_config.cache_config is None: |
| 90 | raise ValueError("QSA requires a paged KV cache") |
| 91 | if vllm_config.model_config.dtype != torch.bfloat16: |
| 92 | raise NotImplementedError("Qwen4Exp QSA currently requires BF16") |
| 93 | |
| 94 | self.layer_id = int(layer_id) |
| 95 | self.index_n_heads = int(config.indexer_n_heads) |
| 96 | self.index_kv_heads = int(config.indexer_kv_heads) |
| 97 | self.index_head_dim = int(config.indexer_head_dim) |
| 98 | self.token_topk = int(config.indexer_budget) |
| 99 | self.compress_ratio = int(config.indexer_compress_ratio) |
| 100 | self.rotary_emb = rotary_emb |
| 101 | self.prefix = prefix |
| 102 | # MTP step 0 selects the target-aligned rows; later steps reuse them |
| 103 | # while continuing to update the QSA side cache. |
| 104 | self.skip_topk = False |
| 105 | |
| 106 | self.index_qk_proj = ReplicatedLinear( |
| 107 | int(config.hidden_size), |
| 108 | (self.index_n_heads + self.index_kv_heads) * self.index_head_dim, |
| 109 | bias=False, |
| 110 | quant_config=quant_config, |
| 111 | prefix=f"{prefix}.index_qk_proj" if prefix else "index_qk_proj", |
| 112 | ) |
| 113 | self.q_layernorm = GemmaRMSNorm( |
| 114 | self.index_head_dim, |
| 115 | eps=float(getattr(config, "rms_norm_eps", 1e-6)), |
| 116 | ) |
| 117 | self.k_layernorm = GemmaRMSNorm( |
| 118 | self.index_head_dim, |
| 119 | eps=float(getattr(config, "rms_norm_eps", 1e-6)), |
| 120 | ) |
| 121 | |
| 122 | cache_config = vllm_config.cache_config |
| 123 | cache_prefix = f"{prefix}." if prefix else "" |
| 124 | self.raw_key_cache = QSAKeyStateCache( |
| 125 | head_size=self.index_head_dim, |
| 126 | dtype=torch.bfloat16, |
| 127 | cache_rope_positions=vllm_config.model_config.uses_mrope, |
| 128 | prefix=f"{cache_prefix}raw_key_cache", |
| 129 | cache_config=cache_config, |
| 130 | compress_ratio=self.compress_ratio, |
| 131 | vllm_config=vllm_config, |
| 132 | ) |
| 133 | self.compressed_key_cache = QSACompressedKeyCache( |
| 134 | head_size=self.index_head_dim, |
| 135 | dtype=torch.bfloat16, |
| 136 | compress_ratio=self.compress_ratio, |
| 137 | prefix=f"{cache_prefix}compressed_key_cache", |
| 138 | cache_config=cache_config, |
| 139 | vllm_config=vllm_config, |
| 140 | ) |
| 141 | |
| 142 | @property |
| 143 | def output_width(self) -> int: |
| 144 | return self.token_topk + self.compress_ratio - 1 |
| 145 | |
| 146 | def project_qk( |
| 147 | self, |
| 148 | hidden_states: torch.Tensor, |
| 149 | positions: torch.Tensor, |
| 150 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 151 | """Project replicated Q/K, normalize+rotate Q, and preserve raw K.""" |
| 152 | qk, _ = self.index_qk_proj(hidden_states) |
| 153 | q_raw, token_k = qk.split( |
| 154 | ( |
| 155 | self.index_n_heads * self.index_head_dim, |
| 156 | self.index_kv_heads * self.index_head_dim, |
| 157 | ), |
| 158 | dim=-1, |
| 159 | ) |
| 160 | q = q_raw.reshape(-1, self.index_n_heads, self.index_head_dim) |
| 161 | q = apply_qsa_rmsnorm( |
| 162 | self.q_layernorm, |
| 163 | q.reshape(-1, self.index_head_dim), |
| 164 | ).reshape_as(q) |
| 165 | q = apply_qsa_rope(self.rotary_emb, positions, q) |
| 166 | return q, token_k.reshape(-1, 1, self.index_head_dim) |
| 167 | |
| 168 | def normalize_compressed_keys( |
| 169 | self, |
| 170 | compressed_keys: torch.Tensor, |
| 171 | first_rope_positions: torch.Tensor, |
| 172 | ) -> torch.Tensor: |
| 173 | """Normalize pooled K and apply the first token's exact group position.""" |
| 174 | keys = compressed_keys.reshape(-1, self.index_head_dim) |
| 175 | keys = apply_qsa_rmsnorm(self.k_layernorm, keys).reshape( |
| 176 | -1, 1, self.index_head_dim |
| 177 | ) |
| 178 | if getattr(self.rotary_emb, "mrope_section", None): |
| 179 | positions = first_rope_positions.transpose(0, 1) |
| 180 | else: |
| 181 | positions = first_rope_positions[:, 0] |
| 182 | return apply_qsa_rope(self.rotary_emb, positions, keys) |
| 183 | |
| 184 | def _metadata( |
| 185 | self, |
| 186 | ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None: |
| 187 | metadata = get_forward_context().attn_metadata |
| 188 | if isinstance(metadata, list): |
| 189 | metadata = metadata[0] |
| 190 | if not isinstance(metadata, dict): |
| 191 | return None |
| 192 | raw = cast(QSAForwardMetadata, metadata[self.raw_key_cache.prefix]) |
| 193 | compressed = cast( |
| 194 | QSAForwardMetadata, metadata[self.compressed_key_cache.prefix] |
| 195 | ) |
| 196 | if raw.num_actual_tokens != compressed.num_actual_tokens: |
| 197 | raise RuntimeError("QSA side-cache metadata token counts disagree") |
| 198 | if not raw.logical_positions.is_cuda and ( |
| 199 | not torch.equal(raw.logical_positions, compressed.logical_positions) |
| 200 | ): |
| 201 | raise RuntimeError("QSA side-cache metadata positions disagree") |
| 202 | return raw, compressed |
| 203 | |
| 204 | def _update_and_compress( |
| 205 | self, |
| 206 | token_k: torch.Tensor, |
| 207 | positions: torch.Tensor, |
| 208 | raw_metadata: QSAForwardMetadata, |
| 209 | compressed_metadata: QSAForwardMetadata, |
| 210 | ) -> None: |
| 211 | num_tokens = raw_metadata.num_actual_tokens |
| 212 | raw_key_cache = self.raw_key_cache.key_cache |
| 213 | rope_position_cache = self.raw_key_cache.rope_position_cache |
| 214 | from .ops.qsa import qsa_compress_groups_with_ratio, qsa_store_cache_rows |
| 215 | |
| 216 | if rope_position_cache is None: |
| 217 | position_rows = raw_metadata.logical_positions.view(-1, 1, 1).expand( |
| 218 | -1, 1, 3 |
| 219 | ) |
| 220 | else: |
| 221 | position_rows = canonical_qsa_rope_positions(positions)[:num_tokens].to( |
| 222 | device=raw_key_cache.device |
| 223 | ) |
| 224 | pooled, first_positions = qsa_compress_groups_with_ratio( |
| 225 | token_k[:num_tokens], |
| 226 | position_rows, |
| 227 | raw_key_cache, |
| 228 | raw_metadata.block_table, |
| 229 | raw_metadata.token_to_req, |
| 230 | raw_metadata.query_start_loc, |
| 231 | raw_metadata.logical_positions, |
| 232 | compressed_metadata.slot_mapping, |
| 233 | self.compress_ratio, |
| 234 | rope_position_cache, |
| 235 | ) |
| 236 | normalized = self.normalize_compressed_keys(pooled, first_positions) |
| 237 | qsa_store_cache_rows( |
| 238 | self.compressed_key_cache.kv_cache, |
| 239 | compressed_metadata.slot_mapping, |
| 240 | normalized, |
| 241 | ) |
| 242 | qsa_store_cache_rows( |
| 243 | raw_key_cache, |
| 244 | raw_metadata.slot_mapping, |
| 245 | token_k[:num_tokens], |
| 246 | ) |
| 247 | if rope_position_cache is not None: |
| 248 | qsa_store_cache_rows( |
| 249 | rope_position_cache, |
| 250 | raw_metadata.slot_mapping, |
| 251 | position_rows, |
| 252 | ) |
| 253 | |
| 254 | def _select( |
| 255 | self, |
| 256 | q: torch.Tensor, |
| 257 | metadata: QSAForwardMetadata, |
| 258 | out: torch.Tensor | None, |
| 259 | ) -> torch.Tensor: |
| 260 | from .ops.qsa import qsa_select_paged_tokens |
| 261 | |
| 262 | return qsa_select_paged_tokens( |
| 263 | q, |
| 264 | self.compressed_key_cache.kv_cache, |
| 265 | metadata.block_table, |
| 266 | metadata.token_to_req, |
| 267 | metadata.logical_positions, |
| 268 | metadata.seq_lens, |
| 269 | self.token_topk, |
| 270 | self.compress_ratio, |
| 271 | out, |
| 272 | ) |
| 273 | |
| 274 | def forward( |
| 275 | self, |
| 276 | hidden_states: torch.Tensor, |
| 277 | positions: torch.Tensor, |
| 278 | out: torch.Tensor | None = None, |
| 279 | ) -> torch.Tensor: |
| 280 | """Return fixed-width request-relative token indices padded with ``-1``.""" |
| 281 | metadata = self._metadata() |
| 282 | if metadata is None: |
| 283 | # Preserve step-0 indices when later MTP steps reuse the buffer. |
| 284 | if self.skip_topk and out is not None: |
| 285 | return out |
| 286 | result = torch.full( |
| 287 | (hidden_states.shape[0], self.output_width), |
| 288 | -1, |
| 289 | dtype=torch.int32, |
| 290 | device=hidden_states.device, |
| 291 | ) |
| 292 | if out is not None: |
| 293 | out.copy_(result) |
| 294 | return out |
| 295 | return result |
| 296 | raw_metadata, compressed_metadata = metadata |
| 297 | num_tokens = raw_metadata.num_actual_tokens |
| 298 | q, token_k = self.project_qk( |
| 299 | hidden_states[:num_tokens], positions[..., :num_tokens] |
| 300 | ) |
| 301 | self._update_and_compress( |
| 302 | token_k, |
| 303 | positions[..., :num_tokens], |
| 304 | raw_metadata, |
| 305 | compressed_metadata, |
| 306 | ) |
| 307 | if self.skip_topk: |
| 308 | if out is None: |
| 309 | raise RuntimeError("QSA top-k reuse requires an output buffer") |
| 310 | return out |
| 311 | return self._select(q, compressed_metadata, out) |
| 312 | |
| 313 | |
| 314 | __all__ = ["QSAIndexer", "apply_qsa_rope"] |