fix(kv): reserve prefill-activation headroom and size KV pool for the tightest TP rank

[carlushuang]

Motivation

On tensor-parallel runs, ModelRunner.get_num_blocks() can size the KV pool too large on the "heavy" ranks, so they overshoot --gpu-memory-utilization and OOM on the first large prefill. Under P/D disaggregation this is silent: the prefill worker OOM-crashes and the proxy hangs waiting on the dead engine (see #1483 for the full investigation).

Root cause is in the budget formula:

non_kv_overhead = peak_torch + cudagraph_overhead + safety_margin
available_for_kv = min(budget - non_kv_overhead, free)

1. peak_torch is the torch-allocator high-watermark from warmup. It misses this process's non-torch resident memory — NCCL/RCCL comm buffers, HIP context, custom-all-reduce / P2P IPC — which lands unevenly across TP ranks (coordinator ranks carry tens of GB more; measured ~34 GB on rank 0/1 for DeepSeek-V4-Pro, and it is not weights: peak_torch is identical across ranks, and not RCCL: NCCL_DEBUG shows ~10 GB/rank uniform). It can also under-measure the true serving prefill activation when warmup runs a reduced path (DeepSeek-V4 short-circuits its sparse-attention dispatch under is_dummy_run).

2. Because the utilization budget binds uniformly, every rank computes the same block count — sized for the lightest rank. The heavy ranks then overshoot the target (observed post-init: 99% vs an 85% target) and OOM on the first prefill whose activation doesn't fit.

Change

Reserve peak-activation headroom off the actual per-rank free (not off the utilization budget), then all_reduce(MIN) across the TP group so the uniform KV pool is sized for the tightest rank:

prefill_headroom = int(float(os.environ.get("ATOM_KV_PREFILL_HEADROOM_FRAC", "0.08")) * total)
available_for_kv = min(available_for_kv_budget, free - prefill_headroom)
# all_reduce(MIN) available_for_kv across the TP group  (no-op when the budget binds on every rank)

On balanced setups the utilization budget still binds on every rank, so available_for_kv is unchanged and the all_reduce(MIN) is a no-op — num_kvcache_blocks is unchanged and there is no throughput impact. On imbalanced setups the heavy rank's smaller free binds, so the pool is sized conservatively enough to leave real prefill headroom (or, if the request cache genuinely cannot fit, the existing "reduce --max-num-seqs" error fires at init instead of a mid-serving OOM).

The headroom fraction is env-tunable via ATOM_KV_PREFILL_HEADROOM_FRAC (default 0.08).

Test Result

8×MI355X (gfx950), deepseek-ai/DeepSeek-V4-Pro, bf16 KV:

Aggregated TP=8 (balanced) — no perf drop:

• num_kvcache_blocks: 49230 → 49228 (0.004% — the all_reduce(MIN) picking the true-min rank; util budget still binds).
• output throughput c64: 2231 → 2195 tok/s; c256: ~4999 → 5304 tok/s (within run-to-run noise; fewer prompts in the spot re-run).

TP=4 P/D (imbalanced) — OOM fixed:

• Original config (--max-num-seqs 512, util 0.85, default chunk): previously started then OOM-crashed the prefill worker at c=16 (silent hang). Now fails fast at init with the actionable "reduce --max-num-seqs" message.
• With --max-num-seqs 128 (as the message suggests) at default attn-prefill-chunk-size: serves c=16 cleanly at 12.8 req/s, prefill worker stays alive (0 OOM events) — no chunked-prefill workaround needed.

Notes / follow-ups

• This addresses the sizing/robustness side. The complementary fix — warming up DeepSeek-V4 without is_dummy_run (or adding an explicit sparse-attention activation reserve) so peak_torch reflects the true prefill peak — is left for a separate change.
• Draft: the 0.08 default headroom fraction is a heuristic; happy to tune it or derive it from max_num_batched_tokens if preferred.

Refs #1483