A high-throughput and memory-efficient inference and serving engine for LLMs
[Doc][Attention] Fix MLA top-of-file comments (#37047)
Signed-off-by: wineandchord <guoqizhou19@gmail.com>
wineandchord committed
Jun 11, 2026 at 23:14 UTC
c2b4cd39acca972da59e53983cf3ddd3b3d32605
1 file changed
+10
-10
vllm/model_executor/layers/attention/mla_attention.py
+10
-10
@@ -14,7 +14,7 @@ MLA has two possible ways of computing, a data-movement friendly approach and a
14
compute friendly approach. We generally want to use the compute friendly
15
approach for "prefill" (i.e. the ratio Sq / Skv is relatively large, often near
16
1) and the data-movement friendly approach for "decode" (i.e. the ratio
17
-Sq / Skv is small).
17
+Sq / Skv is small, often near 0).
18
19
NOTE what we deem small and large is currently determined by if it is labelled
20
prefill or decode by the scheduler, but this is something we should probably
@@ -28,7 +28,7 @@ Deepseek's MLA attention works the following way:
28
* For decode (i.e. the memory friendly approach) the attention "simulates" a
29
multi-head attention, while the compute is similar to multi-query attention.
30
31
-Below is example of both paths assuming batchsize = 1
31
+Below is an example of both paths assuming batch size = 1
32
33
## More Extent Definitions:
34
@@ -77,13 +77,13 @@ v = (kv_c @ W_UV.view(Lkv, N * V)).view(Skv, N, V)
77
78
// MHA with QK headdim = P + R
79
// V headdim = V
80
-// spda_o shape [Sq, N, V]
81
-spda_o = scaled_dot_product_attention(
80
+// sdpa_o shape [Sq, N, V]
81
+sdpa_o = scaled_dot_product_attention(
82
torch.cat([q_nope, q_pe], dim=-1),
83
torch.cat([k_nope, k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
84
v
85
)
86
-return spda_o @ W_O
86
+return sdpa_o @ W_O
87
88
NOTE: in the actual code,
89
`kv_b_proj` is [W_UK; W_UV] concatenated per head
@@ -105,16 +105,16 @@ k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)
105
106
// MQA with QK headdim = Lkv + R
107
// V headdim = Lkv
108
-// spda_o shape [Sq, N, Lkv]
108
+// sdpa_o shape [Sq, N, Lkv]
109
// NOTE: this is less compute-friendly since Lkv > P
110
// but is more data-movement friendly since its MQA vs MHA
111
-spda_o = scaled_dot_product_attention(
111
+sdpa_o = scaled_dot_product_attention(
112
torch.cat([ql_nope, q_pe], dim=-1),
113
torch.cat([kv_c, k_pe], dim=-1),
114
kv_c
115
)
116
117
-o = einsum("snl,lnv->snv", spda_o.reshape(-1, N, Lkv), W_UV)
117
+o = einsum("snl,lnv->snv", sdpa_o.reshape(-1, N, Lkv), W_UV)
118
return o.view(-1, N * V) @ W_O
119
120
@@ -153,7 +153,7 @@ curr_o, curr_lse = scaled_dot_product_attention(
153
torch.cat([q_nope, q_pe], dim=-1),
154
torch.cat([new_k_nope, new_k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
155
new_v,
156
- casual=True,
156
+ causal=True,
157
return_softmax_lse=True
158
)
159
@@ -173,7 +173,7 @@ for chunk_idx in range(cdiv(C, MCC)):
173
cache_k_pe_chunk.unsqueeze(1).expand(-1, N, -1)],
174
dim=-1),
175
cache_v_chunk,
176
- casual=False,
176
+ causal=False,
177
return_softmax_lse=True
178
)
179