forked from mooncake-track/Mooncake
* feat(efa): auto-split large MR registrations exceeding max_mr_size
Buffers larger than the EFA device's max_mr_size are now transparently
split into chunks, each registered as a separate MR. This fixes the
silent truncation bug where only the first max_mr_size bytes were
registered, causing transfers to unregistered regions to fail at runtime.
Key changes:
- Query EFA device max_mr_size via ibverbs during init (libfabric does
not expose this) and clamp globalConfig accordingly
- Auto-split buffers > (max_mr_size - 1GB) into chunks in
registerLocalMemoryInternal, each with its own BufferDesc metadata
- Track chunk mappings for proper cleanup in unregisterLocalMemory
- Change lkey()/rkey() from exact-match to range lookup (matching
mrDesc() pattern) so key lookups work for any address within a chunk
- Replace silent truncation in EfaContext with a hard error as safety net
- Remove preTouchMemory truncation to touch the full buffer
Tested on P5EN (16 EFA NIC, max_mr_size=192GB): 191GB single-chunk
registration succeeds. Auto-split triggers correctly for larger buffers
(200GB splits into 191GB + 9GB). Total registerable size per buffer is
bounded by system pinned_vm limits (~191GB on 16-NIC P5EN).
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat(efa): per-NIC partition for large buffer MR registration
Each EFA NIC can only register up to max_mr_size total. The previous
auto-split approach registered every chunk on ALL NICs, hitting the
per-NIC limit for buffers > max_mr_size. This change assigns each
chunk to a disjoint subset of NICs, enabling registration of buffers
up to max_mr_size × num_NICs (e.g. ~1.5TB on P5EN with 16 NICs).
Key changes:
- chunk_limit = max_mr_size / 2 (was max_mr_size - 1GB) for headroom
- NIC assignment: chunks distributed evenly across available NICs
- selectDevice(): skips NICs with rkey=0 (unassigned for that chunk)
- Striping path: filters by lkey!=0 to avoid unregistered NICs
- Unregister: only deregisters from assigned NICs per chunk
- ChunkRegistration struct tracks per-chunk NIC assignments
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat(efa): PTE-aware auto-split replaces conservative max_mr_size/2 threshold
The previous chunk_limit of max_mr_size/2 (~96GB on P5EN) caused unnecessary
buffer splitting even when hugepages were available. This detects the actual
backing page size via /proc/self/smaps and computes the PTE-based limit:
- 4KB pages: 22M PTEs × 4KB = 88GB (genuine hardware constraint)
- 2MB hugepages: 22M PTEs × 2MB = 44TB (effectively max_mr_size)
Verified on P5EN (H200, 16 EFA): 100GB pool with hugepages no longer splits,
restoring full 16-NIC throughput (108 GB/s vs 46 GB/s with the old threshold).
MR registration time dropped from 302s to 1.6s.
Adds MC_EFA_MAX_PTE_ENTRIES env var for override (default 22M).
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat(efa): add KV cache prefix transfer benchmark script
Python benchmark to measure EFA transfer performance for LLM prefix
cache hit scenarios. Tests different pool sizes (10GB-500GB) and prefix
lengths (4K-32K tokens) to evaluate per-NIC partition impact on
transfer latency and throughput.
Default KV bytes/token matches GLM-5.1 (754B MoE, MLA attention):
(kv_lora_rank=512 + qk_rope_head_dim=64) * 2 * 78 layers = 88KB/token
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix(bench): add connection warmup before benchmark loop
The first prefix size's measurements were skewed by EFA connection
establishment (openSegment, endpoint creation). Add 3 small transfers
before entering the benchmark loop to warm up the connection.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix(bench): per-offset warmup and p50 throughput reporting
Warmup now exercises all offsets the benchmark will measure, eliminating
first-access TLB/page-fault outliers (4K token p99 dropped from 80ms to
3.5ms). Throughput is reported from p50 latency instead of avg for more
stable numbers.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat(bench): add --threads option for concurrent prefix transfer
Adds optional multi-threaded transfer support. Each thread transfers a
chunk of the prefix in parallel via separate transfer_sync_read calls.
Default is 1 (single transfer, same as before). Testing shows threads=2
matches single-thread throughput (~108 GB/s), while higher values add
scheduling overhead with no benefit since EFA already stripes internally.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat(efa): full NIC coverage for multi-chunk MR registration
When a buffer exceeds max_mr_size and must be split into multiple chunks,
register every chunk on ALL NICs instead of disjoint per-NIC partition,
as long as total PTE usage per NIC fits within the PTE budget.
With hugepages (2MB), 500GB buffer uses only 250K PTE/NIC (budget: 24M),
so all 16 NICs cover every address. Falls back to disjoint partition when
PTE budget is exceeded (e.g. 4KB pages with large buffers).
500GB pool throughput: 35 GB/s (disjoint, 5-6 NIC) → 108 GB/s (full, 16 NIC).
Registration time unchanged (~9.5s) due to parallel MR registration.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat(bench): add batch memory registration benchmark
Test script for registering multiple independent memory blocks
(e.g. multi-tenant KV cache pools). Supports both per-block
register_memory and batch_register_memory APIs. Target mode
allocates and registers N blocks; initiator mode transfers
data and measures throughput.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* style: apply clang-format to EFA transport files
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat(efa): expose discoverTopology C API and Rust bindings
Add discoverTopology() to the C API and discover_topology()/install_transport()
to the Rust bindings, enabling EFA transport initialization from C/Rust callers.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* test(efa): add C API test for discoverTopology and EFA transport
Verify that the new discoverTopology() C API correctly populates the
device list, enabling installTransport("efa") and memory registration
via the pure C interface (used by Rust/Go bindings).
All 4 tests passed on p5en.48xlarge (16 EFA NICs).
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix(efa): O(log n) MR lookup and chunk registration rollback
- Replace unordered_map with std::map for mr_map_ and use upper_bound
for O(log n) range lookups in rkey/lkey/mrDesc instead of O(n) scan
- Add rollbackChunks lambda to unregister already-registered chunks on
failure in registerLocalMemoryInternal, preventing MR leaks
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* style: apply clang-format to efa_c_api_test.cpp
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix: rename thr_tag to thread_tag to pass typos spell check
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix(efa): round-robin multi-chunk MR and large MR registration test
When a buffer splits into more chunks than available NICs, round-robin
assign chunks across NICs with per-NIC PTE budget validation, instead
of hard-failing with "Buffer requires N chunks but only M NICs".
Add efa_single_nic_large_mr_test: tests single-NIC and all-NIC large
MR registration with hugepages. Supports --chunk_gb for multi-buffer
mode (e.g. 200×2GB).
Verified on P5EN (16 EFA NIC, 2MB hugepages):
- 1 NIC × 200×2GB (400GB): 53.7s
- 16 NICs × 200×2GB (400GB): 64.8s
- 1 NIC × 200GB single buffer: 2.7s
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* feat(efa): add cross-node transfer test for multi-buffer MR registration
Tests 200x2GB buffer registration on all NICs with actual data transfer
between two P5EN nodes. Supports target/initiator modes with single-read
and multi-block batch benchmarks.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix(efa): avoid CUDA context leak when built without GPU support
libfabric 2.4's EFA provider dlopens libcudart/libcuda during
fi_getinfo/fi_domain to detect HMEM support, even though the caller
never touches GPU memory. This creates a CUDA primary context on
GPU 0 and permanently holds ~616 MiB of device memory.
When Mooncake is built with USE_CUDA=OFF (and USE_HIP=OFF), set
FI_HMEM=system before fi_getinfo and drop FI_MR_HMEM from the domain
hints so the provider skips GPU hmem initialization entirely.
* feat(efa): multi-thread initiator and wildcard location for 16-NIC coverage
Register target buffers with wildcard location "*" instead of "cpu:0"
so initiator-side remote NIC selection distributes evenly across all 16
NICs (both NUMA nodes). With "cpu:0", selectDevice only picked NUMA-0's
8 NICs, leaving NUMA-1 idle — throughput capped at ~107 GB/s instead of
~149 GB/s on P5EN (16×EFA 200Gbps).
Also convert the initiator from single-threaded to multi-threaded
(--threads flag) and print the actual P2P handshake address.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* docs(efa): add p6-b300 bandwidth results (752 GB/s GPU, 230 GB/s CPU)
- GPU-to-GPU peak 752 GB/s write / 713 GB/s read at ~94% line rate
(16×400 Gbps = 800 GB/s theoretical)
- CPU-to-CPU peak 230 GB/s, DRAM-limited on Xeon 8559C
(NUMA-0 NICs 90 Gbps vs NUMA-1 NICs 53 Gbps)
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* feat(rust): customer_pattern sample for per-NUMA multi-MR registration
Adds a Rust binary mirroring the customer's EFA usage pattern: target
registers many MRs per NUMA (e.g. 16 NICs x N x 10GB via cpu:<numa>
locations); initiator reads/writes a specific (numa, buffer_index)
tuple, using --source-numa to pick which local NIC set is exercised.
Also makes rust/build.rs robust against non-default transfer_engine
build configurations: etcd-cpp-api is opt-in via MOONCAKE_WITH_ETCD=1,
CUDA linking is opt-in via MOONCAKE_WITH_CUDA=1 with CUDART_LIB_DIR
override, and libbase.a + libasio.so + libfabric paths are picked up
whether the repo uses the standalone or top-level CMake layout.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* chore(rust): drop customer_pattern sample from upstream
Revert Cargo.toml [[bin]] additions and remove the bench demo source so
the public tree no longer ships it. Keep it gitignored locally so it can
be iterated on without accidental re-adds.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* feat(efa): eager endpoint warmup for segments to remove first-batch stall
libfabric FI_EP_RDM endpoints resolve peer addresses lazily — the first
submitTransfer() to a new segment serializes fi_av_insert + handshake over
every (local_ctx, peer_nic) pair. On 16-NIC instances this is a ~4 s stall
on the first 100 × 0.5 MB batch (measured on p6-B300).
Add EfaTransport::warmupSegment(name) that pre-connects all pairs
concurrently via std::async, plus C wrapper warmupEfaSegment() (guarded by
USE_EFA) and Rust binding TransferEngine::warmup_efa_segment(). Idempotent,
safe to re-run after metadata changes. RDMA/TCP paths untouched.
Measured on p6-B300 (16 × 16 endpoints, dual-NUMA initiator):
first-batch: 4043 ms -> 13.5 ms (~300x)
steady-state: 141 GB/s -> 230 GB/s
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* fix(efa): enable auto-split when max_mr_size is not configured
When MC_EFA_MAX_MR_SIZE is unset, max_mr was 0 and chunk_limit collapsed
to 0, bypassing the PTE-aware split entirely. Large 4KB-paged buffers
then hit the per-NIC PTE ceiling at registration time.
Fall back to pte_limit so splitting kicks in based on the PTE budget
alone when no explicit max_mr_size is provided.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* style: apply clang-format to EFA transport/tests and PEP8 to kvcache bench
Pure formatting changes to satisfy CI format hook (clang-format-20) and
address Gemini review comment on kvcache_prefix_bench.py indentation.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Ubuntu <ubuntu@ip-172-31-13-185.us-east-2.compute.internal>
Co-authored-by: whn09 <whn09@github.com>
|
||
|---|---|---|
| .. | ||
| http-metadata-server | ||
| http-metadata-server-python | ||
| CMakeLists.txt | ||
| batch_register_bench.py | ||
| efa_latency_bench.png | ||
| efa_latency_bench.py | ||
| efa_per_transfer_latency_bench.py | ||
| kvcache_prefix_bench.py | ||
| memory_pool.cpp | ||
| transfer_engine_ascend_direct_perf.cpp | ||
| transfer_engine_ascend_one_sided.cpp | ||
| transfer_engine_ascend_perf.cpp | ||
| transfer_engine_bench.cpp | ||
| transfer_engine_bench_with_notify.cpp | ||
| transfer_engine_bench_with_retry.cpp | ||
| transfer_engine_heterogeneous_ascend_perf_initiator.cpp | ||
| transfer_engine_validator.cpp | ||