## 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
| `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
| 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
| `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
| `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
| `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 |