| 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 typing import TYPE_CHECKING, cast |
| 6 | |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | from transformers import Qwen4ExpTextConfig |
| 10 | |
| 11 | from vllm.config import VllmConfig |
| 12 | from vllm.forward_context import get_forward_context |
| 13 | from vllm.model_executor.layers.layernorm import GemmaRMSNorm |
| 14 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 15 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 16 | from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope |
| 17 | |
| 18 | from ..common.qsa_cache import ( |
| 19 | QSACompressedKeyCache, |
| 20 | QSAForwardMetadata, |
| 21 | QSAKeyStateCache, |
| 22 | canonical_qsa_rope_positions, |
| 23 | ) |
| 24 | from .ops.qsa_prepare import qsa_prepare |
| 25 | |
| 26 | if TYPE_CHECKING: |
| 27 | from .qsa import Qwen4ExpQSAAttention |
| 28 | |
| 29 | |
| 30 | def apply_qsa_rope( |
| 31 | rotary_emb: nn.Module, |
| 32 | positions: torch.Tensor, |
| 33 | tensor: torch.Tensor, |
| 34 | ) -> torch.Tensor: |
| 35 | """Apply the main attention's exact 1D/MRoPE composition to QSA heads.""" |
| 36 | num_tokens, _, head_dim = tensor.shape |
| 37 | rotary_dim = rotary_emb.rotary_dim |
| 38 | cache = rotary_emb._match_cos_sin_cache_dtype(tensor) # noqa: SLF001 |
| 39 | cos_sin = cache[positions] |
| 40 | cos, sin = cos_sin.chunk(2, dim=-1) |
| 41 | if positions.ndim == 2: |
| 42 | shape = tensor.shape |
| 43 | tensor, _ = triton_mrope( |
| 44 | tensor.reshape(num_tokens, -1), |
| 45 | tensor.new_empty((num_tokens, head_dim)), |
| 46 | cos, |
| 47 | sin, |
| 48 | rotary_emb.mrope_section, |
| 49 | head_dim, |
| 50 | rotary_dim, |
| 51 | rotary_emb.mrope_interleaved, |
| 52 | rotary_emb.is_neox_style, |
| 53 | ) |
| 54 | return tensor.reshape(shape) |
| 55 | |
| 56 | rotated = rotary_emb.apply_rotary_emb.forward_cuda( |
| 57 | tensor[..., :rotary_dim], |
| 58 | cos, |
| 59 | sin, |
| 60 | ) |
| 61 | return torch.cat((rotated, tensor[..., rotary_dim:]), dim=-1) |
| 62 | |
| 63 | |
| 64 | def _supports_fused_pre_indexer( |
| 65 | rotary_emb: nn.Module, |
| 66 | head_dim: int, |
| 67 | num_kv_heads: int, |
| 68 | compress_ratio: int, |
| 69 | ) -> bool: |
| 70 | rotary_dim = int(rotary_emb.rotary_dim) |
| 71 | mrope_section = getattr(rotary_emb, "mrope_section", None) |
| 72 | return ( |
| 73 | bool(getattr(rotary_emb, "is_neox_style", False)) |
| 74 | and ( |
| 75 | not mrope_section |
| 76 | or ( |
| 77 | len(mrope_section) == 3 |
| 78 | and sum(mrope_section) == rotary_dim // 2 |
| 79 | and bool(getattr(rotary_emb, "mrope_interleaved", False)) |
| 80 | ) |
| 81 | ) |
| 82 | and head_dim == 128 |
| 83 | and rotary_dim == 64 |
| 84 | and num_kv_heads == 1 |
| 85 | and compress_ratio > 1 |
| 86 | and compress_ratio & (compress_ratio - 1) == 0 |
| 87 | ) |
| 88 | |
| 89 | |
| 90 | class QSAIndexer(nn.Module): |
| 91 | """QSA projection weights, side caches, and paged, weight-free selection. |
| 92 | |
| 93 | ``prefix`` must be the checkpoint's indexer prefix, normally |
| 94 | ``model.layers.N.self_attn.indexer``. Consequently the trainable names are |
| 95 | ``index_qk_proj``, ``q_layernorm`` and ``k_layernorm`` under that prefix. |
| 96 | """ |
| 97 | |
| 98 | def __init__( |
| 99 | self, |
| 100 | *, |
| 101 | vllm_config: VllmConfig, |
| 102 | config: Qwen4ExpTextConfig, |
| 103 | layer_id: int, |
| 104 | rotary_emb: nn.Module, |
| 105 | quant_config: QuantizationConfig | None = None, |
| 106 | prefix: str = "", |
| 107 | ) -> None: |
| 108 | super().__init__() |
| 109 | if vllm_config.cache_config is None: |
| 110 | raise ValueError("QSA requires a paged KV cache") |
| 111 | if vllm_config.model_config.dtype != torch.bfloat16: |
| 112 | raise NotImplementedError("Qwen4Exp QSA currently requires BF16") |
| 113 | |
| 114 | self.layer_id = int(layer_id) |
| 115 | self.index_n_heads = int(config.indexer_n_heads) |
| 116 | self.index_kv_heads = int(config.indexer_kv_heads) |
| 117 | self.index_head_dim = int(config.indexer_head_dim) |
| 118 | self.token_topk = int(config.indexer_budget) |
| 119 | self.compress_ratio = int(config.indexer_compress_ratio) |
| 120 | self.rotary_emb = rotary_emb |
| 121 | self.use_fused_pre_indexer = _supports_fused_pre_indexer( |
| 122 | rotary_emb, |
| 123 | self.index_head_dim, |
| 124 | self.index_kv_heads, |
| 125 | self.compress_ratio, |
| 126 | ) |
| 127 | self.prefix = prefix |
| 128 | # MTP step 0 selects the target-aligned rows; later steps reuse them |
| 129 | # while continuing to update the QSA side cache. |
| 130 | self.skip_topk = False |
| 131 | |
| 132 | self.index_qk_proj = ReplicatedLinear( |
| 133 | int(config.hidden_size), |
| 134 | (self.index_n_heads + self.index_kv_heads) * self.index_head_dim, |
| 135 | bias=False, |
| 136 | quant_config=quant_config, |
| 137 | prefix=f"{prefix}.index_qk_proj" if prefix else "index_qk_proj", |
| 138 | ) |
| 139 | self.q_layernorm = GemmaRMSNorm( |
| 140 | self.index_head_dim, |
| 141 | eps=float(getattr(config, "rms_norm_eps", 1e-6)), |
| 142 | ) |
| 143 | self.k_layernorm = GemmaRMSNorm( |
| 144 | self.index_head_dim, |
| 145 | eps=float(getattr(config, "rms_norm_eps", 1e-6)), |
| 146 | ) |
| 147 | |
| 148 | cache_config = vllm_config.cache_config |
| 149 | cache_prefix = f"{prefix}." if prefix else "" |
| 150 | # Plain e4m3 without scales: Q and the compressed K are RMSNormed |
| 151 | # before quantization, and the logits kernels dot fp8 x fp8 directly. |
| 152 | self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype( |
| 153 | "bf16" |
| 154 | ) |
| 155 | if self.indexer_kv_dtype == "fp8": |
| 156 | indexer_dtype = torch.float8_e4m3fn |
| 157 | elif self.indexer_kv_dtype == "bf16": |
| 158 | indexer_dtype = torch.bfloat16 |
| 159 | else: |
| 160 | raise NotImplementedError( |
| 161 | f"indexer_kv_dtype={self.indexer_kv_dtype!r} is not supported " |
| 162 | "by the Qwen4Exp QSA indexer (only 'bf16' or 'fp8')." |
| 163 | ) |
| 164 | self.indexer_dtype = indexer_dtype |
| 165 | self.raw_key_cache = QSAKeyStateCache( |
| 166 | head_size=self.index_head_dim, |
| 167 | dtype=torch.bfloat16, |
| 168 | cache_rope_positions=vllm_config.model_config.uses_mrope, |
| 169 | prefix=f"{cache_prefix}raw_key_cache", |
| 170 | cache_config=cache_config, |
| 171 | compress_ratio=self.compress_ratio, |
| 172 | vllm_config=vllm_config, |
| 173 | ) |
| 174 | self.compressed_key_cache = QSACompressedKeyCache( |
| 175 | head_size=self.index_head_dim, |
| 176 | dtype=indexer_dtype, |
| 177 | compress_ratio=self.compress_ratio, |
| 178 | prefix=f"{cache_prefix}compressed_key_cache", |
| 179 | cache_config=cache_config, |
| 180 | vllm_config=vllm_config, |
| 181 | ) |
| 182 | |
| 183 | @property |
| 184 | def output_width(self) -> int: |
| 185 | """Selection (index) columns per row.""" |
| 186 | return self.token_topk + self.compress_ratio - 1 |
| 187 | |
| 188 | @property |
| 189 | def packed_output_width(self) -> int: |
| 190 | """Packed selection-buffer width: output_width + 1. |
| 191 | |
| 192 | The trailing column holds each row's valid-entry count (written by |
| 193 | the expand kernel) — never a token index; the sparse attention |
| 194 | kernel reads it as its tile-loop bound. |
| 195 | """ |
| 196 | return self.output_width + 1 |
| 197 | |
| 198 | def _metadata( |
| 199 | self, |
| 200 | ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None: |
| 201 | metadata = get_forward_context().attn_metadata |
| 202 | if isinstance(metadata, list): |
| 203 | metadata = metadata[0] |
| 204 | if not isinstance(metadata, dict): |
| 205 | return None |
| 206 | raw = cast(QSAForwardMetadata, metadata[self.raw_key_cache.prefix]) |
| 207 | compressed = cast( |
| 208 | QSAForwardMetadata, metadata[self.compressed_key_cache.prefix] |
| 209 | ) |
| 210 | if raw.num_actual_tokens != compressed.num_actual_tokens: |
| 211 | raise RuntimeError("QSA side-cache metadata token counts disagree") |
| 212 | raw_split = ( |
| 213 | raw.num_decodes, |
| 214 | raw.num_decode_tokens, |
| 215 | raw.num_prefills, |
| 216 | raw.num_prefill_tokens, |
| 217 | raw.decode_query_len, |
| 218 | ) |
| 219 | compressed_split = ( |
| 220 | compressed.num_decodes, |
| 221 | compressed.num_decode_tokens, |
| 222 | compressed.num_prefills, |
| 223 | compressed.num_prefill_tokens, |
| 224 | compressed.decode_query_len, |
| 225 | ) |
| 226 | if raw_split != compressed_split: |
| 227 | raise RuntimeError("QSA side-cache metadata batch splits disagree") |
| 228 | if not raw.logical_positions.is_cuda and ( |
| 229 | not torch.equal(raw.logical_positions, compressed.logical_positions) |
| 230 | ): |
| 231 | raise RuntimeError("QSA side-cache metadata positions disagree") |
| 232 | return raw, compressed |
| 233 | |
| 234 | def forward( |
| 235 | self, |
| 236 | projected_qk: torch.Tensor, |
| 237 | positions: torch.Tensor, |
| 238 | out: torch.Tensor | None = None, |
| 239 | *, |
| 240 | attn: "Qwen4ExpQSAAttention", |
| 241 | qkv: torch.Tensor | None = None, |
| 242 | slot_mapping: torch.Tensor | None = None, |
| 243 | ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: |
| 244 | """Update side caches and select token indices from pre-projected Q/K. |
| 245 | |
| 246 | Returns the packed buffer of shape [num_tokens, output_width + 1]: |
| 247 | the leading ``output_width`` columns are ``-1``-padded |
| 248 | request-relative token indices, and the trailing column is the row's |
| 249 | valid-entry count (the attention kernel's loop bound, never a token |
| 250 | index). |
| 251 | |
| 252 | With ``attn.use_fused_qsa_prepare``, the same launch also writes |
| 253 | ``attn``'s K/V into ``attn.kv_cache`` at ``slot_mapping`` and prepares |
| 254 | its Q and gate from ``qkv``, returned as the second element (None |
| 255 | otherwise). ``qkv`` and ``slot_mapping`` are only read in that mode. |
| 256 | """ |
| 257 | metadata = self._metadata() |
| 258 | if metadata is None: |
| 259 | # Preserve step-0 indices when later MTP steps reuse the buffer. |
| 260 | if self.skip_topk and out is not None: |
| 261 | return out, None |
| 262 | result = torch.full( |
| 263 | (projected_qk.shape[0], self.packed_output_width), |
| 264 | -1, |
| 265 | dtype=torch.int32, |
| 266 | device=projected_qk.device, |
| 267 | ) |
| 268 | # Inert rows carry a zero valid count (empty loop bound), not -1. |
| 269 | result[:, -1] = 0 |
| 270 | if out is not None: |
| 271 | out.copy_(result) |
| 272 | return out, None |
| 273 | return result, None |
| 274 | |
| 275 | from .ops.qsa import qsa_compress_groups_with_ratio, qsa_store_cache_rows |
| 276 | from .ops.qsa_indexer import ( |
| 277 | expand_qsa_block_indices, |
| 278 | qsa_select_paged_decode, |
| 279 | qsa_select_paged_prefill, |
| 280 | ) |
| 281 | |
| 282 | raw_metadata, compressed_metadata = metadata |
| 283 | num_tokens = raw_metadata.num_actual_tokens |
| 284 | projected_qk = projected_qk[:num_tokens] |
| 285 | positions = positions[..., :num_tokens] |
| 286 | |
| 287 | projected_q, raw_keys = projected_qk.split( |
| 288 | ( |
| 289 | self.index_n_heads * self.index_head_dim, |
| 290 | self.index_kv_heads * self.index_head_dim, |
| 291 | ), |
| 292 | dim=-1, |
| 293 | ) |
| 294 | raw_key_state_cache = self.raw_key_cache |
| 295 | compressed_key_cache = self.compressed_key_cache.kv_cache |
| 296 | |
| 297 | main_outputs: tuple[torch.Tensor, torch.Tensor] | None = None |
| 298 | if attn.use_fused_qsa_prepare: |
| 299 | if qkv is None or slot_mapping is None: |
| 300 | raise ValueError("fused QSA prepare requires qkv and slot_mapping") |
| 301 | q = projected_q.new_empty( |
| 302 | num_tokens, |
| 303 | self.index_n_heads, |
| 304 | self.index_head_dim, |
| 305 | dtype=self.indexer_dtype, |
| 306 | ) |
| 307 | main_kv_cache = attn.kv_cache.transpose(1, 2) |
| 308 | if attn.kv_cache_dtype in ("fp8", "fp8_e4m3"): |
| 309 | main_kv_cache = main_kv_cache.view(torch.float8_e4m3fn) |
| 310 | main_outputs = qsa_prepare( |
| 311 | projected_q, |
| 312 | raw_keys, |
| 313 | positions, |
| 314 | self.rotary_emb.cos_sin_cache, |
| 315 | self.q_layernorm.weight, |
| 316 | self.k_layernorm.weight, |
| 317 | self.q_layernorm.variance_epsilon, |
| 318 | q, |
| 319 | raw_key_state_cache.kv_cache, |
| 320 | raw_metadata.slot_mapping, |
| 321 | raw_metadata.block_table, |
| 322 | raw_metadata.query_start_loc, |
| 323 | raw_metadata.logical_positions, |
| 324 | compressed_key_cache, |
| 325 | compressed_metadata.slot_mapping, |
| 326 | compressed_metadata.k_work_metadata, |
| 327 | compress_ratio=self.compress_ratio, |
| 328 | mrope_section=getattr(self.rotary_emb, "mrope_section", None), |
| 329 | rope_pos_offset=( |
| 330 | raw_key_state_cache.rope_position_offset |
| 331 | if raw_key_state_cache.rope_position_cache is not None |
| 332 | else None |
| 333 | ), |
| 334 | main_qkv=qkv[:num_tokens], |
| 335 | main_q_norm_weight=attn.q_norm.weight, |
| 336 | main_k_norm_weight=attn.k_norm.weight, |
| 337 | main_eps=attn.q_norm.variance_epsilon, |
| 338 | main_kv_cache=main_kv_cache, |
| 339 | main_slot_mapping=slot_mapping[:num_tokens], |
| 340 | main_k_scale=attn._k_scale_float, |
| 341 | main_v_scale=attn._v_scale_float, |
| 342 | ) |
| 343 | else: |
| 344 | # Unfused reference path |
| 345 | from flashinfer.norm import gemma_rmsnorm |
| 346 | |
| 347 | q = projected_q.reshape(-1, self.index_n_heads, self.index_head_dim) |
| 348 | q = gemma_rmsnorm( |
| 349 | q.reshape(-1, self.index_head_dim), |
| 350 | self.q_layernorm.weight, |
| 351 | self.q_layernorm.variance_epsilon, |
| 352 | ).reshape_as(q) |
| 353 | q = apply_qsa_rope(self.rotary_emb, positions, q) |
| 354 | q = q.to(self.indexer_dtype) |
| 355 | |
| 356 | raw_key_cache = raw_key_state_cache.key_cache |
| 357 | rope_position_cache = raw_key_state_cache.rope_position_cache |
| 358 | if rope_position_cache is None: |
| 359 | position_rows = raw_metadata.logical_positions.view(-1, 1, 1).expand( |
| 360 | -1, 1, 3 |
| 361 | ) |
| 362 | else: |
| 363 | position_rows = canonical_qsa_rope_positions(positions).to( |
| 364 | device=raw_key_cache.device |
| 365 | ) |
| 366 | pooled, first_positions = qsa_compress_groups_with_ratio( |
| 367 | raw_keys.reshape(-1, 1, self.index_head_dim), |
| 368 | position_rows, |
| 369 | raw_key_cache, |
| 370 | raw_metadata.block_table, |
| 371 | raw_metadata.token_to_req, |
| 372 | raw_metadata.query_start_loc, |
| 373 | raw_metadata.logical_positions, |
| 374 | compressed_metadata.slot_mapping, |
| 375 | self.compress_ratio, |
| 376 | rope_position_cache, |
| 377 | ) |
| 378 | compressed_keys = gemma_rmsnorm( |
| 379 | pooled.reshape(-1, self.index_head_dim), |
| 380 | self.k_layernorm.weight, |
| 381 | self.k_layernorm.variance_epsilon, |
| 382 | ).reshape(-1, 1, self.index_head_dim) |
| 383 | if getattr(self.rotary_emb, "mrope_section", None): |
| 384 | first_positions = first_positions.transpose(0, 1) |
| 385 | else: |
| 386 | first_positions = first_positions[:, 0] |
| 387 | compressed_keys = apply_qsa_rope( |
| 388 | self.rotary_emb, |
| 389 | first_positions, |
| 390 | compressed_keys, |
| 391 | ) |
| 392 | qsa_store_cache_rows( |
| 393 | compressed_key_cache, |
| 394 | compressed_metadata.slot_mapping, |
| 395 | compressed_keys, |
| 396 | ) |
| 397 | qsa_store_cache_rows( |
| 398 | raw_key_cache, |
| 399 | raw_metadata.slot_mapping, |
| 400 | raw_keys, |
| 401 | ) |
| 402 | if rope_position_cache is not None: |
| 403 | qsa_store_cache_rows( |
| 404 | rope_position_cache, |
| 405 | raw_metadata.slot_mapping, |
| 406 | position_rows, |
| 407 | ) |
| 408 | |
| 409 | if self.skip_topk: |
| 410 | if out is None: |
| 411 | raise RuntimeError("QSA top-k reuse requires an output buffer") |
| 412 | return out, main_outputs |
| 413 | |
| 414 | if out is None: |
| 415 | out = torch.empty( |
| 416 | num_tokens, |
| 417 | self.packed_output_width, |
| 418 | dtype=torch.int32, |
| 419 | device=q.device, |
| 420 | ) |
| 421 | elif out.shape != (num_tokens, self.packed_output_width): |
| 422 | raise ValueError("QSA selection output has an invalid shape") |
| 423 | |
| 424 | num_decode_tokens = compressed_metadata.num_decode_tokens |
| 425 | decode_query_len = compressed_metadata.decode_query_len |
| 426 | visible_blocks = compressed_metadata.visible_blocks[:num_tokens] |
| 427 | block_indices = torch.empty( |
| 428 | num_tokens, |
| 429 | self.token_topk // self.compress_ratio, |
| 430 | dtype=torch.int32, |
| 431 | device=q.device, |
| 432 | ) |
| 433 | |
| 434 | # Decode requests occupy the leading rows and share one query length. |
| 435 | if num_decode_tokens: |
| 436 | num_decodes = compressed_metadata.num_decodes |
| 437 | if num_decodes * decode_query_len != num_decode_tokens: |
| 438 | raise ValueError("QSA decode rows must form a uniform request batch") |
| 439 | decode_slice = slice(0, num_decode_tokens) |
| 440 | qsa_select_paged_decode( |
| 441 | q[decode_slice], |
| 442 | compressed_key_cache, |
| 443 | compressed_metadata.block_table[:num_decodes], |
| 444 | visible_blocks[decode_slice], |
| 445 | self.token_topk, |
| 446 | self.compress_ratio, |
| 447 | decode_query_len, |
| 448 | block_indices[decode_slice], |
| 449 | ) |
| 450 | |
| 451 | # Prefill requests follow the leading decode rows in the reordered batch. |
| 452 | if num_decode_tokens < num_tokens: |
| 453 | num_decodes = compressed_metadata.num_decodes |
| 454 | prefill_slice = slice(num_decode_tokens, num_tokens) |
| 455 | qsa_select_paged_prefill( |
| 456 | q[prefill_slice], |
| 457 | compressed_key_cache, |
| 458 | compressed_metadata.block_table[num_decodes:], |
| 459 | compressed_metadata.query_start_loc[num_decodes:], |
| 460 | visible_blocks[prefill_slice], |
| 461 | self.token_topk, |
| 462 | self.compress_ratio, |
| 463 | compressed_metadata.max_query_len, |
| 464 | block_indices[prefill_slice], |
| 465 | compressed_metadata.max_seq_len, |
| 466 | ) |
| 467 | expand_qsa_block_indices( |
| 468 | block_indices, |
| 469 | compressed_metadata.logical_positions[:num_tokens], |
| 470 | visible_blocks, |
| 471 | self.compress_ratio, |
| 472 | self.token_topk, |
| 473 | out, |
| 474 | ) |
| 475 | return out, main_outputs |
| 476 | |
| 477 | |
| 478 | __all__ = ["QSAIndexer", "apply_qsa_rope"] |