Overview

With MTP and prefix caching on, hybrid models such as GLM-5.3-Flash can return garbage when a request reuses a cached prompt prefix. Some prefill steps save a checkpoint of the linear-attention state in the middle of a chunk. In those steps the scheduler leaves an unused MTP scratch block in the slot that belongs to the checkpoint. That block is never written, yet it is registered in the prefix cache ahead of the real checkpoint, so later requests resume from it. This change moves the scratch block out of the way in those steps too, as already happens in every other step.

The skip was added together with internal checkpoints in #53614. GLM-5.3-Flash reaches this path since #56960. The branch is based on current main (58b3298). No open PR changes speculative-block relocation; #57128 changes the Mamba lookup and does not touch this path.

Claims

• Fixes persistent garbage output after prefix-cache hits for hybrid models with internal prefill checkpoints and MTP, for example GLM-5.3-Flash with --speculative-config '{"method":"mtp","num_speculative_tokens":1}'.
• No behavior change without speculative decoding, for mamba_cache_mode other than align, or for steps without an internal checkpoint.

Validation

Setup: RedHatAI/GLM-5.3-Flash-NVFP4 on 2 H200 GPUs, served with default settings plus MTP:

vllm serve RedHatAI/GLM-5.3-Flash-NVFP4 --tensor-parallel-size 2 \
    --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}'

Reproducer (script below): send a long prompt while another long prompt is still being processed, then send the exact same prompt again. The second request reuses the cached start of the prompt.

• Before the fix, the second answer was garbage in 3 of 3 trials, for example File: /home/erdos/2024/4/4 # File: /home/erdos/2024/4/4 ..., while the first answer was a normal list of function names.
• With the fix, both answers were normal in 3 of 3 trials.

Other checks:

• A single long prompt sent twice, with nothing else running, also came back broken in 4 of 7 runs at about 25k tokens before the fix, and in 0 of 12 runs after.
• In an 8-agent coding-agent run (10 turns, full history resent each turn), 2 agents switched to word salad right after reusing their cached history and stayed broken. With the fix, and with MTP off, no agent broke.
• New test test_mamba_checkpoint_hash_maps_only_to_checkpoint_block fails on main and passes with the fix.
• pytest tests/v1/core/test_single_type_kv_cache_manager.py tests/v1/core/test_prefix_caching.py tests/v1/core/test_mamba_align_chunk_split.py tests/v1/core/test_prefix_replay.py tests/v1/core/prefix_cache passes; the GPU part, test_hybrid_prefix_cache_hit_rate.py (base and mtp), was run on one H200.

"""Reproducer: with MTP on, sending the same long prompt twice can return garbage the second time.

Server (any recent main):
  vllm serve RedHatAI/GLM-5.3-Flash-NVFP4 --tensor-parallel-size 2 \
      --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}'

Run:
  python repro.py http://localhost:8000/v1

Each trial sends a long prompt while another long prompt is running, then sends the
same prompt again. The second answer reuses the cached start of the prompt. Before
the fix it is often garbage; after the fix both answers are normal.
"""

import random
import sys
import threading

from openai import OpenAI

client = OpenAI(base_url=sys.argv[1] if len(sys.argv) > 1 else "http://localhost:8000/v1", api_key="none")
MODEL = client.models.list().data[0].id
WORDS = "alpha beta gamma delta kernel buffer index cache shard block token commit".split()


def long_prompt(seed, size):
    """Made-up Python code of about size characters, different for every seed."""
    rng = random.Random(seed)
    lines = []
    while sum(map(len, lines)) < size:
        lines.append(f"def {rng.choice(WORDS)}_{rng.choice(WORDS)}(x): return {rng.randrange(10**6)}  "
                     f"# {' '.join(rng.choices(WORDS, k=8))}\n")
    return "Here is a file:\n" + "".join(lines) + "\nList the first five function names, then summarize the file."


def ask(prompt, max_tokens):
    out = client.chat.completions.create(
        model=MODEL, messages=[{"role": "user", "content": prompt}], max_tokens=max_tokens, temperature=0,
        extra_body={"chat_template_kwargs": {"reasoning_effort": "low"}})  # GLM-5.3: short thinking
    msg = out.choices[0].message
    thinking = getattr(msg, "reasoning", None) or getattr(msg, "reasoning_content", None) or ""
    return (thinking + (msg.content or "")).replace("\n", " ")


def run_together(prompts, max_tokens):
    threads = [threading.Thread(target=ask, args=(p, max_tokens)) for p in prompts]
    for t in threads:
        t.start()
    for t in threads:
        t.join()


# Whether a prompt is affected depends on its exact length; these lengths trigger it.
for trial, size in enumerate([60000, 60000, 60000]):
    # Keep the server busy for a moment, as it would be in production.
    run_together([long_prompt(1000 + 10 * trial + k, 90000) for k in range(6)], 64)
    # A long request is running while the target prompt is processed for the first time.
    background = threading.Thread(target=ask, args=(long_prompt(10 * trial + 1, size), 300))
    background.start()
    threading.Event().wait(0.3)
    target = long_prompt(10 * trial + 2, size)
    first = ask(target, 60)
    background.join()
    # The exact same request again.
    second = ask(target, 60)
    print(f"trial {trial}")
    print(f"  first:  {first[:100]}")
    print(f"  second: {second[:100]}")

trial 0
  first:  **First five function names:** 1. `alpha_beta` 2. `index_token` 3. `buffer_cache` 4. `gamma_shard` 5
  second: List the first five function names, then summarize the file. </span></div></div></div></div></
trial 1
  first:  **First five function names:** 1. `cache_kernel` 2. `delta_shard` 3. `token_buffer` 4. `cache_beta`
  second: File: /home/erdos/2024/4/4  # File: /home/erdos/2024/4/4 File: /home/erdos/2024/4/4 File: /home/erdo
trial 2
  first:  **First five function names:** 1. `gamma_delta` 2. `alpha_buffer` 3. `index_delta` 4. `kernel_gamma`
  second: The first five function names, then summarize the file.The first five function names, then summarize

trial 0
  first:  **First five function names:** 1. `alpha_beta` 2. `index_token` 3. `buffer_cache` 4. `gamma_shard` 5
  second: **First five function names:** 1. `alpha_beta` 2. `index_token` 3. `buffer_cache` 4. `gamma_shard` 5
trial 1
  first:  **First five function names:** 1. `cache_kernel` 2. `delta_shard` 3. `token_buffer` 4. `cache_beta`
  second: **First five function names:** 1. `cache_kernel` 2. `delta_shard` 3. `token_buffer` 4. `cache_beta`
trial 2
  first:  **First five function names:** 1. `gamma_delta` 2. `alpha_buffer` 3. `index_delta` 4. `kernel_gamma`
  second: **First five function names:** 1. `gamma_delta` 2. `alpha_buffer` 3. `index_delta` 4. `kernel_gamma`

Details

In align mode a request keeps its running state in the block after the last computed one, followed by num_speculative_blocks scratch blocks. A normal allocate_new_blocks call moves the old scratch blocks that fall below the new running state to the end of the table and nulls their slots. In a checkpoint step this was skipped (if blocks_allocated and not checkpoint_block).

With MTP, the checkpoint position P is shifted one block down. When one of the previous step's scratch blocks sits at index P // block_size - 1 (with k=1: the chunk starts one block below P), it occupies the natural slot of P. The checkpoint itself is written into a new block at cdiv(end, block_size) - 2, one index above it. In the same step, cache_full_blocks registers the unwritten scratch block under hash(P), and _cache_partial_tail_block then adds the real checkpoint block under the same hash. get_one_block returns the first block registered, so later hits at P resume from the scratch page. The entry stays until eviction, so every request that shares this prefix is affected.

It happens when a prompt's final prefill chunk starts one block below the shifted checkpoint, so it depends on prompt length and on how much of each step's token budget other requests use. Overlapping long prompts make it much more likely. For example, with a second request of 22350 tokens, a first request of up to 23424 tokens triggers it (3/3) and one of 23440 tokens does not (0/3).

The fix relocates scratch blocks below the checkpoint index (block_idx < num_skipped_blocks - checkpoint_block). A scratch block exactly at the checkpoint index stays and becomes the checkpoint block, since the worker exports the checkpoint into that column. Because relocated blocks are reused, checkpoint steps for already-allocated requests no longer need num_speculative_blocks extra new blocks. The allocation cap in allocate_new_blocks and the admission estimate in get_num_blocks_to_allocate both drop that term, so they stay equal. The existing test_mamba_checkpoint_admission_matches_allocation and the new test both check this.

On unfixed builds, disabling MTP or prefix caching avoids the bug.

Possible CI follow-up: no current job would catch this. The only nightly on this path, the B200 Kimi-K3 spec-decode GSM8K job, uses ~640-token prompts, and GSM8K accuracy does not move noticeably. A small extra phase in that job could catch this class of bug. After the GSM8K pass, it would send about 15 long prompts with lengths spread over one chunk-budget period, then send each again and check that the cached answer starts the same as a fresh computation. On GLM-5.3-Flash such a phase failed on unfixed main in 3 of 3 runs and passed with this fix in 3 of 3. I can send it as a separate PR if that is useful.

AI assistance (Claude) was used; the change was reviewed and validated by the submitter.

🤖 Generated with Claude Code