| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Triton attention backend with different K/V head dimensions (DiffKV). |
| 4 | |
| 5 | The KV cache layout is identical to ``FlashAttentionDiffKVBackend``: K and V |
| 6 | are packed along the last dim in the logical shape |
| 7 | ``[num_blocks, num_kv_heads, block_size, head_size_qk + head_size_v]``. |
| 8 | """ |
| 9 | |
| 10 | from typing import ClassVar |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | from vllm.config import VllmConfig |
| 15 | from vllm.config.cache import CacheDType |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.utils.math_utils import next_power_of_2 |
| 18 | from vllm.utils.torch_utils import is_quantized_kv_cache |
| 19 | from vllm.v1.attention.backend import AttentionLayer, AttentionType |
| 20 | from vllm.v1.attention.backends.triton_attn import ( |
| 21 | TritonAttentionBackend, |
| 22 | TritonAttentionImpl, |
| 23 | TritonAttentionMetadata, |
| 24 | TritonAttentionMetadataBuilder, |
| 25 | ) |
| 26 | from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( |
| 27 | triton_reshape_and_cache_flash_diffkv, |
| 28 | ) |
| 29 | from vllm.v1.attention.ops.triton_unified_attention_diffkv import ( |
| 30 | unified_attention_diffkv, |
| 31 | ) |
| 32 | from vllm.v1.kv_cache_interface import AttentionSpec |
| 33 | |
| 34 | logger = init_logger(__name__) |
| 35 | |
| 36 | |
| 37 | class TritonAttentionDiffKVMetadataBuilder(TritonAttentionMetadataBuilder): |
| 38 | """Override the parent's softmax buffer last-dim to head_size_v. |
| 39 | |
| 40 | The parent allocates ``softmax_segm_output`` with last-dim sized to |
| 41 | ``next_power_of_2(head_size)`` (== Q/K head size). For DiffKV the |
| 42 | accumulator and per-segment partial outputs are V-shaped, so we |
| 43 | re-allocate with ``next_power_of_2(head_size_v)`` instead. |
| 44 | """ |
| 45 | |
| 46 | def __init__( |
| 47 | self, |
| 48 | kv_cache_spec: AttentionSpec, |
| 49 | layer_names: list[str], |
| 50 | vllm_config: VllmConfig, |
| 51 | device: torch.device, |
| 52 | ): |
| 53 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 54 | |
| 55 | head_size_v = TritonAttentionDiffKVBackend.head_size_v |
| 56 | head_size_v_padded = next_power_of_2(head_size_v) |
| 57 | self.softmax_segm_output = torch.empty( |
| 58 | ( |
| 59 | self.softmax_segm_max.shape[0], |
| 60 | self.num_heads_q, |
| 61 | self.num_par_softmax_segments, |
| 62 | head_size_v_padded, |
| 63 | ), |
| 64 | dtype=torch.float32, |
| 65 | device=device, |
| 66 | ) |
| 67 | |
| 68 | |
| 69 | class TritonAttentionDiffKVBackend(TritonAttentionBackend): |
| 70 | # V head dim — set per layer via ``set_head_size_v`` before instantiation. |
| 71 | head_size_v: int = 128 |
| 72 | |
| 73 | # No FP8 / int8 KV cache for the DiffKV path yet; require fp16/bf16/fp32. |
| 74 | supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ |
| 75 | "auto", |
| 76 | "bfloat16", |
| 77 | ] |
| 78 | |
| 79 | @classmethod |
| 80 | def set_head_size_v(cls, head_size_v: int) -> None: |
| 81 | cls.head_size_v = head_size_v |
| 82 | |
| 83 | @staticmethod |
| 84 | def get_name() -> str: |
| 85 | return "TRITON_ATTN_DIFFKV" |
| 86 | |
| 87 | @staticmethod |
| 88 | def get_impl_cls() -> type["TritonAttentionDiffKVImpl"]: |
| 89 | return TritonAttentionDiffKVImpl |
| 90 | |
| 91 | @staticmethod |
| 92 | def get_builder_cls() -> type["TritonAttentionDiffKVMetadataBuilder"]: |
| 93 | return TritonAttentionDiffKVMetadataBuilder |
| 94 | |
| 95 | @classmethod |
| 96 | def supports_head_size(cls, head_size: int) -> bool: |
| 97 | # DiffKV K head sizes (e.g. 192 for MiMo-V2.5) need to be allowed. |
| 98 | return head_size >= 32 |
| 99 | |
| 100 | @classmethod |
| 101 | def supports_attn_type(cls, attn_type: str) -> bool: |
| 102 | # DiffKV only implements decoder self-attention. Unlike the parent |
| 103 | # TritonAttentionBackend (which advertises all types), encoder |
| 104 | # attention is not supported, so gate it here at backend selection. |
| 105 | return attn_type == AttentionType.DECODER |
| 106 | |
| 107 | |
| 108 | class TritonAttentionDiffKVImpl(TritonAttentionImpl): |
| 109 | """Triton attention impl for the DiffKV packed KV cache layout.""" |
| 110 | |
| 111 | def __init__(self, *args, **kwargs) -> None: |
| 112 | super().__init__(*args, **kwargs) |
| 113 | if is_quantized_kv_cache(self.kv_cache_dtype): |
| 114 | raise NotImplementedError( |
| 115 | "TritonAttentionDiffKVBackend does not yet support quantized " |
| 116 | f"KV cache (got kv_cache_dtype={self.kv_cache_dtype!r})." |
| 117 | ) |
| 118 | if self._is_per_token_head_quant: |
| 119 | raise NotImplementedError( |
| 120 | "TritonAttentionDiffKVBackend does not support per-token-head " |
| 121 | "quantization." |
| 122 | ) |
| 123 | if self.chunk_lookback > -1: |
| 124 | raise NotImplementedError( |
| 125 | "TritonAttentionDiffKVBackend does not support chunked " |
| 126 | "attention with lookback." |
| 127 | ) |
| 128 | |
| 129 | def do_kv_cache_update( |
| 130 | self, |
| 131 | layer: AttentionLayer, |
| 132 | key: torch.Tensor, |
| 133 | value: torch.Tensor, |
| 134 | kv_cache: torch.Tensor, |
| 135 | slot_mapping: torch.Tensor, |
| 136 | ) -> None: |
| 137 | # Cache is logical (B, H, N, C); the diffkv reshape kernel expects |
| 138 | # (B, N, H, C). |
| 139 | triton_reshape_and_cache_flash_diffkv( |
| 140 | key, |
| 141 | value, |
| 142 | kv_cache.transpose(1, 2), |
| 143 | slot_mapping, |
| 144 | self.kv_cache_dtype, |
| 145 | layer._k_scale, |
| 146 | layer._v_scale, |
| 147 | ) |
| 148 | |
| 149 | def fused_rope_kvcache_supported(self): |
| 150 | # The fused rope+cache path assumes the standard 2-tensor layout. |
| 151 | return False |
| 152 | |
| 153 | def forward( |
| 154 | self, |
| 155 | layer: torch.nn.Module, |
| 156 | query: torch.Tensor, |
| 157 | key: torch.Tensor, |
| 158 | value: torch.Tensor, |
| 159 | kv_cache: torch.Tensor, |
| 160 | attn_metadata: TritonAttentionMetadata, |
| 161 | output: torch.Tensor, |
| 162 | output_scale: torch.Tensor | None = None, |
| 163 | output_block_scale: torch.Tensor | None = None, |
| 164 | ) -> torch.Tensor: |
| 165 | """Forward pass. |
| 166 | |
| 167 | Shapes: |
| 168 | query: [num_tokens, num_heads, head_size_qk] |
| 169 | key: [num_tokens, num_kv_heads, head_size_qk] |
| 170 | value: [num_tokens, num_kv_heads, head_size_v] |
| 171 | kv_cache: [num_blocks, num_kv_heads, block_size, |
| 172 | head_size_qk + head_size_v] |
| 173 | output: [num_tokens, num_heads, head_size_v] |
| 174 | """ |
| 175 | if output_scale is not None or output_block_scale is not None: |
| 176 | raise NotImplementedError( |
| 177 | "fused output quantization is not supported for " |
| 178 | "TritonAttentionDiffKVImpl" |
| 179 | ) |
| 180 | |
| 181 | if attn_metadata is None: |
| 182 | return output.fill_(0) |
| 183 | |
| 184 | assert attn_metadata.use_cascade is False, ( |
| 185 | "Cascade attention not supported for TritonAttentionDiffKVImpl" |
| 186 | ) |
| 187 | |
| 188 | num_actual_tokens = attn_metadata.num_actual_tokens |
| 189 | head_size_qk = self.head_size |
| 190 | head_size_v = TritonAttentionDiffKVBackend.head_size_v |
| 191 | |
| 192 | # Triton DiffKV kernels consume (B, N, H, D) cache views. |
| 193 | kv_cache = kv_cache.transpose(1, 2) |
| 194 | key_cache = kv_cache[..., :head_size_qk] |
| 195 | value_cache = kv_cache[..., head_size_qk : head_size_qk + head_size_v] |
| 196 | |
| 197 | unified_attention_diffkv( |
| 198 | q=query[:num_actual_tokens], |
| 199 | k=key_cache, |
| 200 | v=value_cache, |
| 201 | out=output[:num_actual_tokens], |
| 202 | cu_seqlens_q=attn_metadata.query_start_loc, |
| 203 | seqused_k=attn_metadata.seq_lens, |
| 204 | softmax_scale=self.scale, |
| 205 | causal=True, |
| 206 | alibi_slopes=self.alibi_slopes, |
| 207 | use_alibi_sqrt=self.use_alibi_sqrt, |
| 208 | window_size=self.sliding_window, |
| 209 | block_table=attn_metadata.block_table, |
| 210 | softcap=self.logits_soft_cap, |
| 211 | sinks=self.sinks, |
| 212 | max_seqlen_q=attn_metadata.max_query_len, |
| 213 | seq_threshold_3D=attn_metadata.seq_threshold_3D, |
| 214 | num_par_softmax_segments=attn_metadata.num_par_softmax_segments, |
| 215 | softmax_segm_output=attn_metadata.softmax_segm_output, |
| 216 | softmax_segm_max=attn_metadata.softmax_segm_max, |
| 217 | softmax_segm_expsum=attn_metadata.softmax_segm_expsum, |
| 218 | ) |
| 219 | return output |