[AMD][DSV4] perf: retune decode split-K heuristic for MI355X (#36094)

python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py

@@ -59,9 +59,15 @@ import triton.language as tl
 from aiter.ops.triton.utils.device_info import get_num_sms
 
 from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
+from sglang.srt.utils import is_hip
 
 LOG2E = 1.4426950408889634  # log2(e); folded into qk_scale so softmax can use exp2.
-_MAX_KV_SPLITS = 64  # Hard cap on kv_splits (see _kv_splits_heuristic).
+
+# Split-K heuristic constants. The (1.5, 16) pair is tuned for MI355X (see
+# _kv_splits_heuristic); CUDA is unmeasured here and keeps the prior (2.0, 64).
+_is_hip = is_hip()
+_MAX_KV_SPLITS = 16 if _is_hip else 64  # Hard cap on kv_splits.
+_TARGET_WG_PER_CU = 1.5 if _is_hip else 2.0
 
 # FP8 KV cache (1xGROUP_SIZE block-scale quantization).
 #
@@ -141,7 +147,7 @@ def _kv_splits_heuristic(
     H: int,
     block_h: int,
     num_cu: int | None = None,
-    target_wg_per_cu: float = 2.0,
+    target_wg_per_cu: float = _TARGET_WG_PER_CU,
     max_kv_splits: int = _MAX_KV_SPLITS,
 ) -> int:
     """Pick KV_SPLITS to fill the GPU. CUDAGraph-safe: depends ONLY on
@@ -154,15 +160,30 @@ def _kv_splits_heuristic(
     ``T * ceil(H/block_h)`` underfills the device.
 
       base_ctas  = T * ceil(H / block_h)
-      target_wg  = target_wg_per_cu * num_cu     (≈ 1.7x to hide load-imbalance)
+      target_wg  = target_wg_per_cu * num_cu
       if base_ctas >= target_wg:  splits = 1     (grid already saturates GPU)
       else:                       splits = prev_pow2(min(target_wg/base_ctas,
                                                           max_kv_splits))
 
-    ``max_kv_splits`` (default 64) caps the number of split-kernel CTAs per
-    token. Higher values would buy more parallelism for bs=1 long-ctx, but
-    when per-token K is short most splits fall-through and the launch
-    overhead dominates. 64 is the sweet spot for MI300/MI355.
+    On HIP the tuned ``target_wg_per_cu`` is 1.5 (CUDA keeps 2.0, unmeasured
+    here). At 2.0 the rule over-split by exactly one power of two across the
+    whole decode range on MI355X: at H=128/block_h=64 it chose 8/4/2 splits for
+    T=32/64/128 where 4/2/1 measure faster. Split-K only pays while the base
+    grid underfills the device, and each extra split adds a partial-buffer write
+    plus reduce-kernel work that the shrinking per-split K no longer amortizes.
+
+    ``max_kv_splits`` caps the number of split-kernel CTAs per token (16 on HIP,
+    64 on CUDA). Higher values buy more parallelism for bs=1 long-ctx in
+    principle, but measured optima on MI355X never exceed 16 even at T=1: when
+    per-token K is short most splits fall through and the launch plus reduce
+    overhead dominates.
+
+    Measured over T in {1..256} x kv_len in {128,512,1024} at H=128 on MI355X,
+    scoring each candidate by distance from the per-shape optimum: (2.0, 64)
+    leaves 33.5% geomean regret (119% worst case), (1.5, 16) leaves 3.7% (36%
+    worst). The per-shape optimum does depend on per-token K, which is not
+    knowable at capture time, so the residual is the price of CUDAGraph safety
+    rather than a tuning gap.
 
     Rounded DOWN to a power of two — rounding up over-splits when
     splits_to_fill isn't already pow2 (e.g. T=2 → 258 → 512 doubles the wg

test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py

@@ -0,0 +1,130 @@
+"""Contract tests for the DSv4 decode split-K heuristic.
+
+`_kv_splits_heuristic` runs at CUDAGraph capture time, so it may depend only on
+capture-time scalars and must never read a tensor. These tests pin that
+contract, the invariants the split-K reduction relies on, and the specific
+shapes the tuned constants were chosen for.
+
+They assert properties rather than the constants themselves, so retuning
+`target_wg_per_cu` / `_MAX_KV_SPLITS` for a future architecture does not
+require rewriting the suite -- only the one explicitly-marked table does.
+"""
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.paged_decode import (  # noqa: E402
+    _MAX_KV_SPLITS,
+    _kv_splits_heuristic,
+    _prev_pow2,
+)
+from sglang.test.ci.ci_register import register_cpu_ci  # noqa: E402
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+# MI355X. Passed explicitly so the tests are device-independent and run on CPU.
+NUM_CU = 256
+# head_dim 512 = 448 nope + 64 rope; DSv4 decode runs 128 heads in 64-head tiles.
+HEADS = 128
+BLOCK_H = 64
+
+
+@pytest.mark.parametrize("n", [1, 2, 3, 5, 8, 17, 64, 100, 1000])
+def test_prev_pow2_is_largest_power_of_two_not_exceeding(n):
+    got = _prev_pow2(n)
+    assert got <= n
+    assert got & (got - 1) == 0, f"{got} is not a power of two"
+    assert got * 2 > n, f"{got} is not the largest such power of two for {n}"
+
+
+@pytest.mark.parametrize("n", [0, -1, -100])
+def test_prev_pow2_clamps_non_positive(n):
+    assert _prev_pow2(n) == 1
+
+
+@pytest.mark.parametrize("tokens", [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024])
+def test_result_is_positive_power_of_two_within_cap(tokens):
+    """Split-K allocates a partial buffer per split and the reduce kernel
+    indexes it by power-of-two stride, so a non-pow2 or out-of-range count
+    would corrupt the reduction rather than merely run slowly."""
+    splits = _kv_splits_heuristic(tokens, HEADS, BLOCK_H, num_cu=NUM_CU)
+    assert splits >= 1
+    assert splits <= _MAX_KV_SPLITS
+    assert splits & (splits - 1) == 0, f"{splits} is not a power of two"
+
+
+def test_splits_never_increase_with_token_count():
+    """More decode tokens means a larger base grid, so the device fills without
+    splitting. A count that rose with tokens would over-subscribe the GPU."""
+    prev = None
+    for tokens in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]:
+        splits = _kv_splits_heuristic(tokens, HEADS, BLOCK_H, num_cu=NUM_CU)
+        if prev is not None:
+            assert (
+                splits <= prev
+            ), f"splits rose from {prev} to {splits} going to T={tokens}"
+        prev = splits
+
+
+def test_saturated_grid_does_not_split():
+    """Once the base grid alone meets the target occupancy, splitting can only
+    add partial-buffer writes and reduce work."""
+    # base_ctas = T * ceil(H/block_h); target_wg = target_wg_per_cu * num_cu.
+    saturating_tokens = NUM_CU * 4  # 1024 tokens x 2 head blocks = 2048 CTAs
+    assert _kv_splits_heuristic(saturating_tokens, HEADS, BLOCK_H, num_cu=NUM_CU) == 1
+
+
+def test_does_not_read_tensors_only_capture_time_scalars():
+    """CUDAGraph safety: the heuristic must be callable with plain ints and no
+    CUDA context. Reading kv_indices/kv_indptr here would bake a value from the
+    capture step into every replay."""
+    splits = _kv_splits_heuristic(32, HEADS, BLOCK_H, num_cu=NUM_CU)
+    assert isinstance(splits, int)
+
+
+@pytest.mark.parametrize(
+    "tokens,expected",
+    [
+        # The tuned operating points on MI355X (256 CU, H=128, block_h=64).
+        # base_ctas = tokens * 2; target_wg = int(1.5 * 256) = 384.
+        (8, 16),  # 384//16 = 24 -> capped at 16
+        (16, 8),  # 384//32 = 12 -> prev_pow2 = 8
+        (32, 4),  # 384//64 = 6  -> prev_pow2 = 4
+        (64, 2),  # 384//128 = 3 -> prev_pow2 = 2
+        (128, 1),  # 384//256 = 1
+        (256, 1),  # base grid already saturates
+    ],
+)
+def test_tuned_operating_points(tokens, expected):
+    """Values measured fastest on MI355X. The MI355X constants are passed
+    explicitly so the table holds on any CI runner (CUDA keeps 2.0/64); update
+    alongside the constants if the heuristic is re-tuned -- this is the one test
+    that pins numbers."""
+    assert (
+        _kv_splits_heuristic(
+            tokens,
+            HEADS,
+            BLOCK_H,
+            num_cu=NUM_CU,
+            target_wg_per_cu=1.5,
+            max_kv_splits=16,
+        )
+        == expected
+    )
+
+
+@pytest.mark.parametrize("num_cu", [64, 104, 128, 228, 256, 304])
+def test_scales_with_device_width(num_cu):
+    """A wider device should never ask for more splits at fixed work: splitting
+    exists to fill the device, and a wider one fills at a lower split count
+    only if the base grid grew, which it did not."""
+    splits = _kv_splits_heuristic(32, HEADS, BLOCK_H, num_cu=num_cu)
+    assert 1 <= splits <= _MAX_KV_SPLITS
+    assert splits & (splits - 1) == 0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))