merge: integrate PR #4742 - re-introduce kvbm-kernels

Merged PR #4742 (oandreeva/kvbm-kernels) which re-introduces the KVBM kernels
library after it was reverted from main.

Key changes:
- Added lib/kvbm-kernels with CUDA tensor kernel operations
- Integrated kernels module into kvbm bindings (kept alongside v2 module)
- Updated workspace to include kvbm-kernels
- Added prebuilt CUDA binaries and static libraries for x86_64
- Updated vLLM Dockerfile with CUDA dev tools
- Moved vectorized_copy.fatbin to kvbm-kernels library

Conflicts resolved:
- Cargo.toml: Added kvbm-kernels to workspace members
- lib/bindings/kvbm/Cargo.toml: Added kvbm_kernels dependency to block-manager feature
- lib/bindings/kvbm/src/lib.rs: Added kernels module alongside existing v2 module
- lib/bindings/kvbm/python/kvbm/__init__.py: Exported both kernels and v2 modules
- Cargo.lock files: Accepted PR version and will regenerate

Signed-off-by: Ryan Olson <[email protected]>

Author: ryanolson

.gitattributes

@@ -8,6 +8,7 @@
 *.[Zz][Ii][Pp]      binary
 *.[Tt][Gg][Zz]      binary
 *.fatbin            binary
+*.a                 binary
 
 # Exclude test data files from linguist language detection
 lib/llm/tests/data/** linguist-vendored

.github/workflows/pre-merge-rust.yml

@@ -92,7 +92,7 @@ jobs:
     runs-on:
       group: Fastchecker
     strategy:
-      matrix: { dir: ['.', 'lib/bindings/python', 'lib/runtime/examples', 'launch/dynamo-run'] }
+      matrix: { dir: ['.', 'lib/bindings/python', 'lib/runtime/examples', 'launch/dynamo-run', 'lib/kvbm-kernels'] }
     permissions:
       contents: read
     steps:
@@ -117,7 +117,10 @@ jobs:
         key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
         restore-keys: ${{ runner.os }}-cargo-
     - name: Set up Rust Toolchain Components
-      run: rustup component add clippy
+      run: rustup component add clippy rustfmt
+    - name: Verify Code Formatting
+      working-directory: ${{ matrix.dir }}
+      run: cargo fmt -- --check
     - name: Run Clippy Checks
       working-directory: ${{ matrix.dir }}
       run: cargo clippy --no-deps --all-targets -- -D warnings

.sandbox/launch_vllm_with_connector.sh

@@ -67,8 +67,8 @@ kv_transfer_config='{
   "kv_role": "kv_both",
   "kv_connector_module_path": "kvbm.v2.vllm.schedulers.connector",
   "kv_connector_extra_config": {
-    "leader": { 
-      "cache": { "host": { "cache_size_gb": 1.0 } },
+    "leader": {
+      "cache": { "host": { "cache_size_gb": 4.0 } },
       "tokio": { "worker_threads": 2 }
     },
     "worker": {

Cargo.lock

@@ -21,6 +21,7 @@ members = [
     "lib/nova-discovery",
     "lib/kvbm",
     "lib/kvbm-config",
+    "lib/kvbm-kernels",
 ]
 # Exclude certain packages that are slow to build and we don't ship as flagship
 # features from default build, but keep them in workspace for convenience.
@@ -59,6 +60,7 @@ dynamo-nova = { path = "lib/nova", version = "0.7.0" }
 dynamo-nova-backend = { path = "lib/nova-backend", version = "0.7.0" }
 dynamo-nova-discovery = { path = "lib/nova-discovery", version = "0.7.0" }
 dynamo-kvbm = { path = "lib/kvbm", version = "0.7.0" }
+dynamo-kvbm-kernels = { path = "lib/kvbm-kernels", version = "0.7.0" }
 dynamo-memory = { path = "lib/memory", version = "0.7.0" }
 dynamo-runtime = { path = "lib/runtime", version = "0.7.0" }
 dynamo-llm = { path = "lib/llm", version = "0.7.0" }

Cargo.toml

@@ -651,6 +651,13 @@ FROM runtime AS dev
 ARG WORKSPACE_DIR=/workspace
 
 USER root
+
+# Copy additional CUDA tools needed for building kernels from source (cargo build)
+COPY --from=base /usr/local/cuda/bin/nvlink /usr/local/cuda/bin/nvlink
+COPY --from=base /usr/local/cuda/bin/crt /usr/local/cuda/bin/crt
+COPY --from=base /usr/local/cuda/lib64/libcudart_static.a /usr/local/cuda/lib64/
+COPY --from=base /usr/local/cuda/lib64/libculibos.a /usr/local/cuda/lib64/
+
 # Install utilities as root
 RUN apt-get update -y && \
     apt-get install -y --no-install-recommends  \

container/Dockerfile.vllm

@@ -21,7 +21,7 @@ crate-type = ["cdylib", "rlib"]
 
 [features]
 default = []
-block-manager = ["dynamo-llm/block-manager", "dep:dlpark", "dep:cudarc"]
+block-manager = ["dynamo-llm/block-manager", "dep:dlpark", "dep:cudarc", "dep:kvbm_kernels"]
 
 [dependencies]
 dynamo-llm = { path = "../../llm" }
@@ -31,6 +31,7 @@ dynamo-memory = { path = "../../memory" }
 dynamo-nova = { path = "../../nova" }
 dynamo-nova-backend = { path = "../../nova-backend" }
 dynamo-runtime = { path = "../../runtime" }
+kvbm_kernels = { path = "../../kvbm-kernels", package = "dynamo-kvbm-kernels", optional = true }
 
 anyhow = { version = "1" }
 async-stream = { version = "0.3" }
@@ -73,7 +74,7 @@ pyo3-async-runtimes = { version = "0.23.0", default-features = false, features =
 pythonize = "0.23"
 
 dlpark = { version = "0.5", features = ["pyo3", "half"], optional = true }
-cudarc = { version = "0.16.2", features = ["cuda-12020"], optional = true }
+cudarc = { version = "0.17.1", features = ["cuda-12020"], optional = true }
 prometheus = "0.14.0"
 
 [dev-dependencies]

lib/bindings/kvbm/Cargo.lock

@@ -16,6 +16,7 @@
 # from kvbm._core import KvbmWorker as KvbmWorker
 
 from kvbm._core import __version__
+from kvbm._core import kernels as kernels
 from kvbm._core import v2 as v2
 
-__all__ = ["__version__", "v2"]
+__all__ = ["__version__", "kernels", "v2"]