[AMD] Add unified kv attention support in dpsk-v4

[1am9trash]

Co-authored-by: @RolaoDenthu

Motivation

Add a unified-KV attention backend for DeepSeek-V4 on amd code path, porting ATOM's sparse attention kernels, which has great perf. It is gated by SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton; when disabled, the default path is unchanged.

Modifications

Memory layout: change to a new class DeepSeekV4UnifiedKVPool.

• File: python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
• Before: three independent packed-fp8 pools (SWA ring, c4, c128).
• After: one bf16 unified KV pool per layer, laid out as [swa_pages + compress_pages, head_dim]:
  • rows [0, swa_pages) — SWA ring (state_slot * ring + pos % ring)
  • rows [swa_pages, ...) — compressed KV (swa_pages + page_index)

Compressed-KV store with unified kv layout

• File: python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh
• In _forward_compress_all_in_one(), after the compressor finishes, fused_norm_rope_flashmla fuses nope + rope + store in one launch, but only supported the original packed layout. Added a bf16_store parameter and a new store path so it can write the compressed KV into the bf16 unified KV pool.

SWA-KV store with unified kv layout

• File: python/sglang/srt/layers/fused_qk_norm_rope_store.py
• In _forward_prepare(), after KV is computed, fused_qk_norm_rope_swa_store fuses nope + rope + store, but again only supported the original layout. Added a bf16_store parameter and a new store path so it can write the SWA KV into the bf16 unified KV pool.

Attention kernels

• File: python/sglang/srt/layers/attention/dsv4/unified_kv_kernels
• New kernels: vendored ATOM paged decode/prefill attention kernels, plus sglang-side glue for building ragged index streams (build_decode_streams, build_prefill_indices), the SWA ring scatter (store_swa_into_unified), and the per-layer compress-tail fill (fill_compress_tail).

Per-forward metadata

• File: python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py
• Decode: init_forward_metadata_decode calls _attach_unified_kv_decode_streams once per forward to build the ragged index streams the attention kernel needs.
• Prefill: init_forward_metadata_prefill calls _attach_unified_kv_prefill_meta once per forward to build the per-token mapping; chunked prefill, target_verify, and draft_extend all share this single derivation.

Accuracy Tests

GSM200:

• tp8: 0.98
• tp8 + mtp: 0.975
• tp8 + dp8: 0.96
• tp8 + dp8 + mtp: 0.96

Speed Tests and Profiling

Server cmd:

export SGLANG_DEFAULT_THINKING=1
export SGLANG_DSV4_REASONING_EFFORT=max
export SGLANG_OPT_DEEPGEMM_HC_PRENORM=false
export SGLANG_USE_AITER=1
export SGLANG_USE_ROCM700A=1
export SGLANG_OPT_USE_FUSED_COMPRESS=true
# export SGLANG_HACK_FLASHMLA_BACKEND=triton
export SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton
export SGLANG_OPT_FP8_WO_A_GEMM=false
export SGLANG_OPT_USE_JIT_INDEXER_METADATA=false
export SGLANG_OPT_USE_TOPK_V2=false
export SGLANG_OPT_USE_AITER_INDEXER=true
export SGLANG_OPT_USE_TILELANG_INDEXER=false
export SGLANG_OPT_USE_TILELANG_MHC_PRE=false
export SGLANG_OPT_USE_TILELANG_MHC_POST=false
export SGLANG_FP8_PAGED_MQA_LOGITS_TORCH=1
export SGLANG_OPT_USE_FUSED_COMPRESS_TRITON=true

export SGLANG_OPT_USE_MULTI_STREAM_OVERLAP=false
export SGLANG_ROCM_USE_MULTI_STREAM=false

export AITER_BF16_FP8_MOE_BOUND=0
export SGLANG_EAGER_INPUT_NO_COPY=true

# --dp 8 --enable-dp-attention \
# --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-num-draft-tokens 4 --speculative-eagle-topk 1
# --disable-cuda-graph \
MODEL=/data/deepseek-ai/DeepSeek-V4-Pro
sglang serve \
    --model-path ${MODEL} \
    --trust-remote-code \
    --tp 8 \
    --disable-radix-cache \
    --attention-backend dsv4 \
    --max-running-request 256 \
    --page-size 256 \
    --mem-fraction-static 0.90 \
    --swa-full-tokens-ratio 0.1 \
    --chunked-prefill-size 8192 \
    --port 8000 \
    --disable-shared-experts-fusion \
    --tool-call-parser deepseekv4 \
    --reasoning-parser deepseek-v4

Client cmd:

python3 -m sglang.bench_serving --host localhost --port 8000 --dataset-name random --random-input 8192 --random-output 1024 --random-range-ratio 1.0 --num-prompt {cc*4} --max-concurrency {cc}

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): 🚫 Run #27188164054
Latest PR Test (Extra): ❌ Run #27188163931

[1am9trash]

/rerun-test registered/attention/unittests/dsv4/test_deepseek_v4.py

[github-actions]

Results for /rerun-test registered/attention/unittests/dsv4/test_deepseek_v4.py:

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/attention/unittests/dsv4/test_deepseek_v4.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/attention/unittests/dsv4/test_deepseek_v4.py

[1am9trash]

@amd-bot

[1am9trash]

@amd-bot ci-status

[1am9trash]

nv V4 test passed.

• base-c-test-deepep-4-gpu-b200 / base-c-test-deepep-4-gpu-b200 (0)
  
• registered/attention/unittests/dsv4/test_deepseek_v4.py
  

[1am9trash]

@amd-bot ci-status

[amd-bot]

@1am9trash

CI Status for PR #27380

Merge verdict: No failing job is caused by this PR — every red X is infrastructure (missing yaml on the H200 runner, a GitHub API "Server Error" on the AMD gate), a fast-fail/finish cascade, or an unrelated-backend flake (Intel-AMX CPU, NPU). However, this PR's entire payload is gated behind SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton, and no PR-CI test sets that value — so the green AMD run verifies only that the default path is unchanged, not that the new unified-KV feature works. Coverage gap matters more here than any of the red Xs.

Caution

The new unified-KV attention path (the whole point of this PR — deepseek_v4_backend_hip_radix.py, unified_kv_kernels/*, DeepSeekV4UnifiedKVPool) is reachable only when SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton (env_gate.py:is_unified_kv_triton). The PR changes zero test files. The only DeepSeek-V4 AMD tests (test/registered/amd/test_deepseek_v4_{flash,pro}_{fp8,fp4}.py) are nightly=True (suites nightly-amd-8-gpu-mi35x-deepseek-v4-*, do not run on PR CI) and they set SGLANG_HACK_FLASHMLA_BACKEND="triton", i.e. they exercise the old path. Nothing — PR CI or nightly — exercises unified_kv_triton. Before merge, run a DeepSeek-V4 accuracy/serving check on AMD MI35x with SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton, or add a registered test that sets it.

Changed files: 16 files, +2418/−84 — all under python/sglang/srt/layers/attention/dsv4/**, mem_cache/deepseek_v4_memory_pool.py, models/deepseek_v4.py, jit_kernel/. No test/** changes.

AMD: 0 real failures (2 infra/cascade) · Others: 0 related failures (7 unrelated/infra/cascade) — all AMD stage-b test shards (27 jobs) passed; stage-c was gate-skipped.

AMD CI Failures

Job	Test File	Test Function	Error	Related?	Why
wait-for-stage-b-amd	N/A	N/A	HttpError: Server Error from GitHub jobs API	🟢	Infra flake — the gate job's actions/.../jobs API call got a 500; all 27 real stage-b AMD shards passed. Not a test failure.
pr-test-amd-finish	N/A	N/A	wait-for-stage-b-amd: failure → exit 1	🟢	Pure cascade — finish job aggregates needs; fails because the gate above failed. All listed test jobs were success/skipped.

Other CI Failures

Job	Test File	Test Function	Error	Related?	Why
base-c-test-8-gpu-h200 (1)	N/A	N/A (setup step)	ModuleNotFoundError: No module named 'yaml' in scripts/ci/runner_configs.py:18	🟢	Runner-config setup step, runs before any test; PyYAML missing on the H200 runner. PR doesn't touch scripts/ci/. Infra.
base-c-test-8-gpu-h200 (2), base-c-test-4-gpu-h100 (4)	N/A	N/A	Fast-fail: skipping — root cause job(s): base-c-test-8-gpu-h200 (1)	🟢	Fast-fail cascade from the yaml infra failure above — collapsed into one root cause.
build-test (xeon-gnr, base-b-test-cpu)	test/registered/scheduler/test_routing_key_scheduling.py	TestRoutingKeyScheduling.setUpClass	RuntimeError: decode: expect req_lens to be int64, got Int	🟢	Crash is in intel_amx_backend.py:161 forward_decode for qwen3 on Intel-AMX CPU — a code path this AMD-GPU PR never touches. Pre-existing CPU-backend issue.
stage-b-test-1-npu-a2 (0)	NPU perf suite	Qwen2.5-0.5B-w8a8 throughput	AssertionError: 645.8 / 475.0 not >= 700	🟢	NPU throughput-threshold flake (varies 475–645 across retries) on Qwen2.5, unrelated to DeepSeek-V4 AMD code.
multimodal-gen-test-1-npu-a3, -2-npu-a3	multimodal_gen/.../ascend/test_server_1_npu.py	test_diffusion_generation[flux_image_t2i_npu]	Diffusion testcase failed 1 check	🟢	NPU Flux diffusion image-gen — different backend and modality entirely. Unrelated.

Details / what to do before merge

• Coverage (the only blocker): add or modify a registered AMD test to set SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton so the new deepseek_v4_backend_hip_radix + unified_kv_kernels paths actually run. Lacking that, manually verify on MI35x: launch DeepSeek-V4 with that env var and confirm accuracy + the bf16 DeepSeekV4UnifiedKVPool store/load path. The 4 existing nightly tests will not cover it (they use triton and don't run on PR CI).
• None of the red Xs require action from the author — they are runner/infra (yaml, GitHub API 500), fast-fail cascades, and CPU/NPU flakes. The AMD signal that did run is green.

---

Generated by amd-bot using Claude Code CLI

[HaiShaw]

LGTM