| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm._aiter_ops import rocm_aiter_ops |
| 7 | from vllm.compilation.breakable_cudagraph import eager_break_during_capture |
| 8 | from vllm.forward_context import get_forward_context |
| 9 | from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer |
| 10 | from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache |
| 11 | from vllm.models.deepseek_v32.attention import DeepseekV32Attention, DeepseekV32Indexer |
| 12 | from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.utils.torch_utils import is_quantized_kv_cache |
| 15 | from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend |
| 16 | from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( |
| 17 | ROCMAiterMLASparseBackend, |
| 18 | ) |
| 19 | |
| 20 | |
| 21 | class DeepseekV32ROCmIndexerCache(DeepseekV32IndexerCache): |
| 22 | def get_attn_backend(self): |
| 23 | return DeepseekV32IndexerBackend |
| 24 | |
| 25 | @property |
| 26 | def uses_shuffled_layout(self) -> bool: |
| 27 | # aiter's gather/insert pair shuffles the cache above block size 1: |
| 28 | # [n_blocks, blk/16, head_dim/16, 16, 16] instead of [n_blocks, blk, head_dim]. |
| 29 | return self.kv_cache.ndim == 3 and self.kv_cache.shape[1] != 1 |
| 30 | |
| 31 | |
| 32 | class DeepseekV32ROCmIndexer(DeepseekV32Indexer): |
| 33 | indexer_cache_cls = DeepseekV32ROCmIndexerCache |
| 34 | |
| 35 | |
| 36 | class DeepseekV32MLAAttention(DeepseekV32Attention): |
| 37 | indexer_cls = DeepseekV32ROCmIndexer |
| 38 | |
| 39 | def __init__(self, vllm_config, config, prefix, topk_indices_buffer=None): |
| 40 | super().__init__( |
| 41 | vllm_config, |
| 42 | config, |
| 43 | prefix, |
| 44 | topk_indices_buffer, |
| 45 | attn_backend=ROCMAiterMLASparseBackend, |
| 46 | ) |
| 47 | |
| 48 | self.indexer_op: SparseAttnIndexer | None = None |
| 49 | if self.indexer is not None: |
| 50 | self.indexer_op = SparseAttnIndexer( |
| 51 | self.indexer.k_cache, |
| 52 | self.indexer.quant_block_size, |
| 53 | self.indexer.scale_fmt, |
| 54 | self.indexer.topk_tokens, |
| 55 | self.indexer.head_dim, |
| 56 | self.indexer.max_model_len, |
| 57 | self.indexer.max_total_seq_len, |
| 58 | topk_indices_buffer, |
| 59 | skip_k_cache_insert=True, |
| 60 | ) |
| 61 | self._fp8_kv = is_quantized_kv_cache(self.kv_cache_dtype) |
| 62 | self._fp8_kv_needs_view = self._fp8_kv and self.kv_cache_dtype != "fp8_ds_mla" |
| 63 | |
| 64 | def forward( # type: ignore[override] |
| 65 | self, |
| 66 | positions: torch.Tensor, |
| 67 | hidden_states: torch.Tensor, |
| 68 | ) -> torch.Tensor: |
| 69 | qkv_lora = self.fused_qkv_a_proj(hidden_states)[0] |
| 70 | q_c, kv_c, k_pe = qkv_lora.split( |
| 71 | [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 |
| 72 | ) |
| 73 | |
| 74 | if self.indexer is not None and not self.skip_topk: |
| 75 | kw = self.indexer.wk_weights_proj(hidden_states)[0] |
| 76 | index_k = kw[:, : self.indexer.head_dim] |
| 77 | index_weights = kw[:, self.indexer.head_dim :] |
| 78 | else: |
| 79 | index_k = None |
| 80 | index_weights = None |
| 81 | |
| 82 | num_tokens = hidden_states.shape[0] |
| 83 | output = torch.empty( |
| 84 | (num_tokens, self.num_local_heads * self.v_head_dim), |
| 85 | dtype=hidden_states.dtype, |
| 86 | device=hidden_states.device, |
| 87 | ) |
| 88 | self._fused_attention( |
| 89 | positions, q_c, kv_c, k_pe, index_k, index_weights, output |
| 90 | ) |
| 91 | return self.o_proj(output)[0] |
| 92 | |
| 93 | def _compute_ql_nope(self, q_c: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 94 | q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) |
| 95 | q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) |
| 96 | q_nope = q_nope.transpose(0, 1) # (N, tokens, P) |
| 97 | |
| 98 | if self.is_aiter_triton_fp4_bmm_enabled: |
| 99 | ql_nope = rocm_aiter_ops.batched_gemm_a16wfp4( |
| 100 | q_nope, self.W_K, self.W_K_scale, transpose_bm=True, prequant=True |
| 101 | ) |
| 102 | elif self.is_aiter_triton_fp8_bmm_enabled: |
| 103 | ql_nope = rocm_aiter_ops.triton_fp8_bmm( |
| 104 | q_nope, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True |
| 105 | ) |
| 106 | else: |
| 107 | ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1) |
| 108 | |
| 109 | return ql_nope, q_pe |
| 110 | |
| 111 | def _run_indexer( |
| 112 | self, |
| 113 | q_c: torch.Tensor, |
| 114 | index_q_fp8: torch.Tensor | None, |
| 115 | index_weights_out: torch.Tensor | None, |
| 116 | ) -> None: |
| 117 | """Run the ROCm sparse indexer (forward_hip) if this layer has an indexer.""" |
| 118 | if self.indexer_op is not None: |
| 119 | self.indexer_op.forward_hip(q_c, index_q_fp8, None, index_weights_out) |
| 120 | |
| 121 | def _build_q_for_attn( |
| 122 | self, |
| 123 | ql_nope: torch.Tensor, |
| 124 | mqa_q: torch.Tensor, |
| 125 | num_actual: int, |
| 126 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 127 | if self._fp8_kv: |
| 128 | # fp8 KV: mqa_q is the full [ql_nope; q_pe] packed as fp8. |
| 129 | return mqa_q[:num_actual] |
| 130 | |
| 131 | return (ql_nope[:num_actual], mqa_q[:num_actual]) |
| 132 | |
| 133 | def _compute_uv_out( |
| 134 | self, |
| 135 | attn_out: torch.Tensor, |
| 136 | output: torch.Tensor, |
| 137 | num_actual: int, |
| 138 | ) -> None: |
| 139 | x = attn_out.view( |
| 140 | num_actual, self.num_local_heads, self.kv_lora_rank |
| 141 | ).transpose(0, 1) # (N, tokens, L) |
| 142 | out_view = output[:num_actual].view( |
| 143 | num_actual, self.num_local_heads, self.v_head_dim |
| 144 | ) |
| 145 | |
| 146 | if self.is_aiter_triton_fp4_bmm_enabled: |
| 147 | rocm_aiter_ops.batched_gemm_a16wfp4( |
| 148 | x, self.W_V, self.W_V_scale, out_view, transpose_bm=True, prequant=True |
| 149 | ) |
| 150 | elif self.is_aiter_triton_fp8_bmm_enabled: |
| 151 | rocm_aiter_ops.triton_fp8_bmm( |
| 152 | x, |
| 153 | self.W_V, |
| 154 | self.W_V_scale, |
| 155 | group_size=128, |
| 156 | transpose_bm=True, |
| 157 | YQ=out_view, |
| 158 | ) |
| 159 | else: |
| 160 | torch.bmm(x, self.W_UV, out=out_view.transpose(0, 1)) |
| 161 | |
| 162 | @eager_break_during_capture |
| 163 | def _fused_attention( |
| 164 | self, |
| 165 | positions: torch.Tensor, |
| 166 | q_c: torch.Tensor, |
| 167 | kv_c: torch.Tensor, |
| 168 | k_pe: torch.Tensor, |
| 169 | index_k: torch.Tensor | None, |
| 170 | index_weights: torch.Tensor | None, |
| 171 | output: torch.Tensor, |
| 172 | ) -> None: |
| 173 | forward_context = get_forward_context() |
| 174 | attn_metadata_raw = forward_context.attn_metadata |
| 175 | if isinstance(attn_metadata_raw, dict): |
| 176 | attn_metadata = attn_metadata_raw.get(self.layer_name) |
| 177 | elif isinstance(attn_metadata_raw, list): |
| 178 | attn_metadata = attn_metadata_raw[0].get(self.layer_name) |
| 179 | else: |
| 180 | attn_metadata = attn_metadata_raw |
| 181 | |
| 182 | slot_mapping = forward_context.slot_mapping |
| 183 | assert isinstance(slot_mapping, dict) |
| 184 | mla_slot = slot_mapping.get(self.layer_name) |
| 185 | |
| 186 | if self.indexer is not None and not self.skip_topk: |
| 187 | has_indexer = True |
| 188 | indexer_k_norm_w = self.indexer.k_norm.weight |
| 189 | indexer_k_norm_bias = self.indexer.k_norm.bias |
| 190 | indexer_k_norm_eps = self.indexer.k_norm.eps |
| 191 | indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache |
| 192 | indexer_k_cache = self.indexer.k_cache.kv_cache |
| 193 | indexer_cache_shuffled = self.indexer.k_cache.uses_shuffled_layout |
| 194 | indexer_softmax_scale = self.indexer.softmax_scale |
| 195 | indexer_n_head_scale = self.indexer.n_head**-0.5 |
| 196 | else: |
| 197 | has_indexer = False |
| 198 | indexer_k_norm_w = None |
| 199 | indexer_k_norm_bias = None |
| 200 | indexer_k_norm_eps = 1e-6 |
| 201 | indexer_k_rope_cos_sin_cache = None |
| 202 | indexer_k_cache = None |
| 203 | indexer_cache_shuffled = False |
| 204 | indexer_softmax_scale = 0.0 |
| 205 | indexer_n_head_scale = 0.0 |
| 206 | |
| 207 | if attn_metadata is None: |
| 208 | mla_kv_cache = None |
| 209 | mla_k_scale = None |
| 210 | indexer_k_cache = None |
| 211 | mla_slot = None |
| 212 | else: |
| 213 | mla_kv_cache = self.kv_cache |
| 214 | mla_k_scale = self._k_scale |
| 215 | |
| 216 | q_c = fused_norm_rope( |
| 217 | positions, |
| 218 | q_c, |
| 219 | self.q_a_layernorm.weight, |
| 220 | self.q_a_layernorm.variance_epsilon, |
| 221 | kv_c, |
| 222 | self.kv_a_layernorm.weight, |
| 223 | self.kv_a_layernorm.variance_epsilon, |
| 224 | k_pe, |
| 225 | self.rotary_emb.cos_sin_cache, |
| 226 | index_k, |
| 227 | indexer_k_norm_w, |
| 228 | indexer_k_norm_bias, |
| 229 | indexer_k_norm_eps, |
| 230 | indexer_k_rope_cos_sin_cache, |
| 231 | self.topk_indices_buffer, |
| 232 | slot_mapping=mla_slot, |
| 233 | indexer_k_cache=indexer_k_cache, |
| 234 | indexer_cache_shuffled=indexer_cache_shuffled, |
| 235 | mla_kv_cache=mla_kv_cache, |
| 236 | mla_kv_cache_dtype=self.kv_cache_dtype, |
| 237 | mla_k_scale=mla_k_scale, |
| 238 | has_indexer=has_indexer, |
| 239 | index_rope_interleave=self._index_rope_interleave, |
| 240 | ) |
| 241 | |
| 242 | ql_nope, q_pe = self._compute_ql_nope(q_c) |
| 243 | |
| 244 | if self.indexer is not None and not self.skip_topk: |
| 245 | index_q = self.indexer.wq_b(q_c)[0] |
| 246 | index_q = index_q.view(-1, self.indexer.n_head, self.indexer.head_dim) |
| 247 | else: |
| 248 | index_q = None |
| 249 | |
| 250 | index_q_fp8, index_weights_out, mqa_q = fused_q( |
| 251 | positions, |
| 252 | q_pe, |
| 253 | self.rotary_emb.cos_sin_cache, |
| 254 | index_q, |
| 255 | self.indexer_rope_emb.cos_sin_cache if has_indexer else None, |
| 256 | ql_nope, |
| 257 | self._q_scale, |
| 258 | index_weights, |
| 259 | indexer_softmax_scale, |
| 260 | indexer_n_head_scale, |
| 261 | has_indexer=has_indexer, |
| 262 | index_rope_interleave=self._index_rope_interleave, |
| 263 | quantize_mqa=self._fp8_kv, |
| 264 | ) |
| 265 | |
| 266 | if self.indexer is not None and not self.skip_topk: |
| 267 | self._run_indexer(q_c, index_q_fp8, index_weights_out) |
| 268 | |
| 269 | if attn_metadata is None: |
| 270 | output.zero_() |
| 271 | return |
| 272 | |
| 273 | num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined] |
| 274 | kv_cache = self.kv_cache |
| 275 | if self._fp8_kv_needs_view: |
| 276 | kv_cache = kv_cache.view(current_platform.fp8_dtype()) |
| 277 | |
| 278 | q_for_attn = self._build_q_for_attn(ql_nope, mqa_q, num_actual) |
| 279 | attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined] |
| 280 | q_for_attn, kv_cache, attn_metadata, self |
| 281 | ) |
| 282 | |
| 283 | self._compute_uv_out(attn_out, output, num_actual) |