[ganyi1996ppo]

Summary

A speculative draft that caches into a KV pool of its own (an MHA draft: MiniMax-M3.1's DSpark, an MHA EAGLE3 draft) can now be given a KV dtype independent of the target's --kv_cache_dtype: --spec-decode-kv-cache-dtype {bf16,fp8}. Omitted, it has the target's dtype, as before. A draft that caches in the target's pool (MTP, an MLA draft) always has the target's dtype; asking for another one is an error at startup.

Independent of #2962 (MiniMax-M3.1 MXFP4 caches): based on main, touches no M3 code, and its tests pass on their own. Split out of #2962 on valarLip's request. Whichever lands second is rebased on the other (see "Interaction with #2962").

What changes

• SpeculativeConfig.kv_cache_dtype is what was asked (never rewritten). SpeculativeConfig.draft_kv_dtype is what Config.__post_init__ settled (resolve_draft_kv_dtype); the draft's pool, its attention, the DSpark/EAGLE proposers, the calibration identity and the compile hash all read it. (Requested and resolved values were one field in [MiniMax-M3.1] MXFP4 index cache and MXFP4 KV cache (--index-cache-dtype fp4, --kv-cache-dtype fp4) #2962; valarLip asked for two.)
• "Does the draft own a pool" is decided once, in Config, from the attention family (selector.DRAFT_POOL_FAMILIES), not by resolving a backend class, so the API process does not import aiter. AttentionBackend.DRAFT_OWNS_KV_POOL and deployment._config_has_own_pool_draft are replaced by it; EAGLE3 is covered as well as DSpark. test_kv_pool_contract.py pins the table to the backends that define make_kv_pool.
• DeepSeek-V4's fp8 -> bf16 fallback off gfx950/gfx1250 moves from DeepseekV4AttentionMetadataBuilder into Config, before the draft's dtype is resolved. Otherwise a V4 draft (it shares the target's pool) would resolve to fp8 against a bf16 pool, and the DSpark q_out dtype check would warn on every run. Not applied under a plugin (same as before: the builder's probe only ran natively).
• A draft dtype different from the target's is refused with --kv-transfer-config (KV offload included): the offload namespace keys the target's dtype only.
• --spec-decode-kv-cache-dtype without --method or with --num-speculative-tokens 0 is an error, naming both conditions.
• Docs: SpeculativeConfig.kv_cache_dtype, the flag row, and the --method row (it listed only mtp).

Behaviour that is deliberately unchanged

No flag, no change: every draft has the target's dtype. The only default that moves is compute_hash (the spec tuple gains the draft dtype, so every speculative model's compile cache key changes once) and calibration_identity (the draft dtype is added only when it differs from the target's, so existing DSpark profiles still match).

Tests

New tests/test_draft_kv_dtype.py (a real Config over a stubbed HF config: resolution table, requested vs settled, the compile hash, MHA EAGLE3, MTP refusal, a plugin config with a V4.1 DSpark draft never resolving a backend, the offload refusal and its remedy). Changed: test_arg_utils_spec (flag reaches SpeculativeConfig, refusal without speculation, no fp4 choice), test_draft_kv_builder (pool decided from the family; an EAGLE3 draft attention is built at its pool's dtype), test_dspark (draft built at its own dtype, q_out follows it), test_dspark_profile, test_kv_pool_contract, test_v4_indexer_defaults (V4 fp8 fallback settled by Config on gfx950/gfx1250/gfx942; not under a plugin), test_compile_hash_spec, two lmcache stubs.

Validation

CPU only, in a container with torch and aiter importable and no GPU. The CPU subset (the test files above plus tests/test_config*, test_spec*, test_mtp_, test_lmcache, test_v4_, test_deepseek_v4, indexer-cp-gate, paged-attention-dispatch; script: run2.sh): on main f7b8962 25 failed / 1860 passed / 55 errors; on this branch 25 failed / 1904 passed / 55 errors; the failing set is identical (80 ids, missing aiter modules and a torch.finfo import error in this image). tests/test_dspark.py does not collect in that image on main either; its two new tests were run with torch.finfo shimmed and pass.
Not run on a GPU: an MHA EAGLE3 or DSpark draft with fp8 KV under a bf16 target; V4 fp8 -> bf16 fallback on a gfx942 host; the plugin path. The regression test for the EAGLE3 draft dtype builds EagleProposer through __new__ and checks the config it hands the draft model, so it needs no GPU but does not exercise attention.

Interaction with #2962

#2962 needs the draft-dtype rule "an fp4 target has a bf16 draft" (an fp4 draft has no layout). If this PR lands first, #2962 rebases onto it and adds that one clause to resolve_draft_kv_dtype and the MTP-fp4 refusal; its diff drops the draft-pool decision, the proposer/calibration/hash changes, the V4 fallback move and their tests. If #2962 lands first (with its flag-free copy of the same machinery), this PR rebases to the flag, kv_cache_dtype (requested), the offload refusal, docs and tests.

🤖 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 3235 --add-label <label>

[Copilot]

🟡 Changes recommended

The new contract tests incorrectly assume AiterBackend is selected, causing false failures when the supported ATOM_USE_UNIFIED_ATTN=1 setting selects TritonMHABackend.

1 open issue

• Disable unified attention for default-backend contract tests · New

What changed in this PR

Adds independent bf16/fp8 KV-cache selection for speculative drafts that own a cache pool, while keeping shared-pool drafts aligned with the target. This separates draft dtype configuration from the MiniMax-M3.1 FP4 cache work.

Changes:

• Adds --spec-decode-kv-cache-dtype, separates requested and resolved values, and validates incompatible configurations.
• Uses the resolved dtype for draft pools, proposers, calibration identity, and compile hashes; moves DeepSeek-V4 fallback resolution into Config.
• Updates documentation and regression coverage for dtype resolution, pool ownership, and CLI behavior.

File	Description
tests/​test_v4_indexer_defaults.py	Tests Config-level V4 fallback and plugin exclusion.
tests/​test_lmcache_mp.py	Adapts pool-ownership stubs to attention families.
tests/​test_lmcache_mp_native_worker.py	Adds speculative configuration to the worker stub.
tests/​test_kv_pool_contract.py	Checks pool-family declarations against backend factories.
tests/​test_dspark.py	Tests draft construction and query dtype selection.
tests/​test_dspark_profile.py	Tests dtype-sensitive calibration identity.
tests/​test_draft_kv_dtype.py	Adds resolution and configuration validation tests.
tests/​test_draft_kv_builder.py	Tests ownership, scale binding, and EAGLE3 dtype propagation.
tests/​test_compile_hash_spec.py	Adds resolved dtype to speculative test fixtures.
tests/​test_arg_utils_spec.py	Tests flag forwarding and invalid CLI combinations.
docs/​configuration_guide.md	Documents draft dtype options and speculative methods.
atom/​utils/​selector.py	Declares families with independent draft pools.
atom/​spec_decode/​eagle_proposer.py	Constructs draft attention with the resolved dtype.
atom/​spec_decode/​dspark_proposer.py	Applies draft dtype to construction and query selection.
atom/​spec_decode/​draft_kv.py	Centralizes ownership checks and uses draft pool dtype.
atom/​spec_decode/​calibration.py	Distinguishes profiles when draft and target dtypes differ.
atom/​model_ops/​attentions/​deepseek_v4_attn.py	Uses the dtype already settled by Config.
atom/​model_ops/​attentions/​backends.py	Removes the backend pool-ownership flag.
atom/​model_ops/​attentions/​aiter_attention.py	Removes the MHA ownership flag superseded by the family table.
atom/​model_engine/​arg_utils.py	Adds and validates the draft dtype CLI option.
atom/​kv_transfer/​offload/​mp/​deployment.py	Reuses the shared draft-ownership check.
atom/​config.py	Resolves draft dtype, enforces constraints, and updates hashing.

_{Review effort: Balanced}

[valar-cop]

a few small things..

gpu run for fp8 draft kv under bf16 target? and v4 fallback on gfx942?

ci green.

[valarLip]

⚠️ Action required before merge (4)

1. Show a GPU run of the new mode before merge (valar-cop's ask). The PR body says nothing was run on a GPU. Please post a short GSM8K or acceptance-rate check for an MHA draft (M3.1 DSpark or an MHA EAGLE3 draft) with --spec-decode-kv-cache-dtype fp8 under a bf16 target, compared against the same draft at bf16. Please also run one DeepSeek-V4 --kv_cache_dtype fp8 launch on gfx942 to confirm the fallback that moved into Config.__post_init__ (atom/config.py:2830) still produces bf16 pools and starts.
2. Reword the stale comment in deployment.py (valar-cop). At atom/kv_transfer/offload/mp/deployment.py:173, "A DSpark draft whose backend owns a KV pool" is now wrong. draft_owns_kv_pool also covers MHA EAGLE3, and the answer comes from selector.DRAFT_POOL_FAMILIES, not from a backend flag. Suggested wording: "A draft with its own KV pool (an MHA DSpark/EAGLE3 draft, draft_owns_kv_pool) …".
3. Fix the two comments this PR left stale. At atom/spec_decode/eagle_proposer.py:239-241, the comment still says only hf_config and compilation_config.level are changed, but line 251 now also sets kv_cache_dtype. Add it, as dspark_proposer.py:429 already does. At atom/spec_decode/draft_kv.py:52-54, "that backend says whether a draft … caches into a pool of its own" should say that the family table (DRAFT_POOL_FAMILIES) decides this.
4. Make the contract tests independent of ATOM_USE_UNIFIED_ATTN (Copilot). At tests/test_kv_pool_contract.py:194-230, get_attn_backend_cls(family, False, False) returns TritonMHABackend when ATOM_USE_UNIFIED_ATTN=1 (atom/utils/selector.py:182). With that set, builders == listed fails even though nothing is broken. Fix: set the env var to 0 with monkeypatch, or compare against the backend's MRO instead of the class name.

---

_{Reviewed at 1619cf6. Pre Checkin CI is green; heavy model jobs were skipped because the PR is not approved yet. The PR does not depend on #2962 and does not probe any aiter API.}

[valar-cop]

ci triage for a3802ba:

Running Copilot Code Review passed. https://github.com/ROCm/ATOM/actions/runs/43219336085

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/43219317315

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/43219316930

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/43219316664

Pre Checkin passed. https://github.com/ROCm/ATOM/actions/runs/43219317094

[Copilot]

🔵 Needs a closer look

The resolver accepts unsupported kv_cache_dtype values from Python callers, causing a later dtype lookup failure instead of a clear configuration-time error.

No open issues 🎉

1 previously missed issue

In code that hasn't changed since last review

Reject unsupported draft KV cache dtypes during configuration

atom/​config.py:1957

The argparse choices do not protect Python callers: an MHA draft's SpeculativeConfig(kv_cache_dtype="invalid") is accepted here and later fails at the dtype lookup in draft_kv_builder. Reject values other than None, "bf16", and "fp8" in this resolver so unsupported requests fail during configuration with a clear error.

1 resolved since last review

• Disable unified attention for default-backend contract tests

_{Review effort: Balanced}

[ganyi1996ppo]

@valarLip @valar-cop review round at 1619cf6, now 9d7ad4c (comments and a test fixture only, no code change).

Done

• Item 4 (Copilot): tests/test_kv_pool_contract.py pins ATOM_USE_UNIFIED_ATTN=0 with an autouse fixture. Both contract tests failed on the old file with it set to 1; they pass now.
• Items 2 and 3: reworded deployment.py (own-pool draft, draft_owns_kv_pool), eagle_proposer.py (now lists kv_cache_dtype) and draft_kv.py (the family table DRAFT_POOL_FAMILIES decides, not a backend flag).
• CPU suite (tests/, minus tests/plugin and tests/diffusion): same 4 failures as the unmodified parent 1619cf6, none touching this PR. black and ruff are clean.

Item 1, first half: fp8 draft KV under a bf16 target, on a GPU
4x MI355X (gfx950), TP4, MiniMax-M3.1-Flash-Quark-MXFP4, --kv_cache_dtype bf16, M3.1 DSpark draft (own KV pool), --num-speculative-tokens 7, real acceptance, 5-shot GSM8K chat, first 200 problems, T=0. The only difference is --spec-decode-kv-cache-dtype bf16 vs fp8. Each arm was started twice, to see the run-to-run spread.

arm	GSM8K strict / flexible (of 200)	tokens per forward (accept rate)	draft pool at startup
draft bf16	182 / 191	4.695 (0.528)	28040 blocks, 9.19 GB per rank
draft bf16, repeat	189 / 193	4.658 (0.523)	28033 blocks
draft fp8	182 / 191	4.612 (0.516)	28764 blocks, 4.86 GB per rank
draft fp8, repeat	181 / 190	4.646 (0.521)	28771 blocks

Paired bf16 vs fp8, strict: 3 vs 3 discordant, McNemar p = 1 (flexible 1 vs 1). The two bf16 runs differ from each other by 3 vs 10 discordant (p = 0.09), so the repeat spread is as large as the arm gap.

Reading: no accuracy or acceptance difference is detectable at N=200 within the repeat noise (about 1.5 points on GSM8K; acceptance is cumulative over 3 chat prompts plus the 200 problems). That is not a proof of zero loss. The fp8 draft pool is about half the bytes. The server starts and serves with the flag.

Not done

• The DeepSeek-V4 --kv_cache_dtype fp8 launch on gfx942: we have gfx950 only. Can CI or a maintainer with a gfx942 card run it? Only the V4 fp8 fallback that moved into Config.__post_init__ needs checking (bf16 pools, server starts).
• Not covered by the run above: an MHA EAGLE3 draft, contexts beyond GSM8K length, N above 200.

[ganyi1996ppo]

@valarLip @valar-cop status at a3802ba (was 9d7ad4c): one new commit, a code check plus one test.

Changed

• Copilot's resolver point (review 5873229871, "Reject unsupported draft KV cache dtypes during configuration"): fixed in a3802ba. resolve_draft_kv_dtype now raises for a SpeculativeConfig.kv_cache_dtype other than None, "bf16" or "fp8", for an own-pool and a shared-pool draft alike (+5 lines in atom/config.py; new test test_an_unsupported_draft_dtype_is_refused_whoever_asks, +6 lines).

Your 11-12 09:19 review

• Item 1, GPU run of the new mode: the fp8-vs-bf16 draft run on 4x MI355X is in my 11:04 comment above. This commit only adds a check, so it was not re-run. The DeepSeek-V4 --kv_cache_dtype fp8 launch on gfx942 is not done: there is no gfx942 host here. Asking: can a maintainer or CI run one launch and confirm bf16 pools and a clean start?
• Items 2, 3 and 4 (stale comments in deployment.py, eagle_proposer.py, draft_kv.py; the contract tests and ATOM_USE_UNIFIED_ATTN): fixed in 9d7ad4c, as in the 11:04 comment. Not changed again.

What I ran, on a CPU host with aiter importable and no GPU: test_draft_kv_dtype, test_draft_kv_builder, test_kv_pool_contract, test_arg_utils_spec, test_compile_hash_spec, test_dspark_profile, test_v4_indexer_defaults, test_lmcache_mp*: 394 passed. tests/test_dspark.py cannot be collected in that environment (aiter stub), so it was not run there. black and ruff are clean on the changed files. I have not looked at CI for a3802ba yet.

[Copilot]

🔵 Needs a closer look

The tests bypass the real draft KV builder, leaving target/draft dtype mismatch allocations unverified and a regression to the target dtype undetected.

No open issues 🎉

1 previously missed issue

In code that hasn't changed since last review

Add CPU regression test for draft KV pool dtype allocation

atom/​spec_decode/​draft_kv.py:74

Please add a CPU regression test that calls the real draft_kv_builder and checks the kv_dtype received by make_kv_pool for both target/draft dtype mismatches. The new proposer tests replace draft_kv_builder with a lambda returning None, while the resolution tests only check config fields, so reverting this allocation to the target dtype would go undetected even though the draft's attention would still be constructed at a different dtype. Stub the backend lookup and AITER dtype mapping rather than the builder itself.

_{Review effort: Balanced}