[ROCm] Optimize HY V4 graph inference on gfx942 and gfx950

[andyluo7]

[ROCm] Optimize HY V4 graph inference on gfx942 and gfx950

Purpose

This PR remains stacked on #54405 (and therefore #54404). This update adds four
targeted changes to the HY V4 ROCm graph-mode serving path:

• use AITER's fused biased sigmoid top-k for HY V4's one-group routing shape;
• make the AITER MXFP8 routing buffers contiguous before fused MoE consumes
  them;
• fuse the BF16 elementwise MLA output gate into one graph-safe Triton op; and
• apply attention sinks with a grouped, graph-safe kernel that also repairs
  empty graph-padded rows.

Unsupported AITER sink layouts are now rejected during construction instead of
on the first request. The shared NVIDIA attention path retains the same eager
gate expression through a small subclass hook.

Platform isolation

• The direct top-k route is guarded by the one-group sigmoid routing
  configuration and HY V4's measured hidden/expert/top-k shape.
• The fused MLA gate is in the ROCm subclass and requires matching contiguous
  BF16 CUDA tensors with the measured width; all other inputs use the existing
  eager expression.
• The attention-sink kernel remains inside the ROCm AITER sparse-MLA backend.
• NVIDIA still executes attn_out * torch.sigmoid(gate_score).
• Model-class dispatch remains capability/platform based. The only gfx-specific
  branch in this update is the AITER sink compatibility check that rejects the
  unsupported gfx942 BF16 64-head bucket.

Exact candidate and runtime

• Published optimization commit:
  53abd3cac3e06993a8afae86b3246fa27d0ab1b1
• Published commit tree: 93b4fced3bf3dfaf67cf9d5f6217240189ba67f1
• Optimization binary-patch SHA-256:
  bf8b13d562587b6b5a4437f4822985d74e6d350091afffedaf1c1440a0638657
• MI350X source manifest SHA-256:
  b33a8fa5bc83e646f67dfe4c23ec3bec8395c3de4e8ecca0ff0678f409ec0a13
• MI350X extension manifest SHA-256:
  a22a8a03367f53d095a541b79d25a469a48d77aeda75d2c99beaa43acea7430b
• MI350X generated-runtime manifest SHA-256:
  d1875bc16df989dd514620ed24575f8e2fdc6e0a6838f391563be69c003fdc67
• MI350X runtime image:
  sha256:03f385c0cdc48a0031aba05e1d3d7b6fa0b51a6744f9b69e6c8cd7cd4de4b49b
• MI350X runtime parent image:
  sha256:40e19c756e3dc9ffc9117770904d40376c7d3bf529cc76ddc379cde7ac4dae2d
• gfx942 source/extension/generated-runtime manifest SHA-256 values:
  6711d26262846e35aad24d1540985c3b9aceea51dbde3dfe36eba72ce4fc023b,
  83aaee20d27b06931a18f639ba0aadcea268e2262077ffc53dcb2b055abea9e1,
  and d719add875735eefb97258888a92393a48bf0d5bcdcaf6dd89bab6bb09928214.
• gfx942 runtime image:
  sha256:8908b8ab5ba28c3b81f9f42bb72e2421f06a180e001c67c4f10ff7f127c5690b
• Performance-runtime prerequisite: open [HIP] [CK] [FlyDSL] FlyDSL split-K preshuffle decode GEMM for small-M shapes (FP8 + MX, gfx950/gfx942) ROCm/aiter#5148 at
  f8500b8543949ec5c856a9395956fb4927948211
• FlyDSL: 0.2.4, wheel SHA-256
  3ba78e6c1cec37c99b83d71a34c1ed50f1058231bcab1a352a1b431082fb2c4f
• AITER FMoE tuning file SHA-256:
  1f05eadfb3654b7d4adb059f069fe1d0de7d4ce24d9e8cb3e1b6f23699b92534
• Model revision: 4215ec29de873a998e849cee902654490c7ff4d1
• GSM8K dataset SHA-256:
  67e9046d913470477dfb321814d4e3240240a397cfc19c6c610af02dd2cb6456
• sgl-eval revision: a231b7a439b235090ff7baa30778fa2b514309ae

Validation

Static and exact-tree correctness:

• Touched-file pre-commit: passed.
• MI350X gfx950 r163: 221 passed, 2 skipped; source/runtime postflight,
  external-user, kernel-error, and clean teardown checks passed. Artifact
  manifest SHA-256:
  b435d4213b0a0d8baef5b34cf271271a7e9f325eca20de473855bd30ec506514.
• gfx942 r155: 210 passed, 13 skipped; source/runtime postflight,
  external-user, kernel-error, and clean teardown checks passed. Artifact
  manifest SHA-256:
  abae83e91b3304206906c2f37dae0fb6105856728457761a5b62c5beac88f381.

Authoritative MI350X TP8/C1 performance used three fresh-process repeats per
arm, exactly 128 input and 256 output tokens per request, and a maximum allowed
within-arm spread of 0.250 ms:

arm	repeat mean ITL values (ms)	mean (ms)	spread (ms)
control before, r178	20.376864, 20.386123, 20.413625	20.392204	0.036761
candidate, r179	18.980972, 18.986528, 19.003041	18.990180	0.022069
control after, r183	20.340497, 20.360134, 20.358735	20.353122	0.019637

The candidate is 1.382483 ms / 6.785971% below the adjacent-control
average (20.372663 ms), below both controls, and 0.672820 ms below the
19.663 ms acceptance target. r180 was rejected rather than averaged because
its 56.644377 ms spread exceeded the stability limit. The sealed post-hoc
validator summary has SHA-256
71398c1c1a74d19e847648cc37434c002800678225172f6effec411aeddbd19b.
The post-hoc step was needed because the wrapper hit an awk newline-formatting
error only after r183 had sealed; the validator rechecked every immutable arm
artifact rather than rerunning or altering the measurements.

The sink-only precursor experiment was neutral to slightly slower, so the sink
kernel is retained for graph-replay correctness and is not presented as an
independent latency win.

Authoritative MI350X full GSM8K and route validation (r186) passed:

• 1,319/1,319 unique rows, zero malformed/duplicate/missing/unexpected IDs,
  zero request errors, zero empty generations, and zero missing grades;
• 1,232/1,319 = 93.4041% symbolic accuracy;
• 1,203/1,227 = 98.0440% normal-stop accuracy;
• 92 length truncations, 29 correct, and 43 missing extracted predictions;
• 2,166,076 completion tokens in 5,501.155 seconds, or 393.7493 output
  tokens/s; and
• health 200 before and after evaluation, all nine optimized-route checks,
  empty runtime/kernel-fault scans, source postflight, and the full 300-second
  exclusive-idle teardown passed.

The pinned workload used TP8, native MTP with three speculative tokens,
full/piecewise graphs, AITER sparse MLA and MXFP8 MoE, explicit fused shared
experts, disabled prefix caching, all 1,319 examples, four request threads,
temperature 0.9, top-p 1.0, and seed 0. The r186 artifact manifest SHA-256 is
c3ba616f12615bc34a55d60d941f77dbe0c9133ddf7c8b53aaae2b86a2d3e284.

Supplementary MI355X r173 completed 1,319/1,319 unique requests with zero
request errors, 93.9348% accuracy (1,239/1,319), 78 length truncations, and
407.3702 output tokens/s. Its artifact manifest SHA-256 is
7c75b7c2624127c12ed796229d38b541eee683b327fd0f938dec9a1ad16da284.
This confirms a second gfx950 execution lane but is not used for the MI350X
performance acceptance claim.

Test plan

pre-commit run --files \
  $(git diff --name-only 4be7ce114bc8fcf5a7ff8175b57f80d3db5418a4..HEAD)

.venv/bin/python -m pytest -q -p no:cacheprovider \
  tests/compile/passes/test_fuse_mla_dual_rms_norm.py \
  tests/kernels/attention/test_rocm_aiter_mla_sink.py \
  tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py \
  tests/kernels/moe/test_mxfp8_aiter_backend_selection.py \
  tests/kernels/moe/test_rocm_aiter_moe.py \
  tests/kernels/moe/test_rocm_aiter_topk.py \
  tests/kernels/test_minimax_m3_amd_ops.py \
  tests/model_executor/layers/test_fused_shared_expert.py \
  tests/models/test_hy_v4_rocm.py \
  tests/test_config.py::test_rocm_hyv4_defaults_to_mrv2_and_breakable_cudagraph \
  tests/v1/attention/test_indexer_native_next_n.py

The hardware suites also cover the surrounding ROCm AITER MoE, sparse-MLA,
MTP, graph-replay, fused-shared-expert, and model-loading paths rather than only
the newly added unit tests.

Related work and stack

• [ROCm][CI] Add attention-sink support to ROCm AITER sparse MLA #54404 supplies the AITER sparse-MLA sink prerequisite; [ROCm][CI] Enable HY-V4 model initialization on ROCm #54405 supplies the
  initial HY V4 ROCm enablement and is this PR's immediate base.
• [ROCM] Enablement of Hy4-Preview on gfx950 #54432 is a separate overlapping gfx950 HY V4 enablement draft. Its fixes and
  this stack must be coordinated before merge rather than landed twice.
• [ROCm] Convert MXFP8 MoE weights to block FP8 on gfx94x #45726 overlaps the gfx942 MXFP8 conversion mechanics; this candidate adds
  graph/MTP integration and native-gfx950 behavior around that route.
• [Kernel][HY-V4] iHC: head op, custom-op registration, tensor-core path for large batches, warmup and tests #55154 adds fused Triton iHC pre/post/head kernels in the NVIDIA path. It is
  adjacent rather than duplicative; this update's new kernels and routing
  changes are ROCm-specific.
• [Bugfix] Zero-init ROCm MLA attention output buffers for graph padding #37682, [ROCm] Fix sparse MLA invalid token kernel returning 0 instead of -1 #45324, Fix/rocm sparse mla long context seq len guard #46763, [ROCm] Fixing GLM-5.1, DeepSeek-V3.2, DeepSeek-V4 on gfx942 and gfx950 #39168, [Perf] [Feat] [ROCm] Add densemha support to ROCm AITER Sparse MLA #49263, [ROCm][MLA] Guard sparse-MLA persistent-only path for gqa_ratio=64 fp8 (#49649) #49755, [ROCm][Perf] Skip redundant sparse index remap on non-indexer layers #51309, [ROCm][MLA] Enable sparse MLA Gluon kernel from Aiter #53492, [Kernel][Hardware][AMD] Fix unsupported fp8 dtype on gfx942 #43804,
  and [ROCm] [Perf] Optimize rocm aiter sparse mla indexer #46172 remain open changes around ROCm sparse MLA or its lower-level
  kernels. Their independent fixes must be preserved through landing order or
  a later rebase.

#54594 is based on #54405 at
29c7fa2027700f0bbbb26fae07868f308d26780b. #54404 has independently advanced
to 1059697b5316c6de3afe3cde3c8c70613e94fbaf; this update does not rewrite the
other contributor's parent branch. Any parent refresh or conflict resolution
changes the effective tree and invalidates this source-dependent validation.

AI assistance

AI assistance was used to investigate, implement, test, and prepare this
change. Before commit and publication, the human submitter reviewed every
changed line in tree 93b4fced3bf3dfaf67cf9d5f6217240189ba67f1,
confirmed end-to-end understanding, and explicitly approved the DCO signoff.
Commit 53abd3cac3e06993a8afae86b3246fa27d0ab1b1 records both the human
Signed-off-by trailer and the Codex Co-authored-by trailer.

[jamesETsmith]

Hy4-preview GPQA results

Previously from this PR I noticed GPQA results which seemed low (83-85%), but that seems due to a couple of factors/bad decisions on my part: 1) output limits that are too small so the model didn't have enough token budget to reason through the problems and 2) some ambiguity in extracting the answers. I don't think Tencent provided details on how they generated the numbers so the second point makes it tougher to compare to the "trusted" 92.3% provided by Tencent. If we use the gpt_oss extractor to pull the answers from the model response, we get results in the much closer to the 92.3%.

TLDR, I think the accuracy looks good on this implementation (at least from these tests), but it might be helpful to try another benchmark with numbers from Tencent to be totally sure. Any suggestions? If not, I'll just pick one and run it.

Run

• vLLM source: PR [ROCm] Optimize HY V4 graph inference on gfx942 and gfx950 #54594 commit 4be7ce114bc8fcf5a7ff8175b57f80d3db5418a4
• Hardware: 8x AMD Instinct MI355X, gfx950, TP8
• Backend: ROCM_AITER_MLA_SPARSE, graph mode, BF16 KV cache, no MTP
• Model limits: 70,000 tokens for 32K/64K and 140,000 tokens for 128K
• Sampling: high reasoning, temperature 0.9, top-p 1.0, paired per-sample seeds
• Dataset: GPQA Diamond, 198 samples, SHA-256 e5a13a2d2f6af845ef7ff4376845a0a15d818e20b8e13421759e12ab2037e214

Primary conservative extraction

Output limit	Correct	Accuracy	Delta from Tencent 92.3%	Extractable	Truncated	Failures
32K	161/198	81.31%	-10.99 pp	168/198	23	0
64K	176/198	88.89%	-3.41 pp	182/198	4	0
128K	175/198	88.38%	-3.92 pp	186/198	0	0

The matched longer-output run recovered 15 net correct answers from 32K to 64K and reduced truncations from 23 to 4. Extending to 128K removed all remaining truncations but did not improve raw accuracy. The one-answer decline from 64K is consistent with stochastic variation: 15 samples changed from incorrect to correct and 16 changed from correct to incorrect.

Paired primary-score transitions:

• correct to correct: 153
• incorrect to correct: 23
• incorrect to incorrect: 14
• correct to incorrect: 8

Extraction sensitivity

Re-scoring the same generations with the installed gpt_oss GPQA extractor produced:

Output limit	Correct	Accuracy	Extractable
32K	166/198	83.84%	198/198
64K	187/198	94.44%	198/198
128K	186/198	93.94%	198/198

The installed extractor falls back to the first character when no explicit answer pattern matches, so these scores are protocol-sensitive and should not replace the conservative primary score without agreement on the benchmark contract. The 64K and 128K scores nevertheless demonstrate that output length and extraction account for most or all of the apparent gap under this protocol.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: 40d63ba0-f352-4e61-868a-4335b64cae35

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[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.

[mergify]

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

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

[AndreasKaratzas]

I dont know if this PR got closed automatically after the merge of my first one. I am trying to reopen it but i seem not to be able to. @andyluo7 please feel free to reopen the pr or create a new one with the same diff if the changes here are still relevant.