Optimize recsys-FlexKV

[Clebrate]

Description

• Optimize _build_slot_mappings by reducing D2H datasize.
• Optimize as_batch
• support SSD

Checklist

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

[Clebrate]

/build

[JacoCheung]

❌ Pipeline #55646281 -- failed

Job	Status	Log
pre_check	❌ failed	view
train_build_x86	✅ success	view
train_build_arm64	✅ success	view
tritonserver_build_x86	✅ success	view
prepare-jet-b200-smoke	✅ success	view
prepare-jet-b200-inference	✅ success	view
prepare-jet-cw-dfw-e2e-benchmark	✅ success	view
build_whl	✅ success	view
dynamicemb_test_fwd_bwd_8gpus	✅ success	view
dynamicemb_test_load_dump_8gpus	✅ success	view
unit_test_1gpu_a100	❌ failed	view
unit_test_1gpu_h100	❌ failed	view
unit_test_4gpu	✅ success	view
unit_test_tp_4gpu	❌ failed	view
L20_unit_test_1gpu	✅ success	view
inference_unit_test_1gpu	✅ success	view
inference_test_1gpu	❌ failed	view

Result: 12/17 jobs passed

View full pipeline

[greptile-apps]

Greptile Summary

This PR optimizes the FlexKV KV-cache manager by batching the D2H transfer in _build_slot_mappings (one transfer per call instead of one per user), caching the _page_offsets tensor across calls, introducing a put_match+launch batch offload path, and adding SSD-tier config support via a new config_path parameter.

• _build_slot_mappings optimization: kv_indices and kv_indptr are now transferred to CPU in a single contiguous copy, and _page_offsets is cached as an instance attribute — benchmark results show an 8–10× reduction in that function's latency for bs=8.
• Batch offload path (as_batch): When as_batch=True and the batch has more than one user, offload uses put_match+launch instead of per-user put_async, reducing total offload+wait time by up to 31% at short sequence lengths.
• SSD support: A new config_path constructor arg forwards a YAML tier config to FlexKV, enabling CPU+SSD two-tier caching; sample configs and updated README are included.

Confidence Score: 5/5

Safe to merge; all core logic changes are well-scoped and the existing put_async fallback is preserved.

The optimizations to _build_slot_mappings and the new batch offload path are well-validated by benchmark data. The only finding is a stale env-var reference in a comment with no runtime impact.

examples/hstu/inference/configs/example.yml — the header comment references an env var that is not wired up anywhere.

Important Files Changed

Filename	Overview
corelib/recsys_kvcache_manager/recsys_kvcache_manager/flex_kvcache_manager.py	Core optimization: D2H transfer batched once per call in _build_slot_mappings; new put_match+launch batch offload path; as_batch and config_path constructor params added. Logic is sound.
corelib/recsys_kvcache_manager/recsys_kvcache_manager/kvcache_manager.py	Parses flexkv_as_batch (with string/bool normalisation) and flexkv_config_path from extra_configs, forwarding both to FlexKVStorageManager. Clean addition.
corelib/recsys_kvcache_manager/test/test_flexkv.py	Test refactored to accept flexkv_as_batch=1 and optional RECSYS_FLEXKV_CONFIG_PATH env var; SSD tier info printed on startup. Commented-out duplicate put block remains (previously flagged).
corelib/recsys_kvcache_manager/configs/flexkv_ssd.yml	New sample SSD config with placeholder ssd_cache_dir; matches the template shown in README. Clean.
examples/hstu/inference/configs/example.yml	New example FlexKV SSD config; header comment references RECSYS_FLEXKV_ENABLE_LAYERWISE=1 which is not implemented anywhere in Python code.
corelib/recsys_kvcache_manager/README.md	New "FlexKV Backend Options" table documenting flexkv_as_batch and flexkv_config_path; SSD config example added with correct generic placeholder path.
corelib/recsys_kvcache_manager/test/FLEXKV_CPU_BREAKDOWN.md	Updated with new as_batch benchmark section and revised offload_launch timings showing 8–10× reduction in _build_slot_mappings latency.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[offload_kvcache_launch] --> B{as_batch AND\nbatch_size > 1?}
    B -- No --> C[put_async per user\ntask_ids collected]
    B -- Yes --> D[put_match per user\nvalid_len >= page_size]
    D --> E{num_unmatched > 0\nAND end > start?}
    E -- No --> F[client.cancel task_id]
    E -- Yes --> G[Collect task_id +\nsliced slot_mapping]
    G --> H{task_ids non-empty?}
    H -- No --> I[Return handle\nno launch call]
    H -- Yes --> J{len task_ids > 1?}
    J -- Yes --> K[client.launch\nas_batch=True]
    J -- No --> L[client.launch\nas_batch=False]
    C --> M[Return HostKVTaskHandle]
    K --> M
    L --> M
    I --> M

Loading

%%{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"}}}%%
flowchart TD
    A[offload_kvcache_launch] --> B{as_batch AND\nbatch_size > 1?}
    B -- No --> C[put_async per user\ntask_ids collected]
    B -- Yes --> D[put_match per user\nvalid_len >= page_size]
    D --> E{num_unmatched > 0\nAND end > start?}
    E -- No --> F[client.cancel task_id]
    E -- Yes --> G[Collect task_id +\nsliced slot_mapping]
    G --> H{task_ids non-empty?}
    H -- No --> I[Return handle\nno launch call]
    H -- Yes --> J{len task_ids > 1?}
    J -- Yes --> K[client.launch\nas_batch=True]
    J -- No --> L[client.launch\nas_batch=False]
    C --> M[Return HostKVTaskHandle]
    K --> M
    L --> M
    I --> M

Loading

_{Reviews (7): Last reviewed commit: "Clean pr" | Re-trigger Greptile}

[Clebrate]

/build

[JacoCheung]

❌ Pipeline #55763541 -- failed

Job	Status	Log
pre_check	✅ success	view
train_build_x86	✅ success	view
train_build_arm64	✅ success	view
tritonserver_build_x86	✅ success	view
prepare-jet-b200-smoke	✅ success	view
prepare-jet-b200-inference	✅ success	view
prepare-jet-cw-dfw-e2e-benchmark	✅ success	view
build_whl	✅ success	view
dynamicemb_test_fwd_bwd_8gpus	✅ success	view
dynamicemb_test_load_dump_8gpus	✅ success	view
unit_test_1gpu_a100	❌ failed	view
unit_test_1gpu_h100	❌ failed	view
unit_test_4gpu	✅ success	view
unit_test_tp_4gpu	❌ failed	view
L20_unit_test_1gpu	✅ success	view
inference_unit_test_1gpu	✅ success	view
inference_test_1gpu	❌ failed	view

Result: 13/17 jobs passed

View full pipeline

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[Clebrate]

/build

[JacoCheung]

❌ Pipeline #57256931 -- failed

Job	Status	Log
pre_check	❌ failed	view
train_build_x86	✅ success	view
train_build_arm64	✅ success	view
tritonserver_build_x86	✅ success	view
prepare-jet-b200-smoke	✅ success	view
prepare-jet-b200-inference	✅ success	view
prepare-jet-cw-dfw-e2e-benchmark	✅ success	view
build_whl	✅ success	view
dynamicemb_test_fwd_bwd_8gpus	✅ success	view
dynamicemb_test_load_dump_8gpus	✅ success	view
unit_test_1gpu_a100	❌ failed	view
unit_test_1gpu_h100	❌ failed	view
unit_test_4gpu	❌ failed	view
unit_test_tp_4gpu	❌ failed	view
L20_unit_test_1gpu	✅ success	view
inference_unit_test_1gpu	✅ success	view
inference_test_1gpu	❌ failed	view

Result: 11/17 jobs passed

View full pipeline