connector/250801

[ryanolson]

Summary by CodeRabbit

New Features

• Introduced a distributed key-value block manager (KVBM) with leader-worker architecture for scalable, multi-tier KV cache management across device, host, and disk.
• Added vLLM integration with a new KV cache manager, connector protocol, and Python interfaces for leader and worker roles.
• Implemented advanced block layouts, including per-layer separate storage and locality-aware block management.
• Added controller APIs for querying, resetting, and managing cache blocks at multiple levels.
• Provided comprehensive documentation and user guides for KVBM and vLLM integration.

Bug Fixes

• Improved error handling and validation throughout block management and transfer operations.
• Enhanced test coverage and removed outdated tests, ensuring robust cache and block lifecycle management.

Documentation

• Added detailed guides, architecture diagrams, and test plans for block manager, slot management, and distributed operation.

Refactor

• Generalized block and block manager abstractions for locality and logical resource support.
• Streamlined block data, transfer, and layout APIs for extensibility and clarity.

Tests

• Introduced new test suites for KVBM, vLLM integration, and slot/block management, covering cache hits, misses, eviction, and edge cases.

Chores

• Updated dependencies, configuration, and workspace settings to support new features and multi-architecture builds.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Caution

Review failed

Failed to post review comments.

Walkthrough

This update introduces a comprehensive distributed key-value block manager (KVBM) system for language model serving, featuring a leader-worker architecture, locality-aware block abstractions, and integration with vLLM for efficient KV cache management. Major changes include new Rust modules for block management, transfer, controller, and connector protocols, as well as extensive Python bindings and vLLM integration layers. Documentation and detailed test plans are provided, and legacy block manager tests are replaced with new KVBM-focused tests.

Changes

Cohort / File(s)	Change Summary
Containerization & Build Integration container/Dockerfile.kvbm, container/build.sh, container/run.sh, Cargo.toml, lib/llm/Cargo.toml, lib/bindings/python/Cargo.toml	Adds a multi-stage Dockerfile for KVBM, extends build/run scripts for KVBM support, updates Cargo features/dependencies, and adjusts dev profile optimization.
Workspace & Python Analysis dynamo.code-workspace	Adds extra Python analysis paths for improved language server support.
Documentation & Guides docs/guides/run_kvbm_in_vllm.md, lib/llm/src/block_manager.md, lib/llm/src/block_manager/distributed/README.md, lib/bindings/python/rust/llm/block_manager/vllm/slot_manager_test_plan.md, lib/bindings/python/rust/llm/block_manager/vllm/slot_test_plan.md	Introduces new documentation: user guide for KVBM in vLLM, block manager lifecycle, distributed async messaging, and comprehensive test plans.
Rust Block Manager Core lib/llm/src/block_manager.rs, lib/llm/src/block_manager/config.rs, lib/llm/src/block_manager/block.rs, lib/llm/src/block_manager/block/data.rs, lib/llm/src/block_manager/block/data/local.rs, lib/llm/src/block_manager/block/data/logical.rs, lib/llm/src/block_manager/block/data/logical/distributed_leader_worker.rs, lib/llm/src/block_manager/block/data/logical/null.rs, lib/llm/src/block_manager/block/data/view.rs, lib/llm/src/block_manager/block/factory.rs, lib/llm/src/block_manager/block/factory/local.rs, lib/llm/src/block_manager/block/factory/logical.rs, lib/llm/src/block_manager/block/locality.rs, lib/llm/src/block_manager/block/state.rs	Refactors block manager for locality-awareness, introduces logical/local block data, new block factories, state management, and error handling.
Block Transfer & Layout lib/llm/src/block_manager/block/transfer.rs, lib/llm/src/block_manager/block/transfer/cuda.rs, lib/llm/src/block_manager/block/transfer/memcpy.rs, lib/llm/src/block_manager/block/transfer/nixl.rs, lib/llm/src/block_manager/layout.rs, lib/llm/src/block_manager/layout/nixl.rs, lib/llm/src/block_manager/layout/utils.rs	Adds/updates block transfer logic (local, CUDA, NIXL), introduces new layout types (LayerSeparate), utility functions, and serialization support.
Distributed Block Manager lib/llm/src/block_manager/distributed.rs, lib/llm/src/block_manager/distributed/leader.rs, lib/llm/src/block_manager/distributed/worker.rs, lib/llm/src/block_manager/distributed/transfer.rs, lib/llm/src/block_manager/distributed/utils.rs, lib/llm/src/block_manager/distributed/zmq.rs	Implements distributed leader-worker block management, ZMQ-based messaging, block transfer handler, and related utilities.
Controller & Connector lib/llm/src/block_manager/controller.rs, lib/llm/src/block_manager/controller/client.rs, lib/llm/src/block_manager/controller/handler.rs, lib/llm/src/block_manager/connector.rs, lib/llm/src/block_manager/connector/protocol.rs, lib/llm/src/block_manager/connector/scheduler.rs	Adds controller for block pool management, connector protocol for leader-worker coordination, and a scheduler for transfer orchestration.
Python Rust Bindings: Core & Block Manager lib/bindings/python/rust/lib.rs, lib/bindings/python/rust/llm.rs, lib/bindings/python/rust/llm/block_manager.rs, lib/bindings/python/rust/llm/block_manager/block.rs, lib/bindings/python/rust/llm/block_manager/block_list.rs, lib/bindings/python/rust/llm/block_manager/controller.rs, lib/bindings/python/rust/llm/block_manager/distributed.rs, lib/bindings/python/rust/llm/block_manager/distributed/leader.rs, lib/bindings/python/rust/llm/block_manager/distributed/utils.rs, lib/bindings/python/rust/llm/block_manager/distributed/worker.rs, lib/bindings/python/rust/llm/block_manager/dlpack.rs, lib/bindings/python/rust/llm/block_manager/layer.rs, lib/bindings/python/rust/llm/block_manager/vllm.rs, lib/bindings/python/rust/llm/block_manager/vllm/block_list.rs, lib/bindings/python/rust/llm/block_manager/vllm/connector.rs, lib/bindings/python/rust/llm/block_manager/vllm/connector/leader.rs, lib/bindings/python/rust/llm/block_manager/vllm/connector/worker.rs, lib/bindings/python/rust/llm/block_manager/vllm/request.rs, lib/bindings/python/rust/llm/block_manager/vllm/slot.rs	Introduces extensive Python bindings for the new block manager, distributed leader/worker, controller, vLLM integration, and exposes new classes and methods for Python interop.
Python vLLM Integration Layer lib/bindings/python/src/dynamo/_core.pyi, lib/bindings/python/src/dynamo/llm/__init__.py, lib/bindings/python/src/dynamo/llm/vllm_integration/__init__.py, lib/bindings/python/src/dynamo/llm/vllm_integration/connector/__init__.py, lib/bindings/python/src/dynamo/llm/vllm_integration/connector/dynamo_connector.py, lib/bindings/python/src/dynamo/llm/vllm_integration/connector_leader.py, lib/bindings/python/src/dynamo/llm/vllm_integration/connector_worker.py, lib/bindings/python/src/dynamo/llm/vllm_integration/kv_cache_manager.py, lib/bindings/python/src/dynamo/llm/vllm_integration/kv_cache_utils.py, lib/bindings/python/src/dynamo/llm/vllm_integration/rust.py	Implements the Python layer for vLLM-KVBM integration: connector protocol, cache manager, leader/worker logic, utility wrappers, and Rust interop.
Testing: KVBM & vLLM Integration lib/bindings/python/tests/test_kvbm.py, lib/bindings/python/tests/test_kvbm_vllm_integration.py	Adds new pytest-based test suites for KVBM and its integration with vLLM, covering allocation, cache hit/miss, and error handling.
Testing: Legacy Block Manager lib/bindings/python/tests/test_block_manager.py	Removes legacy block manager test suite (now superseded by new KVBM tests).

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant vLLM
    participant Python KVBM Integration
    participant Rust KVBM Leader/Worker
    participant Controller/Scheduler

    User->>vLLM: Submit inference request
    vLLM->>Python KVBM Integration: Request KV cache blocks
    Python KVBM Integration->>Rust KVBM Leader/Worker: Allocate/check blocks
    Rust KVBM Leader/Worker->>Controller/Scheduler: (If needed) Schedule/coordinate transfer
    Controller/Scheduler-->>Rust KVBM Leader/Worker: Grant/deny transfer
    Rust KVBM Leader/Worker-->>Python KVBM Integration: Return block info
    Python KVBM Integration-->>vLLM: Provide block IDs/handles
    vLLM-->>User: Return model output

Loading

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120+ minutes

Complexity: This PR is extremely complex, introducing a distributed block manager system, new abstractions for locality, controller and connector protocols, extensive Python bindings, vLLM integration, new test plans, and documentation. The changes span core Rust, Python, Docker, and test code, requiring careful, multi-domain review.

Possibly related PRs

• feat: Restructure kv manager block registration #1093: Directly related, as both PRs restructure the block registration system in the key-value manager, including changes to block state and registration logic.
• chore(dynamo-run): Remove out=sglang|vllm|trtllm #1920: Related by context, as both PRs involve refactoring/disaggregation of components in the dynamo-run tool, though this PR focuses on block manager architecture.

Poem

In fields of memory, blocks align,
Rust and Python intertwine.
Leaders, workers, ZMQ’s call,
KV caches serve us all.
vLLM now runs with glee,
Distributed, fast, and rabbit-free!
🥕🐇✨

Note

⚡️ Unit Test Generation is now available in beta!

Learn more here, or try it out under "Finishing Touches" below.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.