Repository navigation
Conversation
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
LocalCopyEngine plans a ranged get into host memory the way HbmCopyEngine plans one into a GPU: one TransferItem per page fragment, grouped by (src, dst) buffer and coalesced. Behind the default BufferedPlanBuilder it still held every fragment of the batch, and the composite builder copied each one into it once more, so host-destination restores got slower with the incremental builder instead of faster. Give it the builder HbmCopyEngine has: Plan()'s loop moves into Add(), and Plan() is the builder over a vector, so there is one copy of the merge logic. The pair map grows with distinct endpoint pairs instead of reserving two buckets per item. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…cremental bench_umbp_transfer_plan models the local half of PoolClient::BatchGetRanges restoring a long cached prefix and stops once the plan set exists: items are built the way BuildLocalRangeTransfers and BuildLocalFileRangeTransfers build them, and the engine is the composite PoolClient assembles. Defaults are moonshotai/Kimi-K3 at TP8 with an FP8 KV cache: 24 MLA layers x 576 B per token, 1536-token blocks, a 999,936-token prefix. --dst=gpu|cpu picks HbmCopyEngine or LocalCopyEngine, and --file-pct=N serves N% of keys from SSD-tier files for GdsEngine (a verbatim stand-in where the hipfile headers are absent); below 100 it is a mixed batch. The same source builds with and without TransferEngine::CreatePlanBuilder, so a binary per tree compares the buffered Plan() with the incremental builder. Each iteration prints a CSV row with an order-insensitive plan-set checksum. Built but not registered with ctest, like the other bench_* targets. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 task done
Keep this PR to the LocalCopyEngine builder and its tests. The bench that produced the numbers in the description stays out of tree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
TianDi101
marked this pull request as ready for review
October 15, 2026 01:25
This was referenced Oct 15, 2026
Collaborator
Author
|
migrated into PR: #774 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-up to #758, which gives
HbmCopyEnginean incrementalTransferPlanBuilder.LocalCopyEngineplans the same per-page-fragment items when a ranged get lands in host memory, but it stays on the defaultBufferedPlanBuilder. That builder still holds every fragment of the batch, and the composite builder copies each one into it once more, so with #758 alone a host-destination restore gets 4–12% slower instead of faster. This givesLocalCopyEnginethe same builder; the same restore is then 74–87% faster than before #758.Stacked on #758. The first four commits are #758's. The change is
8f6599743; the two commits after it add and then remove a benchmark, so they cancel out. I'll rebase ontomainonce #758 lands, and it stays a draft until then.Change
LocalCopyEngine::Plan()'s loop moves intoLocalCopyPlanBuilder::Add(), as #758 did forHbmCopyEngine.Plan()is now the builder over a vector, so there is one copy of the merge logic. Grouping by (src, dst) buffer and coalescing are unchanged; the pair map grows with distinct endpoint pairs instead of reserving two buckets per item. +34 / −11 source lines ignoring whitespace, plus 68 lines of tests.Results
Measured with a standalone microbenchmark (not in this PR). It builds items the way
BuildLocalRangeTransfersdoes, runs them through the compositePoolClientassembles, and stops once the plan set exists; addresses are synthetic and nothing is copied. The shape is #758's: moonshotai/Kimi-K3, TP8, FP8 KV cache, 999,936 cached tokens. That is 24 MLA layers × 576 B per token in 1536-token blocks (651 keys × 24 per-layer ranges), one DRAM buffer, one registered region per layer. Each number is build items + plan, the median of 36 samples on one pinned core, with the three trees interleaved. The plans are identical in every tree.Host destination (
LocalCopyEngine):MemoryDescsTransferItems is never built. The merge itself also drops from 5.7 to 2.8 ms at 64 KiB.Test plan
test_transfer_engine: 20/20, the 17 from perf(umbp): build transfer plans incrementally #758 plusLocalCopyEngine.BuilderMergesAdjacentRanges,BuilderRejectsOutOfBoundsRangesandBuilderMatchesPlan. Release build inrocm/mori-dev.tests/cpp/umbpsuite: not run locally; left to CI's UMBP unit-test job.🤖 Generated with Claude Code