Skip to content

[Store] Fix EFA transport auto-discovery - #3207

Open
yyun-cpu wants to merge 3 commits into
kvcache-ai:mainfrom
yyun-cpu:yyun/fix-efa-store-autodiscovery
Open

[Store] Fix EFA transport auto-discovery#3207
yyun-cpu wants to merge 3 commits into
kvcache-ai:mainfrom
yyun-cpu:yyun/fix-efa-store-autodiscovery

Conversation

@yyun-cpu

@yyun-cpu yyun-cpu commented Jul 30, 2026

Copy link
Copy Markdown

Description

When Mooncake Store uses protocol="efa" with an empty device list, generic
discovery finds the EFA devices but installs Mooncake's rdma transport. That
transport then attempts RC queue pairs, which EFA does not support.

This change passes the requested protocol into Transfer Engine auto-discovery.
The Transfer Engine keeps its existing topology-discovery path and selects
efa instead of rdma when EFA was requested. Existing boolean callers retain
their current behavior.

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:

# Linux container: configure and build Mooncake Store, then run the focused test.
cmake -S . -B /tmp/mooncake-build -G Ninja \
  -DWITH_STORE=ON \
  -DWITH_STORE_RUST=OFF \
  -DWITH_STORE_GO=OFF \
  -DBUILD_UNIT_TESTS=ON \
  -DBUILD_EXAMPLES=OFF \
  -DBUILD_BENCHMARK=OFF
cmake --build /tmp/mooncake-build \
  --target mooncake_store transport_uint_test
ctest --test-dir /tmp/mooncake-build \
  -R transport_uint_test --output-on-failure

# Documentation.
python -m sphinx -W --keep-going -b html docs/source /tmp/mooncake-docs

# Formatting and lightweight checks on the changed files.
PATH="/opt/homebrew/opt/llvm@20/bin:$PATH" \
  uvx pre-commit run --files \
  docs/source/design/transfer-engine/efa_transport.md \
  mooncake-store/src/client_service.cpp \
  mooncake-transfer-engine/include/transfer_engine.h \
  mooncake-transfer-engine/include/transfer_engine_impl.h \
  mooncake-transfer-engine/src/transfer_engine.cpp \
  mooncake-transfer-engine/src/transfer_engine_impl.cpp \
  mooncake-transfer-engine/tests/transport_uint_test.cpp

Test results:

  • Unit tests pass
  • Integration tests pass
  • Manual testing done

The focused tests cover protocol-aware EFA selection and compatibility with
existing boolean auto-discovery callers. On an AWS p6-b200.48xlarge node,
the Store discovered all eight EFA devices, installed the efa transport, and
completed a verified put/get with protocol="efa" and an empty device list.
The test did not log RDMA RC queue-pair errors.

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 (not applicable)

All pre-commit hooks pass on the changed files. A repository-wide run still
fails on unrelated existing issues.

AI Assistance Disclosure

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

Codex helped investigate the existing EFA paths and prepare the
implementation, tests, documentation, and validation. The submitter reviewed
the changes before submission.

@github-actions github-actions Bot added documentation Improvements or additions to documentation run-ci Store labels Jul 30, 2026
@yyun-cpu

Copy link
Copy Markdown
Author

The failing checks seem to be an issue with timing/connection in the tests, according to the github agent.

@whn09

whn09 commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the fix — root cause is correct and matches #2070. I validated it on a p5.48xlarge (H100, 32 EFA devices) with a real USE_EFA=ON -DUSE_CUDA=ON build, A/B against the parent commit.

Reproduced. protocol="efa", MC_MS_AUTO_DISC=1, empty device list:

  • Parent: installTransport, type=rdmaFailed to create QP: Operation not supported [95]put rc=-800
  • This PR: Found 32 HCAsprovider: efa (shared endpoint) → put/get PASS

No regressions: efa auto-discover (3× soak), efa + manual devices, MC_MS_FILTERS (correctly 2 of 32), MC_CUSTOM_TOPO_JSON valid + missing-file fallback, MC_FORCE_TCP, protocol=rdma (unchanged), protocol=tcp, client_transport_config_test 4/4 — all pass.

Ordering is safe on both sides: with auto_discover_=false, TransferEngineImpl::init() never touches local_topology_, so the pre-init() discovery can't be clobbered; and Store registers memory only after InitTransferEngine returns, so installing post-init() still picks up all regions.

Please rebase — merge conflict

The PR is CONFLICTING against main. Single trivial hunk in mooncake-store/tests/CMakeLists.txt: #3003 added the object_checksum_client_test block where your add_store_test(client_transport_config_test ...) line goes. Keep both, nothing semantic.

Non-blocking comments

  1. discover_efa_topology() duplicates the MC_CUSTOM_TOPO_JSON logic already in transfer_engine_impl.cpp:259-272 (incl. loadTopologyJsonFile). Two copies of the same env-var contract will drift, and the fallback semantics are subtle (empty file vs unopenable file both fall back). I'd like the TE helper exposed and reused instead — the one item I'd genuinely like folded in.

  2. explicitly_discover_efa vs raw auto_discover is a latent trap. The install branch at :812 tests discovery_plan, the else if at :821 still tests raw auto_discover. Correct today only because explicitly_discover_efa implies auto_discover — worth a comment stating that invariant, or deriving one from the other.

  3. force_tcp in MakeTransportDiscoveryPlan is inert. MC_FORCE_TCP is already handled in TransferEngineImpl::init(), which returns before this runs. Fine as defense-in-depth, but ForceTcpDisablesDiscovery isn't coverage of the real MC_FORCE_TCP guarantee.

Approving the fix. Rebase for the conflict, and comment 1 would be good to fold in.

@Aionw Aionw left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we pass the auto-discovery configuration into setAutoDiscover() and let the Transfer Engine derive the appropriate behavior?

Client::InitTransferEngine() already has the relevant configuration, including auto_discover, protocol, and whitelist filters. Instead of calculating an EFA-specific TransportDiscoveryPlan in Store, it could pass that configuration to the Transfer Engine:

transfer_engine_->setAutoDiscover({
    .enabled = auto_discover,
    .protocol = protocol,
    .filters = filters,
});

@whn09

whn09 commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

I think this reviewer is right on the design, and I'd like to back it up with what I found while validating on a 32-EFA p5.48xlarge — plus one reason I'd still land the narrow fix first.

The premise checks out. In transfer_engine_impl.cpp:257-417, the auto-discovery ladder has zero EFA branches — a USE_EFA build falls straight through to the rdma branch at :366-390. And the TE genuinely can't do better today: init() never receives protocol, and TransferEngineImpl has no protocol_ member. So the TE has no way to honor an EFA request even in principle.

That's why this workaround already exists 5 times in-treetransfer_engine_py.cpp:241, transfer_engine_bench.cpp:512, efa_first_submit_probe.cpp, and three EFA test files all hand-roll the same getLocalTopology()->discover() + installTransport("efa") dance. This PR makes it 6. My earlier comment about the duplicated MC_CUSTOM_TOPO_JSON logic was really a symptom of exactly what @Reviewer is naming: the knowledge lives in the wrong layer.

The sharpest evidence: with protocol="rdma" and auto-discovery on EFA hardware, this PR still installs rdma and fails — I measured put rc=-800 with 24× Failed to create QP: Operation not supported [95]. A TransportDiscoveryPlan computed in Store structurally cannot fix that, because the Store never sees the install decision. Pushing the config into the TE would fix that case coherently. The API risk is also lower than it looks: setAutoDiscover() has only 2 in-tree callers (client_service.cpp:689, pybind_client_test.cpp:705) and is exposed neither to Python nor the C API, so an overload taking a config struct is backward-compatible.

Why I'd still merge this one first. That ladder is 16 nested preprocessor branches shared by Ascend, MACA, MNNVL, BAREX, SUNRISE, HIP and UB. Making the TE protocol-aware means touching the common install path for every platform, and most of those I can't test — I have EFA and TCP hardware, not Ascend or MACA. This PR is contained to Store and I've A/B-verified it end-to-end on real EFA hardware, including no-regression runs for rdma, tcp, MC_FORCE_TCP, MC_MS_FILTERS and MC_CUSTOM_TOPO_JSON. Users hitting -800 today get unblocked without putting seven other platforms at risk.

So my suggestion: land this as the targeted fix, and open a follow-up issue for the TE-side refactor — setAutoDiscover(config) with protocol-aware transport selection, which would let us delete all 6 hand-rolled copies including the one this PR adds. I'm happy to own that follow-up since I maintain the EFA transport and have the hardware to validate it.

That said, this is a maintainer call, not mine. If @Aionw or the maintainers would rather see it done properly in one pass, I won't argue the point — I'd just ask that the TE-side change get tested on at least one non-EFA accelerator before merge, since the blast radius is that shared ladder.

@Aionw

Aionw commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

Thanks to @whn09 for continuing the broader refactoring effort in this issue. Let’s keep this PR focused on fixing EFA auto-discovery.

@yyun-cpu
yyun-cpu force-pushed the yyun/fix-efa-store-autodiscovery branch from c6e29c3 to 9896e95 Compare July 31, 2026 07:06
@yyun-cpu

yyun-cpu commented Jul 31, 2026

Copy link
Copy Markdown
Author

Thanks both. Rebased onto current main.

Aionw commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

To clarify my earlier review: I agree that this PR should remain a targeted EFA fix. I am not asking for a broad refactor of the Transfer Engine's platform-specific transport-selection ladder.

However, I do not think the current implementation is the minimal targeted fix.

The bug exists because Store calls:

transfer_engine_->setAutoDiscover(auto_discover);

without passing the configuration that tells the Transfer Engine that the requested protocol is EFA. The current patch works around that by disabling the Transfer Engine's discovery path, duplicating topology discovery in Store, introducing TransportDiscoveryPlan, and installing EFA after init().

I suggest passing the relevant auto-discovery configuration, including the requested protocol, into setAutoDiscover() instead:

transfer_engine_->setAutoDiscover({
    .enabled = auto_discover,
    .protocol = protocol,
});

The Transfer Engine can then keep its existing topology-discovery path and make one targeted adjustment in the default transport-selection branch:

if (auto_discover_config_.protocol == "efa") {
    installTransport("efa", local_topology_);
} else {
    // Keep the existing RDMA/TCP selection unchanged.
}

The existing setAutoDiscover(bool) behavior can remain as the default for other callers, so this does not require restructuring or changing the platform-specific branches.

This approach would preserve the existing handling of:

  • MC_CUSTOM_TOPO_JSON;
  • whitelist filters;
  • MC_FORCE_TCP; and
  • companion transports such as HIP.

It would also allow this PR to remove:

  • discover_efa_topology();
  • TransportDiscoveryPlan;
  • client_transport_config.h;
  • efa_discovery_filters;
  • the pre-init() discovery path; and
  • the post-init() EFA installation branch.

That should produce a substantially smaller diff while still keeping the change narrowly scoped to EFA auto-discovery.

@codecov-commenter

codecov-commenter commented Jul 31, 2026

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 78.18182% with 12 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...ncake-transfer-engine/src/transfer_engine_impl.cpp 10.00% 9 Missing ⚠️
mooncake-store/src/client_service.cpp 33.33% 2 Missing ⚠️
...ake-transfer-engine/include/transfer_engine_impl.h 91.66% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@yyun-cpu

Copy link
Copy Markdown
Author

Changed the PR to pass the protocol into setAutoDiscover() as suggested.

Keep the explicit BAREX override ahead of protocol-aware EFA selection so existing BAREX configurations remain unchanged.

Co-authored-by: Cursor <cursoragent@cursor.com>

@Aionw Aionw left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM for store side.

@ykwd

ykwd commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

Could you take a look at the failed CI? Thanks

@Aionw

Aionw commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

I took another look at the failed CI. The relevant error is a linker failure involving -lcuda; the subsequent “runner lost communication” message appears to have prevented GitHub from preserving the complete job log.

This does not appear to be introduced by this PR:

  • The latest commit only changes the Transfer Engine auto-discovery logic and its unit test; it does not modify any CUDA, CMake, or linker configuration.
  • The CUDA 13 / Python 3.10 job in the same run passed with the same -L/usr/local/cuda/lib64/stubs -lcuda linker configuration.
  • The CUDA 13 / Python 3.12 job also passed on the previous PR revision with the same linker command.
  • The workflow already adds the CUDA driver stub directories to LIBRARY_PATH.

This therefore looks more like a transient runner/container issue where the CUDA driver stub was unavailable or inaccessible, rather than a regression caused by the PR. I suggest rerunning the failed job. If the same cannot find -lcuda error reproduces consistently, we should then inspect the CUDA 13/Python 3.12 builder image and its stub-library layout.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation run-ci Store Transfer Engine

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants