Refactor KV Cache Manager with a new torch exportable KVCache backend

[geoffreyQiu]

Description

• Refactor KV Cache Manager with a new torch exportable KVCache backend
• Implement KV Cache custom ops for torch export and AOTInductor.
• Add PyTorch & C++ torch runtime example for exported aoti HSTU model inference with KVCache.
• Add triton server deployment example for exported aoti HSTU model inference with KVCache.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[greptile-apps]

Greptile Summary

This PR refactors the KV cache manager with a new pluggable KVCacheBackend ABC, migrates logic from KVCacheManager into DefaultKVCacheBackend, and adds a new torch-exportable C++ ExportKVCacheRuntime backed by custom kvcache_manager_ops torch library ops. It also adds AOTI inference examples for HSTU models with KV caching and a Triton server deployment path.

• Backend abstraction: KVCacheManager is now a thin delegation wrapper over KVCacheBackend; all logic lives in DefaultKVCacheBackend, with the ABC providing the contract for custom backends.
• Torch export path: A new C++ ExportKVCacheRuntime and kvcache_manager_ops torch library expose lookup, allocate, onboard, and offload ops usable inside torch.export/AOTI compiled graphs; Python fake-op shims enable tracing without the runtime library.
• AOTI examples: ExportKVCachedInferenceRankingGR wires the ops into an exportable forward pass, with C++ and Triton-server inference examples under examples/hstu/inference_aoti.

Confidence Score: 5/5

The core backend abstraction, torch op schemas, and AOTI export path are logically correct; previously-flagged gaps have been addressed and the new findings are minor quality issues that do not affect correctness under normal usage.

All three new findings are style or low-impact: a missing sleep in offload_flush that could cause CPU spinning on async backends, a duplicate dead-code guard, and an int64/int32 type mismatch in a from_blob call that is only exercised in the valid_len==0 branch where the value is zero regardless of interpretation.

default_kvcache_backend.py (offload_flush busy-loop) and export_kvcache_runtime.cpp (type mismatch and dead-code guard)

Important Files Changed

Filename	Overview
corelib/recsys_kvcache_manager/recsys_kvcache_manager/kvcache_manager.py	Refactored to thin delegation wrapper over KVCacheBackend; backward-compat properties added; offload_reap_completed and offload_flush now forwarded correctly.
corelib/recsys_kvcache_manager/recsys_kvcache_manager/kvcache_backend.py	New KVCacheBackend ABC defining the full lifecycle contract including offload_reap_completed and offload_flush.
corelib/recsys_kvcache_manager/recsys_kvcache_manager/default_kvcache_backend.py	Extracted from KVCacheManager; offload_flush busy-loops without sleep between poll iterations unlike the C++ counterpart which sleeps 100ms.
corelib/recsys_kvcache_manager/recsys_kvcache_manager/fake_kvcache_manager_ops.py	New shape-only fake-op shims for torch.export tracing; runtime config lazily loaded only on first fake-op invocation, avoiding import-time crashes.
corelib/recsys_kvcache_manager/src/runtime/export_kvcache_runtime.cpp	New C++ runtime; duplicate dead-code guard in onboard_kvcache_wait and int64/int32 type mismatch in from_blob in the valid_len==0 branch.
corelib/recsys_kvcache_manager/src/torch_binding/kvcache_manager_ops.cpp	New torch library implementing all kvcache_manager_ops; offload_flush_impl returns default-constructed tensor when queue is empty — previously flagged and discussed.
corelib/recsys_kvcache_manager/src/torch_binding/kvcache_manager_context.cpp	New context singleton backed by thread_local kvcache_manager_; per-thread runtime initialization issue — previously flagged and discussed.
corelib/recsys_kvcache_manager/recsys_kvcache_manager/init.py	Lazy-import pattern; fake ops registered only when ops present, fixing prior unconditional crash; ExportKVCacheBackend removed from _LAZY_IMPORTS.
examples/hstu/model/export_kvcached_inference_ranking_gr.py	New ExportKVCachedInferenceRankingGR wiring kvcache_manager_ops into AOTI-exportable forward pass; op call sites match C++ schemas correctly.
examples/hstu/modules/paged_hstu_infer_layer.py	Adds export_mode flag at init time; sm12 (Blackwell) added to supported architectures; fusion kernels disabled in export mode.

_{Reviews (19): Last reviewed commit: "Update the readme and doc" | Re-trigger Greptile}

[shijieliu]

1. move GUIDE_TO_RUN_AOTI_MODEL_INFERENCE_E2E.md as README.md
2. add a benchmark folder under inference_aoti and move aoti vs. python and aoti e2e benchmark results into it.

[shijieliu]

about the fbgemm update, i think main branch has some update. make sure it does not reverse the changes on main.

[geoffreyQiu]

/build fbgemm devel

[JacoCheung]

❔ Pipeline #58380880 -- canceling

Job	Status	Log
build_devel_x86	❔ canceling	view
build_devel_arm64	❔ canceling	view
build_tritonserver_devel_x86	❔ canceling	view
build_tritonserver_devel_arm64	❌ failed	view
manifest_devel	🚫 canceled	view
manifest_tritonserver_devel	🚫 canceled	view
pre_check	🚫 canceled	view
train_build_x86	🚫 canceled	view
train_build_arm64	🚫 canceled	view
tritonserver_build_x86	🚫 canceled	view
prepare-jet-b200-smoke	🚫 canceled	view
prepare-jet-b200-inference	🚫 canceled	view
prepare-jet-cw-dfw-e2e-benchmark	🚫 canceled	view
build_whl	🚫 canceled	view
dynamicemb_test_fwd_bwd_8gpus	🚫 canceled	view
dynamicemb_test_load_dump_8gpus	🚫 canceled	view
unit_test_1gpu_a100	🚫 canceled	view
unit_test_1gpu_h100	🚫 canceled	view
unit_test_4gpu	🚫 canceled	view
unit_test_tp_4gpu	🚫 canceled	view
L20_unit_test_1gpu	🚫 canceled	view
inference_unit_test_1gpu	🚫 canceled	view
inference_test_1gpu	🚫 canceled	view

View full pipeline

[geoffreyQiu]

Final CI: pipeline
Re:

1. Renamed and reorged the Readme file in inference_aoti.
2. Add specific benchmark page (inference_aoti/benchmark/README.md) for inference_aoti
3. FBGemm HSTU update：merged into tracking branch recsys-examples-dev.