[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