Refactor KV Cache Manager with a new torch exportable KVCache backend

.pre-commit-config.yaml

@@ -29,6 +29,7 @@ repos:
         args: ["--ignore-missing-imports", 
                 "--explicit-package-bases", 
                 "--exclude=examples/hstu/inference/triton/hstu_.*/[0-9]*/model.py$",
+                "--exclude=examples/hstu/triton_libs/",
                 "--exclude=^examples/sid-gr-inference/"
             ]
         pass_filenames: false

corelib/dynamicemb/CMakeLists.txt

@@ -34,7 +34,7 @@ if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
     if(DEFINED SM)
         set(CMAKE_CUDA_ARCHITECTURES ${SM})
     else()
-        set(CMAKE_CUDA_ARCHITECTURES 70 75 80 89 90 100)
+        set(CMAKE_CUDA_ARCHITECTURES 70 75 80 89 90 100 120)
     endif()
 endif()
 

corelib/dynamicemb/dynamicemb/key_value_table.py

@@ -1116,7 +1116,9 @@ def export_keys_values_iter(
             # (logical) score-column order so both the checkpoint and callers of
             # export_keys_values see scores in the order they configured. Identity
             # when the logical order matches the physical layout.
-            dump_perm = score_dump_permutation(state.options_list[table_id].score_strategy)
+            dump_perm = score_dump_permutation(
+                state.options_list[table_id].score_strategy
+            )
             if dump_perm != list(range(out_scores.size(1))):
                 out_scores = out_scores[:, dump_perm].contiguous()
         else:

corelib/dynamicemb/dynamicemb/scored_hashtable.py

@@ -31,14 +31,14 @@
     device_timestamp,
     table_copy_score_blocks,
     table_count_matched,
-    table_gather_score_blocks,
-    table_scatter_score_blocks,
     table_erase,
     table_export_batch,
+    table_gather_score_blocks,
     table_insert,
     table_insert_and_evict,
     table_lookup,
     table_partition,
+    table_scatter_score_blocks,
     table_update_counter_with_layout,
 )
 

corelib/dynamicemb/setup.py

@@ -127,6 +127,8 @@ def get_extensions():
             "arch=compute_90,code=sm_90",
             "-gencode",
             "arch=compute_100,code=sm_100",
+            "-gencode",
+            "arch=compute_120,code=sm_120",
             "-w",
             "-U__CUDA_NO_HALF_OPERATORS__",
             "-U__CUDA_NO_HALF_CONVERSIONS__",

corelib/dynamicemb/test/unit_tests/table_operation/test_lru_lfu.py

@@ -280,7 +280,9 @@ def _read_score_blocks(storage, keys, tids):
 
 
 @pytest.mark.parametrize(
-    "score_strategy", [LRU_LFU_TS_FIRST, LRU_LFU_LFU_FIRST], ids=["ts_first", "lfu_first"]
+    "score_strategy",
+    [LRU_LFU_TS_FIRST, LRU_LFU_LFU_FIRST],
+    ids=["ts_first", "lfu_first"],
 )
 def test_lru_lfu_storage_dump_load_roundtrip(current_device, tmp_path, score_strategy):
     """DynamicEmbStorage dump -> load preserves BOTH LruLfu score words
@@ -313,7 +315,9 @@ def test_lru_lfu_storage_dump_load_roundtrip(current_device, tmp_path, score_str
     # export_keys_values scan is a separate iteration and need not reproduce the
     # dump's row order). The frequency column must land at the position where the
     # user placed LFU in score_strategy.
-    file_keys = torch.tensor(np.fromfile(f"{d}/keys", dtype=np.int64), dtype=torch.int64)
+    file_keys = torch.tensor(
+        np.fromfile(f"{d}/keys", dtype=np.int64), dtype=torch.int64
+    )
     file_scores = torch.tensor(
         np.fromfile(f"{d}/score", dtype=np.uint64), dtype=torch.uint64
     ).view(-1, 2)
@@ -323,9 +327,9 @@ def test_lru_lfu_storage_dump_load_roundtrip(current_device, tmp_path, score_str
         int(k): int(file_scores[i, freq_col]) for i, k in enumerate(file_keys.tolist())
     }
     for k, f in zip(keys.tolist(), freq_ref.tolist()):
-        assert key_to_freq[k] == f, (
-            f"checkpoint frequency column ({freq_col}) must match user's logical order"
-        )
+        assert (
+            key_to_freq[k] == f
+        ), f"checkpoint frequency column ({freq_col}) must match user's logical order"
 
     dst = _lru_lfu_storage(dim=dim, score_strategy=score_strategy)
     dst.load(0, *paths, include_optim=False)

corelib/dynamicemb/test/unit_tests/test_lfu_scores.py

@@ -308,9 +308,7 @@ def test_lfu_score_validation(
     else:
         print(f"  - Caching: DISABLED")
     print(f"  - Global HBM budget scale: {global_hbm_budget_scale}")
-    print(
-        f"  - compound (TIMESTAMP, LFU) score: {lru_lfu} (LruLfu 2-word score)"
-    )
+    print(f"  - compound (TIMESTAMP, LFU) score: {lru_lfu} (LruLfu 2-word score)")
     if not caching and global_hbm_budget_scale < 1.0:
         print(f"  - Storage path: expect HybridStorage (StorageMode DEFAULT)")
 

corelib/recsys_kvcache_manager/CMakeLists.txt

@@ -0,0 +1,114 @@
+cmake_minimum_required(VERSION 3.18)
+project(KVCacheManagerOps LANGUAGES CXX CUDA)
+
+include(GNUInstallDirs)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CXX_EXTENSIONS OFF)
+set(CMAKE_CUDA_STANDARD 17)
+set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+set(CMAKE_CUDA_EXTENSIONS OFF)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+find_package(Python3 COMPONENTS Interpreter REQUIRED)
+
+find_package(PkgConfig REQUIRED)
+pkg_check_modules(ZMQ REQUIRED libzmq)
+
+if(NOT DEFINED CMAKE_PREFIX_PATH)
+    execute_process(
+        COMMAND "${Python3_EXECUTABLE}" -c "import torch; print(torch.utils.cmake_prefix_path)"
+        OUTPUT_VARIABLE TORCH_CMAKE_PREFIX_PATH
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        ERROR_QUIET
+        RESULT_VARIABLE TORCH_PREFIX_STATUS
+    )
+    if(TORCH_PREFIX_STATUS EQUAL 0 AND TORCH_CMAKE_PREFIX_PATH)
+        list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX_PATH}")
+    endif()
+endif()
+
+find_package(Torch REQUIRED)
+find_package(CUDAToolkit REQUIRED)
+
+
+if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+    if(DEFINED SM)
+        set(CMAKE_CUDA_ARCHITECTURES ${SM})
+    else()
+        set(CMAKE_CUDA_ARCHITECTURES 70 75 80 86 89 90 100 120)
+    endif()
+endif()
+
+message(STATUS "Found Torch ${TORCH_VERSION}")
+message(STATUS "CUDA architectures: ${CMAKE_CUDA_ARCHITECTURES}")
+
+set(KVCACHE_MANAGER_OPS_SOURCES
+    src/torch_binding/kvcache_manager_ops.cpp
+    src/torch_binding/kvcache_manager_context.cpp
+    src/runtime/kvcache_runtime_config.cpp
+    src/runtime/export_kvcache_runtime.cpp
+    src/gpu_kvcache_manager_impl.cpp
+    src/torch_binding/flexkv_cpp_client.cpp
+    src/torch_binding/flexkv_aoti_protocol.cpp
+    src/gather_scatter.cpp
+    src/gather_scatter_kernels.cu
+)
+
+add_library(kvcache_manager_ops SHARED ${KVCACHE_MANAGER_OPS_SOURCES})
+
+separate_arguments(TORCH_CXX_FLAGS_LIST NATIVE_COMMAND "${TORCH_CXX_FLAGS}")
+if(TORCH_CXX_FLAGS_LIST)
+    target_compile_options(kvcache_manager_ops PRIVATE ${TORCH_CXX_FLAGS_LIST})
+endif()
+
+target_include_directories(kvcache_manager_ops PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}/src
+    ${CMAKE_CURRENT_SOURCE_DIR}/src/runtime
+    ${CMAKE_CURRENT_SOURCE_DIR}/src/torch_binding
+    ${TORCH_INCLUDE_DIRS}
+    ${CUDAToolkit_INCLUDE_DIRS}
+)
+
+target_compile_options(kvcache_manager_ops PRIVATE
+    $<$<COMPILE_LANGUAGE:CXX>:
+        -O3
+        -fdiagnostics-color=always
+        -fvisibility=hidden
+        -w
+    >
+    $<$<COMPILE_LANGUAGE:CUDA>:
+        -O3
+        --expt-relaxed-constexpr
+        --expt-extended-lambda
+        --use_fast_math
+        -Xcompiler=-fvisibility=hidden
+        -w
+        -U__CUDA_NO_HALF_OPERATORS__
+        -U__CUDA_NO_HALF_CONVERSIONS__
+        -U__CUDA_NO_HALF2_OPERATORS__
+        -U__CUDA_NO_BFLOAT16_CONVERSIONS__
+    >
+)
+
+target_compile_definitions(kvcache_manager_ops PRIVATE
+    TORCH_EXTENSION_NAME=kvcache_manager_ops
+)
+
+target_link_libraries(kvcache_manager_ops PRIVATE
+    ${TORCH_LIBRARIES}
+    CUDA::cudart
+    zmq
+)
+
+set_target_properties(kvcache_manager_ops PROPERTIES
+    PREFIX ""
+    OUTPUT_NAME "kvcache_manager_ops"
+    CUDA_SEPARABLE_COMPILATION OFF
+)
+
+install(TARGETS kvcache_manager_ops
+    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
+)

corelib/recsys_kvcache_manager/README.md

@@ -52,6 +52,11 @@ Furthermore, the `NativeHostKVCacheManager` backend supports layerwise KV data o
 `HostKVStorageManagerBase` is provided as an interface for other LLM-compatible KV-cache systems for host memory, storage, and remote data pools.
 This can be extended to integrate other KV-cache systems. Currently, we provide integration with [`FlexKV`](https://github.com/taco-project/FlexKV/tree/main) as the lower-tier KV storage backend.
 
+- **[NEW] Compatible with Torch Export and AOTI**: 
+Recsys KVCache Manager now includes an export-compatible backend for KV-cache lookup, allocation, onboarding, offloading, and eviction when using the paged GPU KV cache together with the FlexKV host cache backend.
+The export path keeps the Python API unchanged, but routes KV-cache operations through tensor-only torch custom ops while keeping stateful cache policy and runtime management in C++, which makes the flow compatible with `torch.export` and AOTInductor/AOTI packaging. The current implementation is dependent on a [customized FlexKV version](https://github.com/geoffreyQiu/FlexKV/tree/cpp_client).
+
+
 
 ## Installation
 

corelib/recsys_kvcache_manager/doc/EXPORT_KVCACHE_DESIGN.md

@@ -0,0 +1,138 @@
+# Export KVCache Design
+
+## 1. Scope and Goals
+This document captures the finalized class naming and architecture for an export-friendly KV cache stack.
+
+Goals:
+1. Keep one stable Python user interface.
+2. Support both default and export paths with shared workflow.
+3. Keep stateful runtime logic in C++ for export and AOTI compatibility.
+4. Keep torch custom ops thin and tensor-only.
+5. Use thread-local runtime context instead of explicit context_id in this phase.
+
+## 2. Final Naming
+
+### 2.1 Python layer
+1. `KVCacheManager`: user-facing API (kept for compatibility).
+2. `KVCacheBackend`: abstract backend interface.
+3. `DefaultKVCacheBackend`: default implementation using current Python manager stack.
+4. `DeviceKVCache`: device-side KV cache component.
+5. `NativeHostKVStorage`: native host storage component.
+6. `FlexKVStorage`: FlexKV-backed host storage component.
+
+### 2.2 C++ layer
+1. `KVCacheRuntimeContext`: thread-local context holder for current runtime instance.
+2. `IKVCacheRuntime`: abstract C++ runtime interface.
+3. `ExportKVCacheRuntime`: concrete C++ runtime implementation.
+
+## 3. High-Level Architecture and Ownership
+
+## 3.1 Layer responsibilities
+1. `KVCacheManager`
+- Owns orchestration workflow used by model/inference code.
+- Delegates backend-specific operations to `KVCacheBackend`.
+
+2. `KVCacheBackend`
+- Defines backend contract for operations such as lookup and allocate.
+- Keeps `KVCacheManager` independent of implementation details.
+
+3. `DefaultKVCacheBackend`
+- Uses current Python components (`DeviceKVCache`, `NativeHostKVStorage` or `FlexKVStorage`).
+- Preserves current behavior as the default path.
+
+4. `ExportKVCacheRuntime`
+- Owns runtime state and policy in C++.
+- Owns references to device and host storage runtime components.
+
+5. `KVCacheRuntimeContext`
+- Holds the current `IKVCacheRuntime` pointer in `thread_local` storage.
+- Serves runtime resolution for all kvcache torch ops on the active thread.
+
+## 3.2 Ownership model
+1. One inference instance owns one `ExportKVCacheRuntime`.
+2. The runtime is set into `KVCacheRuntimeContext` on inference entry.
+3. Torch ops read runtime from context and invoke methods.
+4. Runtime owns mutable state and backend handles.
+5. Context is cleared at inference exit.
+
+## 4. Key Relations (Composition + Shared Workflow)
+1. `KVCacheManager` has-a `KVCacheBackend`.
+2. `DefaultKVCacheBackend` has-a `DeviceKVCache` and host storage component.
+3. Export path calls ops from `torch.ops.kvcache_manager_ops`
+4. `ExportKVCacheRuntime` implements `IKVCacheRuntime`.
+5. `KVCacheRuntimeContext` stores `std::shared_ptr<IKVCacheRuntime>` per thread.
+
+## 5. lookup_kvcache: Implementation Outline
+
+## 5.1 Public API contract
+`KVCacheManager.lookup_kvcache(user_ids, sequence_lengths)` returns lookup metadata/results with the same logical schema regardless of backend.
+
+## 5.2 Default path (non-export)
+1. `KVCacheManager.lookup_kvcache` delegates to `DefaultKVCacheBackend.lookup_kvcache`.
+2. Backend calls `DeviceKVCache.lookup(user_ids)`.
+3. Backend builds index metadata and calls host storage lookup.
+4. Backend merges host and device lookup results.
+5. Backend returns merged result to `KVCacheManager`.
+
+## 5.3 Export path (AOTI and torch.export focus)
+1. Inference initialization sets `ExportKVCacheRuntime` into `KVCacheRuntimeContext`.
+2. Lookup in model forward calls `torch.ops.kvcache_manager_ops.lookup_kvcache(...)`.
+3. C++ op shim resolves runtime from `KVCacheRuntimeContext`.
+4. C++ op shim calls `ExportKVCacheRuntime::lookup_kvcache(...)`.
+5. Runtime performs device lookup and host lookup, merges outputs.
+6. Runtime returns tensor-only outputs to op shim.
+7. Op shim returns outputs to Python backend.
+
+## 5.4 Export path pseudocode
+```python
+# Python
+class ExportableModel:
+    def forward(self, user_ids, sequence_lengths, ...):
+        lookup_results = torch.ops.kvcache_manager_ops.lookup_kvcache(user_ids, sequence_lengths, ordering_tensor)
+```
+
+```cpp
+// C++ op shim
+std::vector<at::Tensor> lookup_kvcache_op(const at::Tensor& user_ids,
+                                          const at::Tensor& sequence_lengths) {
+    auto runtime = KVCacheRuntimeContext::instance().runtime();
+    return runtime->lookup_kvcache(user_ids, sequence_lengths);
+}
+```
+
+```cpp
+// C++ runtime
+std::vector<at::Tensor> ExportKVCacheRuntime::lookup_kvcache(
+    const at::Tensor& user_ids,
+    const at::Tensor& sequence_lengths) {
+    auto gpu = device_kvcache_->lookup(user_ids);
+    auto idx = host_kvstorage_->build_index_meta(user_ids, sequence_lengths);
+    auto host = host_kvstorage_->lookup_kvcache(idx);
+    return merge_lookup_outputs(gpu, host);
+}
+```
+
+## 6. Thread-Local Context Decision
+1. Current phase uses thread-local context and does not use explicit `context_id` op arguments.
+2. This is valid under these assumptions:
+- Each inference instance has its own device KV cache table.
+- Host backend handles concurrency outside this layer.
+3. If future execution interleaves multiple inference instances on the same thread, add explicit `context_id` to op schemas.
+
+## 7. Torch Op Design Constraints
+1. Torch ops remain thin adapters only.
+2. Torch op inputs and outputs are tensor-first and export-friendly.
+3. Runtime policy and mutable state live in C++ runtime, not in op wrappers.
+4. Every exported op should provide shape-compatible Meta behavior.
+
+## 8. Migration Notes
+1. Keep existing public API name `KVCacheManager` for compatibility.
+2. Introduce backend interface and two implementations first.
+3. Route `lookup_kvcache` through backend as the first migrated API.
+4. Migrate `allocate` and onboard/offload APIs incrementally.
+5. Keep fallback to `DefaultKVCacheBackend` while validating export path.
+
+## 9. Optional Additions
+1. Add an RAII helper around thread-local context set and restore in C++ implementation.
+2. Add a backend selection config flag to switch default and export implementations.
+3. Add parity tests that compare default and export lookup outputs for identical inputs.