[ASM] [HIP] [MLA v4 nm] Split plan helper and persistent decode kernel with in-kernel split and merge

[yhl-amd]

Summary

Two parts for the gfx950 V4-NM fp8 decode (mla_decode_fwd_v4_nm), used by DeepSeek-V4 decode in ATOM. This PR supersedes #5890: its commits are included here.

1. Host-only split count for callers, from a cost model (formerly #5890)

• get_mla_v4_nm_num_kv_splits(num_seqs, num_heads, kv_len) -> int. Split count for mla_decode_fwd_v4_nm, computed on the host with no device work. Callers pass it with the uniform split_indptr [0, s, 2s, ...]; any prefix of a longer one is valid, so a caller can keep one constant buffer per split count and slice it (ATOM does), and CUDA graphs read constant memory.
  • The count minimizes a per-arch latency model (workgroup waves x per-wave KV time, plus the partial-write/stage-2 cost of splitting). Coefficients live in a per-arch table; gfx950 is fitted. Other archs keep the wrapper's occupancy-only pick.
  • It depends only on num_seqs, num_heads and the caller's kv_len, never on kv_page_indices, whose length is meaningless for capacity-sized graph buffers.
• Wrapper unchanged when no split count is passed. The auto pick inside mla_decode_fwd_v4_nm is the same as main.
• Stage-2 fix. With a caller-supplied split_indptr that gives every seq one split on a grid > 1, the merge took the in-place BF16 copy (keyed on split_indptr[-1] == num_seqs) even though the kernel wrote FP32 partials, returning NaN. The v4 nm merge now always reduces the FP32 partials.

Why callers should not rely on the old auto pick: it maximizes CU occupancy only, so for gqa=128 single-token decode it picks 13-16 splits (several workgroup waves) for batches whose KV already fits one wave, adding partial writes and stage-2 work without adding parallelism.

2. Persistent decode kernel with in-kernel split and merge

A persistent variant of the kernel (ps=1 row in mla_v4_asm.csv). One launch per call: a fixed grid of 2*P workgroups walks per-partition work descriptors, the split plan is computed inside the kernel from the real kv_indptr every call, and split rows are merged by the last-arriving workgroup. Rows that are not split write the final bf16 output directly. Because nothing about the split depends on host-side shapes, it is CUDA-graph safe with varying KV lengths.

• hsa/gfx950/mla_v4/mla_a8w8_qh64_qseqlen1_gqaratio64_nm_ps.co + CSV row (ps=1, lse=1).
• csrc/py_itfs_cu/asm_mla_v4.cu: mla_decode_v4_ps_asm entry (documented packed kernarg struct, shape / workspace checks).
• aiter/mla.py: get_mla_v4_nm_ps_workspace(device, num_partitions=128) and mla_decode_fwd_v4_nm_ps(q_packed, q_rope, kv_packed, kv_rope, kv_indptr, kv_page_indices, sink, workspace, out=None, return_lse=False, lse=None).
• Contract: 128 heads, N <= 32768, row-dense KV pools, rows with K=0 left unwritten (same as the non-persistent kernel). Calls sharing a workspace must be ordered on the GPU; allocate the workspace outside graph capture.

Merged with current main; the op_tests/test_mla_v4_nm.py conflict with #5905 is resolved (the varlen check uses _gated_allclose, which now also returns the mismatch fraction).

Consumer: ROCm/ATOM#2462 (split plan for the decode ASM path, HCA layers on the persistent kernel).

Test plan

• op_tests/test_mla_v4_nm.py on gfx950 (MI355X), after merging main, fresh JIT build: 47 passed
  • split count: test_v4_nm_split_planner_picks, test_v4_nm_num_kv_splits_uses_cost_model, test_v4_nm_split_plan_accuracy, test_v4_nm_split_plan_cudagraph_replay (captured once, replayed with different KV lengths), test_v4_nm_one_split_per_seq_on_wide_grid
  • persistent kernel: accuracy vs reference incl. K=0 rows and long K, lse, offset kv_indptr, workspace reuse across batch sizes, CUDA-graph replay with changing KV lengths, agreement with mla_decode_fwd_v4_nm, guard-page out-of-bounds checks, argument rejection
• ruff (0.16.0) / black on changed files
• End-to-end DeepSeek-V4-Pro serving through ATOM (companion PR) with real speculative acceptance: graph capture OK, accuracy checks pass
• GSM8K 5-shot (1319) on DeepSeek-V4-Pro-DSpark through ATOM, TP8 + DP attention + EP8, real acceptance, together with the MegaMoE changes: 0.9606 strict / 0.9613 flexible (stock image 0.9530)

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6126 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[yhl-amd]

Local accuracy validation with the companion ROCm/ATOM#2462 (GSM8K, ATOM CI accuracy configs)

Setup. 8x MI355X. Image: rocm/atom-dev:latest @ sha256:f6940ab6…, the image the nightly CI used for ATOM main 5b228e925.

• Baseline: the image as-is (ATOM 5b228e925, aiter c99c361).
• PR: the same image with [DSV4] Tuned ASM FP8 decode with split plans; HCA on aiter's persistent V4-NM kernel ATOM#2462 (425f8852d, which is 5b228e925 plus this PR) and [ASM] [HIP] [MLA v4 nm] Split plan helper and persistent decode kernel with in-kernel split and merge #6126 applied to the image's aiter. module_mla_v4_asm was rebuilt; op_tests/test_mla_v4_nm.py passes (47) in that image.
• Run: the server is launched with .github/scripts/atom_test.sh launch and the configs from .github/benchmark/models_accuracy.json. The client is the CI default: lm_eval, gsm8k, 3-shot, num_concurrent=65, 1319 samples.

Config	CI threshold	CI nightly @5b228e925	Local baseline	Local PR
DeepSeek-V4-Pro (TP8)	0.94	0.9515 / 0.9522	0.9484 / 0.9492	0.9477 / 0.9477
DeepSeek-V4-Pro MTP (MTP3)	0.94, accept 0.622	0.9530 / 0.9530	0.9583 / 0.9591, accept 65.92%	0.9507 / 0.9515, accept 65.97%
DeepSeek-V4-Pro DSpark (DPA)	0.93	0.9500 / 0.9500	0.9484 / 0.9500, accept 43.64%	0.9477 / 0.9484, accept 43.91%
DeepSeek-V4-Pro DEP RCCL (DPA + EP)	0.94	0.9560 / 0.9568	0.9522 / 0.9538	0.9575 / 0.9575

Each entry is flexible-extract / strict-match. Acceptance is accepted / total draft tokens summed over every [MTP Stats Interval] line of every rank. For DPA that sum is needed because each rank logs its own stats.

• Every PR result is above the CI threshold.
• PR − baseline ranges from −0.0076 to +0.0053 with no consistent direction. That is within the ±0.006 standard error of a single 1319-sample run. Acceptance is unchanged for MTP3 and slightly higher for DSpark.
• DPA configs (DSpark, DEP RCCL): these run at 128 local heads. The HCA persistent kernel is enabled there; all 8 ranks log V4 HCA persistent decode: device N ready.
• TP8 and MTP3: these run at 16 local heads. The persistent kernel stays off by design (it needs 128 heads), so these configs cover the split-plan path of the decode ASM kernel.

🤖 Generated with Claude Code