[Bugfix][KV Offload] Restore MTP-retained sliding-window history

[jacklin78911-collab]

Purpose

Fix native KV offload restoring MTP-retained sliding-window blocks that lookup never checked and the producer may never have stored. Implements the conservative coverage option discussed in #56701.

For example, with a 64-token external hit, SWA block size 16, window 32, and three extra retained tokens, lookup checks chunks 2–3 while allocation also restores chunk 1. If chunk 1 is absent or still being written, CPUOffloadingManager.prepare_load() asserts after the hit has already been admitted. Hybrid store pruning can itself omit chunk 1.

Include extra_retained_tokens in aligned/partial lookup coverage, the existing in-flight-load check, and producer store reachability. Preserve the base coverage lower bound so retaining one extra token does not unnecessarily request another chunk. The store mask uses a temporary coverage spec; actual attention windows, allocation, cache keys, and transfer ownership are unchanged. Restoring the required older blocks can increase CPU cache usage and transfer volume.

Duplicate checks on 2026-09-13 found no open PR for 56701 in:body or extra_retained_tokens offload. #55818 handles allocation scaling and #53499 page-size normalization; neither fixes this lookup/store mismatch.

Test Plan

Extend the existing scheduler tests to cover retained history at aligned and unaligned hybrid boundaries, EAGLE tail handling, missing/pending/loading chunks, and a producer followed by a cold GPU-cache consumer. The consumer only hits keys actually stored by the producer.

PYTHONPATH="$PWD" .venv/bin/python -m pytest \
  tests/v1/kv_connector/unit/offloading_connector \
  tests/v1/kv_offload/cpu/test_manager.py -q --tb=short

.venv/bin/python -m pre_commit run --files \
  vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py \
  tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

.venv/bin/python -m pre_commit run mypy-3.12 --hook-stage manual --files \
  vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py \
  tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

Test Result

• Related suites: 474 passed locally. Native H20: 378 passed, then the remaining 96 passed with --lf after seeding the missing public facebook/opt-125m config/tokenizer cache and enabling offline mode; no code changes.
• Focused regressions: 102 passed. With only the production file reverted to base de500290ffffeb7bd834872c0dea5331590dd40d: 14 failed, 88 passed.
• Pre-commit, including mypy 3.10, and the explicit mypy 3.12 hook: passed.
• Tiny Step3.5 smoke with two MTP modules, four speculative tokens, SWA window 32 / extra retention 3: both RTX 4060 V1 and native H20 V2 restored 64 cached tokens / 7 chunks, with 16/16 identical cold/warm output tokens.
• H20 V2 instantiated MultiModuleMTPSpeculator; the unmodified proposals exercised 30 decode calls with four rejected drafts, entering the three-token re-prefill path. A separate controlled draft-rejection check also passed the restore/output assertions.

V1 used Python 3.12.13 / PyTorch 2.11.0+cu130; H20 V2 used Python 3.12.14 / PyTorch 2.13.0+cu130 with UVA available. Compiled extensions were reused, and both changed source files on H20 were SHA-256 verified against 1e89fa93a662a321e29860f3f30a25a362070fc8 (the reused environment's package version metadata identifies the older build).

These GPU checks use tiny dummy weights, BF16, eager mode, and disabled optional JIT warmup. They validate the serving/re-prefill paths for this configuration. Trained-model quality, acceptance rate, and performance have not been evaluated; no throughput improvement is claimed.

Reproduce the H20 V2 functional smoke

Save the script below as ../tiny_mtp_v2.py relative to the checkout and run:

PYTHONPATH="$PWD" CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 \
  .venv/bin/python ../tiny_mtp_v2.py
FORCE_DRAFT_REJECTIONS=1 PYTHONPATH="$PWD" CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 \
  .venv/bin/python ../tiny_mtp_v2.py

The temporary model directory below is generalized from the executed script.

"""Tiny Step3.5 two-module MTP, real GPU kernels and CPU offload transfers."""
import os
os.environ['VLLM_ENABLE_V1_MULTIPROCESSING']='0'
os.environ['VLLM_USE_V2_MODEL_RUNNER']='1'
import json
from pathlib import Path
from vllm import LLM, SamplingParams
from vllm.transformers_utils.configs.step3p5 import Step3p5Config
from vllm.v1.kv_cache_interface import SlidingWindowSpec

model_dir=Path('/tmp/vllm-offload-mtp-smoke')
Step3p5Config(hidden_size=128, intermediate_size=256, num_attention_heads=2,
    num_attention_groups=1, head_dim=64, num_hidden_layers=2, vocab_size=256,
    max_position_embeddings=256, max_seq_len=256, sliding_window=32,
    num_nextn_predict_layers=2, moe_layers_enum='1', moe_num_experts=2,
    moe_top_k=1, moe_intermediate_size=128, share_expert_dim=128,
    architectures=['Step3p5ForCausalLM']).save_pretrained(model_dir)
llm=LLM(model=str(model_dir),skip_tokenizer_init=True,load_format='dummy',
    enforce_eager=True,dtype='bfloat16',seed=42,max_model_len=256,
    max_num_seqs=2,max_num_batched_tokens=256,gpu_memory_utilization=0.25,
    kv_cache_memory_bytes=16<<20, block_size=16,enable_prefix_caching=True,
    disable_hybrid_kv_cache_manager=False,compilation_config={'mode':0},
    kernel_config={'enable_jit_warmup':False},
    speculative_config={'method':'mtp','num_speculative_tokens':4},
    kv_transfer_config={'kv_connector':'OffloadingConnector','kv_role':'kv_both',
        'kv_connector_extra_config':{'cpu_bytes_to_use':16<<20}})
try:
    core=llm.llm_engine.engine_core.engine_core
    cfg=core.vllm_config
    assert cfg.speculative_config.use_multi_module_mtp()
    from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import MultiModuleMTPSpeculator
    runner=core.model_executor.driver_worker.worker.model_runner
    speculator=runner.speculator
    assert isinstance(speculator,MultiModuleMTPSpeculator),type(speculator)
    proposal_trace=[]
    force_rejections=os.environ.get('FORCE_DRAFT_REJECTIONS')=='1'
    original_propose=speculator.propose
    def track_propose(*args,**kwargs):
        batch=kwargs.get('input_batch',args[0] if args else None)
        rejected=kwargs.get('num_rejected',args[6] if len(args)>6 else None)
        proposal_trace.append({'has_prefill':bool(batch.has_prefill),
            'num_rejected':rejected[:batch.num_reqs].tolist()})
        drafts=original_propose(*args,**kwargs)
        if force_rejections:
            drafts.fill_(75)
        return drafts
    speculator.propose=track_propose
    connector=core.scheduler.connector.connector_scheduler
    specs=[g.kv_cache_spec for g in connector.config.kv_group_configs]
    assert any(isinstance(s,SlidingWindowSpec) and s.extra_retained_tokens==3 for s in specs)
    load_keys=[]
    original_load=connector.manager.prepare_load
    def track_load(keys,ctx):
        keys=list(keys)
        load_keys.append(len(keys))
        return original_load(keys,ctx)
    connector.manager.prepare_load=track_load
    prompt={'prompt_token_ids':list(range(10,75))}
    params=SamplingParams(temperature=0,max_tokens=16,ignore_eos=True)
    cold=llm.generate([prompt],params,use_tqdm=False)[0]
    for _ in range(1000):
        if not core.scheduler.has_requests(): break
        core.step()
    assert not core.scheduler.has_requests()
    assert llm.reset_prefix_cache(reset_connector=False)
    warm=llm.generate([prompt],params,use_tqdm=False)[0]
    assert cold.outputs[0].token_ids==warm.outputs[0].token_ids
    assert warm.num_cached_tokens>0 and load_keys
    assert any(not call['has_prefill'] for call in proposal_trace)
    if force_rejections:
        assert any(max(call['num_rejected'])>1 for call in proposal_trace)
    result={'proposal_trace':proposal_trace,'force_draft_rejections':force_rejections,
        'speculator_class':type(speculator).__name__,'cold':cold.outputs[0].token_ids,'warm':warm.outputs[0].token_ids,
        'cached_tokens':warm.num_cached_tokens,'load_chunk_counts':load_keys,
        'kv_specs':[str(s) for s in specs],'weights':'dummy',
        'mtp_modules':2,'speculative_tokens':4,'model_runner':'V2'}
    Path('../tiny_mtp_v2_rejection_result.json' if force_rejections else '../tiny_mtp_v2_result.json').write_text(json.dumps(result,indent=2)+'\n')
    print(json.dumps(result))
finally:
    llm.llm_engine.engine_core.shutdown()

AI assistance: OpenAI Codex assisted with investigation, implementation, testing, and this description; the commit includes attribution.

---

PR description checklist

• Purpose and related issue provided.
• Test commands provided.
• Regression and GPU test results, including limitations, provided.
• Documentation impact considered: internal correctness fix; no public option or API change.

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

[orozery]

/ci run

[github-actions]

❌ This PR is 112 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[orozery]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89041 for commit 5a5ad8442336.

[jacklin78911-collab]

Thanks @orozery for updating the branch and running CI. Both V1 KV Connectors jobs hit the upstream NIXL test-stub failure already covered by #56955. The KV Offload test lanes and all KV Offload lm-eval lanes passed. Once #56955 lands, we can update this branch and rerun CI.

[github-actions]

✅ @jacklin78911-collab, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch 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.

[orozery]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89251 for commit ebf29254005a.