[AMD][DSV4] perf: retune decode split-K heuristic for MI355X

[karverma-amd]

Motivation

_kv_splits_heuristic picks how many KV splits DSv4 decode attention launches.
It over-split by exactly one power of two across the whole decode range: 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. Each extra split
adds a partial-buffer write plus reduce-kernel work, and once per-split K gets
short that overhead is no longer amortized.

Modification

Two constants in one file: target_wg_per_cu 2.0 -> 1.5 and the
_MAX_KV_SPLITS cap 64 -> 16. No kernel changes, no new code paths.

Swept 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:

constants	geomean regret	worst case
(2.0, 64) — current	33.5%	119%
(1.5, 16) — this PR	3.7%	36%

The per-shape optimum does depend on per-token K, which is not knowable at
CUDAGraph capture time, so the residual 3.7% is the price of capture-time
safety rather than a tuning gap.

Kernel level

Decode attention at the production shape (T=32, top-k 1024, H=128):
55.7 us -> 44.1 us, a 20.8% improvement. Selection is unchanged, so results
are identical modulo bf16 reassociation of the split reduction.

Performance

DeepSeek-V4-Pro, MI355X (gfx950), rocm/sgl-dev:v0.5.18-rocm720-mi35x-20260822,
8192 in / 1024 out, TP8 + DP8 attention + EP8/MoRI + EAGLE MTP. Both arms are the
same image and the same flags; the only difference is the two constants.

conc	tput base (tok/s)	tput PR (tok/s)	tput	TPOT base (ms)	TPOT PR (ms)	TPOT
4	860.07	873.28	+1.5%	39.79	39.11	-1.7%
8	1508.46	1545.51	+2.5%	42.92	42.17	-1.8%
16	2590.30	2610.88	+0.8%	50.02	49.62	-0.8%
32	4001.04	4073.41	+1.8%	65.30	64.24	-1.6%
64	5539.24	5776.29	+4.3%	96.41	92.41	-4.2%

Throughput rises and TPOT falls at every concurrency, and the two move together
at each point, which is harder to explain as drift than either alone. The gain
grows with concurrency because more tokens per rank makes decode attention a
larger share of the step.

One run per cell. Note the effect is specific to DP attention: with plain TP8,
the 128 heads sharded to 16 per rank, the same build measures flat (all deltas
under 1%, signs alternating). DP attention keeps all 128 heads resident per rank,
which is the shape the constants were tuned on.

Commands

# Server, per arm; $CONC in {4, 8, 16, 32, 64}
SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton \
SGLANG_OPT_FUSE_COMPRESS_NORM_ROPE=1 \
SGLANG_OPT_NATIVE_BPRESHUFFLE_SCALE=1 \
SGLANG_OPT_USE_AITER_BATCHED_GEMM=1 \
AITER_BF16_FP8_MOE_BOUND=0 \
SGLANG_USE_AITER=1 \
SGLANG_SHARED_EXPERT_TP1=1 SGLANG_DP_SHARED_EXPERT_LOCAL=1 \
SGLANG_DP_USE_GATHERV=1 SGLANG_DP_USE_REDUCE_SCATTER=1 \
GPU_MAX_HW_QUEUES=5 MORI_SHMEM_MODE=ISOLATION MORI_SHMEM_HEAP_SIZE=8G \
MORI_EP_LAUNCH_CONFIG_MODE=AUTO \
python3 -m sglang.launch_server \
  --model-path deepseek-ai/DeepSeek-V4-Pro \
  --tensor-parallel-size 8 --dp 8 --enable-dp-attention \
  --ep-size 8 --moe-a2a-backend mori --deepep-mode normal \
  --moe-dense-tp-size 1 --enable-dp-lm-head \
  --attention-backend dsv4 --page-size 256 --kv-cache-dtype fp8_e4m3 \
  --enforce-shared-experts-fusion \
  --speculative-algorithm EAGLE --speculative-num-steps 3 \
  --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 \
  --context-length 16384 --chunked-prefill-size 8192 \
  --mem-fraction-static 0.78 \
  --max-running-requests $CONC --cuda-graph-max-bs $CONC

# Perf client, per concurrency
python3 -m sglang.bench_serving --backend sglang-oai \
  --host 127.0.0.1 --port $PORT --model deepseek-ai/DeepSeek-V4-Pro \
  --dataset-name random --random-input-len 8192 --random-output-len 1024 \
  --random-range-ratio 1.0 --num-prompts $((CONC * 10)) --max-concurrency $CONC

# Accuracy
python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 --parallel 1319

Note: under DP attention --max-running-requests is divided by dp_size, so a
value below the DP degree floors to 0 per rank and trips
assert max_running_request > 0. At $CONC 4 it was raised to 8; the client's
--max-concurrency is what bounds the load either way.

Accuracy

GSM8K, 1319 questions, parallel 1319: 0.958 and 0.946 across two runs,
invalid 0.000 both times.

Unit test

_kv_splits_heuristic had no test. This PR adds
test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py (38 cases),
pinning the contract rather than the constants:

• result is always a positive power of two within the cap — the reduce kernel
  indexes partial buffers by pow2 stride, so a bad value corrupts results rather
  than merely running slowly
• splits never rise as token count rises
• no splitting once the base grid saturates the device
• CUDAGraph safety: depends only on capture-time scalars, never reads a tensor

One clearly-marked table pins the six tuned operating points, so retuning for a
future architecture touches that test and nothing else.

pytest test/registered/unit/layers/test_dsv4_kv_splits_heuristic.py

Checklist

• Kernel-level and end-to-end benchmarked on MI355X (tables above)
• Accuracy verified, 1319 questions, zero invalid
• Unit test added for a previously untested function
• No new code paths, no monkeypatching; two constants in one file
• CUDA path unaffected — this file is the ROCm-only dsv4 unified_kv backend

---

CI States

Latest PR Test (Base): ❌ Run #32750870340
Latest PR Test (Extra): ❌ Run #32750870208
Latest PR Test (AMD ROCm 7.2): ⏳ Run #32750870441

[HaiShaw]

/tag-and-rerun-ci