[feat][DCP] Optimize MLA DCP (Enable Prefix Cache / Chunked Prefill + FP8 KVCache)

[yitingw1]

Expand the initial MLA DCP PR #847 with following:

• Make DCP compatible with Prefix Cache
• Make DCP compatible with Chunked Prefill
• Support fp8 KVCache
• Fix fused RoPE+cache kernel bug
• Add DCP doc

1. Motivation

The base DCP PR (#847) enabled Decode Context Parallel for MLA decode only, bf16, no prefix cache. config.py hard-asserted DCP off whenever prefix caching, chunked prefill, or fp8 KVCache was enabled, so DCP could not be used in any realistic serving config.

This PR removes those restrictions and makes DCP work end-to-end with prefix cache, chunked prefill, and fp8 KVCache, in both ATOM server and vllm-atom plugin modes. It also switches the DCP q/KV write off the slower 2-kernel workaround onto the fixed fused fused_qk_rope_concat_and_cache_mla (see the companion aiter PR), and consolidates the DCP access boilerplate.

Depends on the aiter PR ROCm/aiter#4342 "Fix fused_qk_rope_concat_and_cache_mla for DCP".

2. Usage

No new CLI flags — the existing -dcp / --decode-context-parallel-size (from #847) now composes with prefix cache / chunked prefill / fp8 KVCache, which were previously rejected at startup.

Argument	Required	Description
-dcp / --decode-context-parallel-size	No (default 1)	DCP world size; must divide tp_size
--kv_cache_dtype {bf16,fp8}	No	fp8 now supported together with DCP
--enable-prefix-caching, chunked prefill	No	now supported together with DCP

# ATOM server — DeepSeek-R1, TP8 + DCP8, fp8 KVCache, prefix cache + chunked prefill
HSA_NO_SCRATCH_RECLAIM=1 python -m atom.entrypoints.openai_server \
  --model /data1/models/DeepSeek-R1 -tp 8 -dcp 8 --kv_cache_dtype fp8

# vllm-atom plugin — same config
vllm serve /data1/models/DeepSeek-R1 \
  --tensor-parallel-size 8 --decode-context-parallel-size 8 --kv-cache-dtype fp8 \
  --async-scheduling --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE"}'

Verify with lm_eval --model local-completions --tasks gsm8k --num_fewshot 5 ....

3. Design

3.1 What changes vs the DCP=1 / decode-only baseline

Aspect	Baseline (#847, decode-only bf16)	This PR
Prefix cache	Disabled under DCP (assert)	Enabled — block accounting uses a virtual block of block_size × dcp
Chunked prefill	Disabled under DCP (assert)	Enabled — cached-prefix context read via compressed-KV AllGather
fp8 KVCache	Disabled under DCP (assert)	Enabled — fp8 stays on the wire, dequant to bf16 for kv_b_proj
Prefill cached-context KV	n/a	Sharded, per rank: gather local → AllGather → reorg → kv_b_proj
q / new-token KV write	2-kernel workaround (rotary_emb + concat_and_cache_mla)	fixed fused kernel with compute_all_q_rope=True (decode and prefill)
DCP access (world_size/group/rank)	inline getattr + get_dcp_group scattered	centralized in atom/distributed/dcp_utils.py

3.2 Cached-prefix context under DCP (compressed-KV AllGather)

Under DCP the cached KV is interleaved across ranks (token i → rank i % dcp), so each rank holds only 1/dcp of the prefix. To attend to the full prefix, each rank gathers its local compressed latent KV and AllGathers it (a copy-only collective, so fp8 payloads are safe — no fp8 arithmetic). Per cached-prefix chunk:

 ┌───────────────────────────────────────────────┐
 │ gather local compressed KV (owned interleaved   │  index_select(slot_ids) [toks, d=kv_lora+pe]
 │ slots for this chunk, keeps cache dtype/fp8)    │
 └───────────────────────┬─────────────────────────┘
                         ▼  AllGather(dim=0) across DCP group  → [toks*dcp, d]  (copy-only; fp8 safe)
 ┌───────────────────────────────────────────────┐
 │ (fp8) dequant: ag_kv.to(bf16) * k_scale         │
 └───────────────────────┬─────────────────────────┘
                         ▼  reorg_kvcache: rank-major → per-seq contiguous, drop padding
 ┌───────────────────────────────────────────────┐
 │ kv_b_proj decompress → (k_nope,v); concat k_pe  │
 └───────────────────────┬─────────────────────────┘
                         ▼  flash_attn(causal=False, return_lse) → (ctx_out, ctx_lse)
 ┌───────────────────────────────────────────────┐
 │ LSE-merge across chunks, then LSE-merge with    │
 │ the new-token (suffix) self-attention output    │
 └───────────────────────────────────────────────┘

3.3 Key design decisions

• Virtual-block prefix-cache accounting. Under DCP one block-table entry addresses a virtual block of block_size × dcp global tokens (each rank stores its block_size interleaved share), so block tables shrink by dcp. BlockManager hashes / counts cached tokens at this hash_block_size = block_size × dcp granularity; the scheduler computes num_new_tokens against it. Why: prefix-cache reuse must be counted in whole virtual blocks or the interleaved slot math breaks.
• Compressed-KV AllGather (not full-KV). Gather + AllGather the compressed latent (kv_lora+pe, ~576) before kv_b_proj, not the decompressed KV — minimizes cross-rank traffic and matches the vLLM plugin's chunked-context scheme.
• fp8 stays on the wire. AllGather the fp8 compressed KV, dequant to bf16 after the gather (server) or during the gather via gather_and_maybe_dequant_cache (plugin); AG/RS + LSE merge run in bf16. Why: AllGather is copy-only (fp8 safe), halving wire traffic; only AllReduce would be unsafe for fp8.
• Fixed fused kernel + compute_all_q_rope. DCP passes compute_all_q_rope=True so the kernel RoPEs every query (all ranks need all queries after the head all-gather) while writing KV only for owned slots — replacing the 2-kernel workaround. Non-DCP passes the default False (early-return preserved → no perf regression). Integration: requires the companion aiter kernel change.
• dcp_utils access layer. get_dcp_world_size / dcp_is_enabled / get_dcp_group / get_dcp_rank, mirroring pcp_utils.py, replacing scattered getattr(config, ...) + get_dcp_group().rank_in_group in the attention / metadata-builder inits.

4. Test Plan

• op-level correctness already covered by the aiter PR's test_concat_cache_mla.py.
• End-to-end gsm8k (DeepSeek-R1, 8×MI300X) across the newly-enabled matrix: DCP8 × {bf16, fp8}, server + plugin, plus a non-DCP regression check. Prefix cache + chunked prefill on (default), so every request exercises the cached-prefix context path.
• Command: lm_eval --model local-completions --model_args model=/data1/models/DeepSeek-R1,base_url=… --tasks gsm8k --num_fewshot 5.

5. Test Result

DeepSeek-R1, gsm8k (flexible / strict):

Mode	Config	flexible	strict
ATOM server	TP8 DCP8 bf16	0.9492	0.9492
ATOM server	TP8 DCP8 fp8	0.9568	0.9553
ATOM server	TP8 non-DCP bf16 (regression)	0.9492	0.9477
vllm-atom	TP8 DCP8 bf16 (3-shot)	0.9500	0.9492
vllm-atom	TP8 DCP8 fp8 (3-shot)	0.9462	0.9454

All within the run-to-run baseline band (bf16 ≈ 0.948–0.956, fp8 ≈ 0.955); no accuracy regression, no crashes, cudagraph capture OK, prefix-cache path hit.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[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 1701 --add-label <label>