[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>

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()
+    )

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:

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,