Support FlexKV Layerwise

[Clebrate]

Description

1. Add optional FlexKV layerwise CPU-to-GPU KV onboarding support to the RecSys KV-cache manager.

• updates the FlexKV submodule to main and uses its public multi-group layerwise API;
• adapts the uniform HSTU KV layout to FlexKV's multi-group registration interface;
• enables layerwise operation through flexkv_enable_layerwise, with RECSYS_FLEXKV_ENABLE_LAYERWISE as an environment fallback;
• passes layerwise_transfer and counter_id to FlexKV onboard requests;
• adds RecSys-side eventfd synchronization so each HSTU layer waits for its corresponding KV transfer;
• integrates the layerwise CPU-hit path into inference_benchmark_flexkv.py through --layerwise;
• covers the layerwise onboard path in test_flexkv.py.

2. Result

• configuration: RTX Pro 6000 Blackwell; BS = 1
• Result

seq_len	GPU Hit (ms)	CPU Hit (ms)	CPU Hit + Layerwise CE (ms)
2K	3.3467	4.4259	3.9425
4K	3.8102	6.2841	4.8013
8K	4.9784	10.2798	7.5510
16K	7.3139	18.4760	13.4220

3. Reproduce

• Create /tmp/flexkv_benchmark.yml:

cpu_cache_gb: 8.0
ssd_cache_gb: 0
ssd_cache_dir: /tmp/ssd
enable_gds: false
enable_p2p_cpu: false
enable_p2p_ssd: false
enable_3rd_remote: false

• Run from examples/hstu:

export FLEXKV_ENABLE_MPS=1
export FLEXKV_USE_CE_TRANSFER_H2D=1
# Enable optimized Copy Engine transfer-path selection.
export FLEXKV_CE_PATH_OPT=1
# Enable 2D memory copies for eligible strided KV layouts.
export FLEXKV_ENABLE_CE_MEMCPY2D=1
export RECSYS_FLEXKV_AS_BATCH=1
# CPU Hit + Layerwise CE
python3 inference/benchmark/inference_benchmark_flexkv.py \
  --scenarios cpu_hit \
  --history-len 1024 \
  --append-history-len 64 \
  --batch-size 1 \
  --timed-iters 200 \
  --disable-cudagraph \
  --flexkv-config-path /tmp/flexkv_cpuhit_benchmark.yml \
  --layerwise \
  --only-onboard

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

[greptile-apps]

Greptile Summary

The PR adds optional layerwise FlexKV CPU-to-GPU onboarding and coordinates each HSTU layer with its corresponding transfer event.

• Registers the HSTU cache through FlexKV’s multi-group layerwise API.
• Adds authenticated eventfd handoff, synchronous initialization failure propagation, and per-layer waits.
• Integrates layerwise CPU-hit benchmarking and test coverage while preserving an optional fallback for other host-storage backends.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
corelib/recsys_kvcache_manager/recsys_kvcache_manager/flexkv_layerwise.py	Implements private socket creation, authenticated descriptor transfer, acknowledgment waiting, and synchronous handoff-error propagation.
corelib/recsys_kvcache_manager/recsys_kvcache_manager/flex_kvcache_manager.py	Adds FlexKV layer-group registration, layerwise launch arguments, eventfd-backed task handles, and per-layer wait support.
corelib/recsys_kvcache_manager/recsys_kvcache_manager/host_kvstorage_manager.py	Keeps layerwise waiting optional through a concrete unsupported-backend fallback.
examples/hstu/modules/hstu_block_inference.py	Gates each CUDA-graph HSTU layer on its corresponding layerwise onboarding event.
examples/hstu/modules/inference_dense_module.py	Avoids the whole-task FlexKV wait for layerwise handles and adds benchmark-controlled offload skipping.
examples/hstu/inference/benchmark/inference_benchmark_flexkv.py	Adds a layerwise CPU-hit benchmark path and updates request preparation and latency measurement.
corelib/recsys_kvcache_manager/test/test_flexkv.py	Exercises per-layer waits when FlexKV layerwise onboarding is enabled.

Sequence Diagram

sequenceDiagram
    participant R as RecSys KV manager
    participant S as Eventfd sender
    participant F as FlexKV worker
    participant H as HSTU layers
    R->>S: Create counters and start authenticated handoff
    R->>F: Start client and register layer-group cache
    S->>F: Send per-layer eventfds via SCM_RIGHTS
    F-->>S: Acknowledge descriptor handoff
    S-->>R: Initialization ready
    R->>F: Launch layerwise onboard(counter_id)
    loop Each HSTU layer
        F-->>H: Signal matching transfer event
        H->>H: Wait, then execute attention layer
    end

Loading

_{Reviews (4): Last reviewed commit: "use FlexKV main + add test result" | Re-trigger Greptile}

[Clebrate]

/build

[JacoCheung]

❌ Pipeline #64432092 -- failed

Job	Status	Log
train_build_x86	✅ success	view
train_build_arm64	✅ 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	❌ failed	view
unit_test_1gpu_a100	✅ success	view
unit_test_4gpu	✅ success	view
L20_unit_test_1gpu	✅ success	view
dynamicemb_test_load_dump_8gpus	❌ failed	view
inference_test_1gpu	❌ failed	view
inference_unit_test_1gpu	❌ failed	view
unit_test_tp_4gpu	❌ failed	view
pre_check	❌ failed	view
unit_test_1gpu_h100	✅ success	view

Result: 10/16 jobs passed

View full pipeline