[Kernel] Triton MLA logits workspace (#46819)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Jun 29, 2026 at 16:54 UTC
6149187a4cca41e4e16c6461de39a6c33005c361
1 file changed
+64
-29
vllm/v1/attention/backends/mla/triton_mla.py
+64
-29
@@ -25,13 +25,58 @@ from vllm.v1.attention.backend import (
25
MultipleOf,
26
)
27
from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
28
+from vllm.v1.worker.workspace import (
29
+ current_workspace_manager,
30
+ is_workspace_manager_initialized,
31
+)
32
33
logger = init_logger(__name__)
34
35
+# num_kv_splits selection (shared by forward_mqa and the workspace reservation
36
+# so the two cannot drift). Both are hardware dependent.
37
+_MIN_WORK_PER_SPLIT = 512
38
+_SPLIT_OCCUPANCY_MULTIPLIER = 2
39
+
40
+
41
+def _compute_num_kv_splits(max_seq_len: int, sm_count: int) -> int:
42
+ # Power of 2 to avoid excessive kernel instantiations, capped by an SM-based
43
+ # maximum (occupancy multiplier allows multiple blocks per SM
44
+ # for latency hiding).
45
+ ideal_splits = triton.next_power_of_2(max(1, max_seq_len // _MIN_WORK_PER_SPLIT))
46
+ max_splits = sm_count * _SPLIT_OCCUPANCY_MULTIPLIER
47
+ return min(ideal_splits, max_splits)
48
+
49
50
class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
51
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
52
53
+ def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
54
+ super().__init__(kv_cache_spec, layer_names, vllm_config, device)
55
+ self._reserve_attn_logits_workspace()
56
+
57
+ def _reserve_attn_logits_workspace(self) -> None:
58
+ """Pre-size the shared workspace for the decode split-KV attn logits.
59
+
60
+ Reserving at the worst case (max_model_len -> max num_kv_splits,
61
+ max_num_seqs decode tokens) before warmup/cudagraph capture means the
62
+ per-call ``get_simultaneous`` in ``forward_mqa`` never has to grow the
63
+ buffer at runtime (which would raise once the workspace is locked).
64
+ """
65
+ if not is_workspace_manager_initialized():
66
+ return
67
+ # Decode reorder threshold is 1, so decode tokens <= max_num_seqs.
68
+ B = self.vllm_config.scheduler_config.max_num_seqs
69
+ # DCP all-gathers the query heads before forward_mqa.
70
+ q_num_heads = self.num_heads * self.dcp_world_size
71
+ max_splits = _compute_num_kv_splits(
72
+ self.model_config.max_model_len,
73
+ current_platform.num_compute_units(),
74
+ )
75
+ lse_dim = self.mla_dims.kv_lora_rank + 1
76
+ current_workspace_manager().get_simultaneous(
77
+ ((B, q_num_heads, max_splits, lse_dim), torch.float32),
78
+ )
79
+
80
81
class TritonMLABackend(MLACommonBackend):
82
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
@@ -166,35 +211,25 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
211
if envs.VLLM_BATCH_INVARIANT:
212
num_kv_splits = 1
213
else:
169
- # Minimum work per split
170
- # hardware dependent
171
- min_work_per_split = 512
172
-
173
- ideal_splits = max(1, attn_metadata.max_seq_len // min_work_per_split)
174
-
175
- # use power of 2 to avoid excessive kernel instantiations
176
- ideal_splits = triton.next_power_of_2(ideal_splits)
177
-
178
- # Calculate SM-based maximum splits with occupancy multiplier
179
- # 2-4x allows multiple blocks per SM for latency hiding
180
- # hardware dependent
181
- occupancy_multiplier = 2
182
- max_splits = self._sm_count * occupancy_multiplier
183
- num_kv_splits = min(ideal_splits, max_splits)
184
-
185
- # TODO(lucas) Allocate ahead of time
186
- attn_logits = torch.empty(
187
- (
188
- B,
189
- q_num_heads,
190
- num_kv_splits,
191
- # NOTE: the +1 stores the LogSumExp (LSE) that the stage2
192
- # kernel uses to merge partial attention outputs across splits.
193
- self.kv_lora_rank + 1,
194
- ),
195
- dtype=torch.float32,
196
- device=q.device,
197
- )
214
+ num_kv_splits = _compute_num_kv_splits(
215
+ attn_metadata.max_seq_len, self._sm_count
216
+ )
217
+
218
+ # NOTE: the +1 stores the LogSumExp (LSE) that the stage2 kernel uses to
219
+ # merge partial attention outputs across splits. The scratch is served
220
+ # from the shared workspace (reserved at max in the metadata builder), so
221
+ # there is no per-call allocation on the decode hot path. Fall back to a
222
+ # direct allocation when the workspace manager is not initialized (e.g.
223
+ # unit tests without a GPUModelRunner).
224
+ logits_shape = (B, q_num_heads, num_kv_splits, self.kv_lora_rank + 1)
225
+ if is_workspace_manager_initialized():
226
+ (attn_logits,) = current_workspace_manager().get_simultaneous(
227
+ (logits_shape, torch.float32),
228
+ )
229
+ else:
230
+ attn_logits = torch.empty(
231
+ logits_shape, dtype=torch.float32, device=q.device
232
+ )
233
234
# Add a head dim of 1
235
kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)