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 2, 08:30 UTC
docs: fix MLA attention docstring examples (#44118)
Co-authored-by: nightcityblade <nightcityblade@gmail.com>
nightcityblade committed
Jun 1, 2026 at 03:28 UTC
8b8546da1c3ba65097357523bc24199e36eddf65
1 file changed
+7
-7
vllm/model_executor/layers/attention/mla_attention.py
+7
-7
@@ -11,12 +11,12 @@ Sq as Q sequence length
11
Skv as KV sequence length
12
13
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 "small", is near 1)
16
-and the data-movement friendly approach for "decode" (i.e. the ratio
17
-Sq / Skv is "large").
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).
18
19
-NOTE what we deem small and large is currently determined by if its labelled
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
21
tune.
22
@@ -96,7 +96,7 @@ NOTE: in the actual code,
96
Runtime
97
q_c = h_t @ W_DQ
98
q_nope = (q_c @ W_UQ).view(-1, N, P)
99
-ql_nope = einsum("snh,lnh->snl", q, W_UK)
99
+ql_nope = einsum("snh,lnh->snl", q_nope, W_UK)
100
q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)
101
new_kv_c = h_t @ W_DKV
102
new_k_pe = RoPE(h_t @ W_KR)
@@ -115,7 +115,7 @@ spda_o = scaled_dot_product_attention(
115
)
116
117
o = einsum("snl,lnv->snv", spda_o.reshape(-1, N, Lkv), W_UV)
118
-return o.view(-1, N * V) @ self.num_heads @ W_O
118
+return o.view(-1, N * V) @ W_O
119
120
121
## Chunked Prefill