[isytwu]

Motivation

Follow-up to #589, which cut what a remote ranged get moves. This PR targets
the case where nothing is remote: a single standalone process whose inner
backend is distributed, with every key resident on this node. That is the shape
the sglang direct linker runs, and no RDMA is involved anywhere in this change.

A single-node deployment has a fair question to ask of that mode: what does it
cost over plain local mode, which has no master, no peer and no arena? Before
this series, a real amount — measured on DeepSeek-V4-Pro geometry, distributed
mode took 1.8× the TTFL of local mode at one rank. After it, distributed is
ahead of local at every rank count. The numbers are in Test Result.

The reason it was behind is that a BatchGetRanges call was spending most of
its time on per-descriptor bookkeeping rather than on moving bytes. Measured on
DeepSeek-V4-Pro's smallest pool (1,728 B fragments, 32 keys, served entirely
from the local medium):

total=149.7us
  resolve    14.9us  10.0%   per-key medium lookup
  classify   27.3us  18.2%   hipPointerGetAttributes, once PER RANGE
  build      15.1us  10.1%   TransferItem assembly
  xfer       78.7us  52.6%   plan 33.0 (22.0%)  submit 44.4 (29.7%)
  other      13.8us   9.2%   validation, metrics

The copy is under a third; per-descriptor cost is 60%. It is that high
because a layer of a page is small — a model does not have one uniform pool, and
DeepSeek-V4-Pro's three span 37,440 B, 8,448 B and 1,728 B.

The metric being optimised is TTFL: the wait for the first layer group, which is
the one no forward compute overlaps, so it lands directly in TTFT.

Technical Details

Seven commits, each independently measurable.

1. test(umbp): attribute a ranged call's time finely enough to act on

The existing UMBP_RANGED_CALL_DEBUG buckets could not answer the question they
were reached for. resolve covered the whole per-key loop, mixing the medium's
index lookup (per key) with describing the caller's buffers (per range); xfer
covered Plan, Submit and Wait together, so it could not say whether a batch of
many small ranges was paying to plan the move or to make it.

resolve splits into resolve / classify / build, xfer into plan / submit /
wait via a StepTiming the engine fills in, and the report gains items-per-call.
Sub-phases are timed nested and subtracted at the call site so they stay
disjoint. Every timer is inert when the switch is off — a null sink never reads
the clock.

Also adds --local-only to the bench, which seeds the reading node instead of
the peer so every call is served locally. Without it the bench forces every read
remote and the local path's costs are buried under the wire.

2. perf(umbp): let a ranged copy reach the non-temporal path

HostCopyBlock only used the non-temporal AVX2 copy at 256 KiB and above, a
threshold chosen when one key was one ~MiB page (the comment at the top of that
file says exactly that). A ranged copy moves one layer, so every ranged copy
fell under it and took the ordinary memcpy, paying read-for-ownership on a
destination it completely overwrites. Now 2 KiB.

3. perf(umbp): fan a local ranged copy out across threads

One memcpy stream reached the same rate whether 1 rank or 8 were copying
concurrently, so the ceiling was the stream, not the memory system. Local mode
already knew this — DRAMTier spreads a batch read over
UMBP_DRAM_READ_THREADS and its comment reads ">1 breaks the single-core memcpy
ceiling". LocalCopyEngine never did, because it was written when one key was
one whole page and the batch executors above it supplied the parallelism; ranged
I/O invalidated that by handing a single thread thousands of small copies inside
one Submit.

Two things this cost before it paid, both found by measurement and both recorded
in the commit message:

• Spawning threads per call is worse than not parallelizing at all. Creating and
  joining threads maps and unmaps their stacks, which serializes on the process
  address space, so the damage does not scale down with helper count — capping
  the whole process to four concurrent spawned helpers still cost 25% of TTFL on
  a remote restore. Helpers are now a persistent pool.
• The default is 2, not local mode's 4. Past 2 the local copies outrun the
  NIC's DMA into the scratch arena on a mixed workload, and on the real geometry
  4 does not pay even locally.

The asynchronous locality installer marks itself background and stays inline:
nobody is blocked on it, so borrowing cores for it only slows whoever is.

4. perf(umbp): describe a ranged caller buffer by its registration

The largest single win. BuildLocalRangeTransfers described each range with
ClassifiedUserBytes, which asks hipPointerGetAttributes what the pointer is —
a HIP runtime call per range. It is avoidable: in standalone-process mode the
server registers every client region it imports, so location and device are
already in the registration table. ClassifiedUserBytes' own comment says
registered buffers never reach it; this path just never looked.

Using the registration pays a second and larger time. A registered ref names the
region base, so every range landing in one caller buffer shares a (src, dst)
pair and LocalCopyEngine::Plan folds them into one plan whose adjacent
fragments coalesce. Describing each range by its own address made every range
its own plan, with its own four vectors, and made the coalescing that Plan
exists to do unreachable.

The registration table is now sorted by base and searched with upper_bound,
and a whole batch resolves under one shared lock instead of one exclusive lock
per range. Linear scanning was fine when lookups were per key; against a caller
holding one buffer per layer it is quadratic.

Unregistered memory still goes through ClassifiedUserBytes — for a pointer
nobody declared, interrogating it is the only correct answer.

5. perf(umbp): resolve a ranged batch per backend, report its bytes once

Two batch APIs were being called one element at a time.

BatchResolve ran inside the per-key loop as BatchResolve({keys[i]}): a
temporary vector<string> (copying the key) and a temporary result vector per
key, and one acquisition per key of the backend mutex that Allocate, Commit
and Evict also take — the mutex where every rank sharing this client
serializes. Now once per backend for the whole batch, with only unclaimed keys
carried to the next backend, so registry order and first-hit-wins are unchanged.

The byte counters were updated per key. AddCounter takes a lock and builds a
metric key from name and labels, and the help text is a hundred-odd bytes
materialized on every call. Bytes are now accumulated and reported once; totals
are identical, only the number of updates changes. Flushing is a scope guard,
not a line before each return, because this function has four exits that can
all have moved bytes.

6. perf(umbp): stop reallocating while assembling a ranged batch

local_items grew by doubling, and a TransferItem carries two TransferRefs
each holding a MemoryDesc, so every reallocation relocates a lot of bytes. The
range count is known up front and is a tight lower bound on the item count, so
the vector is reserved. MakeReadRanges returned a fresh vector per key for a
buffer dead by the next iteration; it now fills one the loop reuses.

7. test(umbp): run the ranged bench against local mode too

The bench could only drive a distributed PoolClient, so the question this PR
exists to answer had no measurement behind it. --local-mode runs the whole
bench against StandaloneClient instead — no master, no peer, no arena, no
network. Both go through one RangedClient interface so the measurement loop is
literally the same code; a second copy of the loop would make any difference
between the two unfalsifiable.

Not to be confused with --local-only, which keeps the distributed client and
merely puts every key on the reading node.

The local client is configured DRAM-only, because UMBPConfig enables an SSD
tier by default and the distributed side has DRAM only — otherwise the
comparison is between two storage stacks rather than two deployments. Measured
both ways; it makes no difference on these shapes, but the result should not
depend on that.

Where the time goes now

Same call as in Motivation — DeepSeek-V4-Pro's smallest pool, 32 keys:

              before    after
  resolve     14.9us    5.0us
  classify    27.3us    3.0us
  build       15.1us    8.5us
  plan        33.0us    2.7us
  submit      44.4us   20.7us    <- coalescing, not a faster copy
  other       13.8us    4.3us
  total      149.7us   44.4us

The copy is now the largest item, which is where a data path should be spending
its time.

A note on how this was measured

The first three commits were originally tuned against a uniform 73,728 B layer
shape and reported considerably larger wins. That shape is larger than any
pool the model actually has, and this path is fragment-size bound — 5.4 GB/s
at 1,728 B against 37.5 GB/s at 37,440 B. Every number here was re-measured on
the real geometry after #589 landed --pools, and the commit messages were
rewritten to match. The non-temporal threshold tuned on the synthetic shape
(64 KiB) sat above all three real pools and would never have fired in
production.

Knobs

	default
UMBP_DRAM_NT_COPY_MIN_BYTES	2 KiB	non-temporal copy threshold
UMBP_DRAM_COPY_THREADS	2	best everywhere measured, not a compromise; no reason for a deployment to set it
UMBP_RANGED_CALL_DEBUG	off	the phase breakdown above

Test Plan

Image lmsysorg/sglang-rocm:v0.5.16-rocm720-mi30x-20260730, host mlx5 RoCE.
Baseline is a separate build of the merge base (bb259b50) with the same bench
binary. This machine's run-to-run spread reaches ±20%, so every reported figure
is a median of 3.

cmake -S . -B build -GNinja -DBUILD_UMBP=ON -DBUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Release
cmake --build build --target bench_umbp_pool_client_ranges -j

# this container's RC queues are smaller than mori-io's defaults
export MORI_IO_QP_MAX_SEND_WR=256 MORI_IO_QP_MAX_RECV_WR=256 MORI_IO_QP_MAX_CQE=1024
POOLS='c4:30:37440:2/2;c4_indexer:30:8448:2/2;c128:31:1728:0,1,3,5,7,9,11,13,15,17,19,21,23,25,27,29,31,33,35,37,39,41,43,45,47,49,51,53,55,57,59'
B=./build/tests/cpp/umbp/distributed/bench_umbp_pool_client_ranges

# distributed mode, every key on the reading node
$B --pools "$POOLS" --ranks 1,2,4,8 --pages 32 --requests 2 --offload-ranks 0 --local-only
# local mode, same loop and shapes -- the deployment comparison
$B --pools "$POOLS" --ranks 1,2,4,8 --pages 32 --requests 2 --offload-ranks 0 --local-mode
# must not regress: the remote path
$B --pools "$POOLS" --ranks 1,2,4,8 --pages 32 --requests 2 --offload-ranks 0
# byte check: recomputes every restored page against a position-dependent seed
$B --pools "$POOLS" --ranks 8 --pages 32 --requests 2 --offload-ranks 0 --local-only --verify

Correctness was checked on six shapes (real geometry distributed-local,
distributed-remote and local-mode, uniform local, --layer-group 1, and the
8,448 B small-object pool), and the ranged suite was also run with the
non-temporal path disabled and with the fan-out off, to confirm neither knob is
load-bearing for correctness.

Test Result

Single-machine, DeepSeek-V4-Pro geometry, 32 pages/request, layer groups of 8.
restore includes a fixed 9.15 ms of simulated compute (61 layers × 150 us).

ranks	restore ms		TTFL us
1	11.51 → 9.70	1.19×	1327 → 509	2.61×
2	12.24 → 9.74	1.26×	1395 → 562	2.48×
4	12.70 → 9.78	1.30×	1465 → 597	2.45×
8	13.54 → 9.98	1.36×	1771 → 769	2.30×

Subtracting the fixed compute, the load itself goes 2.36 → 0.55 ms at 1 rank
and 4.39 → 0.83 ms at 8 — 4.3× and 5.3×.

Against local mode

Same loop, same shapes, only the deployment differs (restore ms / TTFL us):

ranks	local mode	distributed, before	distributed, after
1	9.93 / 722	11.51 / 1327	9.70 / 509
2	10.10 / 908	12.24 / 1395	9.74 / 562
4	10.74 / 1220	12.70 / 1465	9.78 / 597
8	17.57 / 1711	13.54 / 1771	9.98 / 769

Distributed mode used to cost 1.8× the TTFL of local mode at one rank. It is now
ahead of local at every rank count, by 1.4× at one rank and 2.2× at eight.

Local mode did not move — it is untouched by this PR — and its own profile
(UMBP_DRAM_PROFILE=1) says why it is now behind. At 8 ranks its DRAMTier
spends 92% of a batch in the copy, at 3.0 GB/s per stream, plus 7% classifying
pointers, with zero lock wait. Those are the same three costs this series
removed from the distributed path — a non-temporal threshold set for whole-page
copies, a HIP pointer query per range, and threads spawned per call. Local mode
still has all three; carrying the fixes across is separate work, and commit 7 is
what makes it measurable.

Worth noting for anyone reading commit 3: local mode is where the idea of
fanning the copy out came from, but its implementation spawns threads per call,
which measured worse than not parallelizing. The distributed path uses a
persistent pool instead, which is why it now beats the thing it borrowed from.

Remote path (not the target; must not regress):

ranks	restore ms	TTFL us
1	16.56 → 12.07	2539 → 2456
2	30.60 → 26.34	3384 → 3289
4	49.27 → 47.37	6569 → 5311
8	85.10 → 86.01	13010 → 12988

Improves or holds. The 8-rank restore reads 1% worse, which is inside this
machine's spread — repeated runs of that cell landed between 80.2 and 86.0 ms
on both sides.

Suites — all pass:

test_umbp_pool_client_ranges	24
test_umbp_range_map	13
test_umbp_ssd_ranged_io	6
test_umbp_cross_node_smoke	19
test_umbp_pool_client_batch_put	9
test_umbp_medium_selection	9
test_umbp_client_metrics	9
test_component_metrics	19
test_standalone_shm_ipc	9
test_umbp_dram_tier_ranges	13

test_umbp_pool_client_ranges additionally passes with
UMBP_DRAM_NT_COPY_MIN_BYTES set past every fragment size and with
UMBP_DRAM_COPY_THREADS=1.

Byte verification: 0 mismatched pages on all six shapes, local mode included.

Note that test_standalone_shm_ipc is not matched by CI's
ctest -R '^(umbp_|test_dummy_ssd_tier|test_prefix_aware_eviction|test_ssd_tier)'
filter, since its registered name does not start with umbp_. It exercises the
registration path commit 4 changes, so it was run by hand; renaming it is
outside this PR.

Not in this PR

Put. BatchPutRanges is untouched and does not benefit: measured, a local
ranged put spends 81% in per-key BatchAllocate and 18% in the unconditional
BatchRoutePut RPC, and 0.2% in the code this PR changes. Fixing it means
batching the backend calls and letting a local allocation answer before the
master is asked — the latter is safe (Router::BatchRoutePut is read-only; the
global index is populated by the commit-time KvEvent::ADD) but gives up
cross-node dedup and placement balancing, so it wants its own PR and a config
gate.

That work needs the bench fixed first: with the default offload load the node's
medium fills, every BatchAllocate returns no-space, and the puts fail
silently — the bench only accumulates successful bytes, so it reports
put_gibps=0 rather than an error. Anyone reading the put half of this bench's
output today is reading nothing.

Local mode's copy path. The three fixes in commits 2–4 apply just as well to
DRAMTier, which still has the 256 KiB non-temporal threshold, a
DetectPointerLocation per copy job, and a spawn-per-call fan-out. That is a
different file with different callers and belongs in its own change.

TransferRef weight. The 8.5 us of build that remains is mostly copying
MemoryDesc (three std::strings plus two 64 B arrays) twice per item. Holding
it by shared_ptr would roughly halve that, but weighted by call count build
is about 12% of the total, so the payoff is ~6% — against changing a struct the
RDMA path shares, whose only coverage is a same-host two-client smoke test. Not
worth it here.

Submission Checklist

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

🤖 Generated with Claude Code

[maning00]

Independently verified this series from the sglang side (we are the consumer of the direct linker this targets). It does what it says — but it does not build the Python bindings, which is how sglang consumes mori, so that needs one fix before we can adopt it.

Blocker: BUILD_UMBP=ON pip wheel fails on this branch

FAILED: src/pybind/CMakeFiles/mori_pybinds.dir/pybind_umbp.cpp.o
pybind11/pybind11.h:535:13: error: static assertion failed due to requirement
  'expected_num_args<...>(sizeof...(Args), ...)':
  The number of argument annotations does not match the number of function arguments

Commit 4 gives IUMBPClient::RegisterMemory a 5th parameter (MemoryRegistration mode, umbp_client.h:162), but git diff --stat <base>..HEAD -- src/pybind/ is empty: the .def still declares four py::args, and MemoryRegistration is not bound as a Python enum at all.

It slipped through because the Test Plan builds only bench_umbp_pool_client_ranges. cmake --build build --target mori_pybinds (or a plain pip wheel) reproduces it immediately.

Two additions fix it:

+  py::enum_<MemoryRegistration>(m, "UMBPMemoryRegistration")
+      .value("Pinned", MemoryRegistration::kPinned)
+      .value("LocalCopyOnly", MemoryRegistration::kLocalCopyOnly)
+      .export_values();
+
   py::class_<UMBPHbmConfig>(m, "UMBPHbmConfig")

       .def("register_memory", &IUMBPClient::RegisterMemory, py::arg("ptr"), py::arg("size"),
            py::arg("loc") = mori::io::MemoryLocationType::CPU, py::arg("device") = -1,
+           py::arg("mode") = MemoryRegistration::kPinned,
            py::call_guard<py::gil_scoped_release>())

With those, the wheel builds clean and every existing 2- and 4-argument register_memory call still works. All numbers below were taken on that build.

Verification, on a different instrument and geometry

One MI355X node (n09-29), image lmsysorg/sglang-rocm:v0.5.17-rocm720-mi35x-20260820, baseline bb259b50 vs this branch (d3be4e4c), both through StandaloneProcessClient (so a gRPC hop your C++ bench does not have).

Model-free probe — replays exactly what the sglang linker puts on the wire, sweeping fragment count at fixed fragment size (8192 B) so the marginal per-fragment cost is separated by construction rather than fitted:

	bb259b50	this branch
local, µs/fragment	0.466	0.459
distributed, µs/fragment	1.227	0.480
distributed − local	+0.761	+0.021
distributed / local	2.63×	1.05×

97% of the distributed-specific per-fragment penalty is gone. Two things corroborate the mechanism you describe rather than something else moving: the local arm does not budge (0.466 → 0.459), matching "local mode did not move"; and the marginal bandwidth is unchanged on both arms (local 9.76 → 9.96, distributed 8.85 → 9.15 GB/s), so the win is bookkeeping removal, not a faster copy.

Real connector, topology 3 — DeepSeek-V4-Pro real weights, TP8, 64 GiB hugepage tier, paired local/distributed arms back-to-back in one session:

tokens	delta (dist − local), bb259b50	delta, this branch
16,384	+18.3 ms	+10.3 ms (−44%)
65,536	+45.7 ms	+13.6 ms (−70%)

device (arm-independent) drifts −6.1% / −8.3% between arms; charging all of that against distributed as a conservative bound puts the two deployments at parity (0.9× at both lengths, from 7.5× and 3.8× before). Correctness also passes: 8/8 linkers report ranged_io=yes, cache_test_len 3/3, and a per-token restore_equivalence probe returns umbp == device at both 16,384 and 32,768.

Your finer buckets confirm the shape you were aiming for — classify 0.6%, plan 1.0%, submit 85.5%. The copy is the dominant item now.

One data point for the follow-up put PR

You note a local ranged put spends 18% in BatchRoutePut, alongside the caveat that the bench's put half is unreadable once the medium fills. On the real connector with this branch applied we measure:

SUMMARY PUT calls=1005 bytes=41.28GiB items_per_call=1793
  | resolve=28.2% classify=0.1% build=1.8% route=50.6% xfer=19.0%(plan=0.4% submit=18.3%)

route is 50.6%, and we independently measured 60.4% on an earlier run with a different tier. So the unconditional BatchRoutePut looks considerably more expensive on real traffic than the bench suggests — worth sizing against the per-key BatchAllocate work before choosing an order for that PR.

Happy to re-run any of the above, and to test a revision on the sglang path.