[Tracker]: DeepSeek-V4 Support and Optimization on ATOM

[jamestangg]

[Tracker]: DeepSeek-V4 Support and Optimization on ATOM (as of 5/22/2026)

High-Level Summary

Recently completed

• ATOM V4 PR1 skeleton (full V4 + CUDAGraph infra + batched compressor/indexer + FP8 CSA cache + ROCm gfx950 fix): feat(deepseek_v4): PR1 skeleton — end-to-end inference with triton MoE #650 ✅ Merged
• KV cache subsystem refactor: refactor: delegate ATOM KV cache subsystem to attention builders #659 ✅ Landed 4/29
• Triton sparse_attn + move attn out of loop: feat(deepseek_v4): use triton sparse attn kernel and move attn kernel out of loop #678 ✅ Landed 5/1
• FP8 batched BMM grouped LoRA wo_a workaround: feat(deepseek_v4): FP8 batched BMM for grouped output LoRA wo_a #676 ✅ Landed 5/1
• A8W4 + SwiGLU + clamp MoE kernel: [MOE DSV4] flydsl a8w4 moe for dsv4 aiter#2951 ✅ Landed
• mHC fused Triton kernel + mhc_pre / hc_split_sinkhorn: [TRITON] mHC-pre: Manifold-constrained Hyper Connection  aiter#2646, fix mhc device aiter#2916 ✅ Landed
• AITER mHC integration into vLLM: [Bugfix] [ROCm] [DSV4] [Perf] Add aiter mhc support vllm-project/vllm#41946 ✅ Landed
• DSv4-Pro full graph + piecewise prefill capture in vLLM: DeepSeekV4-Pro enable cuda graph full and piecewise mode vllm-project/vllm#42604 ✅ Landed 5/17

---

Motivation

This issue tracks the end-to-end enablement and optimization checklist for DeepSeek-V4 on ATOM. DeepSeek-V4 includes optimization and functional opportunities (mHC / HCA / CSA / MoE / MTP), and ATOM readiness depends on both kernel-level work in AITER and ATOM-side integration. We're launching a joint effort to land V4-Pro Base FP4 functionality and parity, then drive performance and concurrency. Please feel free to take on any task, and we'd love to hear more optimization ideas.

Purpose

• Track DeepSeek-V4 functionality and performance readiness on ATOM.
• Keep module-level optimization items visible and actionable.
• Align acceptance criteria for FP4 native release readiness.

Scope today: V4-Pro Base (FP4) native, eager mode, 4-batch coherent on real ckpt, GSM8K ~0.6 (5 P0/P1 gaps tracked in #642). V4-Flash variants and V4-Pro (FP8 cross-quant) are not currently in ATOM scope.

General Checklist

1) Functionality / Bugfix / Feature

• [Functionality] V4 PR1 skeleton: feat(deepseek_v4): PR1 skeleton — end-to-end inference with triton MoE #650 (Lingpeng Jin). ✅ MERGED. Full V4 + CUDAGraph infra + batched compressor/indexer + FP8 CSA cache + ROCm gfx950 fix. Hosts the W4.1 to W4.6 multi-request batch concurrency sequence.
• [Functionality] KV cache subsystem refactor for V4: refactor: delegate ATOM KV cache subsystem to attention builders #659 (Lingpeng Jin). ✅ Landed 4/29.
• [Functionality] AITER MLA decode top-k consumer: Use AITER DSv4 Indexer top-k kernel #684 (Oseltamivir). Consumes the AITER Dsv4 sparse indexer aiter#2998 sparse indexer.
• [Functionality] FP8 batched BMM grouped LoRA wo_a (workaround for grouped o_lora_rank): feat(deepseek_v4): FP8 batched BMM for grouped output LoRA wo_a #676 (zufayu). Landed 5/1. Not a fused replacement.
• [Functionality] G-loop FP8 blockscale grouped LoRA wo_a: feat(deepseek_v4): G-loop FP8 blockscale for grouped output LoRA wo_a #677 (zufayu).
• [Functionality] Triton sparse_attn kernel + move attn out of loop: feat(deepseek_v4): use triton sparse attn kernel and move attn kernel out of loop #678 (junhaha666). Landed 5/1.
• [Functionality] MTP-K skeleton + SWA race/aliasing fixes + per-PR MTP-1 CI: feat(deepseek_v4): MTP-K skeleton + SWA race/aliasing fixes + per-PR MTP-1 CI #817 (valarLip). 🟡 In progress. MTP scaffolding with SWA fixes and CI hookup.
• [Bugfix] AITER P0/P1 kernels with no PR yet that gate ATOM accuracy (full list in [Tracking] Enable DeepSeek-V4 (Flash 284B / Pro 1.6T, 1M context) inference #642):
  • sqrtsoftplus CK kernel for gfx950 perf: add topk_softplus kernel aiter#2995 ✅ Landed. PR is currently Triton-only; CK version still absent.
  • MLA grouped o_lora_rank GEMM (fused replacement, not BMM workaround).
  • MLA sliding_window=128 passthrough (HCA layers go through generic path today).
  • mHC kernel gfx950-tuned path (today: generic gfx942 fallback).
• [Feature] encoding_dsv4 chat template port to ATOM (P3 from [Tracking] Enable DeepSeek-V4 (Flash 284B / Pro 1.6T, 1M context) inference #642). ATOM chat-completions endpoint missing V4 reasoning encoding.
• [Feature] Tool argument parsing for DSML.
• [Feature] Variant scope decision: open V4-Flash variants on ATOM, or keep ATOM as V4-Pro Base FP4 only.

Performance Checklist

1) High-Level Performance / Feature

• Multi-request batch concurrency on V4-Pro Base FP4 (W4.1 through W4.6 sequence under feat(deepseek_v4): PR1 skeleton — end-to-end inference with triton MoE #650).
• HIP graph (full graph) capture for Decode steps.
• Piecewise HIP graph capture for Prefill steps.
• Reasoning-mode chat template + accuracy parity.
• V4-Pro Base GSM8K parity vs B300 baseline (current ~0.6 to target ~0.96).
• Multi-node KV cache (gated by KV cache layout / quant design).
• MTP support for DSv4 on ATOM. Scaffold landing via feat(deepseek_v4): MTP-K skeleton + SWA race/aliasing fixes + per-PR MTP-1 CI #817.
• DSv4 disaggregated serving (ATOM side). Test docker rocm/atom-dev:nightly_202605171131-mooncake_build available; multi-node issues outstanding before official release. Repro release for SA gated. Reference: atom/kv_transfer/disaggregation/README.md.

2) Kernel Fusion

Element-wise fusion

• Inverse RoPE + FP8 quant fusion before wo_a projection. Currently torch fallback. vLLM-proven 2 to 3×. AITER-side in flight: Add fused inv_rope + FP8 block-scaled quantization kernel for DSV4 MLA aiter#3079.
• Silu and Mul fusion. Currently torch native. Owner committed 5/9.
• Q/KV pre-attn fusion (Q norm + KV RoPE + K insert, warp-dispatched). Highest single-fusion ROI: 10 to 20× per vLLM. Open and unclaimed.
• Triton kernel work (per-head Q Norm, fp4_act_quant, rotate_activation, small BF16 matmul, mHC inline norm) giving 20-30% improvement on ATOM path; kernels also becoming available in AITER.

CSA

• Replace torch native sparse_mla path with sparse_attn proper kernel integration (CSA path: index-gather + online softmax + attn_sink bias). Parallel approaches encouraged: ASM, flyDSL, Triton/Gluon, HIP/C++. flyDSL Approach 2 in flight.
• New OPUS path: pa_sparse_prefill_opus kernel for DSv4 prefill on gfx950 [OPUS] [ATOM] Add pa_sparse_prefill_opus kernel for DeepSeek-V4 prefill on gfx950 aiter#3225.
• CSA Indexer (FP4) integration once AITER kernel lands.
• CSA Indexer Compressor (head_dim=128, Hadamard, FP4) integration.
• CSA Attention Compressor (head_dim=512, FP8, overlap=True) integration.
• CSA multi-stream execution in Decode (default + indexer stream overlap). Aligned with DSv4 blog design. Key gap vs ATOM at conc=512 (5/21).
• Full-pipeline CSA compressor fusion. vLLM-proven 1.4 to 3×.
• CSA HIP graph compatibility for Decode.
• AITER DSv4 indexer top-k kernel consumer wiring: Use AITER DSv4 Indexer top-k kernel #684.

HCA

• HCA Compressor (head_dim=512, ratio=128, no overlap). Torch fallback today. At 1M context, HCA is approximately 12× CSA Flash per-query cost. Open.
• MLA sliding_window=128 passthrough so HCA layers stop going through the generic path.
• HCA HIP graph compatibility for Decode.

mHC

• AITER mHC fused Triton kernel integration: [TRITON] mHC-pre: Manifold-constrained Hyper Connection  aiter#2646 ✅ Landed. 315/315 tests passing.
• mhc_pre / hc_split_sinkhorn: fix mhc device aiter#2916 ✅ Landed.
• mHC C++ path wired into ATOM scheduler.
• mHC kernel gfx950-tuned path. Today: generic gfx942 fallback. P1 from [Tracking] Enable DeepSeek-V4 (Flash 284B / Pro 1.6T, 1M context) inference #642.
• Fused full mHC block (Sinkhorn to norm to Attn/MoE to post). Future optimization; eliminates 4 launches per layer.

MoE

• AITER A8W4 + SwiGLU + clamp Approach 1 ([MOE DSV4] flydsl a8w4 moe for dsv4 aiter#2951) integration into ATOM ✅ Landed. First DSv4-correct MoE kernel; closes 9× amplitude-loss bug on gfx950.
• Native A8W4 expert GEMM (top-3) integration once AITER kernel lands.
• FP4 upcast-at-load path (TileLang-style). aiter FusedMoE upcast path works as fallback. 2× weight-bandwidth efficiency.
• Hash routing (layers 0-2) ✅ working in ATOM.
• AITER FP4 bandwidth-optimized expert GEMM integration (when native A8W4 kernel ready).
• DSv4 FP4 MoE tuned config for FlyDSL kernels: [Config] Add DeepSeek-V4 FP4 MoE tuned config for FlyDSL kernels aiter#3183 (in progress).
• DSv4-Pro MTP spec-decoding tuned GEMM: [Misc] Add DSv4-Pro MTP spec-decoding tuned GEMM (a8w8 N=7168 K=7168 up to M=65536, bf16 N=384 K=7168) aiter#3238 (in progress).
• MegaMoE Approach 2 (Expert-Parallel All-to-All + FP4 GEMM). Design in flight; cowork on comms once design suggests better perf than current impl. Mid-term; depends on AMD NCCL-GIN-equivalent comms.

3) GEMM

• A8W8 (FP8 act × FP8 weight) via CK GemmMXFP8 ✅ working.
• Grouped o_lora_rank GEMM (BF16 grouped LoRA, fused replacement). einsum("bsgd,grd->bsgr"); G=8 (Flash), G=16 (Pro). 1-2 wk scope. Blocks feat(deepseek_v4): PR1 skeleton — end-to-end inference with triton MoE #650 finalization. Open and unclaimed.
• FP4 upcast-at-load (TileLang-style). 2-3 wk per ROCm/frameworks-internal#16453. Open.

4) KV Cache Codesign

• KV cache layout / quant design proposal. Hard-blocks multi-node work. Needs commit date.
• KV cache infra co-design. Match vLLM 3-bucket page-size design.
• Hybrid KV cache manager: three storage paths (CSA Attention 4:1 overlap, CSA Indexer 4:1 overlap, HCA 128:1) under a single page-pool design.
• KV cache subsystem refactor: refactor: delegate ATOM KV cache subsystem to attention builders #659 ✅ landed 4/29.
• DSv4 disaggregated serving (KV transfer): test docker rocm/atom-dev:nightly_202605171131-mooncake_build released for internal testing; multi-node still has issues; official release pending. Repro release to SA gated. README: atom/kv_transfer/disaggregation/README.md.

5) Quantization & Transform

• act_quant (BF16 to FP8 e4m3, per-1×128) ✅ via aiter.
• DSv4 padding issue fix on AITER: fix deepseek v4 padding issue aiter#3184 (open).
• Indexer multi-stream overlap. Approximately 5-6% E2E.

6) Reasoning / Chat Template

• encoding_dsv4 chat template port (P3 from [Tracking] Enable DeepSeek-V4 (Flash 284B / Pro 1.6T, 1M context) inference #642). ATOM chat-completions endpoint missing V4 reasoning encoding.
• Tool argument parsing for DSML.