| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass |
| 5 | from typing import Any, ClassVar, cast |
| 6 | |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | |
| 10 | from vllm.config import VllmConfig, get_current_vllm_config |
| 11 | from vllm.forward_context import get_forward_context |
| 12 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 13 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 14 | from vllm.model_executor.layers.linear import MergedColumnParallelLinear |
| 15 | from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( |
| 16 | fused_save_compress_norm, |
| 17 | rope_quant_insert, |
| 18 | ) |
| 19 | from vllm.platforms import current_platform |
| 20 | from vllm.triton_utils import tl, triton |
| 21 | from vllm.v1.attention.backend import ( |
| 22 | AttentionBackend, |
| 23 | AttentionCGSupport, |
| 24 | AttentionMetadataBuilder, |
| 25 | CommonAttentionMetadata, |
| 26 | MultipleOf, |
| 27 | ) |
| 28 | from vllm.v1.kv_cache_interface import CircularBufferSpec, KVCacheSpec |
| 29 | |
| 30 | |
| 31 | class CompressorBackend(AttentionBackend): |
| 32 | def __init__(self): |
| 33 | super().__init__() |
| 34 | |
| 35 | @staticmethod |
| 36 | def get_name() -> str: |
| 37 | return "CompressorBackend" |
| 38 | |
| 39 | @staticmethod |
| 40 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 41 | return [MultipleOf(1)] |
| 42 | |
| 43 | @classmethod |
| 44 | def get_supported_head_sizes(cls) -> list[int]: |
| 45 | return [512, 1024] |
| 46 | |
| 47 | @staticmethod |
| 48 | def get_builder_cls() -> type["CompressorMetadataBuilder"]: |
| 49 | return CompressorMetadataBuilder |
| 50 | |
| 51 | |
| 52 | @dataclass |
| 53 | class CompressorMetadata: |
| 54 | # [num_tokens] ring slot of every token: block * capacity + pos % capacity. |
| 55 | slot_mapping: torch.Tensor |
| 56 | query_start_loc: torch.Tensor # [num_reqs + 1] |
| 57 | token_to_req_indices: torch.Tensor # [num_tokens] |
| 58 | |
| 59 | |
| 60 | @triton.jit(do_not_specialize=["block_table_stride", "num_actual_tokens", "num_tokens"]) |
| 61 | def _ring_slot_mapping_kernel( |
| 62 | slot_mapping_ptr, |
| 63 | block_table_ptr, |
| 64 | block_table_stride, |
| 65 | token_to_req_ptr, |
| 66 | positions_ptr, |
| 67 | num_actual_tokens, |
| 68 | num_tokens, |
| 69 | CAPACITY: tl.constexpr, |
| 70 | BLOCK: tl.constexpr, |
| 71 | ): |
| 72 | offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) |
| 73 | valid = offsets < num_actual_tokens |
| 74 | req = tl.load(token_to_req_ptr + offsets, mask=valid, other=0).to(tl.int64) |
| 75 | block = tl.load(block_table_ptr + req * block_table_stride, mask=valid, other=0) |
| 76 | pos = tl.load(positions_ptr + offsets, mask=valid, other=0) |
| 77 | slot = block.to(tl.int64) * CAPACITY + pos % CAPACITY |
| 78 | tl.store( |
| 79 | slot_mapping_ptr + offsets, tl.where(valid, slot, -1), mask=offsets < num_tokens |
| 80 | ) |
| 81 | |
| 82 | |
| 83 | class CompressorMetadataBuilder(AttentionMetadataBuilder): |
| 84 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS |
| 85 | |
| 86 | def __init__(self, *args, **kwargs): |
| 87 | super().__init__(*args, **kwargs) |
| 88 | assert isinstance(self.kv_cache_spec, CircularBufferSpec) |
| 89 | self.capacity = self.kv_cache_spec.block_size |
| 90 | max_num_batched_tokens = ( |
| 91 | self.vllm_config.scheduler_config.max_num_batched_tokens |
| 92 | ) |
| 93 | self.token_to_req_indices = torch.zeros( |
| 94 | max_num_batched_tokens, dtype=torch.int32, device=self.device |
| 95 | ) |
| 96 | self.slot_mapping_buffer = torch.empty( |
| 97 | max_num_batched_tokens, dtype=torch.int64, device=self.device |
| 98 | ) |
| 99 | |
| 100 | def build( |
| 101 | self, |
| 102 | common_prefix_len: int, |
| 103 | common_attn_metadata: CommonAttentionMetadata, |
| 104 | fast_build: bool = False, |
| 105 | ) -> CompressorMetadata: |
| 106 | num_tokens = common_attn_metadata.slot_mapping.numel() |
| 107 | positions = common_attn_metadata.positions |
| 108 | assert positions is not None |
| 109 | token_to_req_indices = common_attn_metadata.token_to_req_indices( |
| 110 | self.token_to_req_indices |
| 111 | ) |
| 112 | slot_mapping = self.slot_mapping_buffer[:num_tokens] |
| 113 | block_table = common_attn_metadata.block_table_tensor |
| 114 | _ring_slot_mapping_kernel[(triton.cdiv(num_tokens, 256),)]( |
| 115 | slot_mapping, |
| 116 | block_table, |
| 117 | block_table.stride(0), |
| 118 | token_to_req_indices, |
| 119 | positions, |
| 120 | common_attn_metadata.num_actual_tokens, |
| 121 | num_tokens, |
| 122 | CAPACITY=self.capacity, |
| 123 | BLOCK=256, |
| 124 | ) |
| 125 | return CompressorMetadata( |
| 126 | slot_mapping=slot_mapping, |
| 127 | query_start_loc=common_attn_metadata.query_start_loc, |
| 128 | token_to_req_indices=token_to_req_indices, |
| 129 | ) |
| 130 | |
| 131 | |
| 132 | class CompressorStateCache(torch.nn.Module, AttentionLayerBase): |
| 133 | """Per-request ring of the open group's [kv, score] rows (ratio > 1).""" |
| 134 | |
| 135 | def __init__(self, state_dim: int, dtype: torch.dtype, prefix: str): |
| 136 | super().__init__() |
| 137 | self.state_dim = state_dim |
| 138 | self.dtype = dtype |
| 139 | self.prefix = prefix |
| 140 | self.kv_cache = torch.tensor([]) |
| 141 | compilation_config = get_current_vllm_config().compilation_config |
| 142 | if prefix in compilation_config.static_forward_context: |
| 143 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 144 | compilation_config.static_forward_context[prefix] = self |
| 145 | |
| 146 | assert self.dtype == torch.float32 |
| 147 | # Drafts + bonus token + the previous row, as a power of two >= 8. |
| 148 | rows_per_step = get_current_vllm_config().num_speculative_tokens + 2 |
| 149 | self.block_size = max(8, 1 << (rows_per_step - 1).bit_length()) |
| 150 | |
| 151 | def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: |
| 152 | # [B, H=1, N, C] -> [B, N, C] |
| 153 | self.kv_cache = kv_cache.squeeze(1) |
| 154 | |
| 155 | def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: |
| 156 | del vllm_config |
| 157 | return CircularBufferSpec( |
| 158 | block_size=self.block_size, |
| 159 | num_kv_heads=1, |
| 160 | head_size=self.state_dim, |
| 161 | head_size_v=0, |
| 162 | dtype=self.dtype, |
| 163 | ) |
| 164 | |
| 165 | def forward(self): ... |
| 166 | |
| 167 | def get_attn_backend(self) -> type[AttentionBackend]: |
| 168 | return CompressorBackend |
| 169 | |
| 170 | |
| 171 | class DeepseekCompressor(nn.Module): |
| 172 | """DeepSeek V4.1 KV/score compressor. |
| 173 | |
| 174 | Pools ``compress_ratio`` consecutive tokens into one KV latent with a |
| 175 | learned softmax gate (ratio 1 has no gate and no pooling). Owns the |
| 176 | linear, norm and state cache. State saving, compression and RMSNorm share |
| 177 | one Triton kernel. The emitted latent feeds independent main-cache and |
| 178 | indexer K kernels, which the attention layer can schedule concurrently. |
| 179 | """ |
| 180 | |
| 181 | def __init__( |
| 182 | self, |
| 183 | vllm_config: VllmConfig, |
| 184 | compress_ratio: int, |
| 185 | hidden_size: int, |
| 186 | head_dim: int, |
| 187 | rotate: bool = False, |
| 188 | prefix: str = "", |
| 189 | k_cache_prefix="", |
| 190 | ): |
| 191 | super().__init__() |
| 192 | if compress_ratio not in (1, 2): |
| 193 | raise NotImplementedError( |
| 194 | "DeepSeek V4.1 compressor supports compress_ratio 1 (full-length " |
| 195 | f"compressed cache) and 2; got {compress_ratio}. The ratio-4/128 " |
| 196 | "CuTe-DSL kernels are v4.0-specific and not wired here." |
| 197 | ) |
| 198 | self.compress_ratio = compress_ratio |
| 199 | self.hidden_size = hidden_size |
| 200 | self.head_dim = head_dim |
| 201 | self.rotate = rotate |
| 202 | self.prefix = prefix |
| 203 | self.k_cache_prefix = k_cache_prefix |
| 204 | # Ratio 1 pools single tokens, so the checkpoint carries no gate. |
| 205 | self.has_gate = compress_ratio > 1 |
| 206 | |
| 207 | config = vllm_config.model_config.hf_config |
| 208 | self.rope_head_dim = config.qk_rope_head_dim |
| 209 | self.nope_head_dim = self.head_dim - self.rope_head_dim |
| 210 | assert self.head_dim == 512 and self.rope_head_dim == 64 |
| 211 | self.rms_norm_eps = config.rms_norm_eps |
| 212 | self.device = current_platform.device_type |
| 213 | self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs |
| 214 | self.max_model_len = vllm_config.model_config.max_model_len |
| 215 | |
| 216 | wkv_wgate_sizes = ( |
| 217 | [self.head_dim, self.head_dim] if self.has_gate else [self.head_dim] |
| 218 | ) |
| 219 | self.fused_wkv_wgate = MergedColumnParallelLinear( |
| 220 | self.hidden_size, |
| 221 | wkv_wgate_sizes, |
| 222 | bias=False, |
| 223 | return_bias=False, |
| 224 | quant_config=None, |
| 225 | disable_tp=True, |
| 226 | prefix=f"{prefix}.fused_wkv_wgate", |
| 227 | ) |
| 228 | self.norm = RMSNorm(self.head_dim, self.rms_norm_eps) |
| 229 | |
| 230 | self.state_cache = ( |
| 231 | CompressorStateCache( |
| 232 | state_dim=2 * self.head_dim, # kv_state + score_state |
| 233 | dtype=torch.float32, |
| 234 | prefix=f"{prefix}.state_cache", |
| 235 | ) |
| 236 | if compress_ratio > 1 |
| 237 | else None |
| 238 | ) |
| 239 | |
| 240 | # Save reference to static_forward_context for forward-time KV cache lookup. |
| 241 | # get_current_vllm_config() is only available during __init__, not forward. |
| 242 | self._static_forward_context = ( |
| 243 | vllm_config.compilation_config.static_forward_context |
| 244 | ) |
| 245 | |
| 246 | def forward( |
| 247 | self, |
| 248 | # [num_tokens, (2 if has_gate else 1) * self.head_dim] |
| 249 | kv_score: torch.Tensor, |
| 250 | # [num_tokens] |
| 251 | positions: torch.Tensor, |
| 252 | ) -> torch.Tensor | None: |
| 253 | """Save states and return the BF16 latent for cache insertion and indexing. |
| 254 | |
| 255 | Only valid group-boundary rows are written. |
| 256 | """ |
| 257 | attn_metadata = get_forward_context().attn_metadata |
| 258 | if not isinstance(attn_metadata, dict): |
| 259 | return None |
| 260 | |
| 261 | if self.state_cache is None: |
| 262 | state_cache = query_start_loc = token_to_req_indices = None |
| 263 | slot_mapping = cast(Any, attn_metadata[self.k_cache_prefix]).slot_mapping |
| 264 | else: |
| 265 | state_metadata = cast( |
| 266 | CompressorMetadata, attn_metadata[self.state_cache.prefix] |
| 267 | ) |
| 268 | state_cache = self.state_cache.kv_cache |
| 269 | slot_mapping = state_metadata.slot_mapping |
| 270 | query_start_loc = state_metadata.query_start_loc |
| 271 | token_to_req_indices = state_metadata.token_to_req_indices |
| 272 | |
| 273 | latent = torch.empty( |
| 274 | kv_score.shape[0], |
| 275 | self.head_dim, |
| 276 | dtype=torch.bfloat16, |
| 277 | device=kv_score.device, |
| 278 | ) |
| 279 | fused_save_compress_norm( |
| 280 | kv_score, |
| 281 | positions, |
| 282 | state_cache, |
| 283 | slot_mapping, |
| 284 | query_start_loc, |
| 285 | token_to_req_indices, |
| 286 | self.norm.weight, |
| 287 | self.rms_norm_eps, |
| 288 | self.compress_ratio, |
| 289 | latent, |
| 290 | ) |
| 291 | return latent |
| 292 | |
| 293 | def insert_cache( |
| 294 | self, |
| 295 | latent: torch.Tensor | None, |
| 296 | positions: torch.Tensor, |
| 297 | rotary_emb, |
| 298 | ) -> None: |
| 299 | """Publish compressed main-cache rows after the latent becomes ready.""" |
| 300 | if latent is None: |
| 301 | return |
| 302 | attn_metadata = get_forward_context().attn_metadata |
| 303 | assert isinstance(attn_metadata, dict) |
| 304 | k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix]) |
| 305 | k_cache_layer = self._static_forward_context[self.k_cache_prefix] |
| 306 | kv_cache = k_cache_layer.kv_cache |
| 307 | # Plain-row per-tensor fp8 caches (FlashInfer) carry the layer's scale; |
| 308 | # fp8_ds_mla and bf16 rows need none. |
| 309 | fp8_scale = ( |
| 310 | getattr(k_cache_layer, "_flashinfer_fp8_kv_scale", None) |
| 311 | if kv_cache.dtype == torch.float8_e4m3fn |
| 312 | else None |
| 313 | ) |
| 314 | |
| 315 | rope_quant_insert( |
| 316 | latent, |
| 317 | positions, |
| 318 | rotary_emb.cos_sin_cache, |
| 319 | kv_cache, |
| 320 | k_cache_metadata.slot_mapping, |
| 321 | self.compress_ratio, |
| 322 | fp8_scale=fp8_scale, |
| 323 | ) |