feat: add kvbm-kernels crate and upgrade cudarc to 0.19 (#6309)

Signed-off-by: Ryan Olson <rolson@nvidia.com>

Cargo.lock

@@ -982,7 +982,7 @@ checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4"
 dependencies = [
  "glob",
  "libc",
- "libloading",
+ "libloading 0.8.9",
 ]
 
 [[package]]
@@ -1448,11 +1448,12 @@ dependencies = [
 
 [[package]]
 name = "cudarc"
-version = "0.17.8"
+version = "0.19.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0bf99ab37ee7072d64d906aa2dada9a3422f1d975cdf8c8055a573bc84897ed8"
+checksum = "aed81f178e780f3d5d354d12b4c5c5a484c4a9c329ecd037ac57f2a0e0648397"
 dependencies = [
- "libloading",
+ "half 2.7.1",
+ "libloading 0.9.0",
 ]
 
 [[package]]
@@ -1959,7 +1960,7 @@ dependencies = [
  "mockito",
  "modelexpress-client",
  "modelexpress-common",
- "ndarray",
+ "ndarray 0.16.1",
  "nix 0.26.4",
  "nixl-sys",
  "object_store",
@@ -2030,7 +2031,7 @@ dependencies = [
  "dynamo-kv-router",
  "dynamo-runtime",
  "dynamo-tokens",
- "ndarray",
+ "ndarray 0.16.1",
  "ndarray-interp",
  "ndarray-npy",
  "rand 0.9.2",
@@ -2848,6 +2849,9 @@ checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b"
 dependencies = [
  "cfg-if 1.0.4",
  "crunchy",
+ "num-traits",
+ "rand 0.9.2",
+ "rand_distr",
  "zerocopy",
 ]
 
@@ -3613,9 +3617,9 @@ dependencies = [
 
 [[package]]
 name = "js-sys"
-version = "0.3.86"
+version = "0.3.87"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d36139f1c97c42c0c86a411910b04e48d4939a0376e6e0f989420cbdee0120e5"
+checksum = "93f0862381daaec758576dcc22eb7bbf4d7efd67328553f3b45a412a51a3fb21"
 dependencies = [
  "once_cell",
  "wasm-bindgen",
@@ -3868,6 +3872,17 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "kvbm-kernels"
+version = "0.9.0"
+dependencies = [
+ "clap 4.5.60",
+ "cudarc",
+ "half 2.7.1",
+ "ndarray 0.17.2",
+ "rand 0.9.2",
+]
+
 [[package]]
 name = "kvbm-logical"
 version = "0.9.0"
@@ -3961,6 +3976,16 @@ dependencies = [
  "windows-link",
 ]
 
+[[package]]
+name = "libloading"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60"
+dependencies = [
+ "cfg-if 1.0.4",
+ "windows-link",
+]
+
 [[package]]
 name = "libm"
 version = "0.2.16"
@@ -4460,13 +4485,28 @@ dependencies = [
  "rawpointer",
 ]
 
+[[package]]
+name = "ndarray"
+version = "0.17.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d"
+dependencies = [
+ "matrixmultiply",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "portable-atomic",
+ "portable-atomic-util",
+ "rawpointer",
+]
+
 [[package]]
 name = "ndarray-interp"
 version = "0.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e43087829efb5ec2736598e88587df286425b59df5a9ce991994cdd2c5855d3f"
 dependencies = [
- "ndarray",
+ "ndarray 0.16.1",
  "num-traits",
  "thiserror 2.0.18",
 ]
@@ -4478,7 +4518,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7b313788c468c49141a9d9b6131fc15f403e6ef4e8446a0b2e18f664ddb278a9"
 dependencies = [
  "byteorder",
- "ndarray",
+ "ndarray 0.16.1",
  "num-complex",
  "num-traits",
  "py_literal",
@@ -4750,6 +4790,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
 dependencies = [
  "autocfg",
+ "libm",
 ]
 
 [[package]]
@@ -5951,6 +5992,16 @@ dependencies = [
  "getrandom 0.3.4",
 ]
 
+[[package]]
+name = "rand_distr"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463"
+dependencies = [
+ "num-traits",
+ "rand 0.9.2",
+]
+
 [[package]]
 name = "rand_xorshift"
 version = "0.4.0"
@@ -8575,9 +8626,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen"
-version = "0.2.109"
+version = "0.2.110"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9ff9c7baef35ac3c0e17d8bfc9ad75eb62f85a2f02bccc906699dadb0aa9c622"
+checksum = "1de241cdc66a9d91bd84f097039eb140cdc6eec47e0cdbaf9d932a1dd6c35866"
 dependencies = [
  "cfg-if 1.0.4",
  "once_cell",
@@ -8588,9 +8639,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-futures"
-version = "0.4.59"
+version = "0.4.60"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d24699cd39db9966cf6e2ef10d2f72779c961ad905911f395ea201c3ec9f545d"
+checksum = "a42e96ea38f49b191e08a1bab66c7ffdba24b06f9995b39a9dd60222e5b6f1da"
 dependencies = [
  "cfg-if 1.0.4",
  "futures-util",
@@ -8602,9 +8653,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro"
-version = "0.2.109"
+version = "0.2.110"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "39455e84ad887a0bbc93c116d72403f1bb0a39e37dd6f235a43e2128a0c7f1fd"
+checksum = "e12fdf6649048f2e3de6d7d5ff3ced779cdedee0e0baffd7dff5cdfa3abc8a52"
 dependencies = [
  "quote",
  "wasm-bindgen-macro-support",
@@ -8612,9 +8663,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro-support"
-version = "0.2.109"
+version = "0.2.110"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dff4761f60b0b51fd13fec8764167b7bbcc34498ce3e52805fe1db6f2d56b6d6"
+checksum = "0e63d1795c565ac3462334c1e396fd46dbf481c40f51f5072c310717bc4fb309"
 dependencies = [
  "bumpalo",
  "proc-macro2",
@@ -8625,9 +8676,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-shared"
-version = "0.2.109"
+version = "0.2.110"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bc6a171c53d98021a93a474c4a4579d76ba97f9517d871bc12e27640f218b6dd"
+checksum = "e9f9cdac23a5ce71f6bf9f8824898a501e511892791ea2a0c6b8568c68b9cb53"
 dependencies = [
  "unicode-ident",
 ]
@@ -8681,9 +8732,9 @@ dependencies = [
 
 [[package]]
 name = "web-sys"
-version = "0.3.86"
+version = "0.3.87"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "668fa5d00434e890a452ab060d24e3904d1be93f7bb01b70e5603baa2b8ab23b"
+checksum = "f2c7c5718134e770ee62af3b6b4a84518ec10101aad610c024b64d6ff29bb1ff"
 dependencies = [
  "js-sys",
  "wasm-bindgen",

Cargo.toml

@@ -10,6 +10,7 @@ members = [
     "lib/mocker",
     "lib/kv-router",
     "lib/memory",
+    "lib/kvbm-kernels",
     "lib/kvbm-logical",
     "lib/async-openai",
     "lib/parsers",
@@ -41,6 +42,7 @@ dynamo-mocker = { path = "lib/mocker", version = "0.9.0" }
 dynamo-kv-router = { path = "lib/kv-router", version = "0.9.0", features = ["metrics"] }
 dynamo-async-openai = { path = "lib/async-openai", version = "0.9.0", features = ["byot"] }
 dynamo-parsers = { path = "lib/parsers", version = "0.9.0" }
+kvbm-kernels = { path = "lib/kvbm-kernels", version = "0.9.0" }
 kvbm-logical = { path = "lib/kvbm-logical", version = "0.9.0" }
 
 # External dependencies
@@ -58,7 +60,7 @@ chrono = { version = "0.4", default-features = false, features = [
     "now",
     "serde",
 ] }
-cudarc = { version = "0.17.8", features = ["cuda-12020"] }
+cudarc = { version = "0.19.2", features = ["cuda-12020"] }
 dashmap = { version = "6.1" }
 derive_builder = { version = "0.20" }
 derive-getters = { version = "0.5" }

lib/bindings/kvbm/Cargo.lock

@@ -808,7 +808,7 @@ checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4"
 dependencies = [
  "glob",
  "libc",
- "libloading",
+ "libloading 0.8.9",
 ]
 
 [[package]]
@@ -1106,20 +1106,11 @@ dependencies = [
 
 [[package]]
 name = "cudarc"
-version = "0.16.6"
+version = "0.19.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "17200eb07e7d85a243aa1bf4569a7aa998385ba98d14833973a817a63cc86e92"
+checksum = "aed81f178e780f3d5d354d12b4c5c5a484c4a9c329ecd037ac57f2a0e0648397"
 dependencies = [
- "libloading",
-]
-
-[[package]]
-name = "cudarc"
-version = "0.17.8"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0bf99ab37ee7072d64d906aa2dada9a3422f1d975cdf8c8055a573bc84897ed8"
-dependencies = [
- "libloading",
+ "libloading 0.9.0",
 ]
 
 [[package]]
@@ -1564,7 +1555,7 @@ dependencies = [
  "bytemuck",
  "bytes",
  "chrono",
- "cudarc 0.17.8",
+ "cudarc",
  "dashmap 5.5.3",
  "derive-getters",
  "derive_builder",
@@ -1632,7 +1623,7 @@ name = "dynamo-memory"
 version = "0.9.0"
 dependencies = [
  "anyhow",
- "cudarc 0.17.8",
+ "cudarc",
  "dynamo-config",
  "libc",
  "nix 0.30.1",
@@ -3018,9 +3009,9 @@ dependencies = [
 
 [[package]]
 name = "js-sys"
-version = "0.3.86"
+version = "0.3.87"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d36139f1c97c42c0c86a411910b04e48d4939a0376e6e0f989420cbdee0120e5"
+checksum = "93f0862381daaec758576dcc22eb7bbf4d7efd67328553f3b45a412a51a3fb21"
 dependencies = [
  "once_cell",
  "wasm-bindgen",
@@ -3248,11 +3239,12 @@ name = "kvbm-py3"
 version = "0.9.0"
 dependencies = [
  "anyhow",
- "cudarc 0.16.6",
+ "cudarc",
  "derive-getters",
  "dlpark",
  "dynamo-llm",
  "dynamo-runtime",
+ "prometheus",
  "pyo3",
  "pyo3-async-runtimes",
  "serde",
@@ -3314,6 +3306,16 @@ dependencies = [
  "windows-link",
 ]
 
+[[package]]
+name = "libloading"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60"
+dependencies = [
+ "cfg-if 1.0.4",
+ "windows-link",
+]
+
 [[package]]
 name = "libm"
 version = "0.2.16"
@@ -7400,9 +7402,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen"
-version = "0.2.109"
+version = "0.2.110"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9ff9c7baef35ac3c0e17d8bfc9ad75eb62f85a2f02bccc906699dadb0aa9c622"
+checksum = "1de241cdc66a9d91bd84f097039eb140cdc6eec47e0cdbaf9d932a1dd6c35866"
 dependencies = [
  "cfg-if 1.0.4",
  "once_cell",
@@ -7413,9 +7415,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-futures"
-version = "0.4.59"
+version = "0.4.60"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d24699cd39db9966cf6e2ef10d2f72779c961ad905911f395ea201c3ec9f545d"
+checksum = "a42e96ea38f49b191e08a1bab66c7ffdba24b06f9995b39a9dd60222e5b6f1da"
 dependencies = [
  "cfg-if 1.0.4",
  "futures-util",
@@ -7427,9 +7429,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro"
-version = "0.2.109"
+version = "0.2.110"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "39455e84ad887a0bbc93c116d72403f1bb0a39e37dd6f235a43e2128a0c7f1fd"
+checksum = "e12fdf6649048f2e3de6d7d5ff3ced779cdedee0e0baffd7dff5cdfa3abc8a52"
 dependencies = [
  "quote",
  "wasm-bindgen-macro-support",
@@ -7437,9 +7439,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro-support"
-version = "0.2.109"
+version = "0.2.110"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dff4761f60b0b51fd13fec8764167b7bbcc34498ce3e52805fe1db6f2d56b6d6"
+checksum = "0e63d1795c565ac3462334c1e396fd46dbf481c40f51f5072c310717bc4fb309"
 dependencies = [
  "bumpalo",
  "proc-macro2",
@@ -7450,9 +7452,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-shared"
-version = "0.2.109"
+version = "0.2.110"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bc6a171c53d98021a93a474c4a4579d76ba97f9517d871bc12e27640f218b6dd"
+checksum = "e9f9cdac23a5ce71f6bf9f8824898a501e511892791ea2a0c6b8568c68b9cb53"
 dependencies = [
  "unicode-ident",
 ]
@@ -7506,9 +7508,9 @@ dependencies = [
 
 [[package]]
 name = "web-sys"
-version = "0.3.86"
+version = "0.3.87"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "668fa5d00434e890a452ab060d24e3904d1be93f7bb01b70e5603baa2b8ab23b"
+checksum = "f2c7c5718134e770ee62af3b6b4a84518ec10101aad610c024b64d6ff29bb1ff"
 dependencies = [
  "js-sys",
  "wasm-bindgen",

lib/bindings/kvbm/Cargo.toml

@@ -56,6 +56,7 @@ pyo3-async-runtimes = { version = "0.23.0", default-features = false, features =
 ] }
 
 dlpark = { version = "0.5", features = ["pyo3", "half"], optional = true }
-cudarc = { version = "0.16.2", features = ["cuda-12020"], optional = true }
+cudarc = { version = "0.19.2", features = ["cuda-12020"], optional = true }
+prometheus = "0.14.0"
 
 [dev-dependencies]

lib/bindings/python/Cargo.lock

@@ -826,7 +826,7 @@ checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4"
 dependencies = [
  "glob",
  "libc",
- "libloading",
+ "libloading 0.8.9",
 ]
 
 [[package]]
@@ -1124,11 +1124,11 @@ dependencies = [
 
 [[package]]
 name = "cudarc"
-version = "0.17.8"
+version = "0.19.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0bf99ab37ee7072d64d906aa2dada9a3422f1d975cdf8c8055a573bc84897ed8"
+checksum = "aed81f178e780f3d5d354d12b4c5c5a484c4a9c329ecd037ac57f2a0e0648397"
 dependencies = [
- "libloading",
+ "libloading 0.9.0",
 ]
 
 [[package]]
@@ -3344,6 +3344,16 @@ dependencies = [
  "windows-link",
 ]
 
+[[package]]
+name = "libloading"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60"
+dependencies = [
+ "cfg-if 1.0.4",
+ "windows-link",
+]
+
 [[package]]
 name = "libm"
 version = "0.2.16"

lib/kvbm-kernels/CLAUDE.md

+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## What This Is
+
+kvbm-kernels is a high-performance CUDA transfer library for batched H2D, D2H, and D2D block copies used by the Dynamo KV cache system. The core API (`vectorized_copy`, `memcpy_batch`) is always available and handles the common case of moving KV cache blocks between host and device without layout changes. Fused permute-and-copy kernels for layout conversion between **Block Stack** (vLLM) and **Universal** (Dynamo storage) formats are feature-gated behind `permute_kernels`.
+
+## Build Commands
+
+```bash
+# Default build (auto-detects nvcc -> source build; no nvcc -> stubs)
+cargo build
+
+# Build from source with custom GPU architectures
+CUDA_ARCHS="80,86,89,90,100" cargo build
+
+# Static linking (embed kernels into binary instead of .so)
+cargo build --features static-kernels
+
+# Check compilation without linking
+cargo check
+
+# Run CUDA integration tests for core transfer APIs (requires GPU + nvcc)
+cargo test --features testing-cuda
+
+# Run all CUDA integration tests including permute kernels
+cargo test --features testing-cuda,permute_kernels
+
+# Run a specific test
+cargo test --features testing-cuda,permute_kernels fused_copy_roundtrip -- --nocapture --test-threads=1
+
+# Run benchmarks (Llama 3.1 70B KV cache profile)
+cargo run --example kvbench --features kvbench
+```
+
+**Environment variables**: `CUDA_ARCHS` (comma-separated SM versions), `CUDA_PTX_ARCHS` (PTX targets), `KVBM_REQUIRE_CUDA` (fail if nvcc missing), `CUDA_PATH`/`CUDA_HOME`.
+
+## Architecture
+
+### Two-tier build system (`build.rs`)
+
+The build script selects one of two modes: **FromSource** (nvcc available, compiles CUDA, requires CUDA >= 12.0) or **Stubs** (no nvcc, C stubs that abort on call). Stubs set the `stub_kernels` cfg flag so tests can be conditionally skipped.
+
+### Core transfer API (always available)
+
+These live in `src/tensor_kernels.rs` and work on any device-visible memory (device allocations or pinned host via unified addressing):
+
+- **`vectorized_copy`** — Batched copy of `(src, dst)` pointer pairs. Per-pair runtime alignment detection selects the widest safe vector width (int4/int2/int/char for 16/8/4/1-byte loads).
+- **`memcpy_batch`** — Takes HOST arrays of src/dst pointers. Dispatches to `cudaMemcpyBatchAsync` (CUDA 12.9+) with fallback to individual `cudaMemcpyAsync` loop. Three modes: `BatchedWithFallback`, `FallbackOnly`, `BatchWithoutFallback`.
+- **`is_using_stubs`** / **`is_memcpy_batch_available`** — Runtime capability queries.
+
+### Permute kernels (feature-gated: `permute_kernels`)
+
+These fuse layout permutation with copy for non-standard transfer paths:
+
+- **`universal_from_block`** / **`block_from_universal`** — Permute between block stack layout (`nl*no` separate allocations, each `[nt, nh, hd]` NHD or `[nh, nt, hd]` HND) and universal layout (contiguous `[nh, nl, no, nt, hd]`).
+
+### Source organization
+
+- `cuda/tensor_kernels.cu` — All CUDA kernels. C++ templates on dtype (F16/BF16/F32/F64) and layout (NHD/HND), exposed via `extern "C"` functions prefixed `kvbm_kernels_launch_*` / `kvbm_kernels_memcpy_batch`.
+- `cuda/stubs.c` — Abort-on-call fallbacks for all `extern "C"` symbols.
+- `src/tensor_kernels.rs` — Rust FFI wrappers, enums (`TensorDataType`, `BlockLayout`, `MemcpyBatchMode`), and integration tests.
+- `examples/kvbench.rs` — Benchmark harness (Llama 3.1 70B profile, CSV output).
+- `scripts/plot_roofline.py` — Roofline bandwidth plots from kvbench output.
+
+### Dimension conventions
+
+`nl` = layers, `no` = outer chunks (2: K and V), `nh` = attention heads, `nt` = tokens per block, `hd` = head dimension.
+
+### Pointer conventions
+
+All pointer-list parameters (e.g. `universal_ptrs`, `src_ptrs`) must be device-accessible: allocated via `cudaMalloc` (device memory) or `cudaMallocHost` / `cuMemHostRegister` (pinned/registered/page-locked host memory).
+
+### Cargo features
+
+| Feature | Purpose |
+|---------|---------|
+| `permute_kernels` | Enable fused permute-and-copy kernels (block<->universal) |
+| `testing-cuda` | Enable CUDA integration tests |
+| `static-kernels` | Link as `.a` instead of `.so` |
+| `kvbench` | Enable benchmark example (pulls in `clap`) |
+
+### Test organization
+
+- `tests/stub_build.rs` — Verifies stub behavior (gated on `stub_kernels`).
+- `tests/memcpy_batch.rs` — Core transfer API roundtrip tests (H2D + D2H via pinned host memory). Gated on `testing-cuda`.
+- `tests/kernel_roundtrip.rs` — Permute kernel roundtrip tests across all dtypes and layouts. Gated on `testing-cuda` + `permute_kernels`.
+- Inline tests in `src/tensor_kernels.rs` — Integration tests including `universal_roundtrip`. Gated on `testing-cuda` + `permute_kernels`.

lib/kvbm-kernels/Cargo.toml

+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+[package]
+name = "kvbm-kernels"
+version.workspace = true
+edition.workspace = true
+authors.workspace = true
+license.workspace = true
+repository = "https://github.com/ai-dynamo/dynamo.git"
+build = "build.rs"
+
+[lib]
+name = "kvbm_kernels"
+crate-type = ["rlib", "cdylib"]
+
+[features]
+default = []
+# Build kernels as a static archive (.a) instead of shared library (.so).
+# When enabled, the kernel code is embedded directly into the consuming crate,
+# eliminating the runtime dependency on libkvbm_kernels.so.
+# Note: This only affects real CUDA builds; stubs always remain dynamic.
+static-kernels = []
+
+# Enable CUDA tests - only works when real CUDA kernels are built (not stubs)
+# Tests are gated with #[cfg(all(test, feature = "testing-cuda", not(stub_kernels)))]
+testing-cuda = []
+
+# Enable operational_copy, universal_from_block, block_from_universal kernels.
+# These kernels perform data layout permutation and are only needed for
+# non-standard transfer paths. The default vectorized_copy kernel handles
+# most FC↔LW transfers efficiently without permutation.
+permute_kernels = []
+
+# Enable kvbench example (pulls in clap for CLI)
+kvbench = ["dep:clap"]
+
+[[example]]
+name = "kvbench"
+required-features = ["kvbench"]
+
+[dependencies]
+cudarc = { workspace = true }
+
+# kvbench
+clap = { version = "4", features = ["derive"], optional = true }
+
+[dev-dependencies]
+ndarray = "0.17.2"
+half = "2"
+rand = { workspace = true }
+cudarc = { workspace = true, features = ["f16"] }
+
+[build-dependencies]

lib/kvbm-kernels/README.md

+## Dynamo KV Block Manager Kernels
+
+GPU kernels for converting KV cache blocks between three memory layouts used by LLM inference frameworks. All conversions run entirely on-device via fused CUDA kernels.
+
+### Dimensions
+
+| Symbol | Meaning                        | Example          |
+|--------|--------------------------------|------------------|
+| `nb`   | Number of blocks in the batch  | 1–128            |
+| `nl`   | Number of layers               | 32 (Llama-70B)   |
+| `no`   | Outer chunks (K and V)         | 2                |
+| `nh`   | Number of attention heads      | 32 or 64         |
+| `nt`   | Tokens per block               | 128 or 256       |
+| `hd`   | Head dimension                 | 128              |
+
+### Layouts
+
+#### Block Stack (NHD or HND)
+
+`nl * no` separate GPU allocations per block. Each allocation holds one layer's keys or values.
+
+- **NHD shape**: `[nt, nh, hd]` — index: `(nt_idx * nh + nh_idx) * hd + hd_idx`
+- **HND shape**: `[nh, nt, hd]` — index: `(nh_idx * nt + nt_idx) * hd + hd_idx`
+
+Passed to kernels as a flat pointer table of length `nb * nl * no`.
+
+#### Operational
+
+Single contiguous buffer per block: `[nl, no, inner]` where `inner = nt * nh * hd`.
+
+The three innermost dimensions (`nt`, `nh`, `hd`) are fused into one `inner` dimension. When no layout permutation is needed (same TP config, same head layout), block-to-operational is a flat copy — the cheapest conversion. Transforming to/from other layouts requires knowing the constituent dimensions.
+
+#### Universal
+
+Single contiguous buffer per block: `[nh, nl, no, nt, hd]`.
+
+Heads are the outermost dimension so that tensor-parallelism resharding is a contiguous slice along `nh`. A block saved from a TP=4 deployment can be loaded into TP=8 by slicing the head dimension differently.
+
+### Layout Cheat Sheet
+
+| Layout              | Logical Shape              | Stored As                          | Notes                         |
+|---------------------|----------------------------|------------------------------------|-------------------------------|
+| NHD block stack     | `[nl][no][nt, nh, hd]`     | list of `nl * no` pointers         | Inner layout = NHD            |
+| HND block stack     | `[nl][no][nh, nt, hd]`     | list of `nl * no` pointers         | Inner layout = HND            |
+| Operational block   | `[nl, no, inner]`          | contiguous buffer per block        | `inner = nt * nh * hd`        |
+| Universal block     | `[nh, nl, no, nt, hd]`     | contiguous buffer per block        | Heads outermost for TP slicing |
+
+### Kernel Functions
+
+All kernels are batched: a single launch processes `nb` blocks from flat pointer tables prepared by host code.
+
+#### Layout permutation kernels
+
+| C API                                        | Conversion                  |
+|----------------------------------------------|-----------------------------|
+| `kvbm_kernels_launch_universal_from_block`   | Block stack → Universal     |
+| `kvbm_kernels_launch_block_from_universal`   | Universal → Block stack     |
+
+Both accept `layout_value` (NHD=0, HND=1) and `dtype_value` (F16=0, BF16=1, F32=2, F64=3). Internally dispatched to C++ template kernels specialized on dtype and layout.
+
+#### Standalone copy utilities
+
+| C API                                    | Description                                              |
+|------------------------------------------|----------------------------------------------------------|
+| `kvbm_kernels_launch_vectorized_copy`    | Adaptive vectorized copy (16/8/4-byte or scalar) across `num_pairs` pointer pairs |
+| `kvbm_kernels_memcpy_batch`              | Batched `cudaMemcpyAsync` from host pointer arrays       |
+| `kvbm_kernels_has_memcpy_batch_async`    | Returns `true` if `cudaMemcpyBatchAsync` is available    |
+| `kvbm_kernels_is_stub_build`             | Returns `true` if built without CUDA (stub mode)         |
+
+
+### Python Bindings (Planned)
+
+Python kernel bindings are not yet implemented. The `lib/bindings/kvbm/` crate currently exposes block manager functionality only. Future work will add Python wrappers for the permute and copy kernels.
+
+### Development
+
+```bash
+# Default build (auto-detects nvcc → source; no nvcc → stubs)
+cargo build
+
+# Custom GPU architectures
+CUDA_ARCHS="80,86,89,90,100" cargo build
+
+# Static linking
+cargo build --features static-kernels
+
+# Run CUDA integration tests (requires GPU + nvcc)
+cargo test --features testing-cuda,permute_kernels
+
+# Specific test with output
+cargo test --features testing-cuda,permute_kernels fused_copy_roundtrip -- --nocapture
+
+# Python bindings
+cd lib/bindings/kvbm
+uv pip install -e ".[dev]"
+pytest tests/
+```
+
+**Environment variables**: `CUDA_ARCHS` (comma-separated SM versions, default `80,86,89,90,100,120`), `CUDA_PATH`/`CUDA_HOME` (toolkit root), `KVBM_REQUIRE_CUDA` (fail build if nvcc missing).
+
+### Benchmarking
+
+```text
+root@9eb240f7ded8:/workspace/lib/kvbm-kernels# cargo run --release --example kvbench --features testing-cuda,kvbench -- --num-blocks=1,128 --tokens-per-block=16,64 --
+backend vectorized,batched --direction h2d
+...
+     Running `/workspace/target/release/examples/kvbench --num-blocks=1,128 --tokens-per-block=16,64 --backend vectorized,batched --direction h2d`
+KV Cache Transfer Benchmark
+  Model: Llama 3.1 70B (bf16)
+  Layers: 80, KV heads: 8, Head dim: 128, Outer dim: 2
+  Warmup: 10, Timed: 100
+  Batch API available: true
+  tokens_per_block: [16, 64]
+  num_blocks: [1, 128]
+  directions: [h2d]
+  patterns: [fc_to_fc, lw_to_fc]
+  backends: [vectorized, batched]
+  Total tests: 16
+
+tokens_per_block,num_blocks,pattern,direction,backend,total_bytes,inner_bytes,copy_size,num_copies,median_ms,bandwidth_gbps
+--- tokens_per_block=16, inner=32768 bytes (32 KB), block=5242880 bytes (5.0 MB) ---
+  [1/16] tpb=16 N=  1 fc_to_fc h2d    vectorized   ... 16,1,fc_to_fc,h2d,vectorized,5242880,32768,5242880,1,1.8686,2.81
+2.81 GB/s (1.8686 ms)
+  [2/16] tpb=16 N=  1 fc_to_fc h2d    batched      ... 16,1,fc_to_fc,h2d,batched,5242880,32768,5242880,1,0.2105,24.91
+24.91 GB/s (0.2105 ms)
+  [3/16] tpb=16 N=  1 lw_to_fc h2d    vectorized   ... 16,1,lw_to_fc,h2d,vectorized,5242880,32768,32768,160,0.2171,24.15
+24.15 GB/s (0.2171 ms)
+  [4/16] tpb=16 N=  1 lw_to_fc h2d    batched      ... 16,1,lw_to_fc,h2d,batched,5242880,32768,32768,160,0.2775,18.89
+18.89 GB/s (0.2775 ms)
+  [5/16] tpb=16 N=128 fc_to_fc h2d    vectorized   ... 16,128,fc_to_fc,h2d,vectorized,671088640,32768,5242880,128,26.6097,25.22
+25.22 GB/s (26.6097 ms)
+  [6/16] tpb=16 N=128 fc_to_fc h2d    batched      ... 16,128,fc_to_fc,h2d,batched,671088640,32768,5242880,128,26.6180,25.21
+25.21 GB/s (26.6180 ms)
+  [7/16] tpb=16 N=128 lw_to_fc h2d    vectorized   ... 16,128,lw_to_fc,h2d,vectorized,671088640,32768,32768,20480,26.6034,25.23
+25.23 GB/s (26.6034 ms)
+  [8/16] tpb=16 N=128 lw_to_fc h2d    batched      ... 16,128,lw_to_fc,h2d,batched,671088640,32768,32768,20480,30.3346,22.12
+22.12 GB/s (30.3346 ms)
+--- tokens_per_block=64, inner=131072 bytes (128 KB), block=20971520 bytes (20.0 MB) ---
+  [9/16] tpb=64 N=  1 fc_to_fc h2d    vectorized   ... 64,1,fc_to_fc,h2d,vectorized,20971520,131072,20971520,1,7.5837,2.77
+2.77 GB/s (7.5837 ms)
+  [10/16] tpb=64 N=  1 fc_to_fc h2d    batched      ... 64,1,fc_to_fc,h2d,batched,20971520,131072,20971520,1,0.8334,25.16
+25.16 GB/s (0.8334 ms)
+  [11/16] tpb=64 N=  1 lw_to_fc h2d    vectorized   ... 64,1,lw_to_fc,h2d,vectorized,20971520,131072,131072,160,0.8407,24.95
+24.95 GB/s (0.8407 ms)
+  [12/16] tpb=64 N=  1 lw_to_fc h2d    batched      ... 64,1,lw_to_fc,h2d,batched,20971520,131072,131072,160,0.9020,23.25
+23.25 GB/s (0.9020 ms)
+  [13/16] tpb=64 N=128 fc_to_fc h2d    vectorized   ... 64,128,fc_to_fc,h2d,vectorized,2684354560,131072,20971520,128,106.3677,25.24
+25.24 GB/s (106.3677 ms)
+  [14/16] tpb=64 N=128 fc_to_fc h2d    batched      ... 64,128,fc_to_fc,h2d,batched,2684354560,131072,20971520,128,106.3199,25.25
+25.25 GB/s (106.3199 ms)
+  [15/16] tpb=64 N=128 lw_to_fc h2d    vectorized   ... 64,128,lw_to_fc,h2d,vectorized,2684354560,131072,131072,20480,106.3158,25.25
+25.25 GB/s (106.3158 ms)
+  [16/16] tpb=64 N=128 lw_to_fc h2d    batched      ... 64,128,lw_to_fc,h2d,batched,2684354560,131072,131072,20480,110.0665,24.39
+24.39 GB/s (110.0665 ms)
+
+Done.
+```
+
+### Troubleshooting
+
+| Symptom                               | Likely Cause / Fix                                                 |
+|---------------------------------------|--------------------------------------------------------------------|
+| `cudaErrorInvalidValue` on launch     | Pointer counts mismatch (`nb`, `nl`, `no`) or non-contiguous input |
+| Wrong values when using HND layout    | Inner tensors not shaped as `[nh, nt, hd]` before passing in       |
+| Python bindings complain about dtype  | Mixed precision in a batch; convert tensors to a common dtype      |
+| Kernels take unexpected time          | Verify that `CUDA_ARCHS` matches your GPU to avoid JIT at runtime  |

lib/kvbm-kernels/build.rs

+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::env;
+use std::fs;
+use std::path::Path;
+use std::path::PathBuf;
+use std::process::Command;
+
+fn main() {
+    // Declare the stub_kernels cfg so Rust knows it's a valid cfg option
+    println!("cargo:rustc-check-cfg=cfg(stub_kernels)");
+
+    let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
+    let out_dir = env::var("OUT_DIR").unwrap();
+
+    // Track file changes
+    let cu_files = discover_cuda_files();
+    for file in &cu_files {
+        println!("cargo:rerun-if-changed={}", file.display());
+    }
+    println!(
+        "cargo:rerun-if-changed={}",
+        Path::new(&manifest_dir).join("cuda/stubs.c").display()
+    );
+    println!("cargo:rerun-if-env-changed=CUDA_ARCHS");
+    println!("cargo:rerun-if-env-changed=CUDA_PTX_ARCHS");
+    println!("cargo:rerun-if-env-changed=KVBM_REQUIRE_CUDA");
+    println!("cargo:rerun-if-env-changed=CUDA_PATH");
+    println!("cargo:rerun-if-env-changed=CUDA_HOME");
+
+    // Check if CUDA is required (set by Python bindings build)
+    let require_cuda = env::var("KVBM_REQUIRE_CUDA").is_ok();
+    let nvcc_available = is_nvcc_available();
+
+    // Fail early if CUDA required but not available
+    if require_cuda && !nvcc_available {
+        panic!(
+            "\n\n\
+            ╔════════════════════════════════════════════════════════════════════════╗\n\
+            ║  KVBM_REQUIRE_CUDA is set but nvcc is not available!                   ║\n\
+            ║                                                                        ║\n\
+            ║  Python bindings require real CUDA kernels. Please:                    ║\n\
+            ║    1. Install CUDA toolkit with nvcc, or                               ║\n\
+            ║    2. Unset KVBM_REQUIRE_CUDA for stub-only build                      ║\n\
+            ╚════════════════════════════════════════════════════════════════════════╝\n\
+            "
+        );
+    }
+
+    // Determine build mode
+    let build_mode = determine_build_mode(nvcc_available);
+
+    // Check if static linking is requested (only applies to CUDA builds, not stubs)
+    #[cfg(feature = "static-kernels")]
+    let use_static = true;
+    #[cfg(not(feature = "static-kernels"))]
+    let use_static = false;
+
+    match build_mode {
+        BuildMode::FromSource => {
+            if use_static {
+                println!("cargo:warning=Building CUDA kernels from source (static linking)");
+            } else {
+                println!("cargo:warning=Building CUDA kernels from source (dynamic linking)");
+            }
+            build_cuda_library(&cu_files, &out_dir, use_static);
+        }
+        BuildMode::Stubs => {
+            // Stubs always use dynamic linking regardless of static-kernels feature
+            println!("cargo:warning=Building stub kernels (no CUDA available, dynamic linking)");
+            build_stub_shared_library(&manifest_dir, &out_dir);
+            // Set cfg flag so tests can be skipped
+            println!("cargo:rustc-cfg=stub_kernels");
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+enum BuildMode {
+    FromSource,
+    Stubs,
+}
+
+/// Determine the build mode based on nvcc availability.
+fn determine_build_mode(nvcc_available: bool) -> BuildMode {
+    if nvcc_available {
+        BuildMode::FromSource
+    } else {
+        BuildMode::Stubs
+    }
+}
+
+fn is_nvcc_available() -> bool {
+    Command::new("nvcc").arg("--version").output().is_ok()
+}
+
+/// Build CUDA kernels from source.
+/// If `use_static` is true, builds a static archive (.a); otherwise builds a shared library (.so).
+fn build_cuda_library(cu_files: &[PathBuf], out_dir: &str, use_static: bool) {
+    let arch_flags = get_cuda_arch_flags();
+
+    // Only build tensor_kernels.cu into the library (it has the extern "C" functions)
+    let tensor_kernels_path = cu_files
+        .iter()
+        .find(|p| p.file_stem().unwrap() == "tensor_kernels")
+        .expect("tensor_kernels.cu not found");
+
+    let obj_path = Path::new(out_dir).join("kvbm_kernels.o");
+
+    // Step 1: Compile to object file
+    let mut nvcc_cmd = Command::new("nvcc");
+    nvcc_cmd
+        .arg("-m64")
+        .arg("-c")
+        .arg("-std=c++17")
+        .arg("-O3")
+        .arg("-Xcompiler")
+        .arg("-fPIC")
+        .arg(tensor_kernels_path)
+        .arg("-o")
+        .arg(&obj_path);
+
+    for flag in &arch_flags {
+        nvcc_cmd.arg(flag);
+    }
+
+    println!("cargo:warning=Compiling tensor_kernels.cu to object file...");
+    let status = nvcc_cmd
+        .status()
+        .expect("Failed to execute nvcc for object file");
+
+    if !status.success() {
+        panic!("nvcc failed to compile tensor_kernels.cu");
+    }
+
+    if use_static {
+        // Step 2a: Create static archive
+        let ar_path = Path::new(out_dir).join("libkvbm_kernels.a");
+
+        let mut ar_cmd = Command::new("ar");
+        ar_cmd.arg("crus").arg(&ar_path).arg(&obj_path);
+
+        println!("cargo:warning=Creating static archive libkvbm_kernels.a...");
+        let status = ar_cmd
+            .status()
+            .expect("Failed to execute ar for static archive");
+
+        if !status.success() {
+            panic!("ar failed to create static archive");
+        }
+
+        // Set up static linking
+        println!("cargo:rustc-link-search=native={}", out_dir);
+        println!("cargo:rustc-link-lib=static=kvbm_kernels");
+
+        // Add CUDA runtime library paths and link cudart dynamically
+        add_cuda_library_paths();
+        println!("cargo:rustc-link-lib=cudart");
+
+        // CUDA object code compiled by nvcc contains C++ runtime symbols
+        // (operator new/delete, __gxx_personality_v0, etc.)
+        println!("cargo:rustc-link-lib=stdc++");
+    } else {
+        // Step 2b: Link into shared library
+        let so_path = Path::new(out_dir).join("libkvbm_kernels.so");
+
+        let mut link_cmd = Command::new("nvcc");
+        link_cmd
+            .arg("-shared")
+            .arg("-o")
+            .arg(&so_path)
+            .arg(&obj_path)
+            .arg("-lcudart");
+
+        println!("cargo:warning=Linking kvbm_kernels into shared library...");
+        let status = link_cmd
+            .status()
+            .expect("Failed to execute nvcc for linking");
+
+        if !status.success() {
+            panic!("nvcc failed to link shared library");
+        }
+
+        // Set up dynamic linking
+        println!("cargo:rustc-link-search=native={}", out_dir);
+        println!("cargo:rustc-link-lib=dylib=kvbm_kernels");
+
+        // Add CUDA runtime library paths
+        add_cuda_library_paths();
+        println!("cargo:rustc-link-lib=cudart");
+    }
+}
+
+/// Build stub shared library from stubs.c when CUDA is not available.
+fn build_stub_shared_library(manifest_dir: &str, out_dir: &str) {
+    let stubs_path = Path::new(manifest_dir).join("cuda/stubs.c");
+
+    if !stubs_path.exists() {
+        panic!(
+            "Stub source file not found at {}. Cannot build without CUDA.",
+            stubs_path.display()
+        );
+    }
+
+    // Build shared library from stubs.c using the system C compiler
+    let so_path = Path::new(out_dir).join("libkvbm_kernels.so");
+    let obj_path = Path::new(out_dir).join("stubs.o");
+
+    // Compile to object file
+    let mut gcc_compile = Command::new("cc");
+    gcc_compile
+        .arg("-c")
+        .arg("-fPIC")
+        .arg("-O2")
+        .arg(&stubs_path)
+        .arg("-o")
+        .arg(&obj_path);
+
+    println!("cargo:warning=Compiling stubs.c...");
+    let status = gcc_compile
+        .status()
+        .expect("Failed to execute cc for stubs");
+
+    if !status.success() {
+        panic!("Failed to compile stubs.c");
+    }
+
+    // Link into shared library
+    let mut gcc_link = Command::new("cc");
+    gcc_link
+        .arg("-shared")
+        .arg("-o")
+        .arg(&so_path)
+        .arg(&obj_path);
+
+    println!("cargo:warning=Linking stub shared library...");
+    let status = gcc_link.status().expect("Failed to link stub library");
+
+    if !status.success() {
+        panic!("Failed to link stub shared library");
+    }
+
+    // Set up linking
+    println!("cargo:rustc-link-search=native={}", out_dir);
+    println!("cargo:rustc-link-lib=dylib=kvbm_kernels");
+}
+
+fn discover_cuda_files() -> Vec<PathBuf> {
+    let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
+    let cuda_dir = Path::new(&manifest_dir).join("cuda");
+    let mut cu_files = Vec::new();
+
+    for entry in fs::read_dir(cuda_dir).expect("Failed to read cuda directory") {
+        let entry = entry.expect("Failed to read entry");
+        let path = entry.path();
+        if path.extension().is_some_and(|ext| ext == "cu") {
+            cu_files.push(path);
+        }
+    }
+    cu_files
+}
+
+/// Parse CUDA toolkit version from `nvcc --version` output.
+/// Returns (major, minor) tuple, e.g. (12, 8) for CUDA 12.8.
+fn parse_cuda_version() -> Option<(u32, u32)> {
+    let output = Command::new("nvcc").arg("--version").output().ok()?;
+
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    // nvcc output contains a line like: "Cuda compilation tools, release 12.8, V12.8.89"
+    for line in stdout.lines() {
+        if let Some(pos) = line.find("release ") {
+            let after = &line[pos + "release ".len()..];
+            let version_str = after.split(',').next().unwrap_or("").trim();
+            let mut parts = version_str.split('.');
+            let major = parts.next()?.parse::<u32>().ok()?;
+            let minor = parts.next()?.parse::<u32>().ok()?;
+            return Some((major, minor));
+        }
+    }
+    None
+}
+
+/// Return the maximum supported compute capability for a given CUDA toolkit version.
+///
+/// Panics if the CUDA version is below 12.0.
+fn max_supported_compute(cuda_version: (u32, u32)) -> u32 {
+    match cuda_version {
+        (major, _) if major < 12 => {
+            panic!("CUDA {major}.x is not supported; CUDA 12.0 or newer is required")
+        }
+        (12, minor) if minor >= 8 => 120,
+        (major, _) if major >= 13 => 120,
+        _ => 90,
+    }
+}
+
+fn get_cuda_arch_flags() -> Vec<String> {
+    let mut flags = Vec::new();
+
+    let cuda_version = parse_cuda_version();
+    let max_compute = cuda_version.map(max_supported_compute);
+
+    if let Some((major, minor)) = cuda_version {
+        println!(
+            "cargo:warning=Detected CUDA {}.{}, max supported compute: sm_{}",
+            major,
+            minor,
+            max_compute.unwrap()
+        );
+    } else {
+        println!("cargo:warning=Could not detect CUDA version, including all architectures");
+    }
+
+    let explicit_archs = env::var("CUDA_ARCHS").ok();
+    let arch_list = explicit_archs.as_deref().unwrap_or("80,86,89,90,100,120");
+
+    for arch in arch_list.split(',') {
+        let arch = arch.trim();
+        if arch.is_empty() {
+            continue;
+        }
+        let arch_num: u32 = match arch.parse() {
+            Ok(n) => n,
+            Err(_) => {
+                println!("cargo:warning=Skipping invalid CUDA_ARCHS entry: {}", arch);
+                continue;
+            }
+        };
+        if let Some(max) = max_compute
+            && arch_num > max
+        {
+            println!(
+                "cargo:warning=Skipping sm_{} (unsupported by detected CUDA toolkit, max: sm_{})",
+                arch_num, max
+            );
+            continue;
+        }
+        flags.push(format!("-gencode=arch=compute_{},code=sm_{}", arch, arch));
+    }
+
+    // Generate forward-compatible PTX for each major architecture family that is
+    // both present in the arch list and supported by the detected CUDA toolkit.
+    let ptx_archs_env = env::var("CUDA_PTX_ARCHS").ok();
+    let ptx_candidates: Vec<u32> = if let Some(ref ptx_env) = ptx_archs_env {
+        ptx_env
+            .split(',')
+            .filter_map(|s| s.trim().parse::<u32>().ok())
+            .collect()
+    } else {
+        vec![90, 100, 120]
+    };
+
+    for &ptx_arch in &ptx_candidates {
+        if let Some(max) = max_compute
+            && ptx_arch > max
+        {
+            continue;
+        }
+        flags.push(format!(
+            "-gencode=arch=compute_{},code=compute_{}",
+            ptx_arch, ptx_arch
+        ));
+    }
+
+    flags
+}
+
+fn add_cuda_library_paths() {
+    if let Ok(cuda_path) = env::var("CUDA_PATH") {
+        println!("cargo:rustc-link-search=native={}/lib64", cuda_path);
+        println!("cargo:rustc-link-search=native={}/lib", cuda_path);
+    } else if let Ok(cuda_home) = env::var("CUDA_HOME") {
+        println!("cargo:rustc-link-search=native={}/lib64", cuda_home);
+        println!("cargo:rustc-link-search=native={}/lib", cuda_home);
+    } else {
+        // Try standard paths
+        println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64");
+        println!("cargo:rustc-link-search=native=/usr/local/cuda/lib");
+    }
+}

lib/kvbm-kernels/cuda/stubs.c

+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+// Stub implementations for CUDA kernel functions.
+// These are used when nvcc is not available, allowing the library to be built
+// without CUDA. The stubs abort() when called, but the binary can be moved to
+// an environment with the real .so and work correctly via LD_LIBRARY_PATH.
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+// cudaError_t equivalent - cudaSuccess = 0
+typedef int cudaError_t;
+
+// cudaStream_t is an opaque pointer
+typedef void* cudaStream_t;
+
+#define STUB_ABORT(name)                                                 \
+  do {                                                                   \
+    fprintf(                                                             \
+        stderr,                                                          \
+        "FATAL: %s called but CUDA kernels not available.\n"             \
+        "This binary was built with stub kernels. To use CUDA:\n"        \
+        "  1. Build with nvcc available, or\n"                           \
+        "  2. Set LD_LIBRARY_PATH to include real libkvbm_kernels.so\n", \
+        name);                                                           \
+    abort();                                                             \
+  } while (0)
+
+cudaError_t
+kvbm_kernels_launch_universal_from_block(
+    void* const* universal_ptrs, const void* const* block_ptrs, size_t num_blocks, size_t nh, size_t nl, size_t no,
+    size_t nt, size_t hd, int dtype_value, int layout_value, cudaStream_t stream)
+{
+  (void)universal_ptrs;
+  (void)block_ptrs;
+  (void)num_blocks;
+  (void)nh;
+  (void)nl;
+  (void)no;
+  (void)nt;
+  (void)hd;
+  (void)dtype_value;
+  (void)layout_value;
+  (void)stream;
+  STUB_ABORT("kvbm_kernels_launch_universal_from_block");
+  return 1;  // Unreachable, but silences compiler warning
+}
+
+cudaError_t
+kvbm_kernels_launch_block_from_universal(
+    const void* const* universal_ptrs, void* const* block_ptrs, size_t num_blocks, size_t nh, size_t nl, size_t no,
+    size_t nt, size_t hd, int dtype_value, int layout_value, cudaStream_t stream)
+{
+  (void)universal_ptrs;
+  (void)block_ptrs;
+  (void)num_blocks;
+  (void)nh;
+  (void)nl;
+  (void)no;
+  (void)nt;
+  (void)hd;
+  (void)dtype_value;
+  (void)layout_value;
+  (void)stream;
+  STUB_ABORT("kvbm_kernels_launch_block_from_universal");
+  return 1;  // Unreachable
+}
+
+cudaError_t
+kvbm_kernels_launch_vectorized_copy(
+    void** src_ptrs, void** dst_ptrs, size_t copy_size_bytes, int num_pairs, cudaStream_t stream)
+{
+  (void)src_ptrs;
+  (void)dst_ptrs;
+  (void)copy_size_bytes;
+  (void)num_pairs;
+  (void)stream;
+  STUB_ABORT("kvbm_kernels_launch_vectorized_copy");
+  return 1;  // Unreachable
+}
+
+// This function is safe to call even with stubs - it just returns false
+// indicating that batch async is not available.
+bool
+kvbm_kernels_has_memcpy_batch_async(void)
+{
+  return false;
+}
+
+// Stub for memcpy_batch - returns not supported since we can't do CUDA ops
+cudaError_t
+kvbm_kernels_memcpy_batch(
+    const void* const* src_ptrs, void* const* dst_ptrs, size_t size_per_copy, size_t num_copies, int mode,
+    cudaStream_t stream)
+{
+  (void)src_ptrs;
+  (void)dst_ptrs;
+  (void)size_per_copy;
+  (void)num_copies;
+  (void)mode;
+  (void)stream;
+  STUB_ABORT("kvbm_kernels_memcpy_batch");
+  return 1;  // Unreachable
+}
+
+// Returns true if this is the stub library (no real CUDA kernels).
+// Downstream crates can use this to skip CUDA tests at runtime.
+bool
+kvbm_kernels_is_stub_build(void)
+{
+  return true;
+}

lib/kvbm-kernels/src/lib.rs

+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+pub mod tensor_kernels;
+
+// Always available - core transfer functionality
+pub use tensor_kernels::{
+    MemcpyBatchMode, is_memcpy_batch_available, is_using_stubs, memcpy_batch, vectorized_copy,
+};
+
+// Permute kernels - data layout transformation (requires permute_kernels feature)
+#[cfg(feature = "permute_kernels")]
+pub use tensor_kernels::{BlockLayout, TensorDataType, block_from_universal, universal_from_block};

lib/kvbm-kernels/tests/kernel_roundtrip.rs

+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Integration tests for CUDA tensor packing kernel roundtrips.
+//!
+//! Mirrors the Python tests in `lib/bindings/kvbm/tests/test_tensor_kernels.py`
+//! using ndarray for reference permutations and cudarc for GPU memory management.
+
+#![cfg(all(
+    feature = "testing-cuda",
+    feature = "permute_kernels",
+    not(stub_kernels)
+))]
+
+use std::ffi::c_void;
+use std::fmt::Debug;
+use std::sync::Arc;
+
+use cudarc::driver::result::memset_d8_async;
+use cudarc::driver::{
+    CudaContext, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, DeviceRepr, DriverError,
+    ValidAsZeroBits,
+};
+use cudarc::runtime::sys as cuda_runtime;
+use half::{bf16, f16};
+use kvbm_kernels::{BlockLayout, TensorDataType, block_from_universal, universal_from_block};
+use ndarray::{Array5, s};
+use rand::Rng;
+
+// ---------------------------------------------------------------------------
+// TestDtype trait — bridges Rust types to kernel enums + tolerances
+// ---------------------------------------------------------------------------
+
+trait TestDtype: Clone + Debug + DeviceRepr + ValidAsZeroBits + 'static {
+    const DTYPE: TensorDataType;
+    const ATOL: f64;
+    const RTOL: f64;
+
+    fn from_f64(v: f64) -> Self;
+    fn to_f64(self) -> f64;
+}
+
+impl TestDtype for f16 {
+    const DTYPE: TensorDataType = TensorDataType::F16;
+    const ATOL: f64 = 1e-2;
+    const RTOL: f64 = 1e-2;
+
+    fn from_f64(v: f64) -> Self {
+        f16::from_f64(v)
+    }
+    fn to_f64(self) -> f64 {
+        f16::to_f64(self)
+    }
+}
+
+impl TestDtype for bf16 {
+    const DTYPE: TensorDataType = TensorDataType::BF16;
+    const ATOL: f64 = 1e-2;
+    const RTOL: f64 = 1e-2;
+
+    fn from_f64(v: f64) -> Self {
+        bf16::from_f64(v)
+    }
+    fn to_f64(self) -> f64 {
+        bf16::to_f64(self)
+    }
+}
+
+impl TestDtype for f32 {
+    const DTYPE: TensorDataType = TensorDataType::F32;
+    const ATOL: f64 = 1e-5;
+    const RTOL: f64 = 1e-5;
+
+    fn from_f64(v: f64) -> Self {
+        v as f32
+    }
+    fn to_f64(self) -> f64 {
+        self as f64
+    }
+}
+
+impl TestDtype for f64 {
+    const DTYPE: TensorDataType = TensorDataType::F64;
+    const ATOL: f64 = 1e-12;
+    const RTOL: f64 = 1e-12;
+
+    fn from_f64(v: f64) -> Self {
+        v
+    }
+    fn to_f64(self) -> f64 {
+        self
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/// Reference permutation using ndarray, mirrors the Python `_make_blocks()`.
+///
+/// Takes a `[nh, nl, no, nt, hd]` universal tensor and produces `nl * no` flat
+/// block chunks, each with layout-dependent axis ordering.
+fn make_blocks<T: TestDtype>(universal: &Array5<T>, layout: BlockLayout) -> Vec<Vec<T>> {
+    let (_nh, nl, no, _nt, _hd) = universal.dim();
+    let mut blocks = Vec::with_capacity(nl * no);
+    for l in 0..nl {
+        for o in 0..no {
+            // Slice out [nh, nt, hd] for this (layer, outer) pair.
+            let chunk = universal.slice(s![.., l, o, .., ..]);
+            let flat = match layout {
+                BlockLayout::NHD => {
+                    // [nh, nt, hd] -> [nt, nh, hd]
+                    let permuted = chunk.permuted_axes([1, 0, 2]);
+                    permuted.as_standard_layout().as_slice().unwrap().to_vec()
+                }
+                BlockLayout::HND => {
+                    // [nh, nt, hd] — identity permutation
+                    chunk.as_standard_layout().as_slice().unwrap().to_vec()
+                }
+            };
+            blocks.push(flat);
+        }
+    }
+    blocks
+}
+
+/// Element-wise comparison with dtype-aware tolerance (mirrors `torch.allclose`).
+fn assert_close<T: TestDtype>(actual: &[T], expected: &[T], context: &str) {
+    assert_eq!(
+        actual.len(),
+        expected.len(),
+        "{context}: length mismatch ({} vs {})",
+        actual.len(),
+        expected.len()
+    );
+    for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() {
+        let a_f64 = a.clone().to_f64();
+        let e_f64 = e.clone().to_f64();
+        let diff = (a_f64 - e_f64).abs();
+        let tol = T::ATOL + T::RTOL * e_f64.abs();
+        assert!(
+            diff <= tol,
+            "{context}[{i}]: {a_f64} vs {e_f64} (diff={diff}, tol={tol})"
+        );
+    }
+}
+
+/// Set up a CUDA context and stream.  Returns `None` if no GPU is available.
+fn cuda_setup() -> Option<(Arc<CudaStream>, cuda_runtime::cudaStream_t)> {
+    let count = CudaContext::device_count().ok()?;
+    if count == 0 {
+        return None;
+    }
+    let ctx = CudaContext::new(0).ok()?;
+    let stream = ctx.default_stream();
+    let raw = stream.cu_stream() as cuda_runtime::cudaStream_t;
+    Some((stream, raw))
+}
+
+// ---------------------------------------------------------------------------
+// GPU allocation helpers
+// ---------------------------------------------------------------------------
+
+/// Upload block chunks to GPU, returning the slices (kept alive) and a device
+/// pointer table suitable for the kernel FFI.
+fn upload_blocks<T: TestDtype>(
+    stream: &Arc<CudaStream>,
+    ref_blocks: &[Vec<Vec<T>>],
+) -> Result<(Vec<Vec<CudaSlice<T>>>, CudaSlice<usize>), DriverError> {
+    let nb = ref_blocks.len();
+    let chunks_per_batch = ref_blocks.first().map_or(0, |b| b.len());
+    let mut all_slices: Vec<Vec<CudaSlice<T>>> = Vec::with_capacity(nb);
+    let mut ptr_values: Vec<usize> = Vec::with_capacity(nb * chunks_per_batch);
+
+    for batch in ref_blocks {
+        let mut slices = Vec::with_capacity(batch.len());
+        for chunk in batch {
+            let slice = stream.clone_htod(chunk)?;
+            {
+                let (ptr, _guard) = slice.device_ptr(stream);
+                ptr_values.push(ptr as usize);
+            }
+            slices.push(slice);
+        }
+        all_slices.push(slices);
+    }
+
+    let ptrs_device = stream.clone_htod(ptr_values.as_slice())?;
+    Ok((all_slices, ptrs_device))
+}
+
+/// Allocate `count` poison-filled (0xDE) device buffers of `volume` elements each.
+/// Returns the slices and a device pointer table.
+fn alloc_buffers<T: TestDtype>(
+    stream: &Arc<CudaStream>,
+    count: usize,
+    volume: usize,
+) -> Result<(Vec<CudaSlice<T>>, CudaSlice<usize>), DriverError> {
+    let mut slices: Vec<CudaSlice<T>> = Vec::with_capacity(count);
+    let mut ptr_values: Vec<usize> = Vec::with_capacity(count);
+    let byte_count = volume * std::mem::size_of::<T>();
+
+    for _ in 0..count {
+        let mut slice = unsafe { stream.alloc::<T>(volume)? };
+        {
+            let (ptr, _guard) = slice.device_ptr_mut(stream);
+            ptr_values.push(ptr as usize);
+            unsafe {
+                memset_d8_async(ptr, 0xDE, byte_count, stream.cu_stream())?;
+            }
+        }
+        slices.push(slice);
+    }
+
+    let ptrs_device = stream.clone_htod(ptr_values.as_slice())?;
+    Ok((slices, ptrs_device))
+}
+
+/// Poison-fill (0xDE) all block chunk slices. `chunk_volume` is the element count per chunk.
+fn poison_fill_blocks<T: TestDtype>(
+    stream: &Arc<CudaStream>,
+    block_slices: &mut [Vec<CudaSlice<T>>],
+    chunk_volume: usize,
+) -> Result<(), DriverError> {
+    let byte_count = chunk_volume * std::mem::size_of::<T>();
+    for batch in block_slices.iter_mut() {
+        for slice in batch.iter_mut() {
+            let (dptr, _guard) = slice.device_ptr_mut(stream);
+            unsafe {
+                memset_d8_async(dptr, 0xDE, byte_count, stream.cu_stream())?;
+            }
+        }
+    }
+    Ok(())
+}
+
+// ---------------------------------------------------------------------------
+// block <-> universal roundtrip
+// ---------------------------------------------------------------------------
+
+fn block_universal_roundtrip_inner<T: TestDtype>(layout: BlockLayout) -> Result<(), DriverError> {
+    let (stream, stream_raw) = match cuda_setup() {
+        Some(s) => s,
+        None => return Ok(()),
+    };
+
+    // Dimensions matching the Python test.
+    let nh = 3usize;
+    let nl = 2usize;
+    let no = 2usize;
+    let nt = 4usize;
+    let hd = 5usize;
+    let nb = 3usize;
+    let universal_volume = nh * nl * no * nt * hd;
+
+    // Generate random universal tensors and compute reference blocks.
+    let mut rng = rand::rng();
+    let universals: Vec<Array5<T>> = (0..nb)
+        .map(|_| {
+            Array5::from_shape_fn((nh, nl, no, nt, hd), |_| {
+                T::from_f64(rng.random::<f64>() * 2.0 - 1.0)
+            })
+        })
+        .collect();
+
+    let ref_blocks: Vec<Vec<Vec<T>>> = universals.iter().map(|u| make_blocks(u, layout)).collect();
+
+    // Upload reference blocks to GPU.
+    let (mut block_slices, block_ptrs) = upload_blocks(&stream, &ref_blocks)?;
+
+    // Allocate universal output buffers on GPU.
+    let (universal_slices, universal_ptrs) = alloc_buffers::<T>(&stream, nb, universal_volume)?;
+
+    // --- Forward: blocks -> universal ---
+    {
+        let (bp, _g1) = block_ptrs.device_ptr(&stream);
+        let (up, _g2) = universal_ptrs.device_ptr(&stream);
+        let status = unsafe {
+            universal_from_block(
+                up as usize as *const *mut c_void,
+                bp as usize as *const *const c_void,
+                nb,
+                nh,
+                nl,
+                no,
+                nt,
+                hd,
+                T::DTYPE,
+                layout,
+                stream_raw,
+            )
+        };
+        assert_eq!(status, cuda_runtime::cudaError::cudaSuccess);
+    }
+    stream.synchronize()?;
+
+    // Verify each universal buffer matches the original tensor.
+    for (i, (slice, expected)) in universal_slices.iter().zip(universals.iter()).enumerate() {
+        let host = stream.clone_dtoh(slice)?;
+        let expected_flat: Vec<T> = expected.as_standard_layout().as_slice().unwrap().to_vec();
+        assert_close::<T>(&host, &expected_flat, &format!("universal batch {i}"));
+    }
+
+    // --- Reverse: poison-fill blocks, then universal -> blocks ---
+    poison_fill_blocks(&stream, &mut block_slices, nh * nt * hd)?;
+    stream.synchronize()?;
+
+    {
+        let (bp, _g1) = block_ptrs.device_ptr(&stream);
+        let (up, _g2) = universal_ptrs.device_ptr(&stream);
+        let status = unsafe {
+            block_from_universal(
+                up as usize as *const *const c_void,
+                bp as usize as *const *mut c_void,
+                nb,
+                nh,
+                nl,
+                no,
+                nt,
+                hd,
+                T::DTYPE,
+                layout,
+                stream_raw,
+            )
+        };
+        assert_eq!(status, cuda_runtime::cudaError::cudaSuccess);
+    }
+    stream.synchronize()?;
+
+    for (bi, (batch, ref_batch)) in block_slices.iter().zip(ref_blocks.iter()).enumerate() {
+        for (ci, (slice, expected)) in batch.iter().zip(ref_batch.iter()).enumerate() {
+            let host = stream.clone_dtoh(slice)?;
+            assert_close::<T>(&host, expected, &format!("block batch {bi} chunk {ci}"));
+        }
+    }
+
+    Ok(())
+}
+
+macro_rules! block_universal_test {
+    ($name:ident, $ty:ty, $layout:expr) => {
+        #[test]
+        fn $name() -> Result<(), DriverError> {
+            block_universal_roundtrip_inner::<$ty>($layout)
+        }
+    };
+}
+
+block_universal_test!(block_universal_roundtrip_nhd_f16, f16, BlockLayout::NHD);
+block_universal_test!(block_universal_roundtrip_nhd_bf16, bf16, BlockLayout::NHD);
+block_universal_test!(block_universal_roundtrip_nhd_f32, f32, BlockLayout::NHD);
+block_universal_test!(block_universal_roundtrip_nhd_f64, f64, BlockLayout::NHD);
+block_universal_test!(block_universal_roundtrip_hnd_f16, f16, BlockLayout::HND);
+block_universal_test!(block_universal_roundtrip_hnd_bf16, bf16, BlockLayout::HND);
+block_universal_test!(block_universal_roundtrip_hnd_f32, f32, BlockLayout::HND);
+block_universal_test!(block_universal_roundtrip_hnd_f64, f64, BlockLayout::HND);
+
+// ---------------------------------------------------------------------------
+// Edge cases
+// ---------------------------------------------------------------------------
+
+/// All kernel functions with num_blocks=0 should be a noop returning cudaSuccess.
+#[test]
+fn empty_batch_noop() -> Result<(), DriverError> {
+    let (_stream, stream_raw) = match cuda_setup() {
+        Some(s) => s,
+        None => return Ok(()),
+    };
+
+    let null_mut = std::ptr::null() as *const *mut c_void;
+    let null_const = std::ptr::null() as *const *const c_void;
+
+    // universal_from_block
+    let status = unsafe {
+        universal_from_block(
+            null_mut,
+            null_const,
+            0,
+            1,
+            1,
+            1,
+            1,
+            1,
+            TensorDataType::F32,
+            BlockLayout::NHD,
+            stream_raw,
+        )
+    };
+    assert_eq!(status, cuda_runtime::cudaError::cudaSuccess);
+
+    // block_from_universal
+    let status = unsafe {
+        block_from_universal(
+            null_const,
+            null_mut,
+            0,
+            1,
+            1,
+            1,
+            1,
+            1,
+            TensorDataType::F32,
+            BlockLayout::NHD,
+            stream_raw,
+        )
+    };
+    assert_eq!(status, cuda_runtime::cudaError::cudaSuccess);
+
+    Ok(())
+}
+
+// ---------------------------------------------------------------------------
+// CPU-only validation of make_blocks reference implementation
+// ---------------------------------------------------------------------------
+
+/// Verify `make_blocks` for NHD layout against first-principles index arithmetic.
+/// Uses deterministic position-encoded values so each element maps to a unique expected value.
+#[test]
+fn make_blocks_reference_nhd() {
+    let nh = 3usize;
+    let nl = 2usize;
+    let no = 2usize;
+    let nt = 4usize;
+    let hd = 5usize;
+
+    let universal =
+        Array5::from_shape_fn((nh, nl, no, nt, hd), |(nh_i, nl_i, no_i, nt_i, hd_i)| {
+            ((((nh_i * nl + nl_i) * no + no_i) * nt + nt_i) * hd + hd_i) as f32
+        });
+
+    let blocks = make_blocks(&universal, BlockLayout::NHD);
+    assert_eq!(blocks.len(), nl * no);
+
+    for nl_i in 0..nl {
+        for no_i in 0..no {
+            let block = &blocks[nl_i * no + no_i];
+            assert_eq!(block.len(), nt * nh * hd);
+            for nt_i in 0..nt {
+                for nh_i in 0..nh {
+                    for hd_i in 0..hd {
+                        // NHD block offset: [nt, nh, hd]
+                        let offset = (nt_i * nh + nh_i) * hd + hd_i;
+                        let expected =
+                            ((((nh_i * nl + nl_i) * no + no_i) * nt + nt_i) * hd + hd_i) as f32;
+                        assert_eq!(
+                            block[offset], expected,
+                            "NHD mismatch at nl={nl_i} no={no_i} nt={nt_i} nh={nh_i} hd={hd_i}"
+                        );
+                    }
+                }
+            }
+        }
+    }
+}
+
+/// Verify `make_blocks` for HND layout against first-principles index arithmetic.
+#[test]
+fn make_blocks_reference_hnd() {
+    let nh = 3usize;
+    let nl = 2usize;
+    let no = 2usize;
+    let nt = 4usize;
+    let hd = 5usize;
+
+    let universal =
+        Array5::from_shape_fn((nh, nl, no, nt, hd), |(nh_i, nl_i, no_i, nt_i, hd_i)| {
+            ((((nh_i * nl + nl_i) * no + no_i) * nt + nt_i) * hd + hd_i) as f32
+        });
+
+    let blocks = make_blocks(&universal, BlockLayout::HND);
+    assert_eq!(blocks.len(), nl * no);
+
+    for nl_i in 0..nl {
+        for no_i in 0..no {
+            let block = &blocks[nl_i * no + no_i];
+            assert_eq!(block.len(), nh * nt * hd);
+            for nh_i in 0..nh {
+                for nt_i in 0..nt {
+                    for hd_i in 0..hd {
+                        // HND block offset: [nh, nt, hd]
+                        let offset = (nh_i * nt + nt_i) * hd + hd_i;
+                        let expected =
+                            ((((nh_i * nl + nl_i) * no + no_i) * nt + nt_i) * hd + hd_i) as f32;
+                        assert_eq!(
+                            block[offset], expected,
+                            "HND mismatch at nl={nl_i} no={no_i} nh={nh_i} nt={nt_i} hd={hd_i}"
+                        );
+                    }
+                }
+            }
+        }
+    }
+}

lib/kvbm-kernels/tests/memcpy_batch.rs

+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Integration tests for `memcpy_batch` and the always-available query helpers
+//! (`is_memcpy_batch_available`, `is_using_stubs`).
+//!
+//! These don't require `permute_kernels` — the functions are unconditionally
+//! linked regardless of feature flags.
+//!
+//! Functional tests use pinned-host -> device -> pinned-host roundtrips (H2D + D2H)
+//! to match the transfer patterns that `cudaMemcpyBatchAsync` is designed for.
+
+#![cfg(all(feature = "testing-cuda", not(stub_kernels)))]
+
+use std::ffi::c_void;
+use std::sync::Arc;
+
+use cudarc::driver::{CudaContext, CudaSlice, CudaStream, DevicePtr, DriverError};
+use cudarc::runtime::sys as cuda_runtime;
+use kvbm_kernels::{MemcpyBatchMode, is_memcpy_batch_available, is_using_stubs, memcpy_batch};
+
+// Direct FFI for cudaMallocHost / cudaFreeHost.
+// We bypass cudarc's runtime::sys because cudarc eagerly resolves ALL runtime
+// symbols on first use, and CUDA 13.x removed `cudaGetDeviceProperties_v2`
+// which causes a panic.  Our test binary links against libcudart directly
+// (through kvbm-kernels' build.rs), so these symbols are always available.
+unsafe extern "C" {
+    fn cudaMallocHost(ptr: *mut *mut c_void, size: usize) -> u32;
+    fn cudaFreeHost(ptr: *mut c_void) -> u32;
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+fn cuda_setup() -> Option<(Arc<CudaStream>, cuda_runtime::cudaStream_t)> {
+    let count = CudaContext::device_count().ok()?;
+    if count == 0 {
+        return None;
+    }
+    let ctx = CudaContext::new(0).ok()?;
+    // Use a non-default stream — cudaMemcpyBatchAsync does not accept the
+    // NULL (default) stream.  A real CUstream from the driver API works fine.
+    let stream = ctx.new_stream().ok()?;
+    let raw = stream.cu_stream() as cuda_runtime::cudaStream_t;
+    Some((stream, raw))
+}
+
+/// Allocate `len` zero bytes on device, return slice + raw device address.
+fn alloc_device_zeroed(
+    stream: &Arc<CudaStream>,
+    len: usize,
+) -> Result<(CudaSlice<u8>, usize), DriverError> {
+    let slice = stream.alloc_zeros::<u8>(len)?;
+    let addr = {
+        let (ptr, _guard) = slice.device_ptr(stream);
+        ptr as usize
+    };
+    Ok((slice, addr))
+}
+
+/// RAII wrapper around pinned host memory allocated with `cudaMallocHost`.
+struct PinnedBuffer {
+    ptr: *mut c_void,
+    len: usize,
+}
+
+impl PinnedBuffer {
+    /// Allocate `len` bytes of pinned host memory, zeroed.
+    fn new_zeroed(len: usize) -> Self {
+        let mut ptr: *mut c_void = std::ptr::null_mut();
+        let err = unsafe { cudaMallocHost(&mut ptr, len) };
+        assert_eq!(err, 0, "cudaMallocHost failed with error {err}");
+        // Zero the buffer
+        unsafe { std::ptr::write_bytes(ptr as *mut u8, 0, len) };
+        Self { ptr, len }
+    }
+
+    /// Allocate pinned host memory and fill it from `data`.
+    fn from_data(data: &[u8]) -> Self {
+        let buf = Self::new_zeroed(data.len());
+        unsafe {
+            std::ptr::copy_nonoverlapping(data.as_ptr(), buf.ptr as *mut u8, data.len());
+        }
+        buf
+    }
+
+    fn as_ptr(&self) -> *mut c_void {
+        self.ptr
+    }
+
+    fn as_const_ptr(&self) -> *const c_void {
+        self.ptr as *const c_void
+    }
+
+    /// Read contents back as a `Vec<u8>`.
+    fn to_vec(&self) -> Vec<u8> {
+        let mut v = vec![0u8; self.len];
+        unsafe {
+            std::ptr::copy_nonoverlapping(self.ptr as *const u8, v.as_mut_ptr(), self.len);
+        }
+        v
+    }
+}
+
+impl Drop for PinnedBuffer {
+    fn drop(&mut self) {
+        if !self.ptr.is_null() {
+            unsafe {
+                cudaFreeHost(self.ptr);
+            }
+        }
+    }
+}
+
+/// Run a pinned-host -> device -> pinned-host roundtrip via two `memcpy_batch` calls.
+///
+/// 1. Batch H2D: copy from `src_pinned` buffers to `device` buffers
+/// 2. Batch D2H: copy from `device` buffers to `dst_pinned` (zeroed) buffers
+/// 3. Assert `dst_pinned` contents match original data
+///
+/// Returns both batch statuses for the caller to inspect.
+fn h2d_d2h_roundtrip(
+    stream: &Arc<CudaStream>,
+    raw: cuda_runtime::cudaStream_t,
+    data_sets: &[Vec<u8>],
+    copy_size: usize,
+    mode: MemcpyBatchMode,
+) -> Result<(cuda_runtime::cudaError, cuda_runtime::cudaError), DriverError> {
+    let num_pairs = data_sets.len();
+
+    // Source: pinned host buffers filled with known data
+    let src_pinned: Vec<PinnedBuffer> = data_sets
+        .iter()
+        .map(|d| PinnedBuffer::from_data(d))
+        .collect();
+
+    // Device buffers (zeroed)
+    let mut dev_slices = Vec::with_capacity(num_pairs);
+    let mut dev_addrs = Vec::with_capacity(num_pairs);
+    for _ in 0..num_pairs {
+        let (s, a) = alloc_device_zeroed(stream, copy_size)?;
+        dev_slices.push(s);
+        dev_addrs.push(a);
+    }
+
+    // Destination: zeroed pinned host buffers
+    let dst_pinned: Vec<PinnedBuffer> = (0..num_pairs)
+        .map(|_| PinnedBuffer::new_zeroed(copy_size))
+        .collect();
+
+    // Build pointer arrays for H2D: src = pinned host, dst = device
+    let h2d_src_ptrs: Vec<*const c_void> = src_pinned.iter().map(|b| b.as_const_ptr()).collect();
+    let h2d_dst_ptrs: Vec<*mut c_void> = dev_addrs.iter().map(|&a| a as *mut c_void).collect();
+
+    let h2d_status = unsafe {
+        memcpy_batch(
+            h2d_src_ptrs.as_ptr() as *const *const c_void,
+            h2d_dst_ptrs.as_ptr() as *const *mut c_void,
+            copy_size,
+            num_pairs,
+            mode,
+            raw,
+        )
+    };
+
+    if h2d_status != cuda_runtime::cudaError::cudaSuccess {
+        return Ok((h2d_status, cuda_runtime::cudaError::cudaSuccess));
+    }
+
+    // Build pointer arrays for D2H: src = device, dst = pinned host
+    let d2h_src_ptrs: Vec<*const c_void> = dev_addrs.iter().map(|&a| a as *const c_void).collect();
+    let d2h_dst_ptrs: Vec<*mut c_void> = dst_pinned.iter().map(|b| b.as_ptr()).collect();
+
+    let d2h_status = unsafe {
+        memcpy_batch(
+            d2h_src_ptrs.as_ptr() as *const *const c_void,
+            d2h_dst_ptrs.as_ptr() as *const *mut c_void,
+            copy_size,
+            num_pairs,
+            mode,
+            raw,
+        )
+    };
+
+    if d2h_status != cuda_runtime::cudaError::cudaSuccess {
+        return Ok((h2d_status, d2h_status));
+    }
+
+    // Synchronize before reading back
+    stream.synchronize()?;
+
+    // Verify roundtrip: dst_pinned should match original data
+    for (i, (dst, expected)) in dst_pinned.iter().zip(data_sets.iter()).enumerate() {
+        let result = dst.to_vec();
+        assert_eq!(
+            result,
+            *expected,
+            "roundtrip mismatch at pair {i}: first differing byte at position {}",
+            result
+                .iter()
+                .zip(expected.iter())
+                .position(|(a, b)| a != b)
+                .unwrap_or(result.len())
+        );
+    }
+
+    Ok((h2d_status, d2h_status))
+}
+
+// ---------------------------------------------------------------------------
+// Query function tests
+// ---------------------------------------------------------------------------
+
+#[test]
+fn stubs_not_active() {
+    // Since the file is gated on not(stub_kernels), this must be false.
+    assert!(!is_using_stubs());
+}
+
+#[test]
+fn availability_is_consistent() {
+    // Just ensure it doesn't crash and returns a stable value.
+    let a = is_memcpy_batch_available();
+    let b = is_memcpy_batch_available();
+    assert_eq!(a, b);
+    eprintln!(
+        "cudaMemcpyBatchAsync available: {} (CUDA {}12.9)",
+        a,
+        if a { ">=" } else { "<" }
+    );
+}
+
+// ---------------------------------------------------------------------------
+// memcpy_batch edge cases (work regardless of CUDA version)
+// ---------------------------------------------------------------------------
+
+#[test]
+fn memcpy_batch_zero_copies_noop() -> Result<(), DriverError> {
+    let (_stream, raw) = match cuda_setup() {
+        Some(s) => s,
+        None => return Ok(()),
+    };
+
+    // All modes should treat zero copies as a no-op
+    for mode in [
+        MemcpyBatchMode::BatchedWithFallback,
+        MemcpyBatchMode::FallbackOnly,
+        MemcpyBatchMode::BatchWithoutFallback,
+    ] {
+        let status = unsafe {
+            memcpy_batch(
+                std::ptr::null(),
+                std::ptr::null(),
+                128,
+                0, // num_copies = 0
+                mode,
+                raw,
+            )
+        };
+        assert_eq!(
+            status,
+            cuda_runtime::cudaError::cudaSuccess,
+            "mode={mode:?}"
+        );
+    }
+    Ok(())
+}
+
+#[test]
+fn memcpy_batch_zero_size_noop() -> Result<(), DriverError> {
+    let (_stream, raw) = match cuda_setup() {
+        Some(s) => s,
+        None => return Ok(()),
+    };
+
+    // All modes should treat zero size as a no-op
+    for mode in [
+        MemcpyBatchMode::BatchedWithFallback,
+        MemcpyBatchMode::FallbackOnly,
+        MemcpyBatchMode::BatchWithoutFallback,
+    ] {
+        let status = unsafe {
+            memcpy_batch(
+                std::ptr::null(),
+                std::ptr::null(),
+                0, // size_per_copy = 0
+                5,
+                mode,
+                raw,
+            )
+        };
+        assert_eq!(
+            status,
+            cuda_runtime::cudaError::cudaSuccess,
+            "mode={mode:?}"
+        );
+    }
+    Ok(())
+}
+
+// ---------------------------------------------------------------------------
+// memcpy_batch functional tests — H2D + D2H roundtrip with pinned memory
+//
+// Each test runs across all three MemcpyBatchMode variants:
+// - BatchedWithFallback: always works (batch or fallback)
+// - FallbackOnly: always works (individual cudaMemcpyAsync)
+// - BatchWithoutFallback: only works when batch API is available (CUDA 12.9+)
+// ---------------------------------------------------------------------------
+
+/// Run a roundtrip test across all modes, handling BatchWithoutFallback gracefully
+/// when the batch API is not available.
+fn run_all_modes(
+    stream: &Arc<CudaStream>,
+    raw: cuda_runtime::cudaStream_t,
+    data_sets: &[Vec<u8>],
+    copy_size: usize,
+) -> Result<(), DriverError> {
+    let batch_available = is_memcpy_batch_available();
+
+    for mode in [
+        MemcpyBatchMode::BatchedWithFallback,
+        MemcpyBatchMode::FallbackOnly,
+        MemcpyBatchMode::BatchWithoutFallback,
+    ] {
+        let (h2d, d2h) = h2d_d2h_roundtrip(stream, raw, data_sets, copy_size, mode)?;
+
+        if mode == MemcpyBatchMode::BatchWithoutFallback && !batch_available {
+            // Expected to fail when batch API is not available
+            eprintln!("  {mode:?}: batch API unavailable, got h2d={h2d:?} (expected non-success)");
+            continue;
+        }
+
+        assert_eq!(
+            h2d,
+            cuda_runtime::cudaError::cudaSuccess,
+            "H2D failed with mode={mode:?}"
+        );
+        assert_eq!(
+            d2h,
+            cuda_runtime::cudaError::cudaSuccess,
+            "D2H failed with mode={mode:?}"
+        );
+        eprintln!("  {mode:?}: OK");
+    }
+    Ok(())
+}
+
+/// Single H2D + D2H roundtrip via `memcpy_batch` (all modes).
+#[test]
+fn memcpy_batch_single_copy() -> Result<(), DriverError> {
+    let (stream, raw) = match cuda_setup() {
+        Some(s) => s,
+        None => return Ok(()),
+    };
+
+    let copy_size = 256;
+    let data: Vec<u8> = (0..copy_size as u16).map(|i| (i % 256) as u8).collect();
+    run_all_modes(&stream, raw, &[data], copy_size)
+}
+
+/// Multiple independent H2D + D2H roundtrips in one batch call (all modes).
+#[test]
+fn memcpy_batch_multiple_copies() -> Result<(), DriverError> {
+    let (stream, raw) = match cuda_setup() {
+        Some(s) => s,
+        None => return Ok(()),
+    };
+
+    let num_pairs = 8;
+    let copy_size = 512;
+    let data_sets: Vec<Vec<u8>> = (0..num_pairs)
+        .map(|i| {
+            (0..copy_size)
+                .map(|j| ((i * 31 + j * 7) % 256) as u8)
+                .collect()
+        })
+        .collect();
+
+    run_all_modes(&stream, raw, &data_sets, copy_size)
+}
+
+/// Large copy (1 MiB per pair) to exercise alignment paths (all modes).
+#[test]
+fn memcpy_batch_large_copy() -> Result<(), DriverError> {
+    let (stream, raw) = match cuda_setup() {
+        Some(s) => s,
+        None => return Ok(()),
+    };
+
+    let copy_size = 1 << 20; // 1 MiB
+    let num_pairs = 3;
+    let data_sets: Vec<Vec<u8>> = (0..num_pairs)
+        .map(|i| (0..copy_size).map(|j| ((i + j) % 251) as u8).collect())
+        .collect();
+
+    run_all_modes(&stream, raw, &data_sets, copy_size)
+}
+
+/// Non-power-of-two copy size (regression guard for alignment assumptions, all modes).
+#[test]
+fn memcpy_batch_odd_size() -> Result<(), DriverError> {
+    let (stream, raw) = match cuda_setup() {
+        Some(s) => s,
+        None => return Ok(()),
+    };
+
+    let copy_size = 999; // not aligned to anything useful
+    let num_pairs = 4;
+    let data_sets: Vec<Vec<u8>> = (0..num_pairs)
+        .map(|i| (0..copy_size).map(|j| ((i * 13 + j) % 256) as u8).collect())
+        .collect();
+
+    run_all_modes(&stream, raw, &data_sets, copy_size)
+}
+
+/// Many small pairs to stress the batch dispatch path (all modes).
+#[test]
+fn memcpy_batch_many_pairs() -> Result<(), DriverError> {
+    let (stream, raw) = match cuda_setup() {
+        Some(s) => s,
+        None => return Ok(()),
+    };
+
+    let num_pairs = 256;
+    let copy_size = 64;
+    let data_sets: Vec<Vec<u8>> = (0..num_pairs)
+        .map(|i| (0..copy_size).map(|j| ((i + j) % 256) as u8).collect())
+        .collect();
+
+    run_all_modes(&stream, raw, &data_sets, copy_size)
+}
+
+// ---------------------------------------------------------------------------
+// Diagnostic: mirrors NVIDIA benchmark calling pattern exactly
+// ---------------------------------------------------------------------------

lib/kvbm-kernels/tests/stub_build.rs

+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Tests that only compile when stub kernels are in use (no CUDA available).
+//!
+//! Complementary to the `memcpy_batch::stubs_not_active` test which asserts
+//! the opposite under `not(stub_kernels)`.
+
+#![cfg(stub_kernels)]
+
+use kvbm_kernels::{is_memcpy_batch_available, is_using_stubs};
+
+#[test]
+fn stubs_active() {
+    assert!(
+        is_using_stubs(),
+        "expected is_using_stubs() == true under stub build"
+    );
+}
+
+#[test]
+fn memcpy_batch_unavailable_under_stubs() {
+    assert!(
+        !is_memcpy_batch_available(),
+        "expected is_memcpy_batch_available() == false under stub build"
+    );
+}