Refactor KV Cache Manager with a new torch exportable KVCache backend#400
Refactor KV Cache Manager with a new torch exportable KVCache backend#400geoffreyQiu wants to merge 9 commits into
Conversation
Introduce KVCacheBackend, split the default Python backend from the public KVCacheManager facade, and add the initial export lookup path.
| 1. Building required custom operators and runtime libraries | ||
| 2. Exporting the HSTU ranking model with KV-cache support through `torch.export` and AOTInductor | ||
| 3. Starting the FlexKV-based KV-cache runtime service | ||
| 4. Running the C++ replay executable against exported artifacts and dumped tensors | ||
| 5. Running the Triton Server demo path for the exported KV-cache AOTI model |
There was a problem hiding this comment.
There is already GUIDE_TO_RUN_CPP_INFERENCE_DEMO.md, i am wondering if we can consolidated them.
Does step 1, 2, 4 already covered in CPP_INFERENCE_DEMO?
| from model.inference_ranking_gr import _STRIP_CACHED_TOKENS_OP | ||
|
|
||
|
|
||
| class ExportKVCachedInferenceRankingGR(torch.nn.Module): |
There was a problem hiding this comment.
Why do we need a sepeate ExportKVCAchedRankingGR instead of reusing existing python kvcache ranking GR?
| @@ -0,0 +1,151 @@ | |||
| # Export KVCache Design | |||
There was a problem hiding this comment.
i would prefer to have a directory corelib/recsys_kvcache_manager/doc and put those under it
* Support PyTorch export and AOTI compile/package with test coverage. * Support C++ Torch runtime validation for AOTI-compiled KV-cache graphs. * Support Triton Inference Server validation using the PyTorch AOTI backend. * Add required KV-cache/runtime/custom-op plumbing and export fake-op shims.
* extract and simplify HSTU KV-cache export wrapper * remove stale debug instrumentation and export bring-up prints * standardize macro-gated logging in KV-cache runtime/binding/client code * rename AOTI ordering placeholder tensors consistently * clean C++ replay demo logging and CUDA check behavior * document end-to-end HSTU KV-cache AOTI C++ and Triton Server workflows
geoffreyQiu
force-pushed
the
aoti_kvcache
branch
from
e4cdaf1 to
d070ab3
Compare
Greptile SummaryThis PR refactors the KV cache manager by introducing a
Confidence Score: 2/5The export backend is not functional in its current state; multiple method calls will fail at runtime even after fixing op names, and the HSTU layer's KV-cache path crashes on any GPU newer than sm80. ExportKVCacheBackend.lookup_kvcache omits the required third ordering_tensor argument to the lookup op, and onboard_launch passes 9 individual tensors where the C++ schema expects 5 arguments with lookup_results as a flat list — both will raise TypeError on first use. On top of previously flagged defects (wrong op names, undefined classes, out-of-scope variables), the backend has no path to a working invocation. Additionally, PagedHSTUInferLayer now unconditionally routes any KV-cache request through hstu_attn_export_impl, which raises RuntimeError on H100/sm89+ hardware regardless of the export_mode setting. export_kvcache_backend.py has the most concentrated defects and needs a full rewrite of every method body. paged_hstu_infer_layer.py needs a hardware-agnostic fallback for the KV-cache attention path. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Model as ExportKVCachedInferenceRankingGR
participant Ops as kvcache_manager_ops (C++)
participant Layer as PagedHSTUInferLayer
participant Backend as ExportKVCacheBackend
Note over Model,Backend: Correct path (ExportKVCachedInferenceRankingGR)
Model->>Ops: offload_reap_completed(user_ids) → completed_offloads
Model->>Ops: lookup(user_ids, seqlens, completed_offloads) → lookup_res[7]
Model->>Ops: allocate(user_ids, seqlens, lookup_res[1], lookup_res[5]) → alloc_result[13]
Model->>Ops: onboard_launch(user_ids, seqlens, lookup_res, alloc_result[0], alloc_result[2]) → (mappings, offsets, task_ids)
Model->>Ops: onboard_wait(task_ids, ordering) → kv_cache_table[]
Model->>Layer: forward(batch, kvcache_metadata)
Layer->>Layer: append_kvcache(...) → updated kv_cache_table
Layer->>Layer: hstu_attn_export_impl(...) [sm80 only]
Model->>Ops: offload_launch(..., slot_mappings, ordering) → task_ids
Note over Backend: Broken path (ExportKVCacheBackend)
Backend->>Ops: lookup_kvcache(user_ids, seqlens) wrong name + missing arg
Backend->>Ops: allocate_kvcache(...) wrong name + wrong unpack
Backend->>Ops: onboard_kvcache_launch(u, s, t1, t2, t3, t4, t5, idx, ptr) wrong name + 9 args vs 5
Backend->>Ops: onboard_kvcache_wait(...) wrong name + returns None
Backend->>Ops: offload_launch(task_handle.task_ids) task_handle out of scope + 1 arg vs 10
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant Model as ExportKVCachedInferenceRankingGR
participant Ops as kvcache_manager_ops (C++)
participant Layer as PagedHSTUInferLayer
participant Backend as ExportKVCacheBackend
Note over Model,Backend: Correct path (ExportKVCachedInferenceRankingGR)
Model->>Ops: offload_reap_completed(user_ids) → completed_offloads
Model->>Ops: lookup(user_ids, seqlens, completed_offloads) → lookup_res[7]
Model->>Ops: allocate(user_ids, seqlens, lookup_res[1], lookup_res[5]) → alloc_result[13]
Model->>Ops: onboard_launch(user_ids, seqlens, lookup_res, alloc_result[0], alloc_result[2]) → (mappings, offsets, task_ids)
Model->>Ops: onboard_wait(task_ids, ordering) → kv_cache_table[]
Model->>Layer: forward(batch, kvcache_metadata)
Layer->>Layer: append_kvcache(...) → updated kv_cache_table
Layer->>Layer: hstu_attn_export_impl(...) [sm80 only]
Model->>Ops: offload_launch(..., slot_mappings, ordering) → task_ids
Note over Backend: Broken path (ExportKVCacheBackend)
Backend->>Ops: lookup_kvcache(user_ids, seqlens) wrong name + missing arg
Backend->>Ops: allocate_kvcache(...) wrong name + wrong unpack
Backend->>Ops: onboard_kvcache_launch(u, s, t1, t2, t3, t4, t5, idx, ptr) wrong name + 9 args vs 5
Backend->>Ops: onboard_kvcache_wait(...) wrong name + returns None
Backend->>Ops: offload_launch(task_handle.task_ids) task_handle out of scope + 1 arg vs 10
Reviews (4): Last reviewed commit: "Fix Docker" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def lookup_kvcache( | ||
| self, | ||
| user_ids: torch.Tensor, | ||
| sequence_lengths: torch.Tensor, | ||
| ) -> Tuple[KVIndexMeta, KVLookupResult]: | ||
| ( | ||
| cached_start_indices, | ||
| cached_lengths, | ||
| gpu_cached_start_indices, | ||
| gpu_cached_lengths, | ||
| host_cached_start_indices, | ||
| host_cached_lengths, | ||
| task_ids, | ||
| ) = torch.ops.kvcache_manager_ops.lookup_kvcache(user_ids, sequence_lengths) | ||
|
|
||
| index_meta = KVIndexMeta( | ||
| user_ids=user_ids, | ||
| seq_lengths=sequence_lengths, | ||
| ) | ||
| lookup_result = KVLookupResult( | ||
| user_ids=user_ids, | ||
| cached_start_indices=cached_start_indices, | ||
| cached_lengths=cached_lengths, | ||
| gpu_cached_start_indices=gpu_cached_start_indices, | ||
| gpu_cached_lengths=gpu_cached_lengths, | ||
| host_cached_start_indices=host_cached_start_indices, | ||
| host_cached_lengths=host_cached_lengths, | ||
| extra={ | ||
| "task_ids": task_ids | ||
| } | ||
| ) | ||
| return index_meta, lookup_result | ||
|
|
||
| def allocate_kvcache( | ||
| self, | ||
| index_meta: KVIndexMeta, | ||
| lookup_results: KVLookupResult, | ||
| output_kvcache_metadata: Optional[KVCacheMetadata] = None, | ||
| ) -> KVCacheMetadata: | ||
| # assert output_kvcache_metadata is None, "Pre-allocated KVCacheMetadata is not supported in ExportKVCacheBackend yet." | ||
|
|
||
| (metadata_buffer, metadata_tensors) = torch.ops.kvcache_manager_ops.allocate_kvcache( | ||
| index_meta.user_ids, | ||
| index_meta.seq_lengths, | ||
| lookup_results.cached_lengths, | ||
| lookup_results.host_cached_lengths, | ||
| ) | ||
|
|
||
| return KVCacheMetadata( | ||
| page_ids_gpu_buffer=metadata_buffer[0], | ||
| metadata_gpu_buffer=metadata_buffer[1], | ||
| kv_indices=metadata_buffer[0], | ||
|
|
||
| kv_indptr=metadata_tensors[0], | ||
| kv_last_page_len=metadata_tensors[1], | ||
| total_history_lengths=metadata_tensors[2], | ||
| total_history_offsets=metadata_tensors[3], | ||
| new_history_offsets=metadata_tensors[4], | ||
| batch_indices=metadata_tensors[5], | ||
| position=metadata_tensors[6], | ||
|
|
||
| new_history_nnz=metadata_tensors[7], | ||
| new_history_nnz_cuda=metadata_tensors[8], | ||
|
|
||
| kv_seqlens=metadata_tensors[9], | ||
| kv_seqlen_offsets=metadata_tensors[10], | ||
| kv_onload_handle=None, | ||
| ) | ||
|
|
||
| def onboard_launch( | ||
| self, | ||
| index_meta: KVIndexMeta, | ||
| lookup_result: KVLookupResult, | ||
| kvcache_metadata: KVCacheMetadata, | ||
| ) -> HostKVTaskHandle: | ||
| slot_mappings = torch.ops.kvcache_manager_ops.onboard_kvcache_launch( | ||
| index_meta.user_ids, | ||
| index_meta.seq_lengths, | ||
| lookup_result.cached_lengths, | ||
| lookup_result.host_cached_lengths, | ||
| lookup_result.gpu_cached_start_indices, | ||
| lookup_result.gpu_cached_lengths, | ||
| lookup_result.extra["task_ids"], | ||
| kvcache_metadata.kv_indices, | ||
| kvcache_metadata.kv_indptr, | ||
| ) | ||
|
|
||
| # In export backend, all user_ids and slot mappings are recorded. User ids with no cache to onboard will have task_id of -1, and the corresponding slot mapping can be ignored in the downstream processing. | ||
| onload_handle = _FlexKVOnloadHandle( | ||
| task_ids=lookup_result.extra["task_ids"], | ||
| uids=index_meta.user_ids, | ||
| slot_mappings=slot_mappings, | ||
| ) | ||
| return HostKVTaskHandle( | ||
| backend="flexkv", | ||
| user_ids=onload_handle.uids, | ||
| handle=onload_handle, | ||
| status=HostKVTaskStatus.LAUNCHED, | ||
| # metadata={ | ||
| # "onboard_start_indices": torch.tensor( | ||
| # onboard_start_indices, dtype=torch.int32 | ||
| # ), | ||
| # "onboard_lengths": torch.tensor(onboard_lengths, dtype=torch.int32), | ||
| # }, | ||
| ) | ||
|
|
||
| def onboard_try_wait( | ||
| self, | ||
| kv_index_meta: KVIndexMeta, | ||
| task_handle: Optional[HostKVTaskHandle], | ||
| ) -> Optional[HostKVWaitResult]: | ||
| raise NotImplementedError("ExportKVCacheBackend.onboard_try_wait is not implemented yet.") | ||
|
|
||
| def onboard_wait( | ||
| self, | ||
| kv_index_meta: KVIndexMeta, | ||
| task_handle: Optional[HostKVTaskHandle], | ||
| ) -> Optional[HostKVWaitResult]: | ||
| torch.ops.kvcache_manager_ops.onboard_kvcache_wait( | ||
| task_handle.handle.task_ids, | ||
| ) | ||
|
|
||
| def offload_launch( | ||
| self, | ||
| index_meta: KVIndexMeta, | ||
| kvcache_metadata: Optional[KVCacheMetadata] = None, | ||
| ): | ||
| torch.ops.kvcache_manager_ops.offload_launch( | ||
| task_handle.handle.task_ids, | ||
| ) |
There was a problem hiding this comment.
The file has several bugs that will each raise an exception the first time the backend is exercised:
-
Wrong op name (line 40):
torch.ops.kvcache_manager_ops.lookup_kvcachedoes not exist — the C++ library registers the op aslookup(m.def("lookup(...)")). Same pattern at line 102 (onboard_kvcache_launch→onboard_launch) and line 145 (onboard_kvcache_wait→onboard_wait). -
_FlexKVOnloadHandleis undefined (line 115): the class is defined inflex_kvcache_manager.pybut is never imported here, causing aNameErrorwheneveronboard_launchis called. -
HostKVTaskStatusis undefined (line 124): the import at line 10 brings inHostKVStorageBase,HostKVTaskHandle, andHostKVWaitResult, but notHostKVTaskStatus(which lives inhost_kvstorage_manager). -
task_handleis out of scope (line 155):offload_launch(self, index_meta, kvcache_metadata)referencestask_handle.handle.task_ids, buttask_handleis not a parameter of this method; this is aNameErroron every call.
|
|
||
| return parallel_input | ||
|
|
||
| def hstu_attn_export_impl( | ||
| self, | ||
| query, | ||
| key, | ||
| value, | ||
| jd: JaggedData, | ||
| kv_cache_metadata, | ||
| kv_cache_table, | ||
| batch_size: int, | ||
| ): | ||
| sm_major_version = torch.cuda.get_device_properties(0).major | ||
| if sm_major_version != 8: | ||
| raise RuntimeError( | ||
| "Export-mode paged-KV HSTU attention currently calls " | ||
| "torch.ops.fbgemm.hstu_varlen_fwd_80 directly. Add the matching " | ||
| "direct dispatcher path before exporting this on non-sm80 GPUs." | ||
| ) | ||
| jagged_attn_output, _ = torch.ops.fbgemm.hstu_varlen_fwd_80( | ||
| query, | ||
| key, | ||
| value, | ||
| jd.seqlen_offsets[: batch_size + 1], | ||
| kv_cache_metadata.kv_seqlen_offsets[: batch_size + 1], | ||
| None, | ||
| None, # seqused_q, seqused_k | ||
| jd.max_seqlen, | ||
| jd.max_seqlen, | ||
| jd.scaling_seqlen, | ||
| None, # num_contexts | ||
| jd.num_candidates[:batch_size], | ||
| self._target_group_size, | ||
| -1, | ||
| 0, | ||
| self._alpha, | ||
| None, |
There was a problem hiding this comment.
hstu_attn_export_impl, breaking inference on H100/sm90+
After the refactor, every call that has kv_cache_metadata is not None unconditionally calls self.hstu_attn_export_impl(...) — there is no branch that falls back to the original hstu_attn_varlen_func for the non-export case. Inside hstu_attn_export_impl, line ~259 raises a RuntimeError whenever sm_major_version != 8. This means any deployment on H100 (sm90), RTX 4090 (sm89), or later hardware that uses KV caching will immediately crash — regardless of whether export_mode is True or False. The previous hstu_attn_varlen_func call, which used a proper kv_cache= path and had no such restriction, was completely removed from this branch.
| item_offsets = _offsets(item_lengths).cpu() | ||
| action_offsets = _offsets(item_lengths).cpu() |
There was a problem hiding this comment.
action_offsets from item_lengths instead of action_lengths. The test only passes today because the test fixture sets action_lengths = item_lengths.clone(). If a future test case uses different lengths for actions and items, the reference will silently produce incorrect offsets and the test would pass against a wrong expected value.
| item_offsets = _offsets(item_lengths).cpu() | |
| action_offsets = _offsets(item_lengths).cpu() | |
| item_offsets = _offsets(item_lengths).cpu() | |
| action_offsets = _offsets(action_lengths).cpu() |
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| (metadata_buffer, metadata_tensors) = torch.ops.kvcache_manager_ops.allocate_kvcache( | ||
| index_meta.user_ids, | ||
| index_meta.seq_lengths, | ||
| lookup_results.cached_lengths, | ||
| lookup_results.host_cached_lengths, | ||
| ) |
There was a problem hiding this comment.
allocate_kvcache op name used here doesn't exist in the registered C++ library. The TORCH_LIBRARY_FRAGMENT(kvcache_manager_ops) in kvcache_manager_ops.cpp registers it as allocate (line 17: m.def("allocate(...)")). Calling allocate_kvcache will raise AttributeError at runtime on the first request that hits this backend.
| (metadata_buffer, metadata_tensors) = torch.ops.kvcache_manager_ops.allocate_kvcache( | |
| index_meta.user_ids, | |
| index_meta.seq_lengths, | |
| lookup_results.cached_lengths, | |
| lookup_results.host_cached_lengths, | |
| ) | |
| (metadata_buffer, metadata_tensors) = torch.ops.kvcache_manager_ops.allocate( | |
| index_meta.user_ids, | |
| index_meta.seq_lengths, | |
| lookup_results.cached_lengths, | |
| lookup_results.host_cached_lengths, | |
| ) |
| def offload_try_wait(self) -> None: | ||
| raise NotImplementedError("ExportKVCacheBackend.offload_try_wait is not implemented yet.") | ||
|
|
||
| def offload_reap_completed(self) -> None: | ||
| torch.ops.kvcache_manager_ops.offload_reap_completed() |
There was a problem hiding this comment.
offload_reap_completed missing required ordering_tensor argument
The C++ op is registered as m.def("offload_reap_completed(Tensor ordering_tensor) -> Tensor") — it requires exactly one tensor argument. Calling torch.ops.kvcache_manager_ops.offload_reap_completed() with no arguments will raise a TypeError at runtime every time this method is invoked. The fake implementation (_offload_reap_completed_fake) also documents the expected dummy tensor parameter.
| def onboard_wait( | ||
| self, | ||
| kv_index_meta: KVIndexMeta, | ||
| task_handle: Optional[HostKVTaskHandle], | ||
| ) -> Optional[HostKVWaitResult]: | ||
| torch.ops.kvcache_manager_ops.onboard_kvcache_wait( | ||
| task_handle.handle.task_ids, | ||
| ) |
There was a problem hiding this comment.
onboard_wait discards the KV cache table result and returns None
torch.ops.kvcache_manager_ops.onboard_wait returns Tensor[] — the list of per-layer KV cache tables. The C++ implementation calls runtime->get_kvcache_tables() and returns those tensors as the op result. The Python method here never captures or returns that result; it falls off the end and returns None. Any caller of KVCacheManager.onboard_wait(...) will receive None instead of the cache tables, silently corrupting downstream attention computation.
| def allocate_kvcache( | ||
| self, | ||
| index_meta: KVIndexMeta, | ||
| lookup_results: KVLookupResult, | ||
| output_kvcache_metadata: Optional[KVCacheMetadata] = None, | ||
| ) -> KVCacheMetadata: | ||
| # assert output_kvcache_metadata is None, "Pre-allocated KVCacheMetadata is not supported in ExportKVCacheBackend yet." | ||
|
|
||
| (metadata_buffer, metadata_tensors) = torch.ops.kvcache_manager_ops.allocate_kvcache( | ||
| index_meta.user_ids, | ||
| index_meta.seq_lengths, | ||
| lookup_results.cached_lengths, | ||
| lookup_results.host_cached_lengths, | ||
| ) | ||
|
|
||
| return KVCacheMetadata( | ||
| page_ids_gpu_buffer=metadata_buffer[0], | ||
| metadata_gpu_buffer=metadata_buffer[1], | ||
| kv_indices=metadata_buffer[0], | ||
|
|
||
| kv_indptr=metadata_tensors[0], | ||
| kv_last_page_len=metadata_tensors[1], | ||
| total_history_lengths=metadata_tensors[2], | ||
| total_history_offsets=metadata_tensors[3], | ||
| new_history_offsets=metadata_tensors[4], | ||
| batch_indices=metadata_tensors[5], | ||
| position=metadata_tensors[6], | ||
|
|
||
| new_history_nnz=metadata_tensors[7], | ||
| new_history_nnz_cuda=metadata_tensors[8], | ||
|
|
||
| kv_seqlens=metadata_tensors[9], | ||
| kv_seqlen_offsets=metadata_tensors[10], | ||
| kv_onload_handle=None, |
There was a problem hiding this comment.
Tensor[] return cannot be unpacked into a 2-tuple
After fixing the op name (allocate_kvcache → allocate), this code will still fail. The C++ allocate op returns a flat list of 13 tensors (confirmed by _allocate_fake returning 13 torch.empty calls). The line (metadata_buffer, metadata_tensors) = torch.ops.kvcache_manager_ops.allocate(...) attempts to unpack that 13-element list into exactly two variables, raising ValueError: too many values to unpack. The correct approach — as used in ExportKVCachedInferenceRankingGR._kvcache_metadata_from_ops — is to index the flat result directly: result[0], result[1], etc.
| host_cached_lengths, | ||
| task_ids, | ||
| ) = torch.ops.kvcache_manager_ops.lookup_kvcache(user_ids, sequence_lengths) | ||
|
|
||
| index_meta = KVIndexMeta( | ||
| user_ids=user_ids, | ||
| seq_lengths=sequence_lengths, | ||
| ) | ||
| lookup_result = KVLookupResult( |
There was a problem hiding this comment.
lookup called with 2 arguments — C++ schema requires 3
The C++ op schema registers lookup(Tensor user_ids, Tensor seqlens, Tensor ordering_tensor) -> Tensor[] (three arguments). This call passes only user_ids and sequence_lengths, omitting the required ordering_tensor (a dependency tensor for ordering/sync). The correct three-argument form is demonstrated in ExportKVCachedInferenceRankingGR.forward() where completed_offloads is passed as the third argument. After fixing the op name (separately flagged), this call would still raise a TypeError at runtime due to the missing argument.
| lookup_result.gpu_cached_lengths, | ||
| lookup_result.extra["task_ids"], | ||
| kvcache_metadata.kv_indices, | ||
| kvcache_metadata.kv_indptr, | ||
| ) | ||
|
|
||
| # In export backend, all user_ids and slot mappings are recorded. User ids with no cache to onboard will have task_id of -1, and the corresponding slot mapping can be ignored in the downstream processing. | ||
| onload_handle = _FlexKVOnloadHandle( | ||
| task_ids=lookup_result.extra["task_ids"], | ||
| uids=index_meta.user_ids, | ||
| slot_mappings=slot_mappings, |
There was a problem hiding this comment.
onboard_launch called with 9 individual tensors — C++ schema expects 5 arguments with a list
The C++ schema is onboard_launch(Tensor user_ids, Tensor seqlens, Tensor[] lookup_results, Tensor kv_page_indices, Tensor kv_page_indptr). The lookup_results parameter is a Tensor[] that must receive the full 7-tensor output of the lookup op as a single list, not the individual fields broken out as separate positional arguments. ExportKVCachedInferenceRankingGR.forward() demonstrates the correct 5-argument call: onboard_launch(user_ids, total_history_lengths, lookup_res, alloc_result[0], alloc_result[2]). Even after fixing the op name, passing 9 arguments instead of 5 will raise a TypeError.
2 participants
Description
Checklist