[ROCm][DSV4][Perf] Optimize Triton sparse-MLA decode on gfx950

[Fangzhou-Ai]

Summary

Optimize the existing in-tree Triton DeepSeek-V4 sparse-MLA decode path for
gfx950/MI355X.

The change adds a dedicated gfx950 partial kernel, workload-aware split
selection, graph-safe adaptive work selection, and gfx950-only cp-gather
de-specialization. It adds no AITER/Gluon runtime dependency or implementation
switch; corrected AITER/Gluon is used only as a performance comparator.

The PR remains draft while refreshed CI runs and the human submitter completes
the final line-by-line review.

gfx950-only scope

All new runtime behavior is explicitly gated by _ON_GFX950:

• the specialized sparse-decode partial kernel and up-to-32 split policy;
• adaptive captured-graph split selection and narrowed ragged views;
• direct BF16 output;
• the compressed-cache NaN-free fast path; and
• cp-gather token/batch de-specialization.

gfx942 retains its existing 1-16 split selector, partial/ragged kernels,
cp-gather specialization, reducer behavior, and output path. A read-only gfx942
AOT audit found byte-identical HSACO text versus the clean baseline.

Design

The gfx950 kernel:

• combines 448 NoPE and 64 RoPE dimensions for one 512-wide QK dot;
• uses four 128-wide PV accumulators to avoid spilling;
• processes main/SWA work in guarded BK32 tiles;
• peels full 64-entry compressed-cache groups into two ordered BK32 updates;
• preserves masked tails, attention sinks, and invalid-index semantics; and
• uses a fixed captured maximum grid while choosing live per-row work.

The split selector targets device occupancy, caps short workloads at four
splits, and uses up to 32 splits for low-batch long-context rows.

Compressed-cache scrub elision requires an explicit, false-by-default
provenance bit. On gfx950, both writers cap UE8M0 codes at 254 and canonicalize
RoPE NaNs. Main-cache values are still scrubbed, invalid indices contribute
zero probability, and cache block address multiplication uses int64.

Performance

Kernel acceptance

On one 256-CU MI355X, the exact 40-shape matrix
T={1,2,4,8,16,32,48,64} x context={4K,32K,100K,500K,1M} passed eager and
warmed-graph correctness and the <=1.03 performance gate:

Comparator	Accepted shapes	Worst measured ratio
Clean in-tree Triton	40/40	1.000000
Corrected AITER/Gluon	40/40	1.023918

All 1,920 correctness comparisons passed with atol=rtol=0.02.

The final source-only cleanup was rechecked on four representative graph
shapes. Outputs were bitwise identical, p50 changed by -0.15% to +0.18%, and
the S8/S32 executable text and resource usage were unchanged.

Canonical 8k/1k serving A/B

Configuration: 8x MI355X, TP8/DP1, concurrency 64, prefix caching disabled,
128 warmups followed by 640 scored requests, and fresh isolated compile caches.
The reference used the exact old sparse-decode source plus the same cp-gather
fix as the candidate. Input/output length arrays were byte-identical.

Metric	Old sparse + common cp fix	Candidate	Delta
Output throughput	1,155.242 tok/s	1,187.054 tok/s	+2.754%
Duration	510.652 s	496.967 s	-2.680%
Median TTFT	451.565 ms	445.356 ms	-1.375%
P90 TTFT	1,508.041 ms	1,569.744 ms	+4.092%
P99 TTFT	21,089.724 ms	19,419.572 ms	-7.919%
P99.9 TTFT	25,961.827 ms	24,709.868 ms	-4.822%

Both arms completed all requests without errors, preemption, or prefix-cache
activity. Every reported P99/P99.9 latency improved; P90 TTFT was the sole
regressed percentile. This is a positive paired no-regression result, not a
claim of statistical consistency from one run per arm.

Full results and methodology are in the
8k/1k follow-up comment.

Validation

ROCR_VISIBLE_DEVICES=6 HIP_VISIBLE_DEVICES=0 \
  .venv/bin/python -m pytest \
  tests/kernels/attention/test_rocm_triton_attn_dsv4.py -vv

ROCR_VISIBLE_DEVICES=6 HIP_VISIBLE_DEVICES=0 \
  .venv/bin/python -m pytest tests/kernels/test_compressor_kv_cache.py \
  -k 'cp_gather_despecialized or gfx950_compressed_cache or \
      indexer_gather_accepts_upper_bound_output' -vv

pre-commit run --files <changed files>
pre-commit run mypy-3.12 --all-files --hook-stage manual

Results on the current head: 57 attention tests passed, four focused
compressor/cp tests passed, all changed-file pre-commit hooks passed, and
Python 3.12 mypy passed.

A full GSM8K run on the preceding pure-Triton source completed all 1,319
examples with strict accuracy 0.960576 and flexible accuracy 0.959818.

Duplicate-work check

Issue #41820 and open PRs were checked for 41820 in:body and
DSV4 sparse MLA Triton. No open PR duplicates this change. #50566 focuses on
reusable split-K scratch and output-copy removal; #51714 adds an optional
external AITER/Gluon path. This PR instead optimizes the default in-tree Triton
implementation on gfx950.

AI assistance and accountability

OpenAI Codex assisted with investigation, implementation, test development,
benchmarking, auditing, and drafting. The human submitter remains responsible
for reviewing every changed line, understanding and defending the change,
rerunning the relevant tests, and completing final CI review.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[AndreasKaratzas]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83840 for commit f3818657a2f0.

[AndreasKaratzas]

LGTM

[jiacao-amd]

@Fangzhou-Ai

Metric	nightly-f8d03e774	AITER Gluon	Triton
Total throughput/GPU	10,089.64 tok/s	11,531.40 tok/s	11,560.29 tok/s
Change vs nightly	—	+14.29%	+14.58%
Input throughput/GPU	10,023.86 tok/s	11,450.33 tok/s	11,478.62 tok/s
Output throughput/GPU	65.78 tok/s	81.06 tok/s	81.67 tok/s
Request rate	0.836 req/s	0.903 req/s	0.905 req/s
Profiled requests	1,025	1,104	1,106
GPU cache hit rate	90.07%	90.586%	90.578%
Mean TTFT	1.632 s	1.653 s	1.760 s
Mean ITL/TPOT	45.38 ms	35.07 ms	34.67 ms
Mean E2E latency	30.09 s	26.99 s	26.87 s
Errors	0	0	0

Under the same 8× MI355X TP8/EP1, c32, 1200s, warmup=10, and FSE-enabled setup, Triton and AITER Gluon are effectively tied, with Triton delivering 0.25% higher throughput.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Fangzhou-Ai.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Fangzhou-Ai]

8k/1k TP8 follow-up

I reran the canonical 8k/1k serving workload after addressing the two issues
found in the high-concurrency investigation:

1. The gfx950 split selector saw the graph-stable padded C128 backing capacity
   rather than live per-row work. The follow-up uses a fixed captured maximum
   grid/workspace but selects 4, one-wave, or maximum active splits per row.
2. _cp_gather_indexer_quant_cache_kernel specialized on the changing token
   count/batch count. That produced hundreds of new Triton binaries during the
   scored run. The follow-up removes the redundant token-count specialization
   and makes the batch count a runtime scalar.

The follow-up is now pushed in 5dfdac01ed, after merging current upstream
main in 464c87dcde. The serving A/B below was run on the pre-merge local
candidate (d98282de plus worktree diff SHA-256 0ee35d4b...); the final
follow-up preserves that gfx950 runtime path while restoring legacy gfx942
selector/kernel behavior. The clean reference is b2e748a plus the common
cp-gather-only diff 320e669c....

Configuration

• 8x MI355X (gfx950), TP=8, DP=1
• no MTP/speculative decoding
• prefix caching disabled (0 queries/hits in both arms)
• max batched tokens 8192, max sequences 1024
• FULL_AND_PIECEWISE, async scheduling, FP8 KV
• random ISL 8192 / OSL 1024, range ratio 0.8, seed 0
• concurrency 64, request rate unlimited
• 128 warmups followed by 640 scored requests
• fresh isolated compile caches for each arm

The 640-element input/output length arrays were byte-identical between arms
(4,727,544 input tokens and 589,927 output tokens).

For a fair comparison, the reference is exact old source b2e748a plus the
same cp-gather de-specialization. Therefore the large per-shape JIT issue is
held constant while comparing old sparse decode against the local follow-up.

Metric	Old sparse + common cp fix	Local follow-up	Delta
Duration	510.652 s	496.967 s	-2.680%
Aggregate output throughput	1,155.242 tok/s	1,187.054 tok/s	+2.754%
Aggregate total-token throughput	10,413.095 tok/s	10,699.843 tok/s	+2.754%
Median TTFT	451.565 ms	445.356 ms	-1.375%
P90 TTFT	1,508.041 ms	1,569.744 ms	+4.092%
P99 TTFT	21,089.724 ms	19,419.572 ms	-7.919%
P99.9 TTFT	25,961.827 ms	24,709.868 ms	-4.822%
Median TPOT	52.911 ms	51.688 ms	-2.313%
P99 TPOT	70.503 ms	68.057 ms	-3.468%
P99.9 TPOT	74.218 ms	71.555 ms	-3.588%
Median ITL	29.708 ms	28.420 ms	-4.336%
P99 ITL	385.875 ms	382.761 ms	-0.807%
P99.9 ITL	397.282 ms	392.829 ms	-1.121%
Median E2E	49,496.371 ms	48,372.896 ms	-2.270%
P99 E2E	80,603.964 ms	77,481.556 ms	-3.874%
P99.9 E2E	83,289.984 ms	80,166.423 ms	-3.750%

Both arms completed 640/640 scored requests and 768/768 total requests
including warmups, with zero errors, aborts, preemptions, or prefix-cache
activity. Each arm compiled exactly one cp-gather binary during warmup and none
during scoring. The old arm had three additional scored Triton variants
(partial/reduce/logits), accounting for about 2 seconds versus the 13.685-second
duration gap, so they do not explain the result.

Result JSON SHA-256: reference a5e03f24...; candidate 3c287b65....

The local follow-up is also +0.808% in output throughput versus the retained
historical Reduced_padding_i384 result (1,187.054 vs 1,177.541 tok/s), although
the common-fix A/B above is the cleaner cache/provenance comparison.

Kernel/correctness validation

• Production-shaped adaptive split sweeps: 64 broad cases plus 17 dense
  q16..32/64K cases, 1,000 graph replays per case; all outputs matched the
  reference and there were no old-relative regression-threshold breaches.
• Isolated graph replay at 8K: q16 improved about 16.0%, q64 about 11.6%; q8
  remained on the guarded static path and 1M contexts remained unchanged.
• CP-gather fresh-cache sweep: 24 old binaries -> 1 patched binary; aggregate
  warm p50/p99/p99.9 changed -1.45%/-0.64%/-2.82%.
• test_rocm_triton_attn_dsv4.py: 64 passed on the final merged head.
• Focused cp-gather correctness test passed; pre-commit and Python 3.12 mypy
  passed.

The full-model numbers above are one fresh paired run per arm, so I am treating
them as a positive no-regression result rather than a statistical consistency
claim. The only regressed reported percentile is P90 TTFT (+61.7 ms); throughput
and every P99/P99.9 metric improved.

AI assistance was used for the patch implementation, test development, and
benchmark analysis. The submitting human remains responsible for reviewing
every changed line before the draft is marked ready.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[AndreasKaratzas]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84060 for commit 39ed5620c36f.