[CUDA] Tune hardware-aware split-K decode scheduling

backends/cuda/tests/test_triton_sdpa_splitk.py

@@ -314,7 +314,7 @@ def test_qwen35_config(self):
     def test_large_positive_logits_stable(self):
         """Large logits stay finite in both split-K kernel families."""
         B, H_q, H_kv, Lk, D = 1, 32, 8, 512, 128
-        num_splits, chunk_size = self.splitk_config(Lk)
+        num_splits, chunk_size = self.splitk_config(Lk, B * H_kv, torch.device("cuda"))
         self.assertGreater(num_splits, 1)
         self.assertNotEqual(0 // chunk_size, 300 // chunk_size)
 
@@ -391,8 +391,8 @@ def test_large_negative_logits_do_not_underflow_to_zero(self):
     def test_non_power_of_two_split_count(self):
         """Reduction masks lanes beyond the runtime split count."""
         B, H_q, H_kv, Lk, D = 1, 8, 2, 768, 128
-        num_splits, _ = self.splitk_config(Lk)
-        self.assertEqual(num_splits, 3)
+        num_splits, _ = self.splitk_config(Lk, B * H_kv, torch.device("cuda"))
+        self.assertEqual(num_splits, 6)
 
         torch.manual_seed(42)
         k = torch.randn(B, H_kv, Lk, D, dtype=torch.bfloat16, device="cuda")

backends/cuda/triton/kernels/sdpa.py

@@ -31,6 +31,7 @@
 and improves tile utilization, especially during decode (seqlen_q=1).
 """
 
+import functools
 import math
 from typing import Optional
 
@@ -51,8 +52,71 @@ def _is_power_of_2(n: int) -> bool:
 _SPLITK_LKV_THRESHOLD = 256
 
 
-def _decode_splitk_config(L_kv: int) -> tuple[int, int]:
-    num_splits = min(max(triton.cdiv(L_kv, _SPLITK_LKV_THRESHOLD), 1), 128)
+# Decode split-K occupancy target. A sweep across both production attention
+# families showed that targeting 16/9 waves gives a good balance between split
+# kernel occupancy and reduction/empty-CTA overhead. Keep the ratio integral so
+# the launch policy is deterministic and does not depend on floating-point
+# rounding.
+_SPLITK_TARGET_WAVES_NUMERATOR = 16
+_SPLITK_TARGET_WAVES_DENOMINATOR = 9
+_SPLITK_MAX_SPLITS = 128
+
+# Do not create more static split CTAs than one per 128 elements in the KV
+# buffer. This permits 16 splits for the 2048-element sliding-window family.
+# A 40-state, >L2 working-set A/B showed gains at the production prompt lengths;
+# single-state measurements are cache-hot and are not representative here.
+# For the 128K global family with grid_y=4 on a 170-SM device, the occupancy
+# target selects 76 splits instead of reaching the merge-base buffer cap of 128.
+_SPLITK_BUFFER_ELEMENTS_PER_SPLIT = 128
+
+
+@functools.lru_cache(maxsize=None)
+def _device_sm_count(device_index: int) -> int:
+    """SM (multiprocessor) count of a CUDA device.
+
+    Read once per device from torch device properties and bake it into the
+    exported artifact before cuda-graph capture. The supported deployment flow
+    exports a native artifact per target GPU and may merge those variants into
+    one PTE; moving one native artifact to a same-architecture GPU with a
+    different SM count remains correct, but may not retain optimal scheduling.
+    """
+    return torch.cuda.get_device_properties(device_index).multi_processor_count
+
+
+def _decode_splitk_config(
+    L_kv: int, grid_y: int, device: torch.device
+) -> tuple[int, int]:
+    """Hardware-derived split count for the flash-decoding decode path.
+
+    The split kernel launches a (num_splits, grid_y) grid where
+    grid_y = B * H_kv. We size num_splits so the total launched CTAs
+    (num_splits * grid_y) target 16/9 waves of this device's SM array, capped by
+    cdiv(L_kv, 128) so a short KV buffer is not needlessly over-split and
+    clamped to >= 1.
+
+    num_splits depends only on host-side constants (L_kv buffer size, the
+    static grid_y, and the device SM count), so it is fixed before cuda-graph
+    capture and identical on every replay.
+    """
+    device_index = (
+        device.index if device.index is not None else torch.cuda.current_device()
+    )
+    n_sm = _device_sm_count(device_index)
+    # Cap so num_splits * grid_y ~= 16/9 * n_sm (>= 1). Use ceil division: a
+    # partial final wave is preferable to leaving an SM idle.
+    sm_cap = max(
+        triton.cdiv(
+            _SPLITK_TARGET_WAVES_NUMERATOR * n_sm,
+            _SPLITK_TARGET_WAVES_DENOMINATOR * max(grid_y, 1),
+        ),
+        1,
+    )
+    buffer_cap = max(triton.cdiv(L_kv, _SPLITK_BUFFER_ELEMENTS_PER_SPLIT), 1)
+    # Avoid letting very high-SM GPUs over-split long global-attention buffers:
+    # beyond 128 partitions the extra reduction and launch work outweighed the
+    # occupancy gain in CUDA-graph measurements. Lower-SM devices continue to
+    # use their hardware-derived cap.
+    num_splits = min(buffer_cap, sm_cap, _SPLITK_MAX_SPLITS)
     return num_splits, triton.cdiv(L_kv, num_splits)
 
 
@@ -1256,6 +1320,7 @@ def _sdpa_decode_splitk_kernel(
     NUM_GROUPS: tl.constexpr,
     BLOCK_G: tl.constexpr,
 ):
+    sm_scale = sm_scale.to(tl.float32)
     split_id = tl.program_id(axis=0)
     pid_bh = tl.program_id(axis=1)
     b = pid_bh // H_kv
@@ -1437,7 +1502,7 @@ def _launch_decode_splitk(
     kv_len_ptr: Optional[torch.Tensor] = None,
     HAS_KV_LEN: bool = False,
 ) -> None:
-    num_splits, chunk_size = _decode_splitk_config(L_kv)
+    num_splits, chunk_size = _decode_splitk_config(L_kv, B * H_kv, query.device)
 
     # Each valid (split, batch, query head) slot belongs to exactly one program,
     # and empty KV ranges still reach the stores with m=-inf and l=acc=0.
@@ -1717,6 +1782,7 @@ def _sdpa_small_query_splitk_kernel(
     BLOCK_G: tl.constexpr,
     BLOCK_M: tl.constexpr,
 ):
+    sm_scale = sm_scale.to(tl.float32)
     split_id = tl.program_id(axis=0)
     pid_bh = tl.program_id(axis=1)
     b = pid_bh // H_kv
@@ -1955,7 +2021,7 @@ def _launch_small_query_splitk(
     HAS_KV_LEN: bool = False,
     mask_is_causal: bool = False,
 ) -> None:
-    num_splits, chunk_size = _decode_splitk_config(L_kv)
+    num_splits, chunk_size = _decode_splitk_config(L_kv, B * H_kv, query.device)
 
     # Each valid (split, batch, query head, query row) slot belongs to exactly
     # one program, and empty KV ranges still store m=-inf and l=acc=0.