[rebel-jinhwan]

🚀 Summary of Changes

Fixes a livelock in which a rank stops generating once the KV cache is full: the scheduler keeps cycling on one prefill and no request ever finishes. This is the cause of the stalls seen in long-context benchmarks.

The scheduler raises num_lookahead_tokens to num_spec_tokens even for ngram (#1128), and the final prefill chunk allocates those slots. The full-sequence admission gate (full_sequence_must_fit) counts request.num_tokens alone. A prompt whose last block has fewer free slots than the lookahead (in particular, any block-aligned prompt) therefore needs one block more on its final chunk than it was admitted with. When the rest of the pool is held by decodes:

1. The final chunk fails to allocate, and FCFS preemption pops the last running request, which is the prefill itself.
2. Next step it is admitted again (it just returned its blocks), and with no mixed batching its admission evicts the decode batch.
3. It prefills alone up to the final chunk and repeats 1.

No decode ever runs, so no block is freed.

The same happens one step later, which review of the upstream fix (vllm-project/vllm#63288) pointed out: the first decode step after the prefill schedules the sampled token and the draft tokens on top of the lookahead, num_tokens + 1 + num_spec_tokens + num_lookahead_tokens slots. A prompt that fits with the lookahead alone is admitted and preempts itself on that step after its prefill ran.

The gate now covers the sequence through that step. The scheduler passes the step's main-model tokens, 1 + num_spec_tokens, to allocate_slots as num_decode_step_tokens (RBLNKVCacheManager forwards it), and a registry patch, patches/admission_lookahead.py, does the counting on two upstream methods:

• KVCacheManager.allocate_slots takes num_decode_step_tokens, which vllm 0.31 does not have, and records it with num_lookahead_tokens in a ContextVar for the duration of a call made with full_sequence_must_fit=True.
• KVCacheCoordinator.get_num_blocks_to_allocate, on the gate's call only (apply_admission_cap=True), adds the decode step to num_tokens_main_model and the lookahead on top of it to num_tokens, both capped at max_model_len. That restores the coordinator's own contract, in which num_tokens includes the lookahead and num_tokens_main_model excludes it; the gate breaks it by passing the prompt for both. Upstream's block arithmetic still does the counting.

Outside a gated call the variable is unset, so upstream's other admission-cap caller (Scheduler._request_remaining_blocks) counts what it did. The waiting loop already passes effective_lookahead_tokens, so the async P/D + EAGLE case that zeroes the lookahead keeps zero. Unlike upstream, the decode step is counted for async KV loads too: this scheduler has no separate spec decode reservation for them. The upstream fix is vllm-project/vllm#63288, which adds the same num_decode_step_tokens argument; once the vllm pin includes it, the patch goes and the scheduler's call stays.

Split out of #1315. #1315 now carries only the sub-block match fix and is stacked on this PR. This one has to land first: the retry in #1315 removes an accidental escape that dev has from this livelock, so #1315 alone would turn a configuration that finishes today into one that stalls.

Model path: vllm model path only. The optimum model path uses optimum_scheduler.py, which does not have this admission path.

---

📌 Related Issues / Tickets

• Related to fix(scheduler): retry admission without the sub-block match when allocation fails #1315 (stacked on this PR) and [Bugfix][Core] Cover the first spec decode step in the full-sequence admission gate vllm-project/vllm#63288 (the upstream fix). No issue filed; found as a rank that stopped generating mid-run in a long-context LMCache benchmark (DP + EP, ngram spec decode, sub-block prefix caching) on vllm-rbln 0.12.1a3.

---

✅ Type of Change

• 🚀 Release (release)
• ✨ Feature (feature)
• 🧠 Model support (model)
• 🧬 Core engine changes (core)
• 🛠 Bug fix (fix)
• ⚙️ Performance improvement (perf)
• 🔁 Refactor or code cleanup (refactor)
• 📄 Documentation (docs)
• ❓ Other (other): please describe

---

🧪 How to Test

1. uv run --no-sync pytest "tests/vllm/v1/core/test_rbln_scheduler.py::TestAdmissionCoversStepAfterPrefill"
   Expected: a prompt admitted with two free blocks, while two drafting decodes hold the rest of the pool, finishes without being preempted. Two prompts, both two chunks long: 32 tokens, whose last chunk's lookahead needs a third block, and 26 tokens, which fits with the lookahead but whose first decode step needs 26 + 1 + 3 + 3 slots. Each runs with sub-block caching off, at sub-block size 4 and at 8, with and without a KV connector.
   On dev all 12 cases fail; with only the lookahead counted, the six 26-token cases still fail.

---

📸 Screenshots / Logs (if applicable)

This host cannot install the locked environment (glibc 2.35; dev pins vllm 0.31.0+cpu), so this was checked on vllm 0.30.0+cpu with the same changes applied on top of 44c3ece, CPU only:

TestAdmissionCoversStepAfterPrefill, dev:                  12 failed
TestAdmissionCoversStepAfterPrefill, lookahead only:        6 failed, 6 passed  (the 26-token cases fail)
TestAdmissionCoversStepAfterPrefill, this PR:              12 passed
test_rbln_scheduler.py + test_sub_block_prefix_caching.py + test_rbln_kv_cache_manager.py + tests/vllm/patches: 406 passed

Please let CI confirm on 0.31.0.

---

📋 Checklist

• PR title follows Conventional Commits format
• This PR is linked to an existing issue
• The test method is described, and the expected result is clearly stated
• Relevant documentation has been updated (if applicable)

---

💬 Notes

• The patch relies on upstream internals. It depends on the gate calling coordinator.get_num_blocks_to_allocate with keywords and marking that call with apply_admission_cap=True, and on callers passing full_sequence_must_fit by keyword. A positional argument would raise; a renamed flag would silently skip the margin, which TestAdmissionCoversStepAfterPrefill catches. The alternative was duplicating upstream's admission arithmetic (retained-hit trimming, prefill_end, watermark) in the scheduler.
• The scheduler now passes an argument vllm 0.31 lacks. Until the pin includes [Bugfix][Core] Cover the first spec decode step in the full-sequence admission gate vllm-project/vllm#63288, the patched allocate_slots consumes num_decode_step_tokens; the patch applies in every vllm-model-path process.
• The margin assumes every draft slot is used, as upstream's _request_remaining_blocks does, and a max_tokens == 1 request reserves it too although it never decodes.
• Earlier revisions of this PR first made the lookahead change inside the scheduler, by swapping the coordinator's method on the instance for one allocate_slots call, then as the registry patch counting the lookahead only. The decode step came from review of the upstream PR.
• sub_block_size=None does not turn sub-block caching off in create_rbln_scheduler: it takes the prefill chunk. The test therefore sets enable_sub_block_cache explicitly, and sub-block caching off is the case without the accidental escape.
• Not done here: making a lone prefill preempt a decode instead of itself, and surfacing repeated preemption. A request that never emits a token is not counted in vllm:num_preemptions_total, which is why these stalls showed only one preemption in metrics. Both are worth separate changes.

🤖 Generated with Claude Code

[rebel-jinhwan]

Reworked: the gate now counts the lookahead through a registry patch (vllm_rbln/patches/admission_lookahead.py) instead of the scheduler-local wrapper, so rbln_scheduler.py is unchanged. The description is updated to match; #1315 is rebased on top.

🤖 Generated with Claude Code

[rebel-jinhwan]

@rebel-cjlee with high kv cache pressure and speculative decoding, livelock can be occurred (admit -> prefill -> try decode -> preempt itself -> admit -> prefill -> try decode -> ...). This PR fixes admission gate to take account not only prompt tokens but also lookahead and decode step tokens. Could you take a look when you have a time?

[rebel-cjlee]

@rebel-jinhwan How often this show up in practice?
Also, the patch implementation looks a bit hacky to me, so I'll take a closer look at it in more detail.

[rebel-jinhwan]

@rebel-cjlee thanks, it appears on high concurrency & spec dec where every block is used. It is not very often but it makes livelock, so i think that it is the issue that must be fixed