@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 58d1b2aa77
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 21:00 UTC

[Attention] MLA support for V1 (#13789)

Signed-off-by: Yang Chen <yangche@fb.com>

Yang Chen committed Feb 27, 2025 at 10:14 UTC 58d1b2aa772deb166355423997fbf5c1b6b186a1
10 files changed +1340 -59
vllm/attention/layer.py
+24 -11
@@ -89,6 +89,7 @@ class Attention(nn.Module):
89 self._k_scale_float = 1.0
90 self._v_scale_float = 1.0
91
92 + self.use_mla = use_mla
93 self.num_heads = num_heads
94 self.head_size = head_size
95 self.num_kv_heads = num_kv_heads
@@ -158,6 +159,10 @@ class Attention(nn.Module):
159 query: torch.Tensor,
160 key: torch.Tensor,
161 value: torch.Tensor,
162 + # For some alternate attention backends like MLA the attention output
163 + # shape does not match the query shape, so we optionally let the model
164 + # definition specify the output tensor shape.
165 + output_shape: Optional[torch.Size] = None,
166 ) -> torch.Tensor:
167 """
168 The KV cache is stored inside this class and is accessed via
@@ -173,17 +178,25 @@ class Attention(nn.Module):
178 if attn_metadata.enable_kv_scales_calculation:
179 self.calc_kv_scales(key, value)
180 if self.use_output:
176 - output = torch.empty_like(query)
177 - hidden_size = query.size(-1)
178 - # Reshape the query, key, and value tensors.
179 - # NOTE(woosuk): We do this outside the custom op to minimize the
180 - # CPU overheads from the non-CUDA-graph regions.
181 - query = query.view(-1, self.num_heads, self.head_size)
182 - output = output.view(-1, self.num_heads, self.head_size)
183 - if key is not None:
184 - key = key.view(-1, self.num_kv_heads, self.head_size)
185 - if value is not None:
186 - value = value.view(-1, self.num_kv_heads, self.head_size)
181 + output_shape = (output_shape
182 + if output_shape is not None else query.shape)
183 + output = torch.empty(output_shape,
184 + dtype=query.dtype,
185 + device=query.device)
186 + hidden_size = output_shape[-1]
187 + # We skip reshaping query, key and value tensors for the MLA
188 + # backend since these tensors have different semantics and are
189 + # processed differently.
190 + if not self.use_mla:
191 + # Reshape the query, key, and value tensors.
192 + # NOTE(woosuk): We do this outside the custom op to minimize the
193 + # CPU overheads from the non-CUDA-graph regions.
194 + query = query.view(-1, self.num_heads, self.head_size)
195 + output = output.view(-1, self.num_heads, self.head_size)
196 + if key is not None:
197 + key = key.view(-1, self.num_kv_heads, self.head_size)
198 + if value is not None:
199 + value = value.view(-1, self.num_kv_heads, self.head_size)
200 if self.use_direct_call:
201 forward_context: ForwardContext = get_forward_context()
202 attn_metadata = forward_context.attn_metadata
vllm/model_executor/models/deepseek_v2.py
+11 -2
@@ -420,9 +420,15 @@ class DeepseekV2MLAAttention(nn.Module):
420 mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
421 self.scaling = self.scaling * mscale * mscale
422
423 + # In the MLA backend, kv_cache includes both k_c and
424 + # pe (i.e. decoupled position embeddings). In particular,
425 + # the concat_and_cache_mla op requires
426 + # k_c.size(1) + k_pe.size(1) == kv_cache.size(2)
427 + # i.e.
428 + # kv_lora_rank + qk_rope_head_dim == head_size
429 self.mla_attn = Attention(
430 num_heads=self.num_local_heads,
425 - head_size=self.kv_lora_rank,
431 + head_size=self.kv_lora_rank + self.qk_rope_head_dim,
432 scale=self.scaling,
433 num_kv_heads=1,
434 cache_config=cache_config,
@@ -458,7 +464,10 @@ class DeepseekV2MLAAttention(nn.Module):
464 kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states)[0].split(
465 [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
466 kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
461 - return self.mla_attn(hidden_states_or_q_c, kv_c_normed, k_pe)
467 + return self.mla_attn(hidden_states_or_q_c,
468 + kv_c_normed,
469 + k_pe,
470 + output_shape=hidden_states.shape)
471
472
473 class DeepseekV2DecoderLayer(nn.Module):
vllm/platforms/cuda.py
+7 -2
@@ -162,8 +162,13 @@ class CudaPlatformBase(Platform):
162 kv_cache_dtype, block_size, use_v1,
163 use_mla) -> str:
164 if use_v1:
165 - logger.info("Using Flash Attention backend on V1 engine.")
166 - return "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"
165 + if use_mla:
166 + logger.info("Using Triton MLA backend on V1 engine.")
167 + return "vllm.v1.attention.backends.triton_mla.TritonMLABackend"
168 + else:
169 + logger.info("Using Flash Attention backend on V1 engine.")
170 + return ("vllm.v1.attention.backends.flash_attn."
171 + "FlashAttentionBackend")
172 if use_mla:
173 if selected_backend == _Backend.FLASHMLA:
174 from vllm.attention.backends.flashmla import (
vllm/platforms/interface.py
+1
@@ -35,6 +35,7 @@ class _Backend(enum.Enum):
35 OPENVINO = enum.auto()
36 FLASHINFER = enum.auto()
37 TRITON_MLA = enum.auto()
38 + TRITON_MLA_VLLM_V1 = enum.auto()
39 FLASHMLA = enum.auto()
40 HPU_ATTN = enum.auto()
41 PALLAS = enum.auto()
vllm/v1/attention/backends/flash_attn.py
+67 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """Attention layer with FlashAttention."""
3 from dataclasses import dataclass
4 -from typing import Any, Dict, List, Optional, Tuple, Type
4 +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
5
6 import numpy as np
7 import torch
@@ -14,6 +14,11 @@ from vllm.logger import init_logger
14 from vllm.platforms import current_platform
15 from vllm.utils import cdiv
16
17 +if TYPE_CHECKING:
18 + from vllm.v1.core.scheduler_output import SchedulerOutput
19 + from vllm.v1.worker.gpu_input_batch import InputBatch
20 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner
21 +
22 if current_platform.is_cuda():
23 from vllm.vllm_flash_attn import flash_attn_varlen_func
24
@@ -40,6 +45,10 @@ class FlashAttentionBackend(AttentionBackend):
45 def get_metadata_cls() -> Type["AttentionMetadata"]:
46 return FlashAttentionMetadata
47
48 + @staticmethod
49 + def get_builder_cls() -> Type["FlashAttentionMetadataBuilder"]:
50 + return FlashAttentionMetadataBuilder
51 +
52 @staticmethod
53 def get_kv_cache_shape(
54 num_blocks: int,
@@ -85,6 +94,62 @@ class FlashAttentionMetadata:
94 num_input_tokens: int = 0 # Number of tokens including padding.
95
96
97 +class FlashAttentionMetadataBuilder:
98 +
99 + def __init__(self, runner: "GPUModelRunner"):
100 + self.runner = runner
101 +
102 + def reorder_batch(self, input_batch: "InputBatch",
103 + scheduler_output: "SchedulerOutput"):
104 + pass
105 +
106 + def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
107 + common_prefix_len: int):
108 + max_seq_len = self.runner.seq_lens_np[:num_reqs].max()
109 + query_start_loc = self.runner.query_start_loc_cpu[:num_reqs + 1].to(
110 + self.runner.device, non_blocking=True)
111 + seq_lens = self.runner.seq_lens_cpu[:num_reqs].to(self.runner.device,
112 + non_blocking=True)
113 + block_table = (
114 + self.runner.input_batch.block_table.get_device_tensor()[:num_reqs])
115 + slot_mapping = self.runner.slot_mapping_cpu[:num_actual_tokens].to(
116 + self.runner.device, non_blocking=True).long()
117 +
118 + use_cascade = common_prefix_len > 0
119 + if use_cascade:
120 + # TODO: Optimize.
121 + cu_prefix_query_lens = torch.tensor([0, num_actual_tokens],
122 + dtype=torch.int32,
123 + device=self.runner.device)
124 + prefix_kv_lens = torch.tensor([common_prefix_len],
125 + dtype=torch.int32,
126 + device=self.runner.device)
127 + suffix_kv_lens = (self.runner.seq_lens_np[:num_reqs] -
128 + common_prefix_len)
129 + suffix_kv_lens = torch.from_numpy(suffix_kv_lens).to(
130 + self.runner.device)
131 + else:
132 + cu_prefix_query_lens = None
133 + prefix_kv_lens = None
134 + suffix_kv_lens = None
135 +
136 + attn_metadata = FlashAttentionMetadata(
137 + num_actual_tokens=num_actual_tokens,
138 + max_query_len=max_query_len,
139 + query_start_loc=query_start_loc,
140 + max_seq_len=max_seq_len,
141 + seq_lens=seq_lens,
142 + block_table=block_table,
143 + slot_mapping=slot_mapping,
144 + use_cascade=use_cascade,
145 + common_prefix_len=common_prefix_len,
146 + cu_prefix_query_lens=cu_prefix_query_lens,
147 + prefix_kv_lens=prefix_kv_lens,
148 + suffix_kv_lens=suffix_kv_lens,
149 + )
150 + return attn_metadata
151 +
152 +
153 class FlashAttentionImpl(AttentionImpl):
154
155 def __init__(
@@ -371,4 +436,4 @@ def cascade_attention(
436
437 # Merge prefix and suffix outputs, and store the result in output.
438 merge_attn_states(output, prefix_output, prefix_lse, suffix_output,
374 - suffix_lse)
\ No newline at end of file
439 + suffix_lse)
vllm/v1/attention/backends/mla/__init__.py
vllm/v1/attention/backends/mla/common.py new
+1022
@@ -0,0 +1,1022 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +This file implements common components for MLA implementations.
4 +
5 +First we define:
6 +
7 +Sq as Q sequence length
8 +Skv as KV sequence length
9 +
10 +MLA has two possible ways of computing, a data-movement friendly approach and a
11 +compute friendly approach, we generally want to use the compute friendly
12 +approach for "prefill" (i.e. the ratio Sq / Skv is "small", is near 1)
13 +and the data-movement friendly approach for "decode" (i.e. the ratio
14 +Sq / Skv is "large").
15 +
16 +NOTE what we deem small and large is currently determined by if its labelled
17 +prefill or decode by the scheduler, but this is something we should probably
18 +tune.
19 +
20 +Main reference: DeepseekV2 paper, and FlashInfer Implementation
21 +(https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
22 +
23 +Deepseek's MLA attention works the following way:
24 +* Use a single latent vector to represent the per-token entry of the KV cache.
25 +* For decode (i.e. the memory friendly approach) the attention "simulates" a
26 +multi-head attention, while the compute is similar to multi-query attention.
27 +
28 +Below is example of both paths assuming batchsize = 1
29 +
30 +## More Extent Definitions:
31 +
32 +C Context length, `Skv - Sq`
33 +H hidden size
34 +N number of attention heads
35 +Lq latent dimension for Q 1536 in DSV3
36 +Lkv latent dimension for K/V 512 in DSV3
37 +P nope dimension, no rope. 128 in DSV3
38 +R rope dimension, goes through rope. 64 in DSV3
39 +V V head dim. 128 in DSV3
40 +
41 +## Vector/Matrix Definitions
42 +
43 +h_t hidden states (input to attention) shape [Sq, H]
44 +q_c latent/compressed Q shape [Sq, Lq]
45 +q_nope uncompressed Q (no-rope) shape [Sq, N, P]
46 +q_pe uncompressed Q (rope) shape [Sq, N, R]
47 +kv_c latent/compressed KV shape [Skv, Lkv]
48 +k_pe decoupled k position embeddings shape [Skv, R]
49 +new_kv_c new kv_c from current iter shape [Sq, Lkv]
50 +new_k_pe new k_pe from current iter shape [Sq, R]
51 +cache_kv_c cached k_c from previous iters shape [C, Lkv]
52 +cache_k_pe cached k_pe from previous iters shape [C, R]
53 +W_DQ project h_t to q_c shape [H, Lq]
54 +W_UQ project q_c to q_nope shape [Lq, N * P]
55 +W_QR project q_c to q_pe shape [Lq, N * R]
56 +W_DKV project h_t to kv_c shape [H, Lkv]
57 +W_UK project kv_c to k_nope shape [Lkv, N * P]
58 +W_KR project h_t to k_pe shape [H, N * R]
59 +W_UV project kv_c to v shape [Lkv, N * V]
60 +W_O project v to h_t shape [N * V, H]
61 +
62 +
63 +## Compute Friendly Approach (i.e. "_forward_prefill"):
64 +
65 +q_c = h_t @ W_DQ
66 +q_nope = (q_c @ W_UQ).view(Sq, N, P)
67 +q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
68 +new_kv_c = h_t @ W_DKV
69 +new_k_pe = RoPE(h_t @ W_KR)
70 +kv_c = torch.cat([new_kv_c, cache_kv_c], dim=0)
71 +k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)
72 +k_nope = (kv_c @ W_UK).view(Skv, N, P)
73 +v = (kv_c @ W_UV).view(Skv, N, V)
74 +
75 +// MHA with QK headdim = P + R
76 +// V headdim = V
77 +// spda_o shape [Sq, N, V]
78 +spda_o = scaled_dot_product_attention(
79 + torch.cat([q_nope, q_pe], dim=-1),
80 + torch.cat([k_nope, k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
81 + v
82 +)
83 +return spda_o @ W_O
84 +
85 +NOTE: in the actual code,
86 + `kv_b_proj` is [W_UK; W_UV] concatnated per head
87 + `q_b_proj` is [W_UQ; W_QR] concatnated per head
88 + `out_proj` is W_O
89 +
90 +
91 +## Data-Movement Friendly Approach (i.e. "_forward_decode"):
92 +
93 +Ahead of time, compute:
94 +
95 +% this projects from q_c to [Sq, N * Lkv]
96 +W_UQ_UK = einsum("qnp,knp -> qnk"
97 + W_UQ.view(Lq, N, P), W_UK.view(Lkv, N, P)
98 + ).view(Lkv, N * Lkv)
99 +% this projects from attn output [Sq, N * Lkv] to [Sq, H]
100 +W_UV_O = einsum("knv,nvh -> nkh"
101 + W_UV.view(Lkv, N, V), W_O.view(N, V, H)
102 + ).view(N * Lkv, H)
103 +
104 +Runtime
105 +q_c = h_t @ W_DQ
106 +q_latent = q_c @ W_UQ_UK.view(Sq, N, Lkv)
107 +q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
108 +new_kv_c = h_t @ W_DKV
109 +new_k_pe = RoPE(h_t @ W_KR)
110 +kv_c = torch.cat([new_kv_c, cache_kv_c], dim=0)
111 +k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)
112 +
113 +// MQA with QK headdim = Lkv + R
114 +// V headdim = Lkv
115 +// spda_o shape [Sq, N, Lkv]
116 +// NOTE: this is less compute-friendly since Lkv > P
117 +// but is more data-movement friendly since its MQA vs MHA
118 +spda_o = scaled_dot_product_attention(
119 + torch.cat([q_latent, q_pe], dim=-1),
120 + torch.cat([kv_c, k_pe], dim=-1),
121 + kv_c
122 +)
123 +return spda_o.reshape(-1, N * Lkv) @ W_UV_O
124 +
125 +
126 +## Chunked Prefill
127 +
128 +For chunked prefill we want to use the compute friendly algorithm. We are
129 +assuming sufficiently large Sq / Skv ratio, in the future may want to switch to
130 +the data-movement friendly approach if the chunk (i.e. `Sq`) is small.
131 +
132 +However, the compute-friendly approach can potentially run out of memory if Skv
133 +is large due to: `k_nope = (kv_c @ W_UK).view(Skv, N, P)`
134 +
135 +To mitigate this, we chunk the computation of attention with respect to the
136 +current context (i.e. `cache_kv_c` and `cache_k_pe`) so that we can used a
137 +fixed workspace size.
138 +
139 +The chunked prefill approach is as follows:
140 +
141 +MCC Max chunk of context to process per iter, computed dynamically,
142 + used to bound the memory usage
143 +
144 +q_c = h_t @ W_DQ
145 +q_nope = (q_c @ W_UQ).view(Sq, N, P)
146 +q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
147 +new_kv_c = h_t @ W_DKV
148 +new_k_pe = RoPE(h_t @ W_KR)
149 +new_k_nope = (new_kv_c @ W_UK).view(Sq, N, P)
150 +new_v = (new_kv_c @ W_UV).view(Sq, N, V)
151 +
152 +// MHA between queries and new KV
153 +// with QK headdim = P + R
154 +// V headdim = V
155 +// curr_o shape [Sq, N, V]
156 +// curr_lse shape [N, Sq], this is just order FA returns
157 +curr_o, curr_lse = scaled_dot_product_attention(
158 + torch.cat([q_nope, q_pe], dim=-1),
159 + torch.cat([new_k_nope, new_k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
160 + new_v,
161 + casual=True,
162 + return_softmax_lse=True
163 +)
164 +
165 +// Compute attention with the already existing context
166 +for chunk_idx in range(cdiv(C, MCC)):
167 + chunk_start = chunk_idx * MCC
168 + chunk_end = min(chunk_start + MCC, C)
169 + Sc = chunk_end - chunk_start
170 + cache_kv_c_chunk = cache_kv_c[chunk_start:chunk_end]
171 + cache_k_pe_chunk = cache_k_pe[chunk_start:chunk_end]
172 + cache_k_nope_chunk = (cache_kv_c_chunk @ W_UK).view(-1, N, P)
173 + cache_v_chunk = (cache_kv_c_chunk @ W_UV).view(-1, N, V)
174 +
175 + chunk_o, chunk_lse = scaled_dot_product_attention(
176 + torch.cat([q_nope, q_pe], dim=-1),
177 + torch.cat([cache_k_nope_chunk,
178 + cache_k_pe_chunk.unsqueeze(1).expand(-1, N, -1)],
179 + dim=-1),
180 + cache_v_chunk,
181 + casual=False,
182 + return_softmax_lse=True
183 + )
184 +
185 + curr_o, curr_lse = merge_attn_states(
186 + suffix_output=curr_o,
187 + suffix_lse=curr_lse,
188 + prefix_output=chunk_o,
189 + prefix_lse=chunk_lse,
190 + )
191 +
192 +return curr_o @ W_O
193 +"""
194 +
195 +import functools
196 +from abc import abstractmethod
197 +from dataclasses import dataclass
198 +from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, Tuple,
199 + Type, TypeVar)
200 +
201 +import torch
202 +from compressed_tensors.quantization import QuantizationStrategy
203 +
204 +from vllm import _custom_ops as ops
205 +from vllm import envs
206 +from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
207 + AttentionMetadata,
208 + MLAAttentionImpl)
209 +from vllm.attention.backends.utils import get_flash_attn_version
210 +from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
211 +from vllm.distributed import (get_tensor_model_parallel_world_size,
212 + tensor_model_parallel_all_reduce)
213 +from vllm.logger import init_logger
214 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
215 + LinearBase, RowParallelLinear,
216 + UnquantizedLinearMethod)
217 +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
218 + CompressedTensorsLinearMethod)
219 +from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
220 + CompressedTensorsW8A8Fp8)
221 +from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
222 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
223 + apply_fp8_linear_generic, current_platform_fp8_dtype, is_fp8)
224 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
225 + scaled_quantize)
226 +from vllm.model_executor.layers.rotary_embedding import (
227 + DeepseekScalingRotaryEmbedding, RotaryEmbedding)
228 +from vllm.utils import cdiv, round_down
229 +
230 +try:
231 + from vllm.vllm_flash_attn import flash_attn_varlen_func
232 +except ImportError:
233 + # For rocm use upstream flash attention
234 + from flash_attn import flash_attn_varlen_func
235 +
236 +if TYPE_CHECKING:
237 + from vllm.v1.core.scheduler_output import SchedulerOutput
238 + from vllm.v1.worker.gpu_input_batch import InputBatch
239 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner
240 +
241 +logger = init_logger(__name__)
242 +
243 +
244 +class MLACommonBackend(AttentionBackend):
245 +
246 + accept_output_buffer: bool = True
247 +
248 + @staticmethod
249 + def get_name() -> str:
250 + return "TRITON_MLA_VLLM_V1"
251 +
252 + @staticmethod
253 + def get_metadata_cls() -> Type["AttentionMetadata"]:
254 + return MLACommonMetadata
255 +
256 + @staticmethod
257 + def get_builder_cls() -> Type["MLACommonMetadataBuilder"]:
258 + return MLACommonMetadataBuilder
259 +
260 + @staticmethod
261 + def get_kv_cache_shape(
262 + num_blocks: int,
263 + block_size: int,
264 + num_kv_heads: int, # assumed to be 1 for MLA
265 + head_size: int,
266 + ) -> Tuple[int, ...]:
267 + return (num_blocks, block_size, head_size)
268 +
269 + @staticmethod
270 + def get_supported_head_sizes() -> List[int]:
271 + return [576]
272 +
273 + @staticmethod
274 + def use_cascade_attention(*args, **kwargs) -> bool:
275 + return False
276 +
277 +
278 +@dataclass
279 +class MLACommonMetadata:
280 + """Metadata for MLACommon.
281 +
282 + NOTE: Please read the comment at the top of the file before trying to
283 + understand this class
284 + """
285 + # New for MLA (compared to FlashAttention)
286 + # Input positions for rotrary embeddings since for MLA the rotary
287 + # position embeddings are applied inside the attention backend
288 + input_positions: torch.Tensor
289 +
290 + # NOTE(sang): Definition of context_len, query_len, and seq_len.
291 + # |---------- N-1 iteration --------|
292 + # |---------------- N iteration ---------------------|
293 + # |- tokenA -|......................|-- newTokens ---|
294 + # |---------- context_len ----------|
295 + # |-------------------- seq_len ---------------------|
296 + # |-- query_len ---|
297 +
298 + num_actual_tokens: int # Number of tokens excluding padding.
299 + max_query_len: int
300 + query_start_loc: torch.Tensor
301 + max_seq_len: int
302 + seq_lens: torch.Tensor
303 + block_table: torch.Tensor
304 + slot_mapping: torch.Tensor
305 +
306 + # For logging.
307 + num_input_tokens: int = 0 # Number of tokens including padding.
308 +
309 + # The dimension of the attention heads
310 + head_dim: Optional[int] = None
311 +
312 + # New for MLA (compared to FlashAttention)
313 + # For chunked prefill
314 + num_decodes: Optional[int] = None
315 + num_decode_tokens: Optional[int] = None
316 + num_prefills: Optional[int] = None
317 + has_context: bool = False
318 + context_chunk_cu_seq_lens: Optional[torch.Tensor] = None
319 + context_chunk_starts: Optional[torch.Tensor] = None
320 + context_chunk_seq_tot: Optional[List[int]] = None
321 + context_chunk_max_seq_lens: Optional[List[int]] = None
322 + chunked_prefill_workspace: Optional[torch.Tensor] = None
323 +
324 + def __post_init__(self):
325 + supported_head_sizes = MLACommonBackend.get_supported_head_sizes()
326 + if self.head_dim is not None and self.head_dim \
327 + not in supported_head_sizes:
328 + raise ValueError(
329 + f"Only {supported_head_sizes} are supported for head_dim,",
330 + f"received {self.head_dim}.")
331 +
332 +
333 +T = TypeVar("T", bound=MLACommonMetadata)
334 +
335 +
336 +class MLACommonMetadataBuilder:
337 + """
338 + NOTE: Please read the comment at the top of the file before trying to
339 + understand this class
340 + """
341 +
342 + def __init__(self, runner: "GPUModelRunner"):
343 + self.runner = runner
344 + scheduler_config = runner.scheduler_config
345 + model_config = runner.model_config
346 + cache_config = runner.cache_config
347 + self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled
348 +
349 + if self.chunked_prefill_enabled:
350 + self.chunked_prefill_workspace_size = min(
351 + # Max sure there is enough for 8 full length request or at least
352 + # 4 pages of cache per request
353 + max(
354 + 8 * model_config.max_model_len, 4 *
355 + scheduler_config.max_num_seqs * cache_config.block_size),
356 + # For long-context models try not to over-allocate limiting
357 + # kv-cache space, limiting it to 64k tokens,
358 + # which would result in the workspace being:
359 + # 2*(576)*(64*1024) = 144mb
360 + # (assuming 576 MLA head dim, and fp16)
361 + # which would result in up-projected context being
362 + # 2*(192*128)*(64*1024) = 3gb
363 + # (assuming 192 QK head dim, 128 heads, and fp16)
364 + 128 * 1024)
365 + assert self.chunked_prefill_workspace_size >= \
366 + scheduler_config.max_num_seqs * cache_config.block_size
367 + self.chunked_prefill_workspace = torch.empty(
368 + (self.chunked_prefill_workspace_size,
369 + model_config.get_head_size()),
370 + dtype=model_config.dtype,
371 + device=runner.device,
372 + )
373 + self.page_size = self.runner.block_size
374 +
375 + def reorder_batch(self, input_batch: "InputBatch",
376 + scheduler_output: "SchedulerOutput"):
377 + # We now want to reorder the batch so that the "decode" requests are and
378 + # the front and the "prefill" requests are at the using the least amount
379 + # swaps possible. (NOTE for now we loosely use "decode" to mean requests
380 + # where attention is likely memory-bound and "prefill" to mean requests
381 + # where attention is likely compute-bound, TODO(lucas): figure out a
382 + # better naming here)
383 + decodes = []
384 + prefills = []
385 + num_decode_tokens = 0
386 + num_prefill_tokens = 0
387 +
388 + for i, req_id in enumerate(input_batch.req_ids):
389 + num_tokens = scheduler_output.num_scheduled_tokens[req_id]
390 + # for now treat 1 scheduled token as "decode" even if its not,
391 + # we should update this to something like < 8 in the future but
392 + # currently the TritonMLA._forward_decode only supports
393 + # num_tokens = 1
394 + if num_tokens == 1:
395 + decodes.append(i)
396 + num_decode_tokens += num_tokens
397 + else:
398 + prefills.append(i)
399 + num_prefill_tokens += num_tokens
400 +
401 + # We hope that this is fairly minimal since decodes
402 + # should be around for a number of iterations so hopefully they are
403 + # relatively stationary (and new request are generally appended to the
404 + # persistent batch so already should be at the back)
405 + # To achieve this we loop over the decodes in descending order and
406 + # the prefills in ascending order. We swap decodes from the "back"
407 + # i.e. past where the last decode should be in the reodorered with
408 + # prefills from the front of the batch.
409 + # `decodes` and `prefills` are already in ascending order just based on
410 + # the above loop
411 + num_decodes = len(decodes)
412 + num_prefills = len(prefills)
413 + first_prefill = 0
414 +
415 + for i in range(1, min(num_decodes, num_prefills) + 1):
416 + # If the decode is at the "back" of the batch, i, we can swap it
417 + # with the prefill closest to the front of the batch
418 + if decodes[num_decodes - i] >= num_decodes:
419 + input_batch.swap_states(prefills[first_prefill],
420 + decodes[num_decodes - i])
421 + first_prefill += 1
422 + else:
423 + break
424 +
425 + # Save for next `build` call
426 + # TODO(lucas): this is a bit of a hack, we should probably have a
427 + # better way of doing this
428 + self._num_decodes = num_decodes
429 + self._num_prefills = num_prefills
430 + self._num_decode_tokens = num_decode_tokens
431 + self._num_prefill_tokens = num_prefill_tokens
432 +
433 + def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
434 + common_prefix_len: int):
435 + device = self.runner.device
436 + max_seq_len = self.runner.seq_lens_np[:num_reqs].max()
437 + query_start_loc = self.runner.query_start_loc_cpu[:num_reqs + 1].to(
438 + device, non_blocking=True)
439 + seq_lens = self.runner.seq_lens_cpu[:num_reqs].to(device,
440 + non_blocking=True)
441 + block_table = (
442 + self.runner.input_batch.block_table.get_device_tensor()[:num_reqs])
443 + slot_mapping = self.runner.slot_mapping_cpu[:num_actual_tokens].to(
444 + device, non_blocking=True).long()
445 + input_positions = self.runner.positions_cpu[:num_actual_tokens].to(
446 + device, non_blocking=True).long()
447 +
448 + context_chunk_cu_seq_lens = None
449 + context_chunk_starts = None
450 + context_chunk_seq_tot = None
451 + context_chunk_max_seq_lens = None
452 +
453 + num_computed_tokens_cpu_tensor = \
454 + self.runner.input_batch.num_computed_tokens_cpu_tensor[:num_reqs]
455 + context_lens_tensor = \
456 + num_computed_tokens_cpu_tensor.to(device, non_blocking=True)
457 +
458 + if self.chunked_prefill_enabled and self._num_prefills > 0 \
459 + and context_lens_tensor[self._num_decodes:].max() > 0:
460 + # NOTE: it is recommend you read the `Chunked Prefill` section in
461 + # the comment at the top of the file before trying to understand
462 + # the following code
463 +
464 + self.has_context = True
465 +
466 + num_prefills_with_context = \
467 + (context_lens_tensor[self._num_decodes:] > 0).sum().item()
468 +
469 + # currently we allocate an equal amount of workspace for each
470 + # prefill in the batch, we could probably use a more advanced
471 + # algorithm here and allocate more workspace to prefills with
472 + # longer context lengths
473 + max_context_chunk = \
474 + self.chunked_prefill_workspace_size // num_prefills_with_context
475 +
476 + # align max_context_chunk to page_size by rounding down,
477 + # currently the `gather_cache` kernel cannot handle
478 + # `context_chunk_starts` that are not aligned to page_size
479 + max_context_chunk = round_down(max_context_chunk, self.page_size)
480 + assert max_context_chunk > 0
481 + num_chunks = cdiv(context_lens_tensor.max(), max_context_chunk)
482 +
483 + # if `max_context_chunk = 256`, `num_chunks = 3`, and
484 + # `num_prefills_with_context = 4`, create a tensor that looks like
485 + # [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]
486 + context_chunk_starts = \
487 + torch.arange(num_chunks, device=device, dtype=torch.int32) \
488 + .unsqueeze(1).expand(-1, self._num_prefills) \
489 + * max_context_chunk
490 + chunk_ends = torch.min(context_lens_tensor[self._num_decodes:] \
491 + .unsqueeze(0), context_chunk_starts + max_context_chunk)
492 + chunk_seq_lens = (chunk_ends - context_chunk_starts).clamp(min=0)
493 + _context_chunk_cu_seq_lens = chunk_seq_lens.cumsum(dim=1).to(
494 + torch.int32)
495 + zero = torch.zeros(num_chunks, dtype=torch.int32, device=device) \
496 + .unsqueeze(-1)
497 + context_chunk_cu_seq_lens = \
498 + torch.cat([zero, _context_chunk_cu_seq_lens], dim=1)
499 + context_chunk_max_seq_lens = \
500 + chunk_seq_lens.max(dim=1).values.tolist()
501 + context_chunk_seq_tot = chunk_seq_lens.sum(dim=1).tolist()
502 + assert max(context_chunk_seq_tot) <= \
503 + self.chunked_prefill_workspace_size
504 +
505 + return MLACommonMetadata(
506 + input_positions=input_positions,
507 + num_actual_tokens=num_actual_tokens,
508 + max_query_len=max_query_len,
509 + query_start_loc=query_start_loc,
510 + max_seq_len=max_seq_len,
511 + seq_lens=seq_lens,
512 + block_table=block_table,
513 + slot_mapping=slot_mapping,
514 + head_dim=self.runner.model_config.get_head_size(),
515 + # MLACommonMetadata Chunk prefill specific
516 + num_decodes=self._num_decodes,
517 + num_decode_tokens=self._num_decode_tokens,
518 + num_prefills=self._num_prefills,
519 + context_chunk_cu_seq_lens=context_chunk_cu_seq_lens,
520 + context_chunk_starts=context_chunk_starts,
521 + context_chunk_seq_tot=context_chunk_seq_tot,
522 + context_chunk_max_seq_lens=context_chunk_max_seq_lens,
523 + )
524 +
525 +
526 +class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
527 + """
528 + NOTE: Please read the comment at the top of the file before trying to
529 + understand this class
530 + """
531 +
532 + def __init__(
533 + self,
534 + num_heads: int,
535 + head_size: int,
536 + scale: float,
537 + num_kv_heads: int,
538 + alibi_slopes: Optional[List[float]],
539 + sliding_window: Optional[int],
540 + kv_cache_dtype: str,
541 + blocksparse_params: Optional[Dict[str, Any]],
542 + logits_soft_cap: Optional[float],
543 + attn_type: str,
544 + # MLA Specific Arguments
545 + q_lora_rank: Optional[int],
546 + kv_lora_rank: int,
547 + qk_nope_head_dim: int,
548 + qk_rope_head_dim: int,
549 + qk_head_dim: int,
550 + v_head_dim: int,
551 + rotary_emb: RotaryEmbedding,
552 + # q_proj should be q_b_proj if q_lora_rank is not None, but from an
553 + # attention backend perspective we rely on the layer to pass in the
554 + # correct matrix
555 + q_proj: ColumnParallelLinear,
556 + kv_b_proj: ColumnParallelLinear,
557 + o_proj: RowParallelLinear,
558 + ) -> None:
559 + self.num_heads = num_heads
560 + self.head_size = head_size
561 + self.scale = float(scale)
562 + self.num_kv_heads = num_kv_heads
563 + self.kv_cache_dtype = kv_cache_dtype
564 +
565 + self.q_lora_rank = q_lora_rank
566 + self.kv_lora_rank = kv_lora_rank
567 + self.qk_nope_head_dim = qk_nope_head_dim
568 + self.qk_rope_head_dim = qk_rope_head_dim
569 + self.qk_head_dim = qk_head_dim
570 + self.v_head_dim = v_head_dim
571 +
572 + self.rotary_emb = rotary_emb
573 + self.use_yarn_rope = isinstance(rotary_emb,
574 + DeepseekScalingRotaryEmbedding)
575 + self.q_proj = q_proj
576 + self.kv_b_proj = kv_b_proj
577 + self.o_proj = o_proj
578 + self.vllm_flash_attn_version = get_flash_attn_version()
579 +
580 + # Handle the differences between the flash_attn_varlen from flash_attn
581 + # and the one from vllm_flash_attn. The former is used on RoCM and the
582 + # latter has an additional parameter to control FA2 vs FA3
583 + self.flash_attn_varlen_func = flash_attn_varlen_func
584 + if self.vllm_flash_attn_version is not None:
585 + self.flash_attn_varlen_func = \
586 + functools.partial(flash_attn_varlen_func,
587 + fa_version=self.vllm_flash_attn_version)
588 +
589 + def _v_up_proj_and_o_proj(self, x):
590 + if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
591 + if is_fp8(self.W_UV_O):
592 + output_parallel = apply_fp8_linear_generic(
593 + x.flatten(start_dim=1), self.W_UV_O, self.W_UV_O_scales,
594 + self.reqaunt_input_group_shape,
595 + self.reqaunt_weight_group_shape)
596 + else:
597 + output_parallel = torch.matmul(x.flatten(start_dim=1),
598 + self.W_UV_O)
599 + if self.tp_size > 1:
600 + output = tensor_model_parallel_all_reduce(output_parallel)
601 + else:
602 + output = output_parallel
603 + return output
604 + else:
605 + x = torch.einsum("bnl,lnv->bnv", x, self.W_UV)
606 + return self.o_proj(x.reshape(-1,
607 + self.num_heads * self.v_head_dim))[0]
608 +
609 + def _q_proj_and_k_up_proj(self, x):
610 + if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
611 + if is_fp8(self.W_Q_UK):
612 + return apply_fp8_linear_generic(
613 + x, self.W_Q_UK, self.W_Q_UK_scales,
614 + self.reqaunt_input_group_shape,
615 + self.reqaunt_weight_group_shape).view(
616 + -1, self.num_heads, self.kv_lora_rank)
617 + return torch.matmul(x, self.W_Q_UK)\
618 + .view(-1, self.num_heads, self.kv_lora_rank)
619 + else:
620 + x = torch.matmul(x, self.W_Q)\
621 + .view(-1, self.num_heads, self.qk_nope_head_dim)
622 + return torch.einsum("bnp,lnp->bnl", x, self.W_UK)\
623 + .view(-1, self.num_heads, self.kv_lora_rank)
624 +
625 + def process_weights_after_loading(self, act_dtype: torch.dtype):
626 +
627 + # TODO(lucas) This is very gross, we need a more wide scale refactor of
628 + # all the FP8 code with a more standard way of
629 + # defining schemes/group-shapes, we should also potentially force
630 + # quant_methods to support a decompress function
631 + #
632 + # returns input_group_shape, weight_group_shape
633 + def get_scale_group_shapes_for_fp8(layer: LinearBase) -> \
634 + Tuple[Tuple[int, int], Tuple[int, int]]:
635 + if isinstance(layer.quant_method, Fp8LinearMethod):
636 + if layer.quant_method.block_quant:
637 + weight_block_size = \
638 + layer.quant_method.quant_config.weight_block_size
639 + # per-token-group (1, X), block-quantized (X, Y)
640 + return (1, weight_block_size[-1]), weight_block_size
641 + else:
642 + return (-1, -1), (-1, -1) # per-tensor, per-tensor
643 + elif isinstance(layer.quant_method, CompressedTensorsLinearMethod)\
644 + and isinstance(layer.scheme, CompressedTensorsW8A8Fp8):
645 + # this is hacky but we always assume the for
646 + # CompressedTensorsW8A8Fp8 the input is dynamic per-token
647 + # we ignore if it is static-per-tensor since we are going to
648 + # requantize after later anyways
649 + strategy = layer.scheme.strategy
650 + if strategy == QuantizationStrategy.TENSOR:
651 + return (1, -1), (-1, -1) # per-token, per-tensor
652 + elif strategy == QuantizationStrategy.CHANNEL:
653 + return (1, -1), (-1, 1) # per-token, per-channel
654 + else:
655 + raise NotImplementedError(
656 + f"QuantizationStrategy.{strategy} is not supported for "
657 + "fp8 MLA, please run with VLLM_MLA_DISABLE=1")
658 + else:
659 + raise NotImplementedError(
660 + "Can't determine scale group shapes for "
661 + f"{layer.quant_method}, please run with VLLM_MLA_DISABLE=1"
662 + )
663 +
664 + def get_layer_weight(layer):
665 + if hasattr(layer, "weight"):
666 + return layer.weight
667 + elif hasattr(layer, "qweight"):
668 + return layer.qweight
669 + else:
670 + raise AttributeError(
671 + f"Layer '{layer}' has neither weight nor qweight")
672 +
673 + def get_and_maybe_dequant_weights(layer: LinearBase):
674 + if not isinstance(layer.quant_method, UnquantizedLinearMethod):
675 + # NOTE: This should only be used offline, since it's O(N^3)
676 + eye = torch.eye(layer.input_size_per_partition,
677 + dtype=act_dtype,
678 + device=get_layer_weight(layer).device)
679 + dequant_weights = layer.quant_method.apply(layer,
680 + eye,
681 + bias=None)
682 + del eye
683 + # standardize to (output, input)
684 + return dequant_weights.T
685 + return layer.weight
686 +
687 + weight_dtype = get_layer_weight(self.kv_b_proj).dtype
688 + assert get_layer_weight(self.o_proj).dtype == weight_dtype
689 + assert get_layer_weight(self.q_proj).dtype == weight_dtype
690 +
691 + kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
692 + assert kv_b_proj_weight.shape == (
693 + self.kv_lora_rank,
694 + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)), (
695 + f"{kv_b_proj_weight.shape=}, "
696 + f"{self.kv_lora_rank=}, "
697 + f"{self.num_heads=}, "
698 + f"{self.qk_nope_head_dim=}, "
699 + f"{self.v_head_dim=}")
700 + kv_b_proj_weight = kv_b_proj_weight.view(
701 + self.kv_lora_rank,
702 + self.num_heads,
703 + self.qk_nope_head_dim + self.v_head_dim,
704 + )
705 +
706 + W_UK, W_UV = kv_b_proj_weight.split(
707 + [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
708 +
709 + q_proj_weight = get_and_maybe_dequant_weights(self.q_proj).T\
710 + .view(-1, self.num_heads, self.qk_head_dim)
711 +
712 + # can be W_Q or W_UQ depending q_lora_rank, the former if
713 + # q_lora_rank is None, the latter otherwise. From the Attention backend
714 + # perspective though we call these both W_Q and rely on the layer
715 + # to pass in the correct matrix
716 + W_Q = q_proj_weight[..., :self.qk_nope_head_dim]
717 + self.W_QR = q_proj_weight[..., self.qk_nope_head_dim:]\
718 + .flatten(start_dim=1).contiguous()
719 +
720 + # W_QR is small so for simplicity we dont bother requantizing it
721 + self.W_QR = self.W_QR.to(act_dtype)
722 +
723 + if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
724 + requantization_enabled = not envs.VLLM_MLA_DISABLE_REQUANTIZATION
725 + if is_fp8(weight_dtype) and requantization_enabled:
726 + # This assumes it wise to requantize using the same group shapes
727 + # (i.e. strategy, per-tensor, per-channel, block etc.) that the
728 + # weights were originally quantized
729 + requant_input_group_shape, requant_weight_group_shape = \
730 + get_scale_group_shapes_for_fp8(self.q_proj)
731 + assert (requant_input_group_shape, requant_weight_group_shape)\
732 + == get_scale_group_shapes_for_fp8(self.kv_b_proj)
733 + assert (requant_input_group_shape, requant_weight_group_shape)\
734 + == get_scale_group_shapes_for_fp8(self.o_proj)
735 + self.reqaunt_input_group_shape = requant_input_group_shape
736 + self.reqaunt_weight_group_shape = requant_weight_group_shape
737 +
738 + #
739 + # Perform matrix-absorption following
740 + # https://github.com/flashinfer-ai/flashinfer/pull/551
741 + # for decode, as a result we end up with absorbed weights for decode
742 + # and another copy of raw weights for prefill.
743 + #
744 + self.W_UK, self.W_UV = kv_b_proj_weight.split(
745 + [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
746 + # We absorb `W_UK` into `W_Q` resulting in either W_Q_UK or W_UQ_UK
747 + # depending q_lora_rank, the former if q_lora_rank is None, the
748 + # latter otherwise
749 + # basically if q_lora_rank is none we are absorbing into q_proj
750 + # instead of UQ
751 + W_Q_UK = torch.einsum("qnd,lnd -> qnl", W_Q, W_UK)\
752 + .flatten(start_dim=1).contiguous()
753 +
754 + if is_fp8(weight_dtype) and requantization_enabled:
755 + W_Q_UK, W_Q_UK_scales = scaled_quantize(
756 + W_Q_UK,
757 + self.reqaunt_weight_group_shape,
758 + quant_dtype=current_platform_fp8_dtype)
759 + # For FP8 save the transpose so we can use
760 + # `apply_w8a8_block_fp8_linear` directly
761 + self.W_Q_UK = W_Q_UK.T.contiguous()
762 + self.W_Q_UK_scales = W_Q_UK_scales.T.contiguous()
763 + else:
764 + self.W_Q_UK = W_Q_UK.to(act_dtype)
765 +
766 + W_O = get_and_maybe_dequant_weights(self.o_proj)\
767 + .view(-1, self.num_heads, self.v_head_dim)
768 + W_UV_O = torch.einsum("lnd,hnd -> nlh", W_UV, W_O)\
769 + .flatten(start_dim=0, end_dim=1).contiguous()
770 +
771 + if is_fp8(weight_dtype) and requantization_enabled:
772 + W_UV_O, W_UV_O_scales = scaled_quantize(
773 + W_UV_O,
774 + self.reqaunt_weight_group_shape,
775 + quant_dtype=current_platform_fp8_dtype)
776 + # For FP8 save the transpose so we can use
777 + # `apply_w8a8_block_fp8_linear` directly
778 + self.W_UV_O = W_UV_O.T.contiguous()
779 + self.W_UV_O_scales = W_UV_O_scales.T.contiguous()
780 + else:
781 + self.W_UV_O = W_UV_O.to(act_dtype)
782 +
783 + self.tp_size = get_tensor_model_parallel_world_size()
784 + else:
785 + if is_fp8(weight_dtype):
786 + raise NotImplementedError(
787 + "Currently fp8 requires matrix absorption")
788 +
789 + self.W_UV = W_UV
790 + self.W_UK = W_UK
791 + self.W_Q = W_Q.flatten(start_dim=1)
792 +
793 + def _compute_prefill_context(
794 + self,
795 + q: torch.Tensor,
796 + kv_c_and_k_pe_cache: torch.Tensor,
797 + attn_metadata: MLACommonMetadata,
798 + ):
799 + assert attn_metadata.num_prefills is not None
800 + assert attn_metadata.context_chunk_seq_tot is not None
801 + assert attn_metadata.context_chunk_cu_seq_lens is not None
802 + assert attn_metadata.context_chunk_starts is not None
803 + assert attn_metadata.context_chunk_max_seq_lens is not None
804 +
805 + output = None
806 + iters = len(attn_metadata.context_chunk_seq_tot)
807 +
808 + assert attn_metadata.chunked_prefill_workspace is not None
809 + workspace = attn_metadata.chunked_prefill_workspace
810 +
811 + for i in range(iters):
812 + toks = attn_metadata.context_chunk_seq_tot[i]
813 +
814 + ops.gather_cache(
815 + src_cache=kv_c_and_k_pe_cache,
816 + dst=workspace,
817 + block_table=attn_metadata.block_table,
818 + cu_seq_lens=attn_metadata.context_chunk_cu_seq_lens[i],
819 + batch_size=attn_metadata.num_prefills,
820 + seq_starts=attn_metadata.context_chunk_starts[i],
821 + )
822 +
823 + kv_c_normed = workspace[:toks]\
824 + [..., :self.kv_lora_rank].unsqueeze(1)
825 + k_pe = workspace[:toks]\
826 + [..., self.kv_lora_rank:].unsqueeze(1)
827 +
828 + kv_nope = self.kv_b_proj(kv_c_normed)[0].view( \
829 + -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
830 + k_nope, v = kv_nope\
831 + .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
832 +
833 + k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))),
834 + dim=-1)
835 +
836 + # For MLA the v head dim is smaller than qk head dim so we pad
837 + # out v with 0s to match the qk head dim
838 + v_padded = torch.nn.functional.pad(v,
839 + [0, q.shape[-1] - v.shape[-1]],
840 + value=0)
841 +
842 + attn_output, attn_softmax_lse = self.flash_attn_varlen_func(
843 + q=q,
844 + k=k,
845 + v=v_padded,
846 + cu_seqlens_q=attn_metadata.query_start_loc,
847 + cu_seqlens_k=attn_metadata.context_chunk_cu_seq_lens[i],
848 + max_seqlen_q=attn_metadata.max_query_len,
849 + max_seqlen_k=attn_metadata.context_chunk_max_seq_lens[i],
850 + softmax_scale=self.scale,
851 + causal=False, # Context is unmasked
852 + return_softmax_lse=True,
853 + )
854 +
855 + if output is None:
856 + output = attn_output
857 + output_lse = attn_softmax_lse
858 + else:
859 + output_tmp = torch.empty_like(output)
860 + output_lse_tmp = torch.empty_like(output_lse)
861 + merge_attn_states(
862 + output=output_tmp,
863 + output_lse=output_lse_tmp,
864 + prefix_output=output,
865 + prefix_lse=output_lse,
866 + suffix_output=attn_output,
867 + suffix_lse=attn_softmax_lse,
868 + )
869 + output = output_tmp
870 + output_lse = output_lse_tmp
871 +
872 + return output, output_lse
873 +
874 + def _forward_prefill(
875 + self,
876 + q: torch.Tensor,
877 + kv_c_normed: torch.Tensor,
878 + k_pe: torch.Tensor,
879 + kv_c_and_k_pe_cache: torch.Tensor,
880 + attn_metadata: MLACommonMetadata,
881 + ) -> torch.Tensor:
882 + has_context = attn_metadata.has_context
883 + kv_nope = self.kv_b_proj(kv_c_normed)[0].view(\
884 + -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
885 + k_nope, v = kv_nope\
886 + .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
887 +
888 + k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
889 +
890 + # For MLA the v head dim is smaller than qk head dim so we pad out
891 + # v with 0s to match the qk head dim
892 + v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
893 + value=0)
894 +
895 + output = self.flash_attn_varlen_func(
896 + q=q,
897 + k=k,
898 + v=v_padded,
899 + cu_seqlens_q=attn_metadata.query_start_loc,
900 + cu_seqlens_k=attn_metadata.query_start_loc,
901 + max_seqlen_q=attn_metadata.max_query_len,
902 + max_seqlen_k=attn_metadata.max_seq_len,
903 + softmax_scale=self.scale,
904 + causal=True,
905 + return_softmax_lse=has_context,
906 + )
907 +
908 + if has_context:
909 + suffix_output, suffix_lse = output
910 + context_output, context_lse = self._compute_prefill_context( \
911 + q, kv_c_and_k_pe_cache, attn_metadata)
912 +
913 + output = torch.empty_like(suffix_output)
914 + merge_attn_states(
915 + output=output,
916 + prefix_output=context_output,
917 + prefix_lse=context_lse,
918 + suffix_output=suffix_output,
919 + suffix_lse=suffix_lse,
920 + )
921 +
922 + # slice by `:v.shape[-1]` in order to remove v headdim padding
923 + output = output\
924 + .view(-1, self.num_heads, q.shape[-1])[..., :v.shape[-1]]\
925 + .reshape(-1, self.num_heads * v.shape[-1])
926 +
927 + return self.o_proj(output)[0]
928 +
929 + @abstractmethod
930 + def _forward_decode(
931 + self,
932 + q_nope: torch.Tensor,
933 + q_pe: torch.Tensor,
934 + kv_c_and_k_pe_cache: torch.Tensor,
935 + attn_metadata: T,
936 + ) -> torch.Tensor:
937 + raise NotImplementedError
938 +
939 + def forward(
940 + self,
941 + layer: AttentionLayer,
942 + hidden_states_or_q_c: torch.Tensor, # query in unified attn
943 + k_c_normed: torch.Tensor, # key in unified attn
944 + k_pe: torch.Tensor, # value in unified attn
945 + kv_cache: torch.Tensor,
946 + attn_metadata: T,
947 + output: Optional[torch.Tensor] = None,
948 + ) -> torch.Tensor:
949 +
950 + assert output is not None, "Output tensor must be provided."
951 +
952 + if attn_metadata is None:
953 + # Profiling run.
954 + return output
955 +
956 + num_actual_toks = attn_metadata.num_actual_tokens
957 +
958 + # Inputs and outputs may be padded for CUDA graphs
959 + output_padded = output
960 + output = output[:num_actual_toks, ...]
961 + hidden_states_or_q_c = hidden_states_or_q_c[:num_actual_toks, ...]
962 + k_c_normed = k_c_normed[:num_actual_toks, ...]
963 + k_pe = k_pe[:num_actual_toks, ...]
964 +
965 + # Restore head dim (for rotary embedding)
966 + k_pe = k_pe.unsqueeze(1)
967 + assert hasattr(attn_metadata, "input_positions")
968 +
969 + assert attn_metadata.num_decodes is not None and \
970 + attn_metadata.num_prefills is not None and \
971 + attn_metadata.num_decode_tokens is not None
972 +
973 + has_decode = attn_metadata.num_decodes > 0
974 + has_prefill = attn_metadata.num_prefills > 0
975 + num_decode_tokens = attn_metadata.num_decode_tokens
976 +
977 + decode_hs_or_q_c = hidden_states_or_q_c[:num_decode_tokens]
978 + decode_k_pe = k_pe[:num_decode_tokens]
979 + decode_input_positions = \
980 + attn_metadata.input_positions[:num_decode_tokens]
981 +
982 + prefill_hs_or_q_c = hidden_states_or_q_c[num_decode_tokens:]
983 + prefill_k_pe = k_pe[num_decode_tokens:]
984 + prefill_input_positions = \
985 + attn_metadata.input_positions[num_decode_tokens:]
986 + prefill_k_c_normed = k_c_normed[num_decode_tokens:]
987 +
988 + if has_decode:
989 + decode_q_nope = self._q_proj_and_k_up_proj(decode_hs_or_q_c)
990 + decode_q_pe = torch.matmul(decode_hs_or_q_c, self.W_QR)\
991 + .view(-1, self.num_heads, self.qk_rope_head_dim)
992 + decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(
993 + decode_input_positions, decode_q_pe, decode_k_pe)
994 +
995 + if has_prefill:
996 + prefill_q = self.q_proj(prefill_hs_or_q_c)[0]\
997 + .view(-1, self.num_heads, self.qk_head_dim)
998 + prefill_q_pe = prefill_q[..., self.qk_nope_head_dim:]
999 + prefill_q_pe[...], prefill_k_pe[...] = self.rotary_emb(
1000 + prefill_input_positions, prefill_q_pe, prefill_k_pe)
1001 +
1002 + # write the latent and rope to kv cache
1003 + if kv_cache.numel() > 0:
1004 + ops.concat_and_cache_mla(
1005 + k_c_normed,
1006 + k_pe.squeeze(1),
1007 + kv_cache,
1008 + attn_metadata.slot_mapping.flatten(),
1009 + kv_cache_dtype=self.kv_cache_dtype,
1010 + scale=layer._k_scale,
1011 + )
1012 +
1013 + if has_prefill:
1014 + output[num_decode_tokens:] = self._forward_prefill(
1015 + prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,
1016 + attn_metadata)
1017 +
1018 + if has_decode:
1019 + output[:num_decode_tokens] = self._forward_decode(
1020 + decode_q_nope, decode_q_pe, kv_cache, attn_metadata)
1021 +
1022 + return output_padded
vllm/v1/attention/backends/triton_mla.py new
+110
@@ -0,0 +1,110 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from typing import Any, Dict, List, Optional, Type
4 +
5 +import torch
6 +
7 +from vllm.attention.backends.abstract import AttentionType
8 +from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
9 +from vllm.logger import init_logger
10 +from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
11 + MLACommonImpl,
12 + MLACommonMetadata)
13 +
14 +logger = init_logger(__name__)
15 +
16 +
17 +class TritonMLABackend(MLACommonBackend):
18 +
19 + @staticmethod
20 + def get_name() -> str:
21 + return "TRITON_MLA_VLLM_V1"
22 +
23 + @staticmethod
24 + def get_impl_cls() -> Type["TritonMLAImpl"]:
25 + return TritonMLAImpl
26 +
27 +
28 +class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
29 +
30 + def __init__(
31 + self,
32 + num_heads: int,
33 + head_size: int,
34 + scale: float,
35 + num_kv_heads: int,
36 + alibi_slopes: Optional[List[float]],
37 + sliding_window: Optional[int],
38 + kv_cache_dtype: str,
39 + blocksparse_params: Optional[Dict[str, Any]],
40 + logits_soft_cap: Optional[float],
41 + attn_type: str,
42 + # MLA Specific Arguments
43 + **mla_args) -> None:
44 + super().__init__(num_heads, head_size, scale, num_kv_heads,
45 + alibi_slopes, sliding_window, kv_cache_dtype,
46 + blocksparse_params, logits_soft_cap, attn_type,
47 + **mla_args)
48 +
49 + unsupported_features = [
50 + alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
51 + ]
52 + if any(unsupported_features):
53 + raise NotImplementedError(
54 + "TritonMLAImpl does not support one of the following: "
55 + "alibi_slopes, sliding_window, blocksparse_params, "
56 + "logits_soft_cap")
57 +
58 + if attn_type != AttentionType.DECODER:
59 + raise NotImplementedError("Encoder self-attention and "
60 + "encoder/decoder cross-attention "
61 + "are not implemented for "
62 + "TritonMLAImpl")
63 +
64 + def _forward_decode(
65 + self,
66 + q_nope: torch.Tensor,
67 + q_pe: torch.Tensor,
68 + kv_c_and_k_pe_cache: torch.Tensor,
69 + attn_metadata: MLACommonMetadata,
70 + ) -> torch.Tensor:
71 + assert kv_c_and_k_pe_cache.numel() > 0
72 + if self.kv_cache_dtype.startswith("fp8"):
73 + raise NotImplementedError("FP8 Triton MLA not yet supported")
74 +
75 + B = q_nope.shape[0]
76 +
77 + q = torch.cat([q_nope, q_pe], dim=-1)
78 + o = torch.zeros(B,
79 + self.num_heads,
80 + self.kv_lora_rank,
81 + dtype=q.dtype,
82 + device=q.device)
83 +
84 + num_kv_splits = 4 # TODO: heuristic
85 +
86 + # TODO(lucas) Allocate ahead of time
87 + attn_logits = torch.empty(
88 + (
89 + B,
90 + self.num_heads,
91 + num_kv_splits,
92 + # NOTE(lucas) idk why the +1 is here but sglang has it so we
93 + # just mirror that
94 + self.kv_lora_rank + 1,
95 + ),
96 + dtype=torch.float32,
97 + device=q.device,
98 + )
99 +
100 + # Add a head dim of 1
101 + kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)
102 + kv_c_cache = kv_c_and_k_pe_cache[..., :self.kv_lora_rank]
103 + PAGE_SIZE = kv_c_and_k_pe_cache.size(1)
104 +
105 + # Run MQA
106 + decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,
107 + attn_metadata.block_table, attn_metadata.seq_lens,
108 + attn_logits, num_kv_splits, self.scale, PAGE_SIZE)
109 +
110 + return self._v_up_proj_and_o_proj(o)
vllm/v1/worker/gpu_input_batch.py
+63 -1
@@ -80,7 +80,14 @@ class InputBatch:
80 self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)
81 self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
82 self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
83 - self.num_computed_tokens_cpu = np.empty(max_num_reqs, dtype=np.int32)
83 + self.num_computed_tokens_cpu_tensor = torch.zeros(
84 + (max_num_reqs, ),
85 + device="cpu",
86 + dtype=torch.int32,
87 + pin_memory=pin_memory,
88 + )
89 + self.num_computed_tokens_cpu = \
90 + self.num_computed_tokens_cpu_tensor.numpy()
91
92 # Block table.
93 self.block_table = BlockTable(
@@ -356,6 +363,61 @@ class InputBatch:
363 self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
364 return req_index
365
366 + def swap_states(self, i1: int, i2: int) -> None:
367 + old_id_i1 = self._req_ids[i1]
368 + old_id_i2 = self._req_ids[i2]
369 + self._req_ids[i1], self._req_ids[i2] =\
370 + self._req_ids[i2], self._req_ids[i1] # noqa
371 + self.req_output_token_ids[i1], self.req_output_token_ids[i2] =\
372 + self.req_output_token_ids[i2], self.req_output_token_ids[i1]
373 + assert old_id_i1 is not None and old_id_i2 is not None
374 + self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] =\
375 + self.req_id_to_index[old_id_i2], self.req_id_to_index[old_id_i1]
376 + self.num_tokens[i1], self.num_tokens[i2] =\
377 + self.num_tokens[i2], self.num_tokens[i1]
378 + self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\
379 + self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...]
380 + self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] =\
381 + self.num_tokens_no_spec[i2], self.num_tokens_no_spec[i1]
382 + self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] =\
383 + self.num_prompt_tokens[i2], self.num_prompt_tokens[i1]
384 + self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] =\
385 + self.num_computed_tokens_cpu[i2], self.num_computed_tokens_cpu[i1]
386 + self.temperature_cpu[i1], self.temperature_cpu[i2] =\
387 + self.temperature_cpu[i2], self.temperature_cpu[i1]
388 + self.top_p_cpu[i1], self.top_p_cpu[i2] =\
389 + self.top_p_cpu[i2], self.top_p_cpu[i1]
390 + self.top_k_cpu[i1], self.top_k_cpu[i2] =\
391 + self.top_k_cpu[i2], self.top_k_cpu[i1]
392 + self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] =\
393 + self.frequency_penalties_cpu[i2], self.frequency_penalties_cpu[i1]
394 + self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] =\
395 + self.presence_penalties_cpu[i2], self.presence_penalties_cpu[i1]
396 + self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] =\
397 + self.repetition_penalties_cpu[i2], self.repetition_penalties_cpu[i1]
398 + self.min_p_cpu[i1], self.min_p_cpu[i2] =\
399 + self.min_p_cpu[i2], self.min_p_cpu[i1]
400 +
401 + g1 = self.generators.get(i1)
402 + g2 = self.generators.get(i2)
403 + if g1 is not None:
404 + self.generators[i2] = g1
405 + if g2 is not None:
406 + self.generators[i1] = g2
407 +
408 + t1 = self.min_tokens.get(i1)
409 + t2 = self.min_tokens.get(i2)
410 + if t1 is not None:
411 + self.min_tokens[i2] = t1
412 + if t2 is not None:
413 + self.min_tokens[i1] = t2
414 +
415 + self.request_lora_mapping[i1], self.request_lora_mapping[i2] =\
416 + self.request_lora_mapping[i2], self.request_lora_mapping[i1]
417 + self.logit_bias[i1], self.logit_bias[i2] =\
418 + self.logit_bias[i2], self.logit_bias[i1]
419 + self.block_table.swap_row(i1, i2)
420 +
421 def condense(self, empty_req_indices: List[int]) -> None:
422 num_reqs = self.num_reqs
423 if num_reqs == 0:
vllm/v1/worker/gpu_model_runner.py
+35 -41
@@ -2,6 +2,7 @@
2
3 import gc
4 import time
5 +import weakref
6 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
7
8 import numpy as np
@@ -9,7 +10,7 @@ import torch
10 import torch.distributed
11 import torch.nn as nn
12
12 -from vllm.attention.backends.abstract import AttentionType
13 +from vllm.attention import AttentionType, get_attn_backend
14 from vllm.attention.layer import Attention
15 from vllm.config import CompilationLevel, VllmConfig
16 from vllm.distributed.parallel_state import get_pp_group, graph_capture
@@ -24,8 +25,7 @@ from vllm.sampling_params import SamplingType
25 from vllm.sequence import IntermediateTensors
26 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
27 LayerBlockType, cdiv, is_pin_memory_available)
27 -from vllm.v1.attention.backends.flash_attn import (FlashAttentionBackend,
28 - FlashAttentionMetadata)
28 +from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
29 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
30 from vllm.v1.engine.mm_input_cache import MMInputCacheClient
31 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
@@ -92,6 +92,27 @@ class GPUModelRunner(LoRAModelRunnerMixin):
92 self.head_size = model_config.get_head_size()
93 self.hidden_size = model_config.get_hidden_size()
94
95 + self.attn_backend = get_attn_backend(
96 + self.head_size,
97 + self.dtype,
98 + self.kv_cache_dtype,
99 + self.block_size,
100 + self.model_config.is_attention_free,
101 + use_mla=self.model_config.use_mla,
102 + )
103 + if self.attn_backend is None:
104 + error_msg = (
105 + f"Error with get_att_backend: {self.head_size=}, "
106 + f"{self.dtype=}, {self.kv_cache_dtype=}, {self.block_size=}, "
107 + f"{self.model_config.is_attention_free=}, "
108 + f"{self.model_config.use_mla=}")
109 + logger.error(error_msg)
110 + raise NotImplementedError(
111 + "Non-Attention backend is not supported by V1 GPUModelRunner.")
112 +
113 + self.attn_metadata_builder = self.attn_backend.get_builder_cls()(
114 + weakref.proxy(self))
115 +
116 # Multi-modal data support
117 self.input_registry = INPUT_REGISTRY
118 self.mm_registry = MULTIMODAL_REGISTRY
@@ -433,6 +454,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
454 num_reqs = self.input_batch.num_reqs
455 assert num_reqs > 0
456
457 + # Some attention backends (namely MLA) may want to separate requests
458 + # based on if the attention computation will be compute-bound or
459 + # memory-bound. This gives them a hook to do that.
460 + self.attn_metadata_builder.reorder_batch(self.input_batch,
461 + scheduler_output)
462 +
463 # OPTIMIZATION: Start copying the block table first.
464 # This way, we can overlap the copy with the following CPU operations.
465 self.input_batch.block_table.commit(num_reqs)
@@ -515,7 +542,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
542 self.seq_lens_np[:num_reqs] = (
543 self.input_batch.num_computed_tokens_cpu[:num_reqs] +
544 num_scheduled_tokens)
518 - max_seq_len = self.seq_lens_np[:num_reqs].max()
545
546 # Copy the tensors to the GPU.
547 self.input_ids[:total_num_scheduled_tokens].copy_(
@@ -530,49 +556,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
556 self.positions[:total_num_scheduled_tokens].copy_(
557 self.positions_cpu[:total_num_scheduled_tokens],
558 non_blocking=True)
533 - query_start_loc = self.query_start_loc_cpu[:num_reqs + 1].to(
534 - self.device, non_blocking=True)
535 - seq_lens = self.seq_lens_cpu[:num_reqs].to(self.device,
536 - non_blocking=True)
537 - slot_mapping = self.slot_mapping_cpu[:total_num_scheduled_tokens].to(
538 - self.device, non_blocking=True).long()
559
560 # Prepare for cascade attention if needed.
561 common_prefix_len = self._compute_cascade_attn_prefix_len(
562 num_scheduled_tokens,
563 scheduler_output.num_common_prefix_blocks,
564 )
545 - use_cascade = common_prefix_len > 0
546 - if use_cascade:
547 - # TODO: Optimize.
548 - cu_prefix_query_lens = torch.tensor(
549 - [0, total_num_scheduled_tokens],
550 - dtype=torch.int32,
551 - device=self.device)
552 - prefix_kv_lens = torch.tensor([common_prefix_len],
553 - dtype=torch.int32,
554 - device=self.device)
555 - suffix_kv_lens = (self.seq_lens_np[:num_reqs] - common_prefix_len)
556 - suffix_kv_lens = torch.from_numpy(suffix_kv_lens).to(self.device)
557 - else:
558 - cu_prefix_query_lens = None
559 - prefix_kv_lens = None
560 - suffix_kv_lens = None
561 -
562 - attn_metadata = FlashAttentionMetadata(
565 + attn_metadata = self.attn_metadata_builder.build(
566 + num_reqs=num_reqs,
567 num_actual_tokens=total_num_scheduled_tokens,
568 max_query_len=max_num_scheduled_tokens,
565 - query_start_loc=query_start_loc,
566 - max_seq_len=max_seq_len,
567 - seq_lens=seq_lens,
568 - block_table=(
569 - self.input_batch.block_table.get_device_tensor()[:num_reqs]),
570 - slot_mapping=slot_mapping,
571 - use_cascade=use_cascade,
569 common_prefix_len=common_prefix_len,
573 - cu_prefix_query_lens=cu_prefix_query_lens,
574 - prefix_kv_lens=prefix_kv_lens,
575 - suffix_kv_lens=suffix_kv_lens,
570 )
571
572 use_spec_decode = len(
@@ -586,7 +580,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
580 # from these partial requests, we do so for simplicity.
581 # We will ignore the sampled tokens from the partial requests.
582 # TODO: Support prompt logprobs.
589 - logits_indices = query_start_loc[1:] - 1
583 + logits_indices = attn_metadata.query_start_loc[1:] - 1
584
585 # Hot-Swap lora model
586 if self.lora_config:
@@ -667,7 +661,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
661 # common_prefix_len should be a multiple of the block size.
662 common_prefix_len = (common_prefix_len // self.block_size *
663 self.block_size)
670 - use_cascade = FlashAttentionBackend.use_cascade_attention(
664 + use_cascade = self.attn_backend.use_cascade_attention(
665 common_prefix_len=common_prefix_len,
666 query_lens=num_scheduled_tokens,
667 num_query_heads=self.num_query_heads,
@@ -1379,7 +1373,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1373 assert tensor_config.size % layer_spec.page_size_bytes == 0
1374 num_blocks = tensor_config.size // layer_spec.page_size_bytes
1375 if isinstance(layer_spec, FullAttentionSpec):
1382 - kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
1376 + kv_cache_shape = self.attn_backend.get_kv_cache_shape(
1377 num_blocks, layer_spec.block_size, layer_spec.num_kv_heads,
1378 layer_spec.head_size)
1379 dtype = layer_spec.dtype