[akii96]

Overview

Step-5 with FP8 KV cache crashes at startup if --block-size is not passed. The default block size of 16 is too small for the CSA tail ring. #61837 worked around this by telling users to pass --block-size 32. This PR makes vLLM pick a large enough block size on its own.

Claims

• --kv-cache-dtype fp8 works without --block-size.
• When 16 is already big enough, nothing changes. That covers BF16, and FP8 at TP2 and TP1.
• Other models are not affected. Only the Step-5 tail ring backend asks for a larger block.

Validation

Before this change, this command crashed on 8x MI300X after the weights loaded:

VLLM_ROCM_USE_AITER=1 vllm serve stepfun-ai/Step-5-Preview \
  --tensor-parallel-size 8 \
  --max-model-len 262144 \
  --kv-cache-dtype fp8 \
  --gpu-memory-utilization 0.90 \
  --moe-backend triton \
  --mm-encoder-attn-backend ROCM_AITER_FA \
  --mm-encoder-tp-mode data \
  --mm-processor-cache-type shm \
  --limit-mm-per-prompt '{"image":1,"video":0}' \
  --reasoning-parser step5 \
  --tool-call-parser step5 \
  --enable-auto-tool-choice \
  --safetensors-load-strategy prefetch

RuntimeError: Worker failed with error 'Step5 CSA tail ring needs 7 rows of 1024 bytes in one 6144-byte KV page; increase --block-size.'

With this branch, I ran the same command on 8x MI325X with --attention-backend ROCM_AITER_UNIFIED_ATTN added. All 8 workers picked block size 32 and the server started:

Setting kv cache block size to 32 for ROCM_AITER_UNIFIED_ATTN/STEP5_CSA_SUMMARY/STEP5_CSA_TAIL_RING backend(s).

I then ran GSM8K against that server with the same setup as #61837: all 1319 questions, 5-shot. It scored 97.27% strict and 97.50% flexible.

To check more than one KV head per GPU, I started the small Step-5 test model from the test registry with dummy weights, FP8 KV and no --block-size:

TP	KV heads per GPU	Block size picked
1	4	16
2	2	16
8	1	32

All three started and generated text.

The new unit test goes through the same startup code and then builds the tail ring's KV spec, which is the call that fails on main. It covers FP8 at TP8, TP2 and TP1, BF16, a larger speculative ring (moves to 64), and a padded page.

Details

The tail ring has to fit in one KV page. It needs 7 rows of 1024 bytes, plus one more row per speculative token. Step-5 has 4 KV heads, so at TP8 each GPU holds one. With FP8, a 16-token page is then 6144 bytes, which is only 6 rows. At TP2 each GPU holds two heads, so the page is already 12 rows and nothing changes.

vLLM already lets each attention backend say which block sizes it supports, and it picks a size that works for all of them at startup. The tail ring backend now takes part in that. It asks for 32, 64 or 128, whichever is the smallest that fits. If 16 already fits, it accepts any size. The math uses the page size without padding, so padding never counts as room for the ring.

The code is in common/csa_tail_ring.py, which the AMD and NVIDIA Step-5 models both use. If a user passes a --block-size that is too small, they still get the same error as before.

AI assistance was used to write this change. I reviewed the diff, ran the unit test, and ran the serve, GSM8K and TP checks above.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).
• AI assistance was used during the creation of this PR.
• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation

[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.

[simondanielsson]

Thanks for the fix!

Minor suggestion: can we check if this works also with >1 KV heads per rank? I.e. TP2

[akii96]

Thanks for the careful review @simondanielsson 🙏

Yes, it works at TP2. With dummy weights, FP8 KV and no --block-size, TP2 (2 KV heads per GPU) stays at block size 16, starts and generates. TP2 is also a case in the unit test now.

[simondanielsson]

Thanks for those changes 👍 suggesting a regression test but otherwise I think this looks good

[github-actions]

✅ @akii96, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[simondanielsson]

LGTM nice job!

[simondanielsson]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99342 for commit 1de53f39768e.