[BadrBasowid]

Motivation

Large range restores create many TransferItem objects before transfer planning, increasing allocation, copying, and page-fault overhead.

Technical Details

• Add an incremental transfer-plan builder API.
• Build HBM plans as fragments arrive while preserving range merging and validation.
• Keep buffered planning as the default for other transfer engines.
• Preserve per-object rollback when range construction fails.
• Route composite transfers through each engine's builder.

Test Plan

• Run test_transfer_engine.
• Benchmark moonshotai/Kimi-K3 with TP8, FP8 KV cache, concurrency 1, and forced GatherFragmentsKernel (The GatherFragmentsKernel outperforms hipMemcpy by at least 1.5x for all profiled shapes in my setup, see perf(umbp): keep large restore fragments on the gather kernel #717).

Test Result

All 17 transfer-engine unit tests passed.

For a 999,936-token cached prefix on moonshotai/Kimi-K3 running on vllm, using median results from 12 samples:

Page size	API latency	TTFT
64 KiB	405.5 -> 349.7 ms (13.8%)	2645.8 -> 2570.4 ms (2.9%)
32 KiB	609.9 -> 406.0 ms (33.4%)	2835.4 -> 2620.5 ms (7.6%)

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

[TianDi101]

@BadrBasowid Thanks for the work! Could you please explain what does the "API latency" in your table mean? Is that total kvcache load latency or just the latency of building transfer plan?

[BadrBasowid]

@TianDi101 Thanks for the review, API latency refers to the wall time of one complete batch_get_ranges_into_ptr(), i obtained this simply by using perf counters.

[TianDi101]

I benchmarked the plan-building phase on its own: items are built as in BatchGetRanges' local half and planned by PoolClient's composite, with no copies (a standalone microbenchmark). The shape is your Kimi-K3 TP8 FP8 setup with 999,936 tokens; the 1536-token blocks × 24 MLA ranges per block are my assumption. The plans are identical before and after, so this is the PR's whole effect.

Page	build + plan, before → after
64 KiB	69.4 → 11.8 ms (−58 ms)
32 KiB	146.8 → 21.6 ms (−125 ms)

• Where the gain comes from: not building the TransferItem vector (720 B per item, so 172 / 343 MiB and 71k / 151k page faults, now 0). The merge itself only saves 3–7 ms. This matches your −56 ms at 64 KiB and covers about 125–162 ms of your −204 ms at 32 KiB.
• Host-destination restores (LocalCopyEngine) get 4–16% slower. That engine still sits behind BufferedPlanBuilder, which holds every item and copies each one once more. perf(umbp): build LocalCopyEngine transfer plans incrementally #765 gives it the same builder (−85%).
• Mixed DRAM+SSD batches gain the most (−90%); SSD-only reads are unchanged.

I will merge it soon.

[TianDi101]

@BadrBasowid Did you check K3 output correcnetss?

[BadrBasowid]

@BadrBasowid Did you check K3 output correcnetss?

@TianDi101 I ran the following lm_eval task two times:

• moonshotai/Kimi-K3
• GSM8K
• 100 distinct prompts
• 20-shot
• batch size 1

The first run populates DRAM and the second one reads from the cache:

Check	Result
Prompts in each pass	100
Identical prompt sets	100/100
Successful per-request HBM clears	100/100
Pass-2 requests with a DRAM hit	100/100
Cacheable prompt tokens	273,408
Restored prompt tokens	273,408
Cacheable state restored from DRAM	100%
Total prompt tokens	328,751
MORI load objects submitted/completed	3,824 / 3,824
MORI load failures	0
MORI lookup failures/timeouts	0 / 0
MORI load bytes	73.644 GiB

The generated text was byte identical for 93% percent of the prompts, 7% had minor differences