[Bugfix][KV] Zero the never-written tail of async-loaded blocks

[zhewenl]

Purpose

Fix NaN outputs for requests whose KV cache was restored by an async load (any KV connector: NIXL P/D, MooncakeStore, MultiConnector, offloading) on models that need KV zeroing (hybrid Mamba/linear-attention + attention pools, mixed-precision caches).

Problem

Admission exempts the blocks an async load will write from alloc-time zeroing (skip_zeroing_group_ids, #59504), since zeroing could race the out-of-band write. A load fills only the leading tokens of the request's last block, so the rest of that block keeps its previous owner's bytes:

last loaded block (one rank):  [ loaded tokens | never written, never zeroed ]
                                               ^ kernel page holding the next token

1. Stale bytes can be NaN. In a hybrid pool the block may last have held another group's state, e.g. fp32 recurrent state, which read back as fp8/bf16 KV contains NaN/Inf bit patterns.
2. Decode reads them. Prefill never reads past its tokens, but decode attention reads the whole kernel page that holds the current position. Masked slots get softmax weight 0, yet P @ V still multiplies their V, and 0 × NaN = NaN. Every logit of the request goes NaN at its first decode step: degenerate output, or a crash once NaN logits reach the sampler.
3. Every connector is affected. A Store hit leaves the tail of the hit block. A NIXL transfer with equal P/D block sizes copies the producer's whole last block, so decode receives the producer's tail as is ([KV Connector] Support NIXL heterogeneous P/D block sizes for hybrid models #49612 / [Bugfix][NIXL] Pair P/D blocks by the last token with different block sizes #59746 only zero what a transfer with different block sizes leaves uncovered).

Fix

• Scheduler. When an async load completes successfully (_update_waiting_for_remote_kv, the completion point shared by all connectors), queue each zeroed attention group's partially filled last block with its valid token count as a KVBlockTail in SchedulerOutput.kv_block_tails_to_zero. This runs before cache_blocks, which can hash a partial block under a finer hash_block_size; cached, shared and null blocks are never queued. Gated by needs_kv_cache_zeroing, no new flag.
• Worker. KVBlockZeroer.zero_block_ids(block_ids, tails) zeroes new blocks whole and tails past their valid slots in one launch of the existing Triton kernel, before the forward pass (the V1 runner keeps a separate tail call so runners that override _zero_block_ids still zero new blocks). For every group with a layer in it, each zeroer segment (one kernel page of a logical block) records the page index, slots per page and slot stride, so groups that share a KV cache tensor all reach the shared segments. A tail entry skips fully valid pages and masks the partial one by slot index (word // slot_stride % page_slots), which is exact for every KVCacheLayout (token-major and head-major pages, head-group-outer, block-major with interleaved layers) and for kernel-block splitting. Views whose slots don't tile int32 words, and indexer caches (their pages store all values, then all scales, unlike their [B, N, C] view), keep the old behavior with a one-time warning.
• Block-free fence. The zeroing runs with the step that queues the tail, even if that step does not schedule the request. With overlapping batches (defer_block_free), queueing a tail fences the request's blocks like a scheduled request's, so an abort cannot hand them to a new async load before the zeroing step has run.
• DCP. A DCP-sharded block spans block_size × dcp tokens dealt to ranks in runs of cp_kv_cache_interleave_size; the zeroer converts the valid token count to this rank's valid slots (replicated groups use the count as is).
• Steps without tails take the unchanged path (HAS_TAILS is a constexpr), and warmup compiles the tail variant as a no-op.

This replaces the first version of this PR, which zeroed the tail inside MooncakeStoreWorker only. That version did not cover NIXL or other connectors, and its register-time layout assert fails under DCP > 1 (the worker's group block size is DCP-scaled).

Co-authored with @GirasoleY.

Not a duplicate. #59504 decides which blocks are exempt from zeroing and still (correctly) exempts the partial last block. #49612 and #59746 zero the part of a NIXL transfer that mismatched P/D block sizes leave uncovered. #39283 extends zeroing to full-attention-only models. None zeroes the never-written tail of a completed load. Open-PR searches for the area keywords returned nothing else.

Test Plan

On a GB200 node:

pytest tests/v1/worker/test_kv_block_zeroer.py \
       tests/v1/core/test_scheduler.py \
       tests/v1/core/test_prefix_replay.py \
       tests/v1/worker/test_utils.py \
       tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py \
       tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py \
       tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py \
       tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py \
       tests/v1/kv_connector/unit/test_error_propagation.py \
       tests/v1/worker/test_gpu_warmup_blocks.py
pre-commit run ruff-check / ruff-format (v0.14.0) on the touched files
mypy on the touched source files (compared against main)

New tests:

• test_async_load_queues_last_block_tail_once (scheduler): the tail is queued exactly once, in the step that first schedules the request, with the loaded token count; nothing for a block-aligned load; a 12,288-token (DCP8-scaled) block.
• test_received_last_block_tail_queued_once (remote prefill): the same for a NIXL-style load, including a hybrid (Mamba) config with a coarser cache block than the hash block, where the count is taken before the full-hit N−1 adjustment.
• test_tail_local_valid_slots_follow_token_ownership: every prefix of a block against explicit per-rank token ownership (DCP 1/2/8 × interleave 1/4/16, plus a replicated group).
• test_zero_tails_keep_exactly_the_valid_slots: byte-exact against a reference for every KVCacheLayout, with and without kernel-block splitting, tails ending inside the first and the second kernel page.
• test_zero_tails_dcp8_mla: DCP8, 1,536-slot rank-local blocks of 24 64-slot fp8 MLA pages; a 275,089-token load ends in block 22 with 595 valid slots on rank 0 and 594 elsewhere. Exactly the rest of that block is zeroed on every rank.
• test_zero_tails_of_groups_sharing_a_tensor: two groups sharing one tensor; each group's tail reaches it.
• test_async_load_tail_fences_block_free: a promoted request that is not scheduled in the queueing step and is then aborted keeps its blocks out of the pool until that step is processed.
• test_zero_new_blocks_and_tails_in_one_launch, test_tails_skip_layouts_whose_slots_are_not_word_aligned, test_tails_skip_indexer_caches.

Test Result

• 382 passed, 8 skipped (layouts that cannot split blocks into kernel pages), on 3 GPUs (two pipeline-parallel scheduler tests need them).
• Mutation checks: keeping one extra slot in the tail mask fails 25 of the new tail tests; dropping the fence fails the fence test; deduplicating shared segments by their first group fails the shared-tensor test.
• ruff clean; mypy reports the same errors as main on the touched files (none new).
• End to end, first version of this PR (Kimi-K3 P/D with NIXL + MooncakeStore, gsm8k 5-shot, two passes): without zeroing 0.9538 / 0.9484, 20 requests non-finite at their first decode step, all with a non-finite unwritten page tail; with tail zeroing 0.9591 / 0.9651, 0 non-finite (healthy baseline 0.963–0.967). An earlier implementation of this design was also checked against FlashInfer's fp8 MLA decode on GB300 under DCP8: a stale tail gave 100% NaN output, the zeroed tail 0% NaN, with output and LSE bit-identical to a clean tail (not rerun on this version). TODO: rerun the end-to-end eval on this version before marking ready.

Known gaps: indexer caches are skipped (see above); the CPU model runner's own _zero_block_ids does not zero tails.

AI assistance (Claude Code) was used for the port, the tests and this description.

🤖 Generated with Claude Code

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zhewenl.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork