Commit Graph

9 Commits

Author SHA1 Message Date
Jason 5b1436196e
fix(transfer-engine): add missing empty checks for batch methods (#2046) 2026-05-07 16:00:54 +08:00
王鹤男 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
dependabot[bot] 897728ddbf
chore(deps): bump tracing-subscriber in /mooncake-transfer-engine/rust (#794)
Bumps [tracing-subscriber](https://github.com/tokio-rs/tracing) from 0.3.18 to 0.3.20.
- [Release notes](https://github.com/tokio-rs/tracing/releases)
- [Commits](https://github.com/tokio-rs/tracing/compare/tracing-subscriber-0.3.18...tracing-subscriber-0.3.20)

---
updated-dependencies:
- dependency-name: tracing-subscriber
  dependency-version: 0.3.20
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2025-08-30 15:57:52 +08:00
Yuan Luo 20a94fc946
[TransferEngine] Support Status return value (#125)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
2025-03-15 13:02:37 +08:00
Ruoyu Qin 19aafbe8d2
[TransferEngine] fix: add auto_discover in transfer_engine_c (#87)
* fix transfer_engine_c

* fix rust and go examples
2025-01-23 09:54:26 +08:00
doujiang24 ab8021d6f7
[TransferEngine] change: auto discover topology & install transport. (#73)
* [TransferEngine] change: auto discover topology & install transport.

Signed-off-by: doujiang24 <doujiang24@gmail.com>
2025-01-17 21:40:36 +08:00
Feng Ren 64ddda46d4
[TransferEngine] Refactor code to hide transport logics from user APIs (#51)
Co-authored-by: doujiang24 <doujiang24@gmail.com>
2025-01-02 16:38:24 +08:00
qinruoyu e841d943ee fix rust example 2024-12-05 09:14:39 +00:00
Feng Ren 348afd7531 Squashed commits related to transfer engine 2024-11-27 13:02:39 +08:00