| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """AMX-only, high-performance MLA backend for DeepSeek V2/V3/R1 on CPU. |
| 4 | |
| 5 | Built on the AMX decode/extend/bmm kernels vendored under |
| 6 | ``csrc/cpu/sgl-kernels/``, plugged into vLLM's ``MLACommonBackend``/ |
| 7 | ``MLACommonImpl`` abstraction the same way every other concrete MLA backend |
| 8 | (``TritonMLAImpl``, etc.) does. |
| 9 | |
| 10 | This is a separate backend from the reference ``CPUMLABackend`` |
| 11 | (``vllm/v1/attention/backends/mla/cpu_mla.py``): that one targets every CPU |
| 12 | (any dtype, block_size=16, ``mla_decode_kvcache`` decode kernel + inherited |
| 13 | SDPA-style prefill) as a functional/CI reference, not performance. This |
| 14 | backend instead requires AMX (bf16 only, block_size a multiple of 32) and is |
| 15 | selected by the platform layer in preference to the reference backend |
| 16 | whenever the host supports it -- see ``CpuPlatform.get_attn_backend_cls``. |
| 17 | |
| 18 | Two points where this backend differs structurally from the GPU backends, |
| 19 | both explained in the CPU MLA design plan: |
| 20 | |
| 21 | - ``forward_mha`` is fully overridden (not inherited from |
| 22 | ``MLACommonBaseImpl``): the GPU "compute-friendly" prefill path depends on |
| 23 | a pluggable ``MLAPrefillBackend`` and CUDA-only chunked-context gather ops, |
| 24 | neither of which exist on CPU. Instead, this attends directly in |
| 25 | latent-MQA space via the ``extend_attention_cpu`` kernel, which handles |
| 26 | cached-prefix continuation and fresh prefill in one causal pass (the KV |
| 27 | cache already contains the new tokens by the time this runs, since |
| 28 | ``do_kv_cache_update`` executes first). |
| 29 | - Weight absorption for the prefill path is done with this impl's own |
| 30 | VNNI-packed copies of ``W_UK``/``W_UV`` (computed in |
| 31 | ``process_weights_after_loading``, using ``bmm_cpu``), rather than reading |
| 32 | the generic ``layer.W_UK_T``/``layer.W_UV``/``layer._v_up_proj`` the way |
| 33 | GPU backends do, since ``forward_mha``'s abstract signature has no |
| 34 | ``layer`` parameter (only ``forward_mqa`` does) and that signature is left |
| 35 | untouched. The decode path needs no such packing: ``MLAAttention.forward_impl`` |
| 36 | already absorbs/de-absorbs Q around the ``forward_mqa`` call using the |
| 37 | generic (unpacked) ``layer.W_UK_T``/``layer._v_up_proj``, so |
| 38 | ``forward_mqa`` here only has to invoke the decode kernel. |
| 39 | """ |
| 40 | |
| 41 | from __future__ import annotations |
| 42 | |
| 43 | from typing import ClassVar |
| 44 | |
| 45 | import torch |
| 46 | |
| 47 | from vllm import _custom_ops as ops |
| 48 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 49 | MLACommonBackend, |
| 50 | MLACommonImpl, |
| 51 | MLACommonMetadata, |
| 52 | MLACommonMetadataBuilder, |
| 53 | ) |
| 54 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 55 | get_and_maybe_dequant_weights, |
| 56 | ) |
| 57 | from vllm.platforms import current_platform |
| 58 | from vllm.v1.attention.backend import AttentionLayer, AttentionType, MultipleOf |
| 59 | |
| 60 | _MIN_WORK_PER_SPLIT = 512 |
| 61 | _SPLIT_OCCUPANCY_MULTIPLIER = 2 |
| 62 | |
| 63 | |
| 64 | def _block_quant_to_tensor_quant( |
| 65 | weight: torch.Tensor, # FP8 [N, K] block-quantized weight |
| 66 | weight_scale: torch.Tensor, # [ceil(N/bn), ceil(K/bk)] block scales |
| 67 | block_size: list[int], # [block_n, block_k] |
| 68 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 69 | """Re-quantize FP8 block-quant weight to FP8 per-tensor quant. |
| 70 | |
| 71 | Returns (fp8_weight, dequant_scale) where fp8_weight * dequant_scale |
| 72 | approximates the original float weight. dequant_scale is a 0-dim |
| 73 | scalar tensor as expected by bmm_cpu's ``scale`` argument. |
| 74 | |
| 75 | NOTE: block->tensor re-quantization may slightly affect accuracy |
| 76 | (equivalent to SGLang's block_quant_to_tensor_quant for CPU). |
| 77 | """ |
| 78 | block_n, block_k = block_size[0], block_size[1] |
| 79 | n, k = weight.shape |
| 80 | |
| 81 | # Expand block scales [ceil(N/bn), ceil(K/bk)] -> [N, K] |
| 82 | scale = weight_scale.to(torch.float32) |
| 83 | if scale.shape[0] < n: |
| 84 | scale = scale.repeat_interleave(block_n, dim=0)[:n, :] |
| 85 | if scale.shape[1] < k: |
| 86 | scale = scale.repeat_interleave(block_k, dim=1)[:, :k] |
| 87 | |
| 88 | # Dequantize to float32 |
| 89 | x_f32 = weight.to(torch.float32) * scale |
| 90 | |
| 91 | # Re-quantize per-tensor back to FP8 |
| 92 | fp8_max = torch.finfo(weight.dtype).max # 448.0 for e4m3fn |
| 93 | amax = x_f32.abs().amax().clamp(min=1e-12) |
| 94 | quant_scale = fp8_max / amax |
| 95 | x_fp8 = (x_f32 * quant_scale).clamp(-fp8_max, fp8_max).to(weight.dtype) |
| 96 | |
| 97 | # dequant_scale = 1 / quant_scale (0-dim tensor for bmm_cpu) |
| 98 | dequant_scale = (amax / fp8_max).to(torch.float32) |
| 99 | return x_fp8, dequant_scale |
| 100 | |
| 101 | |
| 102 | def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int: |
| 103 | """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread |
| 104 | count in place of SM count.""" |
| 105 | ideal_splits = 1 |
| 106 | while ideal_splits < max(1, max_seq_len // _MIN_WORK_PER_SPLIT): |
| 107 | ideal_splits *= 2 |
| 108 | max_splits = num_threads * _SPLIT_OCCUPANCY_MULTIPLIER |
| 109 | return min(ideal_splits, max_splits) |
| 110 | |
| 111 | |
| 112 | def _expand_block_table(block_table: torch.Tensor, block_size: int) -> torch.Tensor: |
| 113 | """Adapter: vLLM's block-table paging -> the flat per-(request, position) |
| 114 | physical row index the decode/extend kernels expect. Called once per step |
| 115 | from ``AMXMLAMetadataBuilder.build``, not per layer. |
| 116 | """ |
| 117 | offsets = torch.arange( |
| 118 | block_size, device=block_table.device, dtype=block_table.dtype |
| 119 | ) |
| 120 | flat = block_table.unsqueeze(-1) * block_size + offsets |
| 121 | return flat.reshape(block_table.size(0), -1).contiguous() |
| 122 | |
| 123 | |
| 124 | class AMXMLABackend(MLACommonBackend): |
| 125 | supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] |
| 126 | |
| 127 | @classmethod |
| 128 | def get_supported_head_sizes(cls) -> list[int]: |
| 129 | # kv_lora_rank(512) + qk_rope_head_dim(64), DeepSeek V2/V3/R1/V3.2. |
| 130 | return [576] |
| 131 | |
| 132 | @staticmethod |
| 133 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 134 | return [MultipleOf(32)] |
| 135 | |
| 136 | @classmethod |
| 137 | def supports_block_size(cls, block_size: int | None) -> bool: |
| 138 | if block_size is None: |
| 139 | return True |
| 140 | return block_size % 32 == 0 |
| 141 | |
| 142 | @staticmethod |
| 143 | def get_name() -> str: |
| 144 | return "AMX_MLA" |
| 145 | |
| 146 | @staticmethod |
| 147 | def get_impl_cls() -> type[AMXMLAImpl]: |
| 148 | return AMXMLAImpl |
| 149 | |
| 150 | @staticmethod |
| 151 | def get_builder_cls() -> type[AMXMLAMetadataBuilder]: |
| 152 | return AMXMLAMetadataBuilder |
| 153 | |
| 154 | |
| 155 | class AMXMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): |
| 156 | def build(self, common_prefix_len, common_attn_metadata, fast_build: bool = False): |
| 157 | attn_metadata = super().build( |
| 158 | common_prefix_len, common_attn_metadata, fast_build |
| 159 | ) |
| 160 | block_size = self.kv_cache_spec.block_size |
| 161 | if attn_metadata.decode is not None: |
| 162 | # Built once per step and reused by every layer, instead of |
| 163 | # recomputing the same tensor arithmetic in every forward_mqa call. |
| 164 | decode = attn_metadata.decode |
| 165 | decode.req_to_token = _expand_block_table( # type: ignore[attr-defined] |
| 166 | decode.block_table, block_size |
| 167 | ) |
| 168 | decode.seq_lens_i64 = decode.seq_lens.to(torch.int64) # type: ignore[attr-defined] |
| 169 | decode.req_pool_indices = torch.arange( # type: ignore[attr-defined] |
| 170 | decode.block_table.size(0), |
| 171 | dtype=torch.int64, |
| 172 | device=decode.block_table.device, |
| 173 | ) |
| 174 | decode.num_kv_splits = _compute_num_kv_splits( # type: ignore[attr-defined] |
| 175 | attn_metadata.max_seq_len, current_platform.num_compute_units() |
| 176 | ) |
| 177 | if attn_metadata.prefill is not None: |
| 178 | # cpu_seq_lens: total per-request context length (prefix + new |
| 179 | # tokens), the one thing the extend kernel needs that isn't |
| 180 | # already on the shared MLACommonPrefillMetadata. |
| 181 | prefill = attn_metadata.prefill |
| 182 | num_decodes = attn_metadata.num_decodes |
| 183 | num_prefills = attn_metadata.num_prefills |
| 184 | prefill.cpu_seq_lens = common_attn_metadata.seq_lens[ # type: ignore[attr-defined] |
| 185 | num_decodes : num_decodes + num_prefills |
| 186 | ].to(torch.int64) |
| 187 | prefill.req_to_token = _expand_block_table( # type: ignore[attr-defined] |
| 188 | prefill.block_table, block_size |
| 189 | ).to(torch.int64) |
| 190 | prefill.req_pool_indices = torch.arange( # type: ignore[attr-defined] |
| 191 | prefill.block_table.size(0), |
| 192 | dtype=torch.int64, |
| 193 | device=prefill.block_table.device, |
| 194 | ) |
| 195 | query_start_loc_i64 = prefill.query_start_loc.to(torch.int64) |
| 196 | extend_seq_lens = query_start_loc_i64[1:] - query_start_loc_i64[:-1] |
| 197 | prefill.extend_seq_lens = extend_seq_lens # type: ignore[attr-defined] |
| 198 | prefill.extend_start_loc = query_start_loc_i64[:-1] # type: ignore[attr-defined] |
| 199 | prefill.max_len_extend = int(extend_seq_lens.max().item()) # type: ignore[attr-defined] |
| 200 | return attn_metadata |
| 201 | |
| 202 | |
| 203 | class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]): |
| 204 | # Tells MLAAttention (mla_attention.py) to use this impl's own packed |
| 205 | # W_UK/W_UV + bmm_cpu for the shared decode absorb/de-absorb, mirroring |
| 206 | # is_aiter_triton_fp8_bmm_enabled/is_aiter_triton_fp4_bmm_enabled. |
| 207 | uses_amx_bmm = True |
| 208 | |
| 209 | def __init__( |
| 210 | self, |
| 211 | num_heads: int, |
| 212 | head_size: int, |
| 213 | scale: float, |
| 214 | num_kv_heads: int, |
| 215 | alibi_slopes: list[float] | None, |
| 216 | sliding_window: int | None, |
| 217 | kv_cache_dtype: str, |
| 218 | logits_soft_cap: float | None, |
| 219 | attn_type: str, |
| 220 | kv_sharing_target_layer_name: str | None, |
| 221 | **mla_args, |
| 222 | ) -> None: |
| 223 | super().__init__( |
| 224 | num_heads, |
| 225 | head_size, |
| 226 | scale, |
| 227 | num_kv_heads, |
| 228 | alibi_slopes, |
| 229 | sliding_window, |
| 230 | kv_cache_dtype, |
| 231 | logits_soft_cap, |
| 232 | attn_type, |
| 233 | kv_sharing_target_layer_name, |
| 234 | **mla_args, |
| 235 | ) |
| 236 | unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap] |
| 237 | if any(unsupported_features): |
| 238 | raise NotImplementedError( |
| 239 | "AMXMLAImpl does not support one of the following: " |
| 240 | "alibi_slopes, sliding_window, logits_soft_cap" |
| 241 | ) |
| 242 | if attn_type != AttentionType.DECODER: |
| 243 | raise NotImplementedError( |
| 244 | "Encoder self-attention and encoder/decoder cross-attention " |
| 245 | "are not implemented for AMXMLAImpl" |
| 246 | ) |
| 247 | self._w_uk_packed: torch.Tensor | None = None |
| 248 | self._w_uv_packed: torch.Tensor | None = None |
| 249 | self._w_scale: torch.Tensor | None = ( |
| 250 | None # per-tensor dequant scale (FP8 W8A16) |
| 251 | ) |
| 252 | |
| 253 | def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: |
| 254 | # forward_mha (unlike forward_mqa) receives no `layer` argument, so it |
| 255 | # cannot read the generic layer.W_UK_T/W_UV/._v_up_proj that |
| 256 | # MLAAttention.process_weights_after_loading already computes. This |
| 257 | # impl derives its own copies from self.kv_b_proj (already available |
| 258 | # via the constructor) and VNNI-packs them once, ahead of time, for |
| 259 | # use with bmm_cpu in forward_mha. |
| 260 | # |
| 261 | # kv_b_proj carries `_cpu_skip_gemm_dispatch`, so the CPU FP8 kernels |
| 262 | # leave its weight in the checkpoint [N, K] layout. For block-quant |
| 263 | # FP8 checkpoints, re-quantize per-tensor and keep the BMM in FP8 |
| 264 | # (W8A16); everything else goes through the BF16 path. |
| 265 | weight = getattr(self.kv_b_proj, "weight", None) |
| 266 | block_size = getattr(self.kv_b_proj, "weight_block_size", None) |
| 267 | if ( |
| 268 | weight is not None |
| 269 | and weight.dtype == torch.float8_e4m3fn |
| 270 | and block_size is not None |
| 271 | ): |
| 272 | block_scale = getattr(self.kv_b_proj, "weight_scale_inv", None) |
| 273 | if block_scale is None: |
| 274 | block_scale = self.kv_b_proj.weight_scale |
| 275 | fp8_weight, self._w_scale = _block_quant_to_tensor_quant( |
| 276 | weight, block_scale, block_size |
| 277 | ) |
| 278 | # [N, K] -> [K, N] = [kv_lora_rank, heads * (nope + v)] |
| 279 | kv_b_proj_weight = fp8_weight.T.contiguous().view( |
| 280 | self.kv_lora_rank, |
| 281 | self.num_heads, |
| 282 | self.qk_nope_head_dim + self.v_head_dim, |
| 283 | ) |
| 284 | else: |
| 285 | kv_b_proj_weight = get_and_maybe_dequant_weights( |
| 286 | self.kv_b_proj, out_dtype=act_dtype |
| 287 | ).T |
| 288 | kv_b_proj_weight = kv_b_proj_weight.view( |
| 289 | self.kv_lora_rank, |
| 290 | self.num_heads, |
| 291 | self.qk_nope_head_dim + self.v_head_dim, |
| 292 | ) |
| 293 | self._w_scale = None |
| 294 | w_uk, w_uv = kv_b_proj_weight.split( |
| 295 | [self.qk_nope_head_dim, self.v_head_dim], dim=-1 |
| 296 | ) |
| 297 | # bmm_cpu computes out[b] = mat1[b] @ mat2[b]^T (Linear-weight |
| 298 | # convention: mat2[b] is (OUT, IN)). |
| 299 | # absorb: ql_nope(N,B,L) = q_nope(N,B,P) @ w_uk_for_bmm(N,L,P)^T |
| 300 | # de-absorb: out(N,B,V) = attn_out(N,B,L) @ w_uv_for_bmm(N,V,L)^T |
| 301 | w_uk_for_bmm = w_uk.permute(1, 0, 2).contiguous() # (N, L, P) |
| 302 | w_uv_for_bmm = w_uv.permute(1, 2, 0).contiguous() # (N, V, L) |
| 303 | self._w_uk_packed = torch.ops._C.convert_weight_packed(w_uk_for_bmm) |
| 304 | self._w_uv_packed = torch.ops._C.convert_weight_packed(w_uv_for_bmm) |
| 305 | |
| 306 | def do_kv_cache_update( |
| 307 | self, |
| 308 | kv_c_normed: torch.Tensor, |
| 309 | k_pe: torch.Tensor, |
| 310 | kv_cache: torch.Tensor, |
| 311 | slot_mapping: torch.Tensor, |
| 312 | kv_cache_dtype: str, |
| 313 | k_scale: torch.Tensor, |
| 314 | ) -> None: |
| 315 | # Overrides the default (CUDA-only concat_and_cache_mla). |
| 316 | if kv_cache.numel() == 0: |
| 317 | return |
| 318 | assert kv_cache_dtype == "auto", ( |
| 319 | "AMXMLAImpl only supports an unquantized (bf16) KV cache; fp8 " |
| 320 | "KV cache for CPU MLA is not yet implemented." |
| 321 | ) |
| 322 | ops.amx_mla_concat_and_cache( |
| 323 | kv_c_normed, k_pe.squeeze(1), kv_cache, slot_mapping.flatten() |
| 324 | ) |
| 325 | |
| 326 | def forward_mqa( |
| 327 | self, |
| 328 | q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 329 | kv_c_and_k_pe_cache: torch.Tensor, |
| 330 | attn_metadata: MLACommonMetadata, |
| 331 | layer: AttentionLayer, |
| 332 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 333 | # q arrives already absorbed (tensor or (ql_nope, q_pe) tuple) -- |
| 334 | # MLAAttention.forward_impl does the Q-absorption bmm generically |
| 335 | # before calling forward_mqa, using layer.W_UK_T. De-absorption is |
| 336 | # likewise done generically afterwards via layer._v_up_proj. Nothing |
| 337 | # left for this impl to do beyond invoking the decode kernel. |
| 338 | if isinstance(q, tuple): |
| 339 | q = torch.cat(q, dim=-1) |
| 340 | assert attn_metadata.decode is not None |
| 341 | |
| 342 | num_tokens = q.shape[0] |
| 343 | kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size) |
| 344 | v_buffer = kv_cache_flat[..., : self.kv_lora_rank] |
| 345 | |
| 346 | seq_lens = attn_metadata.decode.seq_lens_i64 # type: ignore[attr-defined] |
| 347 | req_to_token = attn_metadata.decode.req_to_token # type: ignore[attr-defined] |
| 348 | req_pool_indices = attn_metadata.decode.req_pool_indices # type: ignore[attr-defined] |
| 349 | |
| 350 | num_kv_splits = attn_metadata.decode.num_kv_splits # type: ignore[attr-defined] |
| 351 | o = torch.zeros( |
| 352 | num_tokens, |
| 353 | self.num_heads, |
| 354 | self.kv_lora_rank, |
| 355 | dtype=q.dtype, |
| 356 | device=q.device, |
| 357 | ) |
| 358 | attn_logits = torch.zeros( |
| 359 | num_tokens, |
| 360 | self.num_heads, |
| 361 | num_kv_splits, |
| 362 | self.kv_lora_rank + 1, |
| 363 | dtype=torch.float32, |
| 364 | device=q.device, |
| 365 | ) |
| 366 | |
| 367 | ops.cpu_mla_decode( |
| 368 | q, |
| 369 | kv_cache_flat, |
| 370 | v_buffer, |
| 371 | o, |
| 372 | None, |
| 373 | None, |
| 374 | None, # loc: only meaningful when key/value are given; we |
| 375 | # already wrote the new K/V via do_kv_cache_update before this. |
| 376 | attn_logits, |
| 377 | req_to_token, |
| 378 | req_pool_indices, |
| 379 | seq_lens, |
| 380 | self.scale, |
| 381 | 0.0, |
| 382 | False, |
| 383 | 0, |
| 384 | None, |
| 385 | None, |
| 386 | ) |
| 387 | return o, None |
| 388 | |
| 389 | def forward_mha( |
| 390 | self, |
| 391 | q: torch.Tensor, |
| 392 | kv_c_normed: torch.Tensor, |
| 393 | k_pe: torch.Tensor, |
| 394 | kv_c_and_k_pe_cache: torch.Tensor, |
| 395 | attn_metadata: MLACommonMetadata, |
| 396 | k_scale: torch.Tensor, |
| 397 | output: torch.Tensor, |
| 398 | output_scale: torch.Tensor | None = None, |
| 399 | ) -> None: |
| 400 | assert output_scale is None, ( |
| 401 | "AMXMLAImpl.forward_mha does not support fused output quantization" |
| 402 | ) |
| 403 | prefill = attn_metadata.prefill |
| 404 | assert prefill is not None |
| 405 | assert self._w_uk_packed is not None and self._w_uv_packed is not None |
| 406 | |
| 407 | num_tokens = q.shape[0] |
| 408 | num_heads = self.num_heads |
| 409 | |
| 410 | q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) |
| 411 | |
| 412 | # Absorb: (N, B, P) x (N, L, P)^T -> (N, B, L) |
| 413 | # bmm_cpu tolerates non-contiguous mat1/out as long as their last dim |
| 414 | # is contiguous, so q_nope's transpose and ql_nope's transposed view |
| 415 | # can be passed directly -- no copy needed for either. |
| 416 | q_nope_t = q_nope.transpose(0, 1) |
| 417 | ql_nope = torch.empty( |
| 418 | num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device |
| 419 | ) |
| 420 | ops.bmm_cpu( |
| 421 | ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, self._w_scale |
| 422 | ) |
| 423 | mqa_q = torch.cat([ql_nope, q_pe], dim=-1) |
| 424 | |
| 425 | kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size) |
| 426 | v_buffer = kv_cache_flat[..., : self.kv_lora_rank] |
| 427 | |
| 428 | req_to_token = prefill.req_to_token # type: ignore[attr-defined] |
| 429 | req_pool_indices = prefill.req_pool_indices # type: ignore[attr-defined] |
| 430 | |
| 431 | seq_lens = prefill.cpu_seq_lens # type: ignore[attr-defined] |
| 432 | extend_seq_lens = prefill.extend_seq_lens # type: ignore[attr-defined] |
| 433 | extend_start_loc = prefill.extend_start_loc # type: ignore[attr-defined] |
| 434 | max_len_extend = prefill.max_len_extend # type: ignore[attr-defined] |
| 435 | |
| 436 | # k_extend/v_extend: the new tokens' own latent K/V, aliased in one |
| 437 | # 576-wide buffer (v_extend is a view of k_extend's first |
| 438 | # kv_lora_rank columns) so the kernel can gather both from one read, |
| 439 | # mirroring the cache's own layout. |
| 440 | k_extend = torch.cat( |
| 441 | [kv_c_normed, k_pe.reshape(num_tokens, self.qk_rope_head_dim)], dim=-1 |
| 442 | ).unsqueeze(1) |
| 443 | v_extend = k_extend[..., : self.kv_lora_rank] |
| 444 | |
| 445 | attn_out = torch.empty( |
| 446 | num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device |
| 447 | ) |
| 448 | ops.cpu_mla_extend( |
| 449 | mqa_q, |
| 450 | k_extend, |
| 451 | v_extend, |
| 452 | attn_out, |
| 453 | kv_cache_flat, |
| 454 | v_buffer, |
| 455 | req_to_token, |
| 456 | req_pool_indices, |
| 457 | seq_lens, |
| 458 | extend_seq_lens, |
| 459 | extend_start_loc, |
| 460 | max_len_extend, |
| 461 | self.scale, |
| 462 | 0.0, |
| 463 | False, |
| 464 | 0, |
| 465 | None, |
| 466 | None, |
| 467 | None, |
| 468 | ) |
| 469 | |
| 470 | # De-absorb directly into `output` (view, no extra copy on either side): |
| 471 | # (N, B, L) x (N, V, L)^T -> (N, B, V) |
| 472 | attn_out_t = attn_out.transpose(0, 1) |
| 473 | output_view = output.view(num_tokens, num_heads, self.v_head_dim).transpose( |
| 474 | 0, 1 |
| 475 | ) |
| 476 | ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, self._w_scale) |