[TE] perf(efa): allow bounding the batch MR registration fan-out

[whn09]

Description

registerLocalMemoryBatch() and unregisterLocalMemoryBatch() in the EFA transport spawn one std::async(std::launch::async) per buffer with no cap, which libstdc++ takes literally — one fresh thread each. Kimi-K3 registers ~180 KV buffers per TP rank in a single batch, and since SGLang creates one TransferEngine per TP rank, an 8-rank node peaks at ~1100 threads all sitting inside fi_mr_reg / fi_mr_regattr.

This replaces both fan-outs with a fixed pool pulling from a shared atomic index, sized by a new MC_MAX_CONCURRENT_REG_MR option.

Default behavior is unchanged. The default of 0 means unbounded: the pool spawns count-1 threads and runs the caller as a worker, matching what std::async did. Error semantics are also unchanged — every item is still attempted and the first non-zero return is propagated.

What the knob does, measured

Measured on p5.48xlarge (32 EFA NICs), replaying K3's KV registration the way SGLang actually issues it: 8 processes, one engine and one GPU each, 182 GPU buffers of 2.5 KB–391 MB per process, barrier-synchronized so all 8 libfabric domains hit the NICs at once. Times are the slowest rank, since nothing serves until every rank has registered.

per-process cap	in flight per node	size-descending	grouped by pool (what SGLang passes)	size-ascending
unset (unbounded)	~1100	108 / 99 / 128 s	112 s	120 s
64	512	184 / 208 s	101 s	60 s
16	128	95 / 98 / 97 s	43 s	36 / 37 s
8	64	107 s	—	44 s
4	32	157 s	—	—

Repeated entries are separate sweeps (~3% run-to-run; unbounded's own spread is ~30%).

• A cap of 16 cuts registration 2.6× on the order SGLang really uses (112 s → 43 s), but only 1.1× if buffers happen to arrive largest-first.
• The input order matters as much as the cap: it moves the capped result by 2.7× (cap 16) to 3.5× (cap 64). Unbounded is insensitive to order, because every buffer gets its own thread and nothing queues — the ordering effect exists only once tasks queue behind a cap.
• The optimum is cap≈16 for both orders, so cap and order look like largely independent effects.

The optimum also tracks the platform:

• p5 host memory (fi_mr_reg, 128 × 2 GiB of 4 KB pages): 274 s unbounded, 120 s @64, 89.6 s @32, 59.1 s @16, 70.9 s @8, 257 s @4. (Caps were swept out of order — 8, 128, 32, 4, 64, 16 — so the curve is not a warming/fragmentation artifact.)
• 2× p6-b300, 16 NICs/rank, full K3 server run: 138.7 s unbounded, 130.1 s @128, 51.0 s @32, 20.7 s @8, 24.8 s @4.

The optimum is cap x nranks ~= cores

The value is not a platform mystery to be swept for: registration threads do CPU-bound page-pinning, so what matters is the global thread count against the core count. Since the cap is per process, that is cap x nranks.

Tested by moving both sides of that relation on the same 192-core p5 node (taskset puts all ranks in one shared cpuset, so they still contend with each other, just with a smaller budget). Optimum in bold:

ranks	cores	cap=4	cap=8	cap=16	cap=32	cap=64
8	192	—	106.7 s	99.8 s	127.8 s	182.6 s
4	192	—	109.1 s	100.4 s	89.2 s	120.6 s
8	64	155.1 s	121.3 s	164.7 s	263.8 s	—

The optimal cap differs in all three rows — 16, 32, 8 — so it is not a per-process constant. Converted to global threads the three optima collapse:

ranks	cores	optimal cap	global threads
8	192	16	128
4	192	32	128
8	64	8	64

• Halving the process count doubled the optimal cap, holding global threads at 128.
• Cutting cores 192 → 64 moved the optimum to 64 threads, i.e. it tracked the cores.

Oversubscription is punished superlinearly: 2.7× oversubscribed (row 1, cap=64) is 1.8× slower than the optimum, and 4× oversubscribed (row 3, cap=32) is 2.2× slower. So a practical rule is cap ~= cores / nranks, giving roughly 0.7–1.0 × cores in flight globally. On a tp8 192-core node that is 16–24, which is what the sweep above found.

This is consistent with the p5 host-memory optimum (cap=16 → also 128 threads). The B300 optimum is lower in global terms (cap=8 → 64 threads); that run was a full K3 server, which leaves far less idle CPU than a standalone replay, but this has not been verified and is only a plausible reading.

Why the knob is still opt-in rather than a compiled-in default

cores / nranks is only computable by something that knows how many peer processes share the node, and a single TransferEngine does not — SGLang creates one per TP rank with no cross-rank coordination. A constant would be wrong by the rank count, and a cores-derived value would oversubscribe by exactly that factor (picking 192 on a tp8 node yields 1536 threads, i.e. the 8× oversubscription this data shows is the worst case). Combined with the ordering penalty below — a cap makes a largest-first batch only 1.1× faster, so an operator whose caller passes buffers in a bad order gets little for it until the companion sort (#3216) lands — the historical unbounded behavior stays the default, and an operator who knows their rank count opts in. What has changed is that they no longer need to sweep: the value is derivable.

The ordering effect: mechanism found

The earlier revision of this PR listed "largest-first is the worst order" as an open mystery. It is now explained, and it is not a scheduling effect at all — which is why longest-processing-time intuition mispredicted it.

Registering device memory on EFA costs roughly k x (bytes of device memory already registered on that domain), k ~= 260 ms/GiB on p5.48xlarge, flat across 0–19 GiB, and nearly independent of the buffer's own size. Measured single-threaded and strictly serial (one buffer per batch_register_memory() call), so no lock and no queue is involved:

	prior bytes	avg cost	n
same 2.79 MB buffer, registered first	0.00 GiB	70 ms	24
same 2.79 MB buffer, registered last	9.17 GiB	2505 ms	24
same 391 MB buffer, first	0.00 GiB	1382 ms	24
same 391 MB buffer, after the smalls	0.06 GiB	1408 ms	24

Control: identical 48 buffers, identical 9.23 GiB total, only the order differs — ascending 35.5 s vs descending 93.3 s, both ending with the same 2421/2422 ms final registration. So cost is a function of what is already in, not of the buffer's own size, and ascending order is a well-defined optimum rather than a heuristic.

This also explains the cap=24 (23.0 s) vs cap=47 (5.5 s) gap noted in the earlier revision, which queue depth alone could not (4.2x where only 2x was available): a worker's second registration happens when the domain is much fuller.

Host memory does not accumulate. The same sweep on mmap'd host buffers is flat (a 391 MB buffer costs ~600 ms whether it is 1st or 48th) and order makes no difference (a control varying only registration order gives 15.3/14.2 s ascending vs 13.2/15.7 s descending — the direction flips between reps). So this is a device-memory property, most likely the CUDA/HMEM path rebuilding per-domain state, and worth a libfabric efa issue separately.

The follow-up this enables — sorting ascending inside registerLocalMemoryBatch() — is a separate PR, #3216, because it changes behavior for every caller rather than only for those who opt in, and because whether the same sort should apply to the other 15 transports is worth a maintainer discussion. Note the two changes are coupled in one direction: the sort only has an effect when a cap is set. Unbounded, every buffer gets its own thread, nothing queues, and the dispatch order never reaches the provider (112 s vs 120 s, i.e. noise).

Known gaps

1. A much larger lever is not addressed by this PR. Mooncake registers every single-chunk buffer on all NICs (nic_assignments[0] gets every context), while NIXL's libfabric plugin registers a VRAM_SEG buffer only on that GPU's topology-local rails (selectRailsForMemory() → getEfaDevicesForPci()). On p5.48xlarge that is 32 NICs vs 4. Restricting Mooncake's device list to emulate it, on 48 x 391 MB GPU buffers serially: 123.7 s at 32 NICs vs 17.6 s at 4 vs 4.4 s at 1 — 7.0x, larger than either the cap or the sort, and it multiplies the accumulation above (the per-domain cost is paid once per domain). That is a throughput trade-off rather than a free win, since it changes which NICs can serve a transfer, so it needs its own PR with bandwidth numbers and not just registration times. Filed as [TE][EFA] GPU buffers are registered on all NICs, costing 7x vs per-GPU rail selection #3217; flagging it here so reviewers know this PR is not claiming to close the gap to NIXL.
2. The B300 numbers predate this harness (16 NICs instead of 32, and being a full server run, a different buffer order), so the two platforms are not yet reconciled. No p6-b300 capacity has been available to re-run them.

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• P2P Store (mooncake-p2p-store)
• Python Wheel (mooncake-wheel)
• Common (mooncake-common)
• Mooncake RL (mooncake-rl)
• CI/CD
• Docs
• Other

Type of Change

• Bug fix
• New feature
• Refactor
• Breaking change
• Documentation update
• Performance improvement
• Other

How Has This Been Tested?

Test commands:

# Unit tests for the new config parsing (8 cases)
./build/mooncake-transfer-engine/tests/config_test --gtest_filter='MaxConcurrentRegMrEnvTest.*'

# Formatting
./scripts/code_format.sh

# Benchmark: 8-process replay of K3's registration shape on p5.48xlarge.
# Sweeps the cap x input-order matrix reported above.
REGBENCH_ORDER=asc CAPS="unset 64 16 8" NRANKS=8 ./capsweep_k3_mp.sh

# Core-budget experiment: vary rank count and core budget to test whether the
# optimum obeys cap x nranks ~= cores. All ranks share one cpuset so they still
# contend with each other.
# (the runner takes an extra arg that taskset-pins all ranks into one cpuset)
hostrun_k3_cores.sh <cap> 4 1.0 <logprefix> 0     # 4 ranks, all 192 cores
hostrun_k3_cores.sh <cap> 8 1.0 <logprefix> 64    # 8 ranks, 64-core budget

Test results:

• Unit tests pass
• Integration tests pass (if applicable)
• Manual testing done (describe below)

config_test covers unset → unbounded, explicit 0 → unbounded, valid values, and rejection of malformed input (non-numeric, negative, trailing garbage, empty) with the default retained.

Manual testing is the measurement matrix above: 8-process GPU registration on p5.48xlarge (32 NICs) across 5 cap values × 3 input orders, plus host-memory and B300 sweeps, plus the 12-cell core-budget experiment (3 rank/core configurations × 4 caps), plus the serial single-buffer sweeps that establish the accumulation mechanism (4 size sequences × GPU and host, 2 order-control reps, and a 32/4/1-NIC fan-out comparison). Unbounded runs reproduce the ~127 s registration wall time seen in a real K3 prefill log, confirming the harness matches production. Verified the default path (MC_MAX_CONCURRENT_REG_MR unset) is unchanged versus the previous std::async implementation.

One caveat on methodology, since it produced a wrong claim in an earlier revision of this PR: this knob cannot be measured single-process. 48 registration threads on a 192-core node never contend, so more threads is unconditionally faster there and the optimum vanishes. An earlier single-process sweep on that basis concluded the bottleneck was a provider lock and that the value should not be core-scaled; both were wrong, and both have been corrected above. Single-process runs are used here only for the per-registration cost curve, which is what they are valid for.

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh
• I have run pre-commit run --all-files and all hooks pass
• I have updated the documentation (if applicable)
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed an RFC issue

pre-commit is not installed in my environment; I ran clang-format (via the repo's 80-column config) on every touched C++ file instead. The change is ~190 LOC, so no RFC.

Checked for overlap: #3174 (idempotent duplicate registration) and #3207 (EFA auto-discovery) touch neither efa_transport.cpp nor the config plumbing, so there is no conflict.

AI Assistance Disclosure

• No AI tools were used
• AI tools were used (specify below)

Claude Code was used to write the benchmark harness, run the measurement sweeps, and draft the code and documentation. Notably, it also caught that an earlier iteration of this work had the wrong process topology — a single-process replay that produced the opposite conclusion (239 s unbounded vs 1051 s capped) and would have justified the wrong default. Every line has been reviewed by me and I can defend the change end-to-end.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!