[Bug] may_append allocates new block one token too late, causing KV cache write to unallocated block

[MrPaoBrother]

Description

BlockManager.may_append allocates a new physical block when len(seq) % block_size == 1, but the correct condition should be == 0. This causes the new block to be allocated after the first token that needs it has already been scheduled, meaning that token's
KV cache entry will be written to an unallocated block.

Root Cause

In block_manager.py:

def can_append(self, seq: Sequence) -> bool:
    return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

def may_append(self, seq: Sequence):
    if len(seq) % self.block_size == 1:
        seq.block_table.append(self._allocate_block())

may_append is called during the schedule phase, before the new token is appended to the sequence. At this point, len(seq) reflects the
number of tokens already stored in KV cache blocks.

For a concrete example with block_size=256:

┌────────────────────────────┬──────────┬───────┬───────────────────────┬───────────────────────────────────┐
│            Step            │ len(seq) │ % 256 │   Current behavior    │         Correct behavior          │
├────────────────────────────┼──────────┼───────┼───────────────────────┼───────────────────────────────────┤
│ Prefill done (256 tokens)  │ 256      │ 0     │ ❌ No block allocated │ Should allocate — block 0 is full │
├────────────────────────────┼──────────┼───────┼───────────────────────┼───────────────────────────────────┤
│ 1st decode token generated │ 257      │ 1     │ Allocates (too late!) │ Already has block allocated       │
├────────────────────────────┼──────────┼───────┼───────────────────────┼───────────────────────────────────┤
│ ...                        │ ...      │ ...   │ ...                   │ ...                               │
├────────────────────────────┼──────────┼───────┼───────────────────────┼───────────────────────────────────┤
│ Block 1 full (512 tokens)  │ 512      │ 0     │ ❌ No block allocated │ Should allocate — block 1 is full │
├────────────────────────────┼──────────┼───────┼───────────────────────┼───────────────────────────────────┤
│ Next decode token          │ 513      │ 1     │ Allocates (too late!) │ Already has block allocated       │
└────────────────────────────┴──────────┴───────┴───────────────────────┴───────────────────────────────────┘

When len(seq) == 256, block 0 is full. The upcoming decode token (the 257th token) needs to be written to block 1, but may_append does
not allocate it because 256 % 256 == 0 ≠ 1. The allocation only happens on the next schedule cycle, after the token has already been
scheduled and presumably written to KV cache — at which point block 1 is not yet in the block table.

Impact

KV cache writes for the first token of every new block will target an unallocated physical block, which can cause:

- Silent data corruption in KV cache
- Memory corruption / out-of-bounds access
- Incorrect generation results

Fix

Change both conditions from == 1 to == 0:

def can_append(self, seq: Sequence) -> bool:
    return len(self.free_block_ids) >= (len(seq) % self.block_size == 0)

def may_append(self, seq: Sequence):
    if len(seq) % self.block_size == 0:
        seq.block_table.append(self._allocate_block())

can_append must also be fixed because it acts as a guard for may_append — if they disagree on whether a new block is needed, the
scheduler could skip preemption and crash with an IndexError when may_append tries to pop from an empty free_block_ids.

[Anai-Guo]

I traced this against scheduler.py and model_runner.py, and I believe the current == 1 is actually correct — the == 0 change would write KV to an unallocated block.

The key is the call ordering: may_append runs inside Scheduler.schedule(), but seq.append_token() runs later in Scheduler.postprocess() (after the forward of the previous step). So at decode-schedule time, len(seq) already includes the token whose KV the upcoming forward will write. In prepare_decode, that forward writes positions = len(seq) - 1 into block_table[-1] at offset last_block_num_tokens - 1.

Trace with block_size=256:

• After prefill: len=256, one block, KV for positions 0–255 written. Prefill postprocess appends T256 → len=257.
• Decode schedule: 257 % 256 == 1 → allocate block 1. prepare_decode writes position 256 → block_table[-1] (block 1), offset 0. ✅
• len=258 … 512: % 256 ∈ {2…0}, no new block; positions 257–511 fill block 1. ✅
• len=513: 513 % 256 == 1 → allocate block 2 for position 512. ✅

With == 0, at len=257 no block would be allocated, yet the forward needs to write position 256 to a second block — block_table[-1] is still block 0, so it would overwrite position 0's KV. Decode may_append never observes len % block_size == 0 right when a block fills, because postprocess always increments len before the next schedule.

Could you double-check the trace on your end? Happy to be corrected if I've misread the flow.

[Anai-Guo]

Thanks for the very detailed writeup — but I believe the current == 1 condition is actually correct, and changing it to == 0 would write the first token of each block into the previous (full) block. The analysis hinges on when the token is appended to seq relative to when may_append runs.

The per-step order (llm_engine) is:

1. scheduler.schedule() → calls block_manager.may_append(seq)
2. model forward (writes KV for the scheduled token)
3. scheduler.postprocess() → calls seq.append_token(token_id)

So the sampled token is appended in postprocess, before the next decode schedule(). That means at decode-schedule time len(seq) already includes the token whose KV this step will write. A decode step writes KV for the token at index len(seq) - 1 (the one appended by the previous postprocess), not for a token that hasn't been generated yet.

A new block is therefore needed exactly when that index is the first slot of a new block:

(len(seq) - 1) % block_size == 0   ⇔   len(seq) % block_size == 1

which is what may_append/can_append already check.

Re-tracing the block_size = 256 example, the key correction is that after the prefill of 256 prompt tokens, postprocess appends the first sampled token, so len(seq) is 257, not 256, by the time the first decode step is scheduled:

Step	len(seq) at schedule	% 256	input token index (KV written this step)	block needed	allocated?
prefill done, 1st token appended in postprocess	257	1	256	1	✅ yes (==1)
next decode	258	2	257	1 (already have it)	no
…	…	…	…	…
token index 511 appended	513	1	512	2	✅ yes (==1)

So block 1 is allocated on the step that writes slot 256 (the first slot of block 1) — just in time, not too late. With == 0, allocation would instead fire when len(seq) is a multiple of block_size, i.e. one step early relative to the token actually being written, leaving the boundary token to be written into the just-filled block.

I ran the repro mentally against scheduler.py postprocess (seq.append_token happens there for both the final prefill token and every decode token), which is what makes len(seq) lead the written index by one. Happy to be corrected if I've misread the engine loop, but I don't think there's KV corruption here. 🙏