Skip to content

[Bugfix][TransferEngine] Honor retry hints in RDMA request preselection - #4023

Draft
adenzhou1350 wants to merge 1 commit into
kvcache-ai:mainfrom
adenzhou1350:fix/rdma-request-retry-hint
Draft

[Bugfix][TransferEngine] Honor retry hints in RDMA request preselection#4023
adenzhou1350 wants to merge 1 commit into
kvcache-ai:mainfrom
adenzhou1350:fix/rdma-request-retry-hint

Conversation

@adenzhou1350

@adenzhou1350 adenzhou1350 commented Sep 11, 2026

Copy link
Copy Markdown

Description

TransferEnginePy::transferSync() retries failed transfers across available
local RDMA contexts by setting TransferRequest::advise_retry_cnt to the outer
retry index. The RDMA request-wide device-selection fast path did not forward
that hint, so a request contained in one registered buffer could preselect the
retry-zero preferred HCA again and bypass the per-slice selection path that
already honors the hint.

Forward request.advise_retry_cnt to the request-wide selectDevice() call.
The default value is zero, so first-attempt selection is unchanged. Requests
that cannot be selected as a whole still use the existing per-slice fallback
behavior.

This is adjacent to the failover work in #2959 / #2960, but narrower: it
restores retry-index selection for an active request-wide preselected context.
It is also complementary to open PR #3696, which adds rail cooldown,
redispatch, chaos coverage, and inactive-context safeguards but leaves the
active request-wide selectDevice() call at retry zero. A merge-tree check
against the observed #3696 head is clean.
It is also complementary to open PR #2524: that PR advances retry hints in the
Python batch wrapper but does not forward the hint through this request-wide
RDMA source-device selection. This change does not modify #2524's batch retry
lifecycle, endpoint recovery, inactive-context handling, or metadata
publication.

vLLM's Mooncake KV connector and SGLang's primary KV/weight transfer paths use
batch transfer APIs. On current Mooncake main those APIs keep
advise_retry_cnt at zero, so this PR alone does not claim a framework KV
batch-failover improvement. If #2524 (or an equivalent producer fix) lands,
this consumer fix is required for a whole-buffer classic RDMA request to honor
the advanced hint instead of reusing retry-zero preselection.

There is also a current single-transfer consumer: SGLang multimodal-generation
disaggregation

calls Mooncake transfer_sync_write() when pushing a staged payload to a peer.
That path reaches this fix when classic RDMA has multiple local contexts. This
is source-level reachability only; no end-to-end SGLang failover or performance
result is claimed here.

Module

  • Transfer Engine (mooncake-transfer-engine)

Type of Change

  • Bug fix

How Has This Been Tested?

The focused unit test uses two bare test RDMA contexts and distinct lkeys. With
advise_retry_cnt = 2, the request-wide fast path must select the fallback
context's lkey.

cmake -S /src -B /build \
  -DWITH_STORE=OFF -DWITH_STORE_RUST=OFF -DWITH_TE=ON \
  -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=OFF -DUSE_CUDA=OFF \
  -DUSE_ETCD=OFF -DUSE_REDIS=OFF -DUSE_HTTP=ON
cmake --build /build --target rdma_transport_submit_task_test -j3
/build/mooncake-transfer-engine/tests/rdma_transport_submit_task_test
pre-commit run --files \
  mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp \
  mooncake-transfer-engine/tests/rdma_transport_submit_task_test.cpp
  • Focused unit tests pass: 5/5.
  • Unchanged-production control is sensitive: 4/5 pass; only
    RequestWideSelectionHonorsRetryHint fails, selecting preferred lkey 11
    instead of fallback lkey 22.
  • Current origin/main merge-tree and PR-scoped pre-commit hooks pass.
  • Forward compatibility screen on the observed [TE] Add RDMA chaos runner and worker-pool failover retries #3696 head passes: the
    combined CPU build succeeds and the same focused suite passes 5/5.
  • Adjacent CPU-only regression suites pass: topology_test 13/13,
    rdma_endpoint_state_test 8/8, and worker_pool_rail_state_test 15/15.
  • Real multi-NIC RDMA transferSync failover and data-integrity test. This
    remains pending; this Draft makes no live-RDMA or performance claim.

Checklist

  • I have performed a self-review of my own code.
  • I have formatted my code using the repository's LLVM 20 changed-line
    formatter.
  • I have run pre-commit on the files changed in this PR and all hooks pass.
  • Documentation is not required because the public API and configured
    behavior are unchanged.
  • I have added a test that fails without the production fix.
  • An RFC is not required; the change is under 500 LOC and restores existing
    retry semantics.

AI Assistance Disclosure

  • AI tools were used.

Codex helped trace the retry path, prepare the focused change and regression
test, find and compare adjacent upstream work, and run paired candidate/control
validation. The human submitter will review every changed line before
publication.

Claim boundary

This is a correctness/reliability fix for classic RDMA request-wide local-device
preselection. It does not claim higher healthy-path throughput, end-to-end
failover success, or results for EFA, CXI, UB, TENT, the batch APIs, GPU/GDR, or
production workloads.

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant