[Bugfix][ROCm] Bound kpool prefill tail-seed stores by NUM_TAIL_BLOCKS

[zzw09773]

Prefill _kpool_tail_seed_kernel stores into tail[blk] using blk = tslot // kpool. After prefix-cache churn, tslot can carry a main-KV-magnitude physical id while the tail tensor is a 1-block-per-request ring (n_tail << physical id). Decode already skipped block >= n_tail; seed did not.

Measured (gfx942, GLM-5.3-Flash TP4)

See #56037 Crash D (2026-09-07).

• Dying step: chunked prefill, computed=211200, scheduled=2338 (2338 % 4 == 2 so the last chunk must seed the tail ring), scheduled_spec_decode_tokens={}, new_block_ids around 3535–3578.
• HIP_LAUNCH_BLOCKING=1 + AMD_LOG_LEVEL=3 on the dying image, GPU 0 only:

hipModuleGetFunction(..., _kpool_tail_seed_kernel)
ShaderName : _kpool_tail_seed_kernel
Memory Fault Error
Memory access fault by GPU node-2

Same probe with if blk >= NUM_TAIL_BLOCKS: return does not fault.

This PR

Pass NUM_TAIL_BLOCKS = tail_kv_cache.shape[0] into the prefill seed kernel (AMD + NVIDIA copies) and return before the store when blk is out of range.

Out of scope: mixed-length MTP decode (Crash C) and equal-length short MTP decode (Crash E) in #56037. Those are separate paths; do not close #56037 on this PR alone.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0b964b0a9d

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Bound the decode tail accesses as well

In the reported prefix-cache scenario where this guard sees blk >= tail_kv_cache.shape[0], it only suppresses the prefill seed fault. The same request retains that tail-block mapping during decode, but both AMD and NVIDIA _kpool_decode_update_batched_kernel only require tail_slot >= 0 before using block = tail_slot // POOL_SIZE for tail reads and stores (AMD lines 499-509 and 600-609; NVIDIA lines 502-512 and 603-612). Consequently, the next decode step can access the same out-of-range block and GPU-fault; the decode path needs the corresponding upper-bound handling rather than relying on the new seed-only check.

Useful? React with 👍 / 👎.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀