From edaf3e7b613599449fbd2f2ac065869e7e4d358a Mon Sep 17 00:00:00 2001 From: Venky <23023424+venkywonka@users.noreply.github.com> Date: Wed, 30 Sep 2026 19:20:33 -0700 Subject: [PATCH] [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) Co-authored-by: vadiklyutiy <156319763+vadiklyutiy@users.noreply.github.com> --- .../v1/attention/test_triton_attn_segments.py | 73 +++++++++++++++++++ vllm/v1/attention/backends/triton_attn.py | 44 ++++++++++- .../attention/backends/triton_attn_diffkv.py | 2 +- 3 files changed, 114 insertions(+), 5 deletions(-) create mode 100644 tests/v1/attention/test_triton_attn_segments.py diff --git a/tests/v1/attention/test_triton_attn_segments.py b/tests/v1/attention/test_triton_attn_segments.py new file mode 100644 index 0000000000..874c61d237 --- /dev/null +++ b/tests/v1/attention/test_triton_attn_segments.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm.config import CUDAGraphMode +from vllm.platforms import current_platform +from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadataBuilder +from vllm.v1.attention.backends.triton_attn_diffkv import ( + TritonAttentionDiffKVMetadataBuilder, +) +from vllm.v1.kv_cache_interface import FullAttentionSpec + + +# 64 segments only while num_seqs * num_kv_heads * 16 < 188 SMs on SM12.0 and +# the 64-segment scratch holds every query token (e.g. multi-query verify). +@pytest.mark.parametrize( + "builder_cls", + [TritonAttentionMetadataBuilder, TritonAttentionDiffKVMetadataBuilder], +) +@pytest.mark.parametrize( + "capability,num_kv_heads,num_seqs,query_len,segments", + [ + ((12, 0), 1, 11, 1, 64), + ((12, 0), 1, 12, 1, 16), + ((12, 0), 8, 1, 1, 64), + ((12, 0), 8, 2, 1, 16), + ((9, 0), 1, 1, 1, 16), + ((12, 0), 1, 2, 4, 64), + ((12, 0), 1, 10, 4, 16), + ], +) +def test_split_k_segments_follow_sm_occupancy( + monkeypatch, builder_cls, capability, num_kv_heads, num_seqs, query_len, segments +): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + current_platform, "is_device_capability", lambda cap: cap == capability + ) + monkeypatch.setattr(current_platform, "num_compute_units", lambda: 188) + config = SimpleNamespace( + model_config=SimpleNamespace( + get_num_attention_heads=lambda _: 8, + get_num_kv_heads=lambda _: num_kv_heads, + get_head_size=lambda: 128, + rswa_window=None, + ), + parallel_config=None, + scheduler_config=SimpleNamespace(max_num_seqs=16), + speculative_config=None, + compilation_config=SimpleNamespace( + cudagraph_mode=CUDAGraphMode.NONE, static_forward_context={} + ), + ) + spec = FullAttentionSpec( + block_size=16, num_kv_heads=num_kv_heads, head_size=128, dtype=torch.bfloat16 + ) + builder = builder_cls(spec, ["layer.0"], config, "cpu") + batch = BatchSpec(seq_lens=[128] * num_seqs, query_lens=[query_len] * num_seqs) + metadata = builder.build(0, create_common_attn_metadata(batch, 16, "cpu")) + assert metadata.num_par_softmax_segments == segments + assert metadata.softmax_segm_output.shape[2] == segments + assert metadata.softmax_segm_max.shape[0] >= num_seqs * query_len + # DiffKV re-allocates the output; it must stay sized like the base buffers. + assert builder.softmax_segm_output.shape[0] == builder.softmax_segm_max.shape[0] + assert ( + metadata.softmax_segm_output.data_ptr() + == builder.softmax_segm_output.data_ptr() + ) diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 8a2fab4cc7..d8577b61f1 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -18,7 +18,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability -from vllm.utils.math_utils import next_power_of_2 +from vllm.utils.math_utils import cdiv, next_power_of_2 from vllm.utils.torch_utils import get_dtype_size, is_quantized_kv_cache from vllm.v1.attention.backend import ( AttentionBackend, @@ -152,10 +152,29 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet ) self.num_par_softmax_segments = NUM_PAR_SOFTMAX_SEGMENTS + # On SM120, batches whose 16-segment grid under-fills the SMs use 64. + self.max_seqs_64_segments = 0 + if current_platform.is_cuda() and current_platform.is_device_capability( + (12, 0) + ): + self.max_seqs_64_segments = (current_platform.num_compute_units() - 1) // ( + NUM_PAR_SOFTMAX_SEGMENTS * self.num_heads_kv + ) + max_num_tokens_3d = self.seq_threshold_3D + if self.max_seqs_64_segments > 0: + self.num_par_softmax_segments = 64 + # build() reuses this scratch at 16 segments with proportionally more rows. + max_num_tokens_3d = max( + min(self.max_seqs_64_segments, max_num_tokens_3d), + cdiv( + max_num_tokens_3d, + self.num_par_softmax_segments // NUM_PAR_SOFTMAX_SEGMENTS, + ), + ) headdim_padded = next_power_of_2(self.headdim) self.softmax_segm_output = torch.empty( ( - self.seq_threshold_3D, + max_num_tokens_3d, self.num_heads_q, self.num_par_softmax_segments, headdim_padded, @@ -164,12 +183,12 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet device=device, ) self.softmax_segm_max = torch.empty( - (self.seq_threshold_3D, self.num_heads_q, self.num_par_softmax_segments), + (max_num_tokens_3d, self.num_heads_q, self.num_par_softmax_segments), dtype=torch.float32, device=device, ) self.softmax_segm_expsum = torch.empty( - (self.seq_threshold_3D, self.num_heads_q, self.num_par_softmax_segments), + (max_num_tokens_3d, self.num_heads_q, self.num_par_softmax_segments), dtype=torch.float32, device=device, ) @@ -246,6 +265,23 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet softmax_segm_max=self.softmax_segm_max, softmax_segm_expsum=self.softmax_segm_expsum, ) + if self.max_seqs_64_segments > 0 and ( + self.max_seqs_64_segments < seq_lens.shape[0] + or self.softmax_segm_max.shape[0] < num_actual_tokens + ): + # The 16-segment grid fills the SMs, or the 64-segment rows cannot + # hold every query token: same scratch, 16-segment layout. + segments = NUM_PAR_SOFTMAX_SEGMENTS + attn_metadata.num_par_softmax_segments = segments + attn_metadata.softmax_segm_output = self.softmax_segm_output.view( + -1, self.num_heads_q, segments, self.softmax_segm_output.shape[-1] + ) + attn_metadata.softmax_segm_max = self.softmax_segm_max.view( + -1, self.num_heads_q, segments + ) + attn_metadata.softmax_segm_expsum = self.softmax_segm_expsum.view( + -1, self.num_heads_q, segments + ) mm_ranges = common_attn_metadata.mm_req_doc_ranges if mm_ranges is not None: diff --git a/vllm/v1/attention/backends/triton_attn_diffkv.py b/vllm/v1/attention/backends/triton_attn_diffkv.py index ddbd066206..4b42838a15 100644 --- a/vllm/v1/attention/backends/triton_attn_diffkv.py +++ b/vllm/v1/attention/backends/triton_attn_diffkv.py @@ -56,7 +56,7 @@ class TritonAttentionDiffKVMetadataBuilder(TritonAttentionMetadataBuilder): head_size_v_padded = next_power_of_2(head_size_v) self.softmax_segm_output = torch.empty( ( - self.seq_threshold_3D, + self.softmax_segm_max.shape[0], self.num_heads_q, self.num_par_softmax_segments, head_size_v_padded,