Stop hardcoding bf16 for the DSV4 unified_kv buffers

[AMD-yanfeiwang]

Problem

DeepSeekV4TokenToKVPool takes --kv-cache-dtype and forwards it to every sub-pool it owns — the SWA / c4 / c128 single pools and the c4 indexer pool — but DeepSeekV4UnifiedKVPool ignored the argument entirely:

torch.zeros(..., dtype=torch.bfloat16, device=device)   # deepseek_v4_memory_pool.py

Under SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton (HIP / MI355X) that makes --kv-cache-dtype fp8_e4m3 a no-op for the main KV buffers, with nothing in the log to say so. The operator asks for the 584 B/token/layer packed layout (nope fp8 448 + rope bf16 128 + fp8 scales 8) that DeepSeekV4SingleKVPool.get_bytes_per_token() implements and that pool_configurator._get_bytes_per_full_token() budgets for, and silently gets head_dim × bf16 = 1024 B/token/layer.

On DeepSeek-V4-Pro (61 layers: 30 × ratio-4, 31 × ratio-128) that is 1.75× the KV bytes the flag implies — both resident and on the wire in PD disaggregation:

	per token per rank
c4 KV, 30 layers, at the implied 584 B/row	4,380 B
c4 KV, 30 layers, as actually allocated	7,680 B
c128 KV, 31 layers, implied / actual	141 B / 248 B

Why the fix is not "make it fp8"

The unified buffers genuinely cannot be fp8 today:

• sparse_attn_v4_paged_prefill rejects unified_kv.dtype != q.dtype, and requires q to be fp16/bf16 — there is no fp8 prefill path at all.
• sparse_attn_v4_paged_decode has an fp8 QUANT_KV path, but it needs a per-slot fp32 kv_scales tensor that this pool neither owns nor populates.

So bf16 is currently required, not merely chosen. This PR keeps the effective result but makes it a decision instead of an accident.

Change

• resolve_unified_kv_dtype(kv_cache_dtype) — passes 16-bit float dtypes through (an fp16 model now gets fp16 buffers instead of bf16 ones) and warns on anything else, naming both the requested dtype and the real byte cost.
• DeepSeekV4UnifiedKVPool takes dtype as a required keyword argument and raises on a dtype its kernels cannot read, so a future caller cannot reintroduce the silent substitution one level up.

Behavior for the existing --kv-cache-dtype fp8_e4m3 DSV4 recipes is unchanged apart from the new warning.

Tests

test/registered/unit/mem_cache/test_dsv4_unified_kv_dtype.py (CPU, 10 tests): dtype pass-through and fp8 fallback-with-warning; buffers allocated at the requested dtype (asserted with fp16, so a reintroduced bf16 hardcode still fails); get_buf_infos() item lengths tracking the element size; construction rejected for fp8/fp32; an AST guard that the single construction site passes dtype=resolve_unified_kv_dtype(dtype); and a premise test that pins why the fallback exists — if the prefill kernel ever accepts a non-q.dtype KV buffer, that test fails and UNIFIED_KV_DTYPES should be widened rather than the fallback quietly kept.

Mutation-checked — each half of the fix reverted in isolation kills exactly its own test and nothing else:

mutant	tests killed
restore dtype=torch.bfloat16 in the allocation	test_buffers_use_the_requested_dtype, test_pool_init_has_no_hardcoded_dtype_literal
caller passes the constant instead of resolving	test_dtype_argument_is_resolved_not_hardcoded
fall back without warning	test_fp8_falls_back_and_says_so
drop the constructor guard	test_unsupported_dtype_is_rejected_at_construction

Not addressed here

_get_bytes_per_full_token() still budgets the unified layout at the packed 584 B/token/layer figure. It happens not to over-commit today only because the same formula also charges a swa_ratio * kv_bytes * num_layers_total SWA term that the unified pool does not allocate per token (its SWA ring is sized by request slots, not by token capacity), and at --swa-full-tokens-ratio 0.1 the two errors roughly cancel. That coincidence deserves its own fix.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #30981717518
Latest PR Test (Extra): ❌ Run #30981717352