Mooncake/docs
王鹤男 952da65651
[TE] PTE-aware auto-split large MR registration for EFA transport (#1912)
* 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>
2026-04-20 15:04:49 +08:00
..
source [TE] PTE-aware auto-split large MR registration for EFA transport (#1912) 2026-04-20 15:04:49 +08:00
Makefile [DOC] feat(docs): build documentation website for Mooncake using Sphinx (#354) 2025-05-12 23:33:18 +08:00
README.md [DOC] feat(docs): build documentation website for Mooncake using Sphinx (#354) 2025-05-12 23:33:18 +08:00
make.bat [DOC] feat(docs): build documentation website for Mooncake using Sphinx (#354) 2025-05-12 23:33:18 +08:00

README.md

Mooncake documents

Build the docs

  • Make sure in docs directory
cd docs
  • Install the dependencies:
pip install -r ../requirements_docs.txt
  • Clean the previous build (optional but recommended):
make clean
  • Generate the HTML documentation:
make html

Be sure to set locale right before building.

Open the docs with your browser

  • Serve the documentation locally:
python -m http.server -d build/html/

This will start a local server at http://localhost:8000. You can now open your browser and view the documentation.

If port 8000 is already in use, you can specify a different port, for example:

python -m http.server 3000 -d build/html/