[Bugfix] Bound KV block zeroing launch geometry

[LucasWilkinson]

Purpose

Fix the KVBlockZeroer launch overflow reproduced on main nightly #83443, at commit 3e372c5ff2:

OverflowError: signed integer is greater than maximum

DeepSeek-V4 combines 181 KV segments with 9,344- and 292-element pages. The old zeroer selected the largest common power-of-two divisor, so the 292-element page forced every segment to use four-element chunks. Zeroing 6,870 blocks therefore flattened to:

6,870 * 181 * 2,336 = 2,904,745,920 programs

That overflows the signed launch dimension passed by the NVIDIA wrapper.

This change:

• maps blocks, segments, and chunks directly onto a 3-D grid, keeping block IDs on the large x-axis;
• uses up to 1,024 elements per program and masks each segment's tail, so small pages no longer shrink all chunks;
• reduces the failing geometry to (6,870, 181, 10), or 12,434,700 programs (233.6x fewer);
• preserves one compiled kernel across different block counts.

The H100/H200 sm90_paged_mqa_logits_metadata failures from the same CI build were separate. They are already fixed on current main by #52035; its exact 4xH100 KV-offload job passed both DeepSeek-V4 cases (2 passed).

Duplicate-work check

Required open PR/issue searches were run for KVBlockZeroer overflow, KV block zeroing grid, and signed integer zeroing.

#50485 also proposes a 3-D grid, but this is materially different:

• [Bugfix][ROCm] Launch KV block zeroing on a 3-D grid #50485 is currently conflicted and predates packed block strides;
• it puts block IDs on the z-axis, which is limited to 65,535 entries on CUDA;
• it retains narrow divisor-based chunks;
• this change puts blocks on x and adds masked wide chunks, addressing both the observed NVIDIA overflow and the 233.6x excess work.

Test plan and results

CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m pytest \
  tests/v1/worker/test_kv_block_zeroer.py \
  tests/v1/worker/test_dsv4_packed_zeroer_geometry.py -q
# 9 passed

CUDA_VISIBLE_DEVICES=0 compute-sanitizer --tool memcheck --error-exitcode 99 \
  .venv/bin/python -m pytest tests/v1/worker/test_kv_block_zeroer.py -q
# 8 passed; ERROR SUMMARY: 0 errors

.venv/bin/pre-commit run --files \
  vllm/v1/worker/utils.py tests/v1/worker/test_kv_block_zeroer.py
# passed

An exact-shape B200 validation launched grid (6870, 181, 10), zeroed all storage, and completed in 0.334 seconds.

No model eval was run because this only changes how already-selected KV-cache bytes are zeroed, not model outputs or scheduling semantics.

Disclosure

AI assistance (OpenAI Codex) was used to investigate the CI history, develop the change, and draft this description. reviewed by the submitter

Breaking PR and overlap

The failure became reachable after #51749 generalized worker-side KV zeroing to every allocating AttentionSpec. That change is needed to prevent stale FP8 sliding-window pages; it exposed a pre-existing launch-geometry scaling bug when DeepSeek-V4 contributes many heterogeneous segments.

#52062 addresses the same observed overflow by reverting #51749. This PR is materially different: it preserves generalized zeroing and bounds the kernel launch using a 3-D grid plus masked wide chunks, fixing the overflow without restoring the stale-cache bug that #51749 corrected.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[njhill]

Thanks @LucasWilkinson

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83624 for commit 5f8ac6872b92.

[zzw09773]

Independently reproduced the engine-killing overflow on 8×MI325X (gfx942) serving
DeepSeek-V4-Flash-0731 at max_model_len=1M: a single 506,470-token prefill raised
OverflowError: signed integer is greater than maximum in _zero_kv_blocks_kernel
and killed the engine (all in-flight requests 500). With this PR backported onto the
08-12 nightly, the same request completes (~90 s prefill, ≈5.6k tok/s). +1 for landing.