Skip to content

[KV Connector] Add scheduler context for lazy block access - #53067

Draft
ivanium wants to merge 1 commit into
vllm-project:mainfrom
ivanium:codex/kv-connector-scheduler-context
Draft

[KV Connector] Add scheduler context for lazy block access#53067
ivanium wants to merge 1 commit into
vllm-project:mainfrom
ivanium:codex/kv-connector-scheduler-context

Conversation

@ivanium

@ivanium ivanium commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add a narrow KVConnectorSchedulerContext that exposes the GPU block pool and temporary access to current request blocks
  • bind the context once during scheduler initialization and forward it through MultiConnector and SimpleCPUOffloadConnector
  • make Mooncake borrow current blocks only when it emits save/load metadata, while serializing block IDs and collecting unique blocks to pin in one pass
  • remove Mooncake RequestTracker.allocated_block_ids and the old connector-level bind_gpu_block_pool() API

Performance

Connectors that do not use the context have no scheduler hot-path work; they receive one context object during initialization. Mooncake no longer snapshots or incrementally maintains request block IDs. It borrows scheduler-owned block lists only for emitted metadata jobs, serializes all KV groups, skips null blocks for pinning, and deduplicates blocks shared with partial-tail offloads.

Relationship to #51358

This is prerequisite groundwork for #51358. It separates the scheduler/connector ownership boundary and lazy Mooncake block access from the Mamba boundary-state expansion and worker batching in that PR.

Duplicate work

I searched open vLLM PRs for KV connector scheduler context and request-block borrowing work and found no PR implementing this interface or Mooncake migration.

Tests

  • /home/yifanqiao/vllm/.venv/bin/python -m pytest tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py -q — 33 passed
  • pre-commit run ruff-check --files <changed files> — passed
  • pre-commit run mypy-3.12 --hook-stage manual --files <changed production files> — passed
  • commit-time pre-commit hooks, including ruff, mypy, forbidden imports, and sign-off checks — passed
  • Mooncake connector, MultiConnector, and scheduler override suites were attempted but could not collect in this CPU environment because the compiled FlashAttention CUDA extension is unavailable

No model evaluation is needed because this change does not affect model outputs or accuracy.

AI assistance

AI assistance was used for code exploration, implementation, test updates, and review. The human submitter reviewed the changes and is responsible for the final contribution.

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant