[DSV4] Tuned ASM FP8 decode with split plans; HCA on aiter's persistent V4-NM kernel

[yhl-amd]

Summary

DeepSeek-V4 native FP8 decode on the gfx950 ASM kernels, in two parts. This PR supersedes #2411: its commits are included here.

1. Adapt to the tuned AITER ASM FP8 decode, with caller-built split plans (formerly #2411)

• Adapt DeepSeek-V4 native FP8 decode to the swapped gfx950 ASM decode kernel from [ASM] [MLA v4 nm] Swap gfx950 asm decode kernel aiter#5195.

• Pass a KV split plan to the decode ASM kernel. Whoever builds a decode CSR builds its split plan next to it, on the host: the split count comes from AITER's get_mla_v4_nm_num_kv_splits ([ASM] [HIP] [MLA v4 nm] Split plan helper and persistent decode kernel with in-kernel split and merge aiter#6126) and split_indptr is a row of a constant [16, T + 1] uniform table (v4_uniform_split_table) built once at init, so no forward writes a plan and every CUDA graph reads constant memory.

  • V4 metadata builder, verify/decode: one plan per layer type (SWA window, CSA window + index_topk, HCA window + max_model_len // 128), sized for the padded row count.
  • V4 metadata builder, MTP draft (prepare_mtp_decode): SWA plan for the draft rows.
  • DSpark draft (DSparkIndexBuffers, its own table): built with the draft CSR at stage 0 and read by every stage (B*T rows, W + T KV).

  Each plan depends only on the row count and config, so CUDA graph capture and replay agree.

• sparse_attn_v4_paged_decode takes a single split_plan; the model selects it in the same ratio branch as kv_indices / kv_indptr.

• v4_decode_split_plan is the only place that checks for AITER's helper. With an AITER build that lacks it, or a pick outside the table, plans are None and AITER keeps its own split pick.

ASM-only: no FlyDSL/Triton decode implementation, route or kernel file is added or changed.

2. HCA layers on aiter's persistent V4-NM kernel

HCA layers (compress_ratio 128) now use mla_decode_fwd_v4_nm_ps instead of the decode ASM + host split plan. The kernel plans the split from the real per-row KV lengths every call, so HCA's growing KV no longer runs with a split fixed at graph-capture time.

• atom/model_ops/v4_kernels/hca_persist.py: gates, one workspace per device, call wrapper. If the installed aiter lacks the API, everything stays on the ASM path (logged once).
• Gates: ATOM_V4_HCA_PERSIST (default 1), fp8 KV, ratio 128, 128 local heads, gfx950, ATOM_V4_HCA_PERSIST_MIN_ROWS (default 15) <= rows <= 32768, row-dense pools. Smaller batches and all SWA/CSA layers keep the ASM + split plan.
• Workspace allocated at model load (HCA DeepseekV4Attention.__init__, all paths incl. plugins) before KV sizing and graph capture; a capture without a workspace stays on the ASM path instead of allocating.
• paged_decode.py takes compress_ratio=; deepseek_v4.py passes it at the decode call site.
• Env vars documented in docs/environment_variables.md.

Also: tests/test_deepseek_v4_transfer_regions.py stubs the three imports these changes add to deepseek_v4_attn (get_tensor_model_parallel_world_size, v4_decode_split_plan, hca_persist), which the CPU unit-test runner was failing on.

Dependencies

• [ASM] [MLA v4 nm] Swap gfx950 asm decode kernel aiter#5195 (merged): gfx950 v4 nm decode code object
• [ASM] [HIP] [MLA v4 nm] Split plan helper and persistent decode kernel with in-kernel split and merge aiter#6126: split plan helper and persistent V4-NM kernel. Without it the split pick stays in AITER and HCA stays on the ASM path.

Test plan

• Non-GPU unit suite as CI runs it (.github/scripts/run_unit_tests.sh, CPU torch, no aiter), merged with current main: 0 failures
• tests/test_v4_prefill_asm_decode_dispatch.py, tests/test_paged_attention_dispatch.py, tests/test_dspark.py: pass with and without the AITER helper
• tests/test_v4_hca_persist_dispatch.py (CPU, mocked) + tests/model_ops/test_v4_hca_persist_gpu.py (GPU: vs ASM path and reference, K=0 padded rows, eager N < T_pad, offset kv_indptr, CUDA-graph replay with changing lengths, capture without workspace)
• ruff / black on changed files
• DeepSeek-V4-Pro-DSpark, FP8 KV, TP8/EP8/DPA8, --cudagraph-mode FULL, real draft acceptance: graph capture, correct outputs, DSpark acceptance unchanged vs. the stock image
• GSM8K 5-shot (1319) on DeepSeek-V4-Pro-DSpark, TP8 + DP attention + EP8, real acceptance, together with the MegaMoE changes: 0.9606 strict / 0.9613 flexible (stock image 0.9530)
• Full-config server smoke (V4-Pro-DSpark, TP8 + DP-attention, FULL graphs, dspark K6, real acceptance): graph capture OK on all ranks, persistent path hit eager + captured, 37*48=1776, arithmetic 480/480

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2462 --add-label <label>

[yhl-amd]

Local accuracy validation (GSM8K, CI accuracy configs)

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

• Baseline: the image as-is (ATOM 5b228e9, aiter c99c361b5).
• PR: the same image with [DSV4] Tuned ASM FP8 decode with split plans; HCA on aiter's persistent V4-NM kernel #2462 (425f885, which is 5b228e9 plus this PR) and [ASM] [HIP] [MLA v4 nm] Split plan helper and persistent decode kernel with in-kernel split and merge aiter#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