Repository navigation
Explore SemBlend-style semantics for GPU-initiated NVMe-KV cache-block retrieval #182
Description
Motivation
PR #177 adds GPU-initiated NVMe Key-Value Store/Retrieve to nvme-ep. That is a natural data-plane primitive for LLM KV-cache reuse: a GPU can retrieve key-addressed KV-cache block payloads from an NVMe-KV target directly into GPU-accessible memory.
SemBlend's (https://github.com/WorldFlowAI/semblend) role is complementary! SemBlend does not need the storage layer to perform semantic search. It produces donor evidence and a segmented prefill plan: fresh spans that must be computed, plus donor spans that may be materialized if the backend accepts them. The question for this issue is how to map those SemBlend-level semantics onto an NVMe-KV/XIO block retrieval path without making the NVMe-KV endpoint responsible for model correctness.
Proposed layering
Keep the layering strict:
- SemBlend / inference runtime: discovers semantic donor spans and emits a segmented plan.
- Backend materializer: validates model/tokenizer/KV-layout compatibility, handles RoPE/source-position mapping, accepts or declines donor spans, and owns fallback.
- ROCm XIO / NVMe-KV: retrieves or stores opaque key-addressed byte payloads efficiently. It should not decide whether semantic reuse is safe.
In this model, SemBlend metadata helps choose which KV blocks to fetch, while XIO provides the GPU-initiated movement primitive.
Candidate object model
A first prototype could treat each NVMe-KV value as a fixed-size KV-cache block shard. The key can be a compact content/block identifier, while richer metadata lives in a sidecar manifest or a separate metadata object.
Candidate key/value shape:
- NVMe-KV key: block id, content hash, or manifest-derived block key.
- NVMe-KV value: fixed-size KV block shard, for example one model/layer/KV/token-block shard or another backend-native block unit.
- Sidecar metadata: model id, tokenizer hash, chat-template/version hash, KV dtype, block size, layer range, tensor-parallel rank, KV-head layout, RoPE base/scaling, source token range, cache generation, tenant/policy namespace, and checksum.
The segmented plan would then refer to a list of block keys for each donor span.
Candidate execution path
recipient request
-> SemBlend donor discovery
-> SegmentedPrefillPlan
F0: fresh span
S1: donor block keys A0..A31, source positions, target positions
F1: fresh span
S2: donor block keys B0..B15, source positions, target positions
-> backend validates compatibility
-> ROCm XIO batched NVMe-KV Retrieve into request-local GPU buffer
-> backend maps/references/copies into paged KV layout
-> backend applies RoPE correction or recomputes unsafe spans
-> exact-prefix cache remains separate
What to measure
This should be evaluated before optimizing the interface too far:
- exact hash KV-block hit rate
- SemBlend candidate block hit rate; ( I have a new and growing dataset for this based on LongBench v1/v2)
- backend-accepted/materialized block hit rate
- XIO NVMe-KV retrieve latency and effective bandwidth/IOPS
- recompute-vs-fetch break-even by context length and block size
- TTFT impact
- quality gates from the inference side, including PPL ratio and structured correctness checks.
Possible MVP
- Store fixed-size synthetic KV-cache block payloads under keys in an NVMe-KV target.
- Have SemBlend or a small harness produce a segmented plan containing donor block keys.
- Use XIO batched Retrieve to fetch those blocks into GPU memory.
- Measure fetch latency and bandwidth against recompute-equivalent tensor sizes.
- Add an inference-runtime integration only after the block movement path is characterized.
This would give us a clean way to connect SemBlend's semantic block-hit-rate lift to an AMD GPU-initiated data path
Activity
zbennett10 commentedon Jun 30, 2026
@sbates130272 Just an update!
I spent some time turning the SemBlend/XIO idea into a small reproducible harness.
HIgh-level overview of what we created on our side:
The semantic index returns a donor record; the storage path retrieves that donor's actual block keys. This keeps provenance and exact identity clean:
What I have working in the harness:
A donor/block manifest shape that carries:
A backend-neutral storage interface with:
An XIO/NVMe-KV command-boundary adapter that maps manifest-derived block keys into the PR nvme-ep: add GPU-initiated NVMe Key-Value Store/Retrieve #177-style
xio-tester nvme-ep --kv-op retrievepath. I built from this PRA multi-span replay shaped like wrapper/tool/RAG prompts:
Separate metrics for:
I also exercised PR #183's containerized Ceph + SPDK + QEMU host-KV path... Immediately wanted to work towards somethinng that's easily reproducible - lots of moving parts here.. But would love some access to some AMD gpus to actually prove real GPU-initiated retrieval of semantic KV blocks.. working on this...
The next concrete test I would like to run on suitable AMD hardware is:
xio-tester nvme-ep --kv-op retrievewith those keys and capture latency/bytes.Using this harness with https://github.com/WorldFlowAI/sembench we've started to see semantic KV reuse across NVMe but want to wait to share much of anything until we get AMD GPUs. I can share a cleaned-up harness or smaller reproducer once it is useful.. but at the moment waiting on those
sbates130272 commentedon Jun 30, 2026
@zbennett10 what is the minimal AMD setup that could work for this?
zbennett10 commentedon Jun 30, 2026
Nothing crazy @sbates130272 - one bare-metal AMD ROCm box with local NVMe
zbennett10 commentedon Jul 16, 2026
@sbates130272 This is slowly becoming a reality within inference engines. Note:
sgl-project/sglang#31057
sbates130272 commentedon Jul 16, 2026
@zbennett10 is there similar work happening in vllm?
zbennett10 commentedon Jul 20, 2026
@sbates130272 With vLLM it's technically already possible through their connector API - although, the way I understand it, there still may need to be some updates in the inference engine to support this deeply there. Most of our focus so far has been on SGLang as we just haven't had the time nor the resources to focus on both deeply
sbates130272 commentedon Sep 16, 2026
@zbennett10 great progress. A few thoughts on where to push next given the current state of main.
The quiesce/ownership infrastructure makes this viable now. PRs #197–#201 (now in main) add:
Together these mean a GPU kernel can now take ownership of an NVMe queue, issue a batch of KV-block Retrieve commands targeting real GPU VRAM, release the queue cleanly, and have the kernel continue normal I/O on it — without the races or data-visibility bugs we saw before these PRs.
Suggested concrete next step for experimentation.
The PR #177 NVMe-KV path (still open) is the right place to start. Before wiring it into an inference runtime, I'd suggest:
Validate that the KV Retrieve path lands data correctly in GPU VRAM on bare metal with modes 1/8/11 (now that PR fix(kernel/rocm-xio): use dynamic dma-buf attach for NVMe P2PDMA passthrough #201 fixes the P2P address). Use the data-integrity checks that r2samsung used in issue [Issue]: xio-tester nvme-ep tests fail for memory-mode (1, 2, 11) but passes for (0, 8) on a MI210 GPU #181 — write a stamped pattern from GPU, flush, Retrieve back into a fresh GPU buffer, compare.
Measure the Retrieve latency and bandwidth per block size (4 KiB, 16 KiB, 128 KiB) on the target NVMe against the
hipMemcpyD2H+H2D round-trip that a recompute path would avoid. This gives the break-even context length.Once step 2 confirms the movement cost, integrate the batched Retrieve call at the SGLang connector boundary (rather than deep in the attention kernel) — the connector API @zbennett10 mentioned for vLLM/SGLang is the right seam. The XIO call becomes
xio_batched_kv_retrieve(block_keys[], n, gpu_buffer, stream)and the connector validates and maps results.The sgl-project/sglang#31057 reference is exactly the right direction. If that connector PR lands, the XIO NVMe-KV Retrieve can plug in there without patching attention kernels.
One note on the layering: the sidecar metadata (model id, tokenizer hash, RoPE base, etc.) should live entirely in the SGLang/vLLM layer, not in the XIO key schema. The NVMe-KV key stays a compact opaque blob derived from the donor block record; XIO doesn't need to know what's in it.