[Attention][SM120] Occupancy-adaptive Triton split-K segment count (#58482)
Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: vadiklyutiy <156319763+vadiklyutiy@users.noreply.github.com>
Venky committed
Sep 30, 2026 at 19:20 UTC
edaf3e7b613599449fbd2f2ac065869e7e4d358a
3 files changed
+114
-5
tests/v1/attention/test_triton_attn_segments.py
new
+73
@@ -0,0 +1,73 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
10
+from vllm.config import CUDAGraphMode
11
+from vllm.platforms import current_platform
12
+from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadataBuilder
13
+from vllm.v1.attention.backends.triton_attn_diffkv import (
14
+ TritonAttentionDiffKVMetadataBuilder,
15
+)
16
+from vllm.v1.kv_cache_interface import FullAttentionSpec
17
+
18
+
19
+# 64 segments only while num_seqs * num_kv_heads * 16 < 188 SMs on SM12.0 and
20
+# the 64-segment scratch holds every query token (e.g. multi-query verify).
21
+@pytest.mark.parametrize(
22
+ "builder_cls",
23
+ [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder],
24
+)
25
+@pytest.mark.parametrize(
26
+ "capability,num_kv_heads,num_seqs,query_len,segments",
27
+ [
28
+ ((12, 0), 1, 11, 1, 64),
29
+ ((12, 0), 1, 12, 1, 16),
30
+ ((12, 0), 8, 1, 1, 64),
31
+ ((12, 0), 8, 2, 1, 16),
32
+ ((9, 0), 1, 1, 1, 16),
33
+ ((12, 0), 1, 2, 4, 64),
34
+ ((12, 0), 1, 10, 4, 16),
35
+ ],
36
+)
37
+def test_split_k_segments_follow_sm_occupancy(
38
+ monkeypatch, builder_cls, capability, num_kv_heads, num_seqs, query_len, segments
39
+):
40
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
41
+ monkeypatch.setattr(
42
+ current_platform, "is_device_capability", lambda cap: cap == capability
43
+ )
44
+ monkeypatch.setattr(current_platform, "num_compute_units", lambda: 188)
45
+ config = SimpleNamespace(
46
+ model_config=SimpleNamespace(
47
+ get_num_attention_heads=lambda _: 8,
48
+ get_num_kv_heads=lambda _: num_kv_heads,
49
+ get_head_size=lambda: 128,
50
+ rswa_window=None,
51
+ ),
52
+ parallel_config=None,
53
+ scheduler_config=SimpleNamespace(max_num_seqs=16),
54
+ speculative_config=None,
55
+ compilation_config=SimpleNamespace(
56
+ cudagraph_mode=CUDAGraphMode.NONE, static_forward_context={}
57
+ ),
58
+ )
59
+ spec = FullAttentionSpec(
60
+ block_size=16, num_kv_heads=num_kv_heads, head_size=128, dtype=torch.bfloat16
61
+ )
62
+ builder = builder_cls(spec, ["layer.0"], config, "cpu")
63
+ batch = BatchSpec(seq_lens=[128] * num_seqs, query_lens=[query_len] * num_seqs)
64
+ metadata = builder.build(0, create_common_attn_metadata(batch, 16, "cpu"))
65
+ assert metadata.num_par_softmax_segments == segments
66
+ assert metadata.softmax_segm_output.shape[2] == segments
67
+ assert metadata.softmax_segm_max.shape[0] >= num_seqs * query_len
68
+ # DiffKV re-allocates the output; it must stay sized like the base buffers.
69
+ assert builder.softmax_segm_output.shape[0] == builder.softmax_segm_max.shape[0]
70
+ assert (
71
+ metadata.softmax_segm_output.data_ptr()
72
+ == builder.softmax_segm_output.data_ptr()
73
+ )
vllm/v1/attention/backends/triton_attn.py
+40
-4
@@ -18,7 +18,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
)
19
from vllm.platforms import current_platform
20
from vllm.platforms.interface import DeviceCapability
21
-from vllm.utils.math_utils import next_power_of_2
21
+from vllm.utils.math_utils import cdiv, next_power_of_2
22
from vllm.utils.torch_utils import get_dtype_size, is_quantized_kv_cache
23
from vllm.v1.attention.backend import (
24
AttentionBackend,
@@ -152,10 +152,29 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
152
)
153
154
self.num_par_softmax_segments = NUM_PAR_SOFTMAX_SEGMENTS
155
+ # On SM120, batches whose 16-segment grid under-fills the SMs use 64.
156
+ self.max_seqs_64_segments = 0
157
+ if current_platform.is_cuda() and current_platform.is_device_capability(
158
+ (12, 0)
159
+ ):
160
+ self.max_seqs_64_segments = (current_platform.num_compute_units() - 1) // (
161
+ NUM_PAR_SOFTMAX_SEGMENTS * self.num_heads_kv
162
+ )
163
+ max_num_tokens_3d = self.seq_threshold_3D
164
+ if self.max_seqs_64_segments > 0:
165
+ self.num_par_softmax_segments = 64
166
+ # build() reuses this scratch at 16 segments with proportionally more rows.
167
+ max_num_tokens_3d = max(
168
+ min(self.max_seqs_64_segments, max_num_tokens_3d),
169
+ cdiv(
170
+ max_num_tokens_3d,
171
+ self.num_par_softmax_segments // NUM_PAR_SOFTMAX_SEGMENTS,
172
+ ),
173
+ )
174
headdim_padded = next_power_of_2(self.headdim)
175
self.softmax_segm_output = torch.empty(
176
(
158
- self.seq_threshold_3D,
177
+ max_num_tokens_3d,
178
self.num_heads_q,
179
self.num_par_softmax_segments,
180
headdim_padded,
@@ -164,12 +183,12 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
183
device=device,
184
)
185
self.softmax_segm_max = torch.empty(
167
- (self.seq_threshold_3D, self.num_heads_q, self.num_par_softmax_segments),
186
+ (max_num_tokens_3d, self.num_heads_q, self.num_par_softmax_segments),
187
dtype=torch.float32,
188
device=device,
189
)
190
self.softmax_segm_expsum = torch.empty(
172
- (self.seq_threshold_3D, self.num_heads_q, self.num_par_softmax_segments),
191
+ (max_num_tokens_3d, self.num_heads_q, self.num_par_softmax_segments),
192
dtype=torch.float32,
193
device=device,
194
)
@@ -246,6 +265,23 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
265
softmax_segm_max=self.softmax_segm_max,
266
softmax_segm_expsum=self.softmax_segm_expsum,
267
)
268
+ if self.max_seqs_64_segments > 0 and (
269
+ self.max_seqs_64_segments < seq_lens.shape[0]
270
+ or self.softmax_segm_max.shape[0] < num_actual_tokens
271
+ ):
272
+ # The 16-segment grid fills the SMs, or the 64-segment rows cannot
273
+ # hold every query token: same scratch, 16-segment layout.
274
+ segments = NUM_PAR_SOFTMAX_SEGMENTS
275
+ attn_metadata.num_par_softmax_segments = segments
276
+ attn_metadata.softmax_segm_output = self.softmax_segm_output.view(
277
+ -1, self.num_heads_q, segments, self.softmax_segm_output.shape[-1]
278
+ )
279
+ attn_metadata.softmax_segm_max = self.softmax_segm_max.view(
280
+ -1, self.num_heads_q, segments
281
+ )
282
+ attn_metadata.softmax_segm_expsum = self.softmax_segm_expsum.view(
283
+ -1, self.num_heads_q, segments
284
+ )
285
286
mm_ranges = common_attn_metadata.mm_req_doc_ranges
287
if mm_ranges is not None:
vllm/v1/attention/backends/triton_attn_diffkv.py
+1
-1
@@ -56,7 +56,7 @@ class TritonAttentionDiffKVMetadataBuilder(TritonAttentionMetadataBuilder):
56
head_size_v_padded = next_power_of_2(head_size_v)
57
self.softmax_segm_output = torch.empty(
58
(
59
- self.seq_threshold_3D,
59
+ self.softmax_segm_max.shape[0],
60
self.num_heads_q,
61
self.num_par_softmax_segments,
62
head_size_v_padded,