diff --git a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py index fc3f83db7..8774d75f4 100644 --- a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py +++ b/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 diff --git a/test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py b/test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py new file mode 100644 index 000000000..0983e6f91 --- /dev/null +++ b/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"]))