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, 19:45 UTC
[Attention] MLA get rid of materialization (#14770)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Mar 14, 2025 at 02:39 UTC
9532c49836ad9b5f2120ebba8caf0c56f998126f
4 files changed
+117
-499
vllm/attention/backends/mla/common.py
+57
-210
@@ -7,22 +7,22 @@ First we define:
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
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
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
@@ -54,9 +54,9 @@ 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]
57
+W_UK project kv_c to k_nope shape [Lkv, N, P]
58
+W_KR project h_t to k_pe shape [H, 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
@@ -69,8 +69,8 @@ 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)
72
+k_nope = (kv_c @ W_UK.view(Lkv, N * P)).view(Skv, N, P)
73
+v = (kv_c @ W_UV.view(Lkv, N * V)).view(Skv, N, V)
74
75
// MHA with QK headdim = P + R
76
// V headdim = V
@@ -90,20 +90,10 @@ NOTE: in the actual code,
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
-
93
Runtime
94
q_c = h_t @ W_DQ
106
-q_latent = q_c @ W_UQ_UK.view(Sq, N, Lkv)
95
+q_nope = (q_c @ W_UQ).view(-1, N, P)
96
+ql_nope = einsum("snh,lnh->snl", q, W_UK)
97
q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
98
new_kv_c = h_t @ W_DKV
99
new_k_pe = RoPE(h_t @ W_KR)
@@ -116,11 +106,13 @@ k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)
106
// NOTE: this is less compute-friendly since Lkv > P
107
// but is more data-movement friendly since its MQA vs MHA
108
spda_o = scaled_dot_product_attention(
119
- torch.cat([q_latent, q_pe], dim=-1),
109
+ torch.cat([ql_nope, q_pe], dim=-1),
110
torch.cat([kv_c, k_pe], dim=-1),
111
kv_c
112
)
123
-return spda_o.reshape(-1, N * Lkv) @ W_UV_O
113
+
114
+o = einsum("snl,lnv->snv", spda_o.reshape(-1, N, Lkv), W_UV)
115
+return o.view(-1, N * V) @ self.num_heads @ W_O
116
117
118
## Chunked Prefill
@@ -146,8 +138,8 @@ q_nope = (q_c @ W_UQ).view(Sq, N, P)
138
q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
139
new_kv_c = h_t @ W_DKV
140
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)
141
+new_k_nope = (new_kv_c @ W_UK.view(Lkv, N * P)).view(Sq, N, P)
142
+new_v = (new_kv_c @ W_UV.view(Lkv, N * V)).view(Sq, N, V)
143
144
// MHA between queries and new KV
145
// with QK headdim = P + R
@@ -171,17 +163,17 @@ for chunk_idx in range(cdiv(C, MCC)):
163
cache_k_pe_chunk = cache_k_pe[chunk_start:chunk_end]
164
cache_k_nope_chunk = (cache_kv_c_chunk @ W_UK).view(-1, N, P)
165
cache_v_chunk = (cache_kv_c_chunk @ W_UV).view(-1, N, V)
174
-
166
+
167
chunk_o, chunk_lse = scaled_dot_product_attention(
168
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)],
169
+ torch.cat([cache_k_nope_chunk,
170
+ cache_k_pe_chunk.unsqueeze(1).expand(-1, N, -1)],
171
dim=-1),
172
cache_v_chunk,
173
casual=False,
174
return_softmax_lse=True
175
)
184
-
176
+
177
curr_o, curr_lse = merge_attn_states(
178
suffix_output=curr_o,
179
suffix_lse=curr_lse,
@@ -202,7 +194,6 @@ from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, Tuple,
194
Type, TypeVar)
195
196
import torch
205
-from compressed_tensors.quantization import QuantizationStrategy
197
198
from vllm import _custom_ops as ops
199
from vllm import envs
@@ -215,20 +206,9 @@ from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
206
get_flash_attn_version,
207
is_block_tables_empty)
208
from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
218
-from vllm.distributed import (get_tensor_model_parallel_world_size,
219
- tensor_model_parallel_all_reduce)
209
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
210
LinearBase, RowParallelLinear,
211
UnquantizedLinearMethod)
223
-from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
224
- CompressedTensorsLinearMethod)
225
-from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
226
- CompressedTensorsW8A8Fp8)
227
-from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
228
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
229
- Fp8LinearGenericOp, is_fp8)
230
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
231
- scaled_quantize)
212
from vllm.model_executor.layers.rotary_embedding import (
213
DeepseekScalingRotaryEmbedding, RotaryEmbedding)
214
from vllm.multimodal import MultiModalPlaceholderMap
@@ -1057,7 +1037,6 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1037
self.kv_b_proj = kv_b_proj
1038
self.o_proj = o_proj
1039
self.triton_fa_func = triton_attention
1060
- self.fp8_linear_generic = Fp8LinearGenericOp()
1040
1041
# Handle the differences between the flash_attn_varlen from flash_attn
1042
# and the one from vllm_flash_attn. The former is used on RoCM and the
@@ -1070,79 +1049,28 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1049
fa_version=self.vllm_flash_attn_version)
1050
1051
def _v_up_proj_and_o_proj(self, x):
1073
- if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
1074
- if is_fp8(self.W_UV_O):
1075
- output_parallel = self.fp8_linear_generic.apply(
1076
- x.flatten(start_dim=1), self.W_UV_O, self.W_UV_O_scales,
1077
- self.reqaunt_input_group_shape,
1078
- self.reqaunt_weight_group_shape)
1079
- else:
1080
- output_parallel = torch.matmul(x.flatten(start_dim=1),
1081
- self.W_UV_O)
1082
- if self.tp_size > 1:
1083
- output = tensor_model_parallel_all_reduce(output_parallel)
1084
- else:
1085
- output = output_parallel
1086
- return output
1087
- else:
1088
- x = torch.einsum("bnl,lnv->bnv", x, self.W_UV)
1089
- return self.o_proj(x.reshape(-1,
1090
- self.num_heads * self.v_head_dim))[0]
1091
-
1052
+ # Convert from (B, N, L) to (N, B, L)
1053
+ x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
1054
+ # Multiply (N, B, L) x (N, L, V) -> (N, B, V)
1055
+ x = torch.bmm(x, self.W_UV)
1056
+ # Convert from (N, B, V) to (B, N * V)
1057
+ x = x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)
1058
+ return self.o_proj(x)[0]
1059
+
1060
+ # Return `ql_nope`, `q_pe`
1061
def _q_proj_and_k_up_proj(self, x):
1093
- if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
1094
- if is_fp8(self.W_Q_UK):
1095
- return self.fp8_linear_generic.apply(
1096
- x, self.W_Q_UK, self.W_Q_UK_scales,
1097
- self.reqaunt_input_group_shape,
1098
- self.reqaunt_weight_group_shape).view(
1099
- -1, self.num_heads, self.kv_lora_rank)
1100
- return torch.matmul(x, self.W_Q_UK)\
1101
- .view(-1, self.num_heads, self.kv_lora_rank)
1102
- else:
1103
- x = torch.matmul(x, self.W_Q)\
1104
- .view(-1, self.num_heads, self.qk_nope_head_dim)
1105
- return torch.einsum("bnp,lnp->bnl", x, self.W_UK)\
1106
- .view(-1, self.num_heads, self.kv_lora_rank)
1062
+ q_nope, q_pe = self.q_proj(x)[0]\
1063
+ .view(-1, self.num_heads, self.qk_head_dim)\
1064
+ .split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
1065
1108
- def process_weights_after_loading(self, act_dtype: torch.dtype):
1066
+ # Convert from (B, N, P) to (N, B, P)
1067
+ q_nope = q_nope.transpose(0, 1)
1068
+ # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
1069
+ ql_nope = torch.bmm(q_nope, self.W_UK_T)
1070
+ # Convert from (N, B, L) to (B, N, L)
1071
+ return ql_nope.transpose(0, 1), q_pe
1072
1110
- # TODO(lucas) This is very gross, we need a more wide scale refactor of
1111
- # all the FP8 code with a more standard way of
1112
- # defining schemes/group-shapes, we should also potentially force
1113
- # quant_methods to support a decompress function
1114
- #
1115
- # returns input_group_shape, weight_group_shape
1116
- def get_scale_group_shapes_for_fp8(layer: LinearBase) -> \
1117
- Tuple[Tuple[int, int], Tuple[int, int]]:
1118
- if isinstance(layer.quant_method, Fp8LinearMethod):
1119
- if layer.quant_method.block_quant:
1120
- weight_block_size = \
1121
- layer.quant_method.quant_config.weight_block_size
1122
- # per-token-group (1, X), block-quantized (X, Y)
1123
- return (1, weight_block_size[-1]), weight_block_size
1124
- else:
1125
- return (-1, -1), (-1, -1) # per-tensor, per-tensor
1126
- elif isinstance(layer.quant_method, CompressedTensorsLinearMethod)\
1127
- and isinstance(layer.scheme, CompressedTensorsW8A8Fp8):
1128
- # this is hacky but we always assume the for
1129
- # CompressedTensorsW8A8Fp8 the input is dynamic per-token
1130
- # we ignore if it is static-per-tensor since we are going to
1131
- # requantize after later anyways
1132
- strategy = layer.scheme.strategy
1133
- if strategy == QuantizationStrategy.TENSOR:
1134
- return (1, -1), (-1, -1) # per-token, per-tensor
1135
- elif strategy == QuantizationStrategy.CHANNEL:
1136
- return (1, -1), (-1, 1) # per-token, per-channel
1137
- else:
1138
- raise NotImplementedError(
1139
- f"QuantizationStrategy.{strategy} is not supported for "
1140
- "fp8 MLA, please run with VLLM_MLA_DISABLE=1")
1141
- else:
1142
- raise NotImplementedError(
1143
- "Can't determine scale group shapes for "
1144
- f"{layer.quant_method}, please run with VLLM_MLA_DISABLE=1"
1145
- )
1073
+ def process_weights_after_loading(self, act_dtype: torch.dtype):
1074
1075
def get_layer_weight(layer):
1076
WEIGHT_NAMES = ("weight", "qweight", "weight_packed")
@@ -1167,10 +1095,9 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1095
return dequant_weights.T
1096
return layer.weight
1097
1170
- weight_dtype = get_layer_weight(self.kv_b_proj).dtype
1171
- assert get_layer_weight(self.o_proj).dtype == weight_dtype
1172
- assert get_layer_weight(self.q_proj).dtype == weight_dtype
1173
-
1098
+ # we currently do not have quantized bmm's which are needed for
1099
+ # `W_UV` and `W_UK_T`, we we just store fp16/bf16 copies and perform
1100
+ # the bmm's in 16-bit, the extra memory overhead of this is fairly low
1101
kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
1102
assert kv_b_proj_weight.shape == (
1103
self.kv_lora_rank,
@@ -1189,89 +1116,10 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1116
W_UK, W_UV = kv_b_proj_weight.split(
1117
[self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1118
1192
- q_proj_weight = get_and_maybe_dequant_weights(self.q_proj).T\
1193
- .view(-1, self.num_heads, self.qk_head_dim)
1194
-
1195
- # can be W_Q or W_UQ depending q_lora_rank, the former if
1196
- # q_lora_rank is None, the latter otherwise. From the Attention backend
1197
- # perspective though we call these both W_Q and rely on the layer
1198
- # to pass in the correct matrix
1199
- W_Q = q_proj_weight[..., :self.qk_nope_head_dim]
1200
- self.W_QR = q_proj_weight[..., self.qk_nope_head_dim:]\
1201
- .flatten(start_dim=1).contiguous()
1202
-
1203
- # W_QR is small so for simplicity we dont bother requantizing it
1204
- self.W_QR = self.W_QR.to(act_dtype)
1205
-
1206
- if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
1207
- requantization_enabled = not envs.VLLM_MLA_DISABLE_REQUANTIZATION
1208
- if is_fp8(weight_dtype) and requantization_enabled:
1209
- # This assumes it wise to requantize using the same group shapes
1210
- # (i.e. strategy, per-tensor, per-channel, block etc.) that the
1211
- # weights were originally quantized
1212
- requant_input_group_shape, requant_weight_group_shape = \
1213
- get_scale_group_shapes_for_fp8(self.q_proj)
1214
- assert (requant_input_group_shape, requant_weight_group_shape)\
1215
- == get_scale_group_shapes_for_fp8(self.kv_b_proj)
1216
- assert (requant_input_group_shape, requant_weight_group_shape)\
1217
- == get_scale_group_shapes_for_fp8(self.o_proj)
1218
- self.reqaunt_input_group_shape = requant_input_group_shape
1219
- self.reqaunt_weight_group_shape = requant_weight_group_shape
1220
-
1221
- #
1222
- # Perform matrix-absorption following
1223
- # https://github.com/flashinfer-ai/flashinfer/pull/551
1224
- # for decode, as a result we end up with absorbed weights for decode
1225
- # and another copy of raw weights for prefill.
1226
- #
1227
- self.W_UK, self.W_UV = kv_b_proj_weight.split(
1228
- [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
1229
- # We absorb `W_UK` into `W_Q` resulting in either W_Q_UK or W_UQ_UK
1230
- # depending q_lora_rank, the former if q_lora_rank is None, the
1231
- # latter otherwise
1232
- # basically if q_lora_rank is none we are absorbing into q_proj
1233
- # instead of UQ
1234
- W_Q_UK = torch.einsum("qnd,lnd -> qnl", W_Q, W_UK)\
1235
- .flatten(start_dim=1).contiguous()
1236
-
1237
- if is_fp8(weight_dtype) and requantization_enabled:
1238
- W_Q_UK, W_Q_UK_scales = scaled_quantize(
1239
- W_Q_UK,
1240
- self.reqaunt_weight_group_shape,
1241
- quant_dtype=current_platform.fp8_dtype())
1242
- # For FP8 save the transpose so we can use
1243
- # `apply_w8a8_block_fp8_linear` directly
1244
- self.W_Q_UK = W_Q_UK.T.contiguous()
1245
- self.W_Q_UK_scales = W_Q_UK_scales.T.contiguous()
1246
- else:
1247
- self.W_Q_UK = W_Q_UK.to(act_dtype)
1248
-
1249
- W_O = get_and_maybe_dequant_weights(self.o_proj)\
1250
- .view(-1, self.num_heads, self.v_head_dim)
1251
- W_UV_O = torch.einsum("lnd,hnd -> nlh", W_UV, W_O)\
1252
- .flatten(start_dim=0, end_dim=1).contiguous()
1253
-
1254
- if is_fp8(weight_dtype) and requantization_enabled:
1255
- W_UV_O, W_UV_O_scales = scaled_quantize(
1256
- W_UV_O,
1257
- self.reqaunt_weight_group_shape,
1258
- quant_dtype=current_platform.fp8_dtype())
1259
- # For FP8 save the transpose so we can use
1260
- # `apply_w8a8_block_fp8_linear` directly
1261
- self.W_UV_O = W_UV_O.T.contiguous()
1262
- self.W_UV_O_scales = W_UV_O_scales.T.contiguous()
1263
- else:
1264
- self.W_UV_O = W_UV_O.to(act_dtype)
1265
-
1266
- self.tp_size = get_tensor_model_parallel_world_size()
1267
- else:
1268
- if is_fp8(weight_dtype):
1269
- raise NotImplementedError(
1270
- "Currently fp8 requires matrix absorption")
1271
-
1272
- self.W_UV = W_UV
1273
- self.W_UK = W_UK
1274
- self.W_Q = W_Q.flatten(start_dim=1)
1119
+ # Convert from (L, N, V) to (N, L, V)
1120
+ self.W_UV = W_UV.transpose(0, 1)
1121
+ # Convert from (L, N, P) to (N, P, L)
1122
+ self.W_UK_T = W_UK.permute(1, 2, 0)
1123
1124
def _compute_prefill_context(
1125
self,
@@ -1471,7 +1319,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1319
@abstractmethod
1320
def _forward_decode(
1321
self,
1474
- q_nope: torch.Tensor,
1322
+ ql_nope: torch.Tensor,
1323
q_pe: torch.Tensor,
1324
kv_c_and_k_pe_cache: torch.Tensor,
1325
attn_metadata: T,
@@ -1525,9 +1373,8 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1373
prefill_k_c_normed = k_c_normed[:num_prefill_tokens]
1374
1375
if has_decode:
1528
- decode_q_nope = self._q_proj_and_k_up_proj(decode_hs_or_q_c)
1529
- decode_q_pe = torch.matmul(decode_hs_or_q_c, self.W_QR)\
1530
- .view(-1, self.num_heads, self.qk_rope_head_dim)
1376
+ decode_ql_nope, decode_q_pe = \
1377
+ self._q_proj_and_k_up_proj(decode_hs_or_q_c)
1378
decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(
1379
decode_input_positions, decode_q_pe, decode_k_pe)
1380
@@ -1561,6 +1408,6 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1408
1409
if has_decode:
1410
output[num_prefill_tokens:] = self._forward_decode(
1564
- decode_q_nope, decode_q_pe, kv_cache, attn_metadata)
1411
+ decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)
1412
1413
return output
vllm/envs.py
-19
@@ -84,8 +84,6 @@ if TYPE_CHECKING:
84
VLLM_SERVER_DEV_MODE: bool = False
85
VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128
86
VLLM_MLA_DISABLE: bool = False
87
- VLLM_MLA_PERFORM_MATRIX_ABSORPTION: bool = True
88
- VLLM_MLA_DISABLE_REQUANTIZATION: bool = False
87
VLLM_MLA_CUDA_MEM_ALIGN_KV_CACHE: bool = True
88
VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON: bool = False
89
VLLM_RAY_PER_WORKER_GPUS: float = 1.0
@@ -563,23 +561,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
561
"VLLM_MLA_DISABLE":
562
lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))),
563
566
- # Flag that can control whether or not we perform matrix-absorption for MLA
567
- # decode, i.e. absorb W_UK into W_Q/W_UK and W_UV into W_O, absorbing the
568
- # matrices reduces the runtime FLOPs needed to compute MLA but requires
569
- # storing more weights, W_Q_UK and W_UV_O, so can increase memory usage,
570
- # the is enabled by default
571
- "VLLM_MLA_PERFORM_MATRIX_ABSORPTION":
572
- lambda: bool(int(os.getenv("VLLM_MLA_PERFORM_MATRIX_ABSORPTION", "1"))),
573
-
574
- # When running MLA with matrix-absorption enabled and fp8 quantized weights
575
- # we perform the matrix-absorption in float32 precision, after the matrices
576
- # are absorbed we requantize the weights back to fp8, this flag can be used
577
- # to disable the requantization step, and instead convert the absorbed
578
- # matrices to match the activation type. This can lead to higher memory and
579
- # compute usage but better preserves the accuracy of the original model.
580
- "VLLM_MLA_DISABLE_REQUANTIZATION":
581
- lambda: bool(int(os.getenv("VLLM_MLA_DISABLE_REQUANTIZATION", "0"))),
582
-
564
# If set, vLLM will use the Triton implementation of moe_align_block_size,
565
# i.e. moe_align_block_size_triton in fused_moe.py.
566
"VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON":
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+2
-57
@@ -13,10 +13,9 @@ import triton.language as tl
13
from vllm import _custom_ops as ops
14
from vllm.logger import init_logger
15
from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
- _normalize_quant_group_shape, scaled_dequantize)
16
+ scaled_dequantize)
17
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
18
- CUTLASS_BLOCK_FP8_SUPPORTED, Fp8LinearOp, cutlass_block_fp8_supported,
19
- cutlass_fp8_supported)
18
+ CUTLASS_BLOCK_FP8_SUPPORTED)
19
from vllm.platforms import current_platform
20
from vllm.utils import direct_register_custom_op
21
@@ -101,60 +100,6 @@ direct_register_custom_op(
100
)
101
102
104
-# Unify the interface between `apply_w8a8_block_fp8_linear` and
105
-# `apply_fp8_linear`
106
-# NOTE(lucas): this is quite messy, we should think through this more formally
107
-# TODO(luka): unify this better
108
-# https://github.com/vllm-project/vllm/issues/14397
109
-class Fp8LinearGenericOp:
110
-
111
- def __init__(
112
- self,
113
- cutlass_fp8_supported: bool = cutlass_fp8_supported(),
114
- cutlass_block_fp8_supported: bool = cutlass_block_fp8_supported(),
115
- ):
116
- self.cutlass_block_fp8_supported = cutlass_block_fp8_supported
117
- self.fp8_linear = Fp8LinearOp(
118
- cutlass_fp8_supported=cutlass_fp8_supported)
119
-
120
- def apply(
121
- self,
122
- input: torch.Tensor,
123
- weight: torch.Tensor,
124
- weight_scale: torch.Tensor,
125
- input_group_shape: Tuple[int, int],
126
- weight_group_shape: Tuple[int, int],
127
- input_scale: Optional[torch.Tensor] = None, # static scale if one
128
- ) -> torch.Tensor:
129
- # View input as 2D matrix for fp8 methods
130
- input = input.view(-1, input.shape[-1])
131
-
132
- weight_group_shape = _normalize_quant_group_shape( \
133
- weight, weight_group_shape)
134
- input_group_shape = _normalize_quant_group_shape(
135
- input, input_group_shape)
136
-
137
- def is_dim_blocked(dim, shape, group_shape):
138
- return group_shape < shape[dim] and group_shape > 1
139
-
140
- if is_dim_blocked(0, weight.shape, weight_group_shape[0])\
141
- and is_dim_blocked(1, weight.shape, weight_group_shape[1]) and\
142
- input_group_shape == (1, weight_group_shape[1]):
143
- return apply_w8a8_block_fp8_linear(
144
- input,
145
- weight,
146
- list(weight_group_shape),
147
- weight_scale,
148
- cutlass_block_fp8_supported=self.cutlass_block_fp8_supported)
149
- else:
150
- # Despite having linear in the name it doesn't conform to
151
- # `torch.nn.functional.linear` which is defined as
152
- # `input @ weight.T` so we explicitly transpose the weight matrix
153
- return self.fp8_linear.apply(input, weight.T, weight_scale.T,
154
- use_per_token_if_dynamic=\
155
- (input_group_shape == (1, input.shape[1])))
156
-
157
-
103
def input_to_float8(
104
x: torch.Tensor,
105
dtype: Optional[torch.dtype] = None
vllm/v1/attention/backends/mla/common.py
+58
-213
@@ -21,7 +21,7 @@ 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.
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
@@ -54,9 +54,9 @@ 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]
57
+W_UK project kv_c to k_nope shape [Lkv, N, P]
58
+W_KR project h_t to k_pe shape [H, 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
@@ -69,8 +69,8 @@ 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)
72
+k_nope = (kv_c @ W_UK.view(Lkv, N * P)).view(Skv, N, P)
73
+v = (kv_c @ W_UV.view(Lkv, N * V)).view(Skv, N, V)
74
75
// MHA with QK headdim = P + R
76
// V headdim = V
@@ -79,7 +79,7 @@ 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
-)
82
+)
83
return spda_o @ W_O
84
85
NOTE: in the actual code,
@@ -90,20 +90,10 @@ NOTE: in the actual code,
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
-
93
Runtime
94
q_c = h_t @ W_DQ
106
-q_latent = q_c @ W_UQ_UK.view(Sq, N, Lkv)
95
+q_nope = (q_c @ W_UQ).view(-1, N, P)
96
+ql_nope = einsum("snh,lnh->snl", q, W_UK)
97
q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
98
new_kv_c = h_t @ W_DKV
99
new_k_pe = RoPE(h_t @ W_KR)
@@ -116,29 +106,31 @@ k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)
106
// NOTE: this is less compute-friendly since Lkv > P
107
// but is more data-movement friendly since its MQA vs MHA
108
spda_o = scaled_dot_product_attention(
119
- torch.cat([q_latent, q_pe], dim=-1),
109
+ torch.cat([ql_nope, q_pe], dim=-1),
110
torch.cat([kv_c, k_pe], dim=-1),
111
kv_c
112
)
123
-return spda_o.reshape(-1, N * Lkv) @ W_UV_O
113
+
114
+o = einsum("snl,lnv->snv", spda_o.reshape(-1, N, Lkv), W_UV)
115
+return o.view(-1, N * V) @ self.num_heads @ W_O
116
117
118
## Chunked Prefill
119
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
120
+For chunked prefill we want to use the compute friendly algorithm. We are
121
+assuming sufficiently large Sq / Skv ratio, in the future may want to switch to
122
the data-movement friendly approach if the chunk (i.e. `Sq`) is small.
123
124
However, the compute-friendly approach can potentially run out of memory if Skv
125
is large due to: `k_nope = (kv_c @ W_UK).view(Skv, N, P)`
126
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
127
+To mitigate this, we chunk the computation of attention with respect to the
128
+current context (i.e. `cache_kv_c` and `cache_k_pe`) so that we can used a
129
fixed workspace size.
130
131
The chunked prefill approach is as follows:
132
141
-MCC Max chunk of context to process per iter, computed dynamically,
133
+MCC Max chunk of context to process per iter, computed dynamically,
134
used to bound the memory usage
135
136
q_c = h_t @ W_DQ
@@ -146,8 +138,8 @@ q_nope = (q_c @ W_UQ).view(Sq, N, P)
138
q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
139
new_kv_c = h_t @ W_DKV
140
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)
141
+new_k_nope = (new_kv_c @ W_UK.view(Lkv, N * P)).view(Sq, N, P)
142
+new_v = (new_kv_c @ W_UV.view(Lkv, N * V)).view(Sq, N, V)
143
144
// MHA between queries and new KV
145
// with QK headdim = P + R
@@ -160,7 +152,7 @@ curr_o, curr_lse = scaled_dot_product_attention(
152
new_v,
153
casual=True,
154
return_softmax_lse=True
163
-)
155
+)
156
157
// Compute attention with the already existing context
158
for chunk_idx in range(cdiv(C, MCC)):
@@ -198,30 +190,17 @@ from dataclasses import dataclass
190
from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
191
192
import torch
201
-from compressed_tensors.quantization import QuantizationStrategy
193
194
from vllm import _custom_ops as ops
204
-from vllm import envs
195
from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
196
AttentionMetadata,
197
MLAAttentionImpl)
198
from vllm.attention.backends.utils import get_flash_attn_version
199
from vllm.attention.ops.triton_merge_attn_states import merge_attn_states
210
-from vllm.distributed import (get_tensor_model_parallel_world_size,
211
- tensor_model_parallel_all_reduce)
200
from vllm.logger import init_logger
201
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
202
LinearBase, RowParallelLinear,
203
UnquantizedLinearMethod)
216
-from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
217
- CompressedTensorsLinearMethod)
218
-from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
219
- CompressedTensorsW8A8Fp8)
220
-from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
221
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
222
- Fp8LinearGenericOp, is_fp8)
223
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
224
- scaled_quantize)
204
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
205
from vllm.platforms import current_platform
206
from vllm.utils import cdiv, round_down
@@ -646,7 +625,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
625
self.kv_b_proj = kv_b_proj
626
self.o_proj = o_proj
627
self.vllm_flash_attn_version = get_flash_attn_version()
649
- self.fp8_linear_generic = Fp8LinearGenericOp()
628
629
# Handle the differences between the flash_attn_varlen from flash_attn
630
# and the one from vllm_flash_attn. The former is used on RoCM and the
@@ -658,88 +636,37 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
636
fa_version=self.vllm_flash_attn_version)
637
638
def _v_up_proj_and_o_proj(self, x):
661
- if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
662
- if is_fp8(self.W_UV_O):
663
- output_parallel = self.fp8_linear_generic.apply(
664
- x.flatten(start_dim=1), self.W_UV_O, self.W_UV_O_scales,
665
- self.reqaunt_input_group_shape,
666
- self.reqaunt_weight_group_shape)
667
- else:
668
- output_parallel = torch.matmul(x.flatten(start_dim=1),
669
- self.W_UV_O)
670
- if self.tp_size > 1:
671
- output = tensor_model_parallel_all_reduce(output_parallel)
672
- else:
673
- output = output_parallel
674
- return output
675
- else:
676
- x = torch.einsum("bnl,lnv->bnv", x, self.W_UV)
677
- return self.o_proj(x.reshape(-1,
678
- self.num_heads * self.v_head_dim))[0]
679
-
639
+ # Convert from (B, N, L) to (N, B, L)
640
+ x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
641
+ # Multiply (N, B, L) x (N, L, V) -> (N, B, V)
642
+ x = torch.bmm(x, self.W_UV)
643
+ # Convert from (N, B, V) to (B, N * V)
644
+ x = x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)
645
+ return self.o_proj(x)[0]
646
+
647
+ # Return `ql_nope`, `q_pe`
648
def _q_proj_and_k_up_proj(self, x):
681
- if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
682
- if is_fp8(self.W_Q_UK):
683
- return self.fp8_linear_generic.apply(
684
- x, self.W_Q_UK, self.W_Q_UK_scales,
685
- self.reqaunt_input_group_shape,
686
- self.reqaunt_weight_group_shape).view(
687
- -1, self.num_heads, self.kv_lora_rank)
688
- return torch.matmul(x, self.W_Q_UK)\
689
- .view(-1, self.num_heads, self.kv_lora_rank)
690
- else:
691
- x = torch.matmul(x, self.W_Q)\
692
- .view(-1, self.num_heads, self.qk_nope_head_dim)
693
- return torch.einsum("bnp,lnp->bnl", x, self.W_UK)\
694
- .view(-1, self.num_heads, self.kv_lora_rank)
649
+ q_nope, q_pe = self.q_proj(x)[0]\
650
+ .view(-1, self.num_heads, self.qk_head_dim)\
651
+ .split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
652
696
- def process_weights_after_loading(self, act_dtype: torch.dtype):
653
+ # Convert from (B, N, P) to (N, B, P)
654
+ q_nope = q_nope.transpose(0, 1)
655
+ # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
656
+ ql_nope = torch.bmm(q_nope, self.W_UK_T)
657
+ # Convert from (N, B, L) to (B, N, L)
658
+ return ql_nope.transpose(0, 1), q_pe
659
698
- # TODO(lucas) This is very gross, we need a more wide scale refactor of
699
- # all the FP8 code with a more standard way of
700
- # defining schemes/group-shapes, we should also potentially force
701
- # quant_methods to support a decompress function
702
- #
703
- # returns input_group_shape, weight_group_shape
704
- def get_scale_group_shapes_for_fp8(layer: LinearBase) -> \
705
- tuple[tuple[int, int], tuple[int, int]]:
706
- if isinstance(layer.quant_method, Fp8LinearMethod):
707
- if layer.quant_method.block_quant:
708
- weight_block_size = \
709
- layer.quant_method.quant_config.weight_block_size
710
- # per-token-group (1, X), block-quantized (X, Y)
711
- return (1, weight_block_size[-1]), weight_block_size
712
- else:
713
- return (-1, -1), (-1, -1) # per-tensor, per-tensor
714
- elif isinstance(layer.quant_method, CompressedTensorsLinearMethod)\
715
- and isinstance(layer.scheme, CompressedTensorsW8A8Fp8):
716
- # this is hacky but we always assume the for
717
- # CompressedTensorsW8A8Fp8 the input is dynamic per-token
718
- # we ignore if it is static-per-tensor since we are going to
719
- # requantize after later anyways
720
- strategy = layer.scheme.strategy
721
- if strategy == QuantizationStrategy.TENSOR:
722
- return (1, -1), (-1, -1) # per-token, per-tensor
723
- elif strategy == QuantizationStrategy.CHANNEL:
724
- return (1, -1), (-1, 1) # per-token, per-channel
725
- else:
726
- raise NotImplementedError(
727
- f"QuantizationStrategy.{strategy} is not supported for "
728
- "fp8 MLA, please run with VLLM_MLA_DISABLE=1")
729
- else:
730
- raise NotImplementedError(
731
- "Can't determine scale group shapes for "
732
- f"{layer.quant_method}, please run with VLLM_MLA_DISABLE=1"
733
- )
660
+ def process_weights_after_loading(self, act_dtype: torch.dtype):
661
662
def get_layer_weight(layer):
736
- if hasattr(layer, "weight"):
737
- return layer.weight
738
- elif hasattr(layer, "qweight"):
739
- return layer.qweight
740
- else:
741
- raise AttributeError(
742
- f"Layer '{layer}' has neither weight nor qweight")
663
+ WEIGHT_NAMES = ("weight", "qweight", "weight_packed")
664
+ for attr in WEIGHT_NAMES:
665
+ if hasattr(layer, attr):
666
+ return getattr(layer, attr)
667
+ raise AttributeError(
668
+ f"Layer '{layer}' has no recognized weight attribute:"
669
+ f" {WEIGHT_NAMES}.")
670
671
def get_and_maybe_dequant_weights(layer: LinearBase):
672
if not isinstance(layer.quant_method, UnquantizedLinearMethod):
@@ -755,10 +682,9 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
682
return dequant_weights.T
683
return layer.weight
684
758
- weight_dtype = get_layer_weight(self.kv_b_proj).dtype
759
- assert get_layer_weight(self.o_proj).dtype == weight_dtype
760
- assert get_layer_weight(self.q_proj).dtype == weight_dtype
761
-
685
+ # we currently do not have quantized bmm's which are needed for
686
+ # `W_UV` and `W_UK_T`, we we just store fp16/bf16 copies and perform
687
+ # the bmm's in 16-bit, the extra memory overhead of this is fairly low
688
kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
689
assert kv_b_proj_weight.shape == (
690
self.kv_lora_rank,
@@ -777,89 +703,10 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
703
W_UK, W_UV = kv_b_proj_weight.split(
704
[self.qk_nope_head_dim, self.v_head_dim], dim=-1)
705
780
- q_proj_weight = get_and_maybe_dequant_weights(self.q_proj).T\
781
- .view(-1, self.num_heads, self.qk_head_dim)
782
-
783
- # can be W_Q or W_UQ depending q_lora_rank, the former if
784
- # q_lora_rank is None, the latter otherwise. From the Attention backend
785
- # perspective though we call these both W_Q and rely on the layer
786
- # to pass in the correct matrix
787
- W_Q = q_proj_weight[..., :self.qk_nope_head_dim]
788
- self.W_QR = q_proj_weight[..., self.qk_nope_head_dim:]\
789
- .flatten(start_dim=1).contiguous()
790
-
791
- # W_QR is small so for simplicity we dont bother requantizing it
792
- self.W_QR = self.W_QR.to(act_dtype)
793
-
794
- if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
795
- requantization_enabled = not envs.VLLM_MLA_DISABLE_REQUANTIZATION
796
- if is_fp8(weight_dtype) and requantization_enabled:
797
- # This assumes it wise to requantize using the same group shapes
798
- # (i.e. strategy, per-tensor, per-channel, block etc.) that the
799
- # weights were originally quantized
800
- requant_input_group_shape, requant_weight_group_shape = \
801
- get_scale_group_shapes_for_fp8(self.q_proj)
802
- assert (requant_input_group_shape, requant_weight_group_shape)\
803
- == get_scale_group_shapes_for_fp8(self.kv_b_proj)
804
- assert (requant_input_group_shape, requant_weight_group_shape)\
805
- == get_scale_group_shapes_for_fp8(self.o_proj)
806
- self.reqaunt_input_group_shape = requant_input_group_shape
807
- self.reqaunt_weight_group_shape = requant_weight_group_shape
808
-
809
- #
810
- # Perform matrix-absorption following
811
- # https://github.com/flashinfer-ai/flashinfer/pull/551
812
- # for decode, as a result we end up with absorbed weights for decode
813
- # and another copy of raw weights for prefill.
814
- #
815
- self.W_UK, self.W_UV = kv_b_proj_weight.split(
816
- [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
817
- # We absorb `W_UK` into `W_Q` resulting in either W_Q_UK or W_UQ_UK
818
- # depending q_lora_rank, the former if q_lora_rank is None, the
819
- # latter otherwise
820
- # basically if q_lora_rank is none we are absorbing into q_proj
821
- # instead of UQ
822
- W_Q_UK = torch.einsum("qnd,lnd -> qnl", W_Q, W_UK)\
823
- .flatten(start_dim=1).contiguous()
824
-
825
- if is_fp8(weight_dtype) and requantization_enabled:
826
- W_Q_UK, W_Q_UK_scales = scaled_quantize(
827
- W_Q_UK,
828
- self.reqaunt_weight_group_shape,
829
- quant_dtype=current_platform.fp8_dtype())
830
- # For FP8 save the transpose so we can use
831
- # `apply_w8a8_block_fp8_linear` directly
832
- self.W_Q_UK = W_Q_UK.T.contiguous()
833
- self.W_Q_UK_scales = W_Q_UK_scales.T.contiguous()
834
- else:
835
- self.W_Q_UK = W_Q_UK.to(act_dtype)
836
-
837
- W_O = get_and_maybe_dequant_weights(self.o_proj)\
838
- .view(-1, self.num_heads, self.v_head_dim)
839
- W_UV_O = torch.einsum("lnd,hnd -> nlh", W_UV, W_O)\
840
- .flatten(start_dim=0, end_dim=1).contiguous()
841
-
842
- if is_fp8(weight_dtype) and requantization_enabled:
843
- W_UV_O, W_UV_O_scales = scaled_quantize(
844
- W_UV_O,
845
- self.reqaunt_weight_group_shape,
846
- quant_dtype=current_platform.fp8_dtype())
847
- # For FP8 save the transpose so we can use
848
- # `apply_w8a8_block_fp8_linear` directly
849
- self.W_UV_O = W_UV_O.T.contiguous()
850
- self.W_UV_O_scales = W_UV_O_scales.T.contiguous()
851
- else:
852
- self.W_UV_O = W_UV_O.to(act_dtype)
853
-
854
- self.tp_size = get_tensor_model_parallel_world_size()
855
- else:
856
- if is_fp8(weight_dtype):
857
- raise NotImplementedError(
858
- "Currently fp8 requires matrix absorption")
859
-
860
- self.W_UV = W_UV
861
- self.W_UK = W_UK
862
- self.W_Q = W_Q.flatten(start_dim=1)
706
+ # Convert from (L, N, V) to (N, L, V)
707
+ self.W_UV = W_UV.transpose(0, 1)
708
+ # Convert from (L, N, P) to (N, P, L)
709
+ self.W_UK_T = W_UK.permute(1, 2, 0)
710
711
def _compute_prefill_context(
712
self,
@@ -998,7 +845,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
845
@abstractmethod
846
def _forward_decode(
847
self,
1001
- q_nope: torch.Tensor,
848
+ ql_nope: torch.Tensor,
849
q_pe: torch.Tensor,
850
kv_c_and_k_pe_cache: torch.Tensor,
851
attn_metadata: M,
@@ -1051,10 +898,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
898
899
if has_decode:
900
assert attn_metadata.decode is not None
1054
- decode_q_nope = self._q_proj_and_k_up_proj(decode_hs_or_q_c)
1055
- decode_q_pe = torch.matmul(decode_hs_or_q_c, self.W_QR)\
1056
- .view(-1, self.num_heads, self.qk_rope_head_dim)
1057
-
901
+ decode_ql_nope, decode_q_pe = \
902
+ self._q_proj_and_k_up_proj(decode_hs_or_q_c)
903
decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(
904
attn_metadata.decode.input_positions, decode_q_pe.contiguous(),
905
decode_k_pe)
@@ -1087,6 +932,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
932
933
if has_decode:
934
output[:num_decode_tokens] = self._forward_decode(
1090
- decode_q_nope, decode_q_pe, kv_cache, attn_metadata)
935
+ decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)
936
937
return output_padded