PLUGINS/UCX: Improve the performance by specifying memory type in sendAm().

[ZhenlongMa]

What?

This PR improves the performance of NIXL by specifying memory type in sendAm().

Why?

Each nixlbench batch triggers a notification (in postXfer function). The notification calls sendAm without providing mh or memory type, in which case UCX must detect the memory type of the notification buffer in ucp_memory_detect_internal.

If memtype cache is NULL, ucp_memory_detect_internal enters the slow path and then queries all possible memory domians, including CUDA memory, which introduces significant overhead.

Below is the experiment result with block size of 1024B:

Batch	no-opt BW (GB/s)	opt BW (GB/s)	Gain
32	0.823	0.844	2.55%
64	0.944	0.961	1.80%
128	1.037	1.059	2.12%
256	1.103	1.121	1.63%

Summary by CodeRabbit

• Bug Fixes
  • Improved compatibility and reliability for host-memory data transfers over UCX.
  • Preserved existing transfer flags and callback behavior.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

👋 Hi ZhenlongMa! Thank you for contributing to ai-dynamo/nixl.

Your PR reviewers will review your contribution then trigger the CI to test your changes.

🚀

[iyastreb]

/build

[brminich]

/ok to test 88a2db3

[brminich]

/build

[iyastreb]

Thanks, nice finding!

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: 73c652fd-d3aa-4bb0-8346-a9affb91e226

📥 Commits

Reviewing files that changed from the base of the PR and between 43efa63 and 88a2db3.

📒 Files selected for processing (1)

• src/plugins/ucx/ucx_utils.cpp

---

📝 Walkthrough

Walkthrough

The UCX active-message send path now declares the request memory type and sets AM payload memory to host memory while preserving existing flags and callback behavior.

Changes

UCX AM memory configuration

Layer / File(s)	Summary
Configure host memory for AM sends src/plugins/ucx/ucx_utils.cpp	sendAm includes UCP_OP_ATTR_FIELD_MEMORY_TYPE and sets the request memory type to UCS_MEMORY_TYPE_HOST.

Estimated code review effort: 1 (Trivial) | ~2 minutes

Possibly related PRs

• ai-dynamo/nixl#2017: Both changes update nixlUcxEp::sendAm request parameters.

Suggested reviewers: colinnv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the UCX performance improvement implemented in sendAm().
Description check	✅ Passed	The description includes the required What and Why sections, with quantitative results; the optional How section is not required.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[svc-nixl]

🤖 CI Triage Agent — nixl-ci-gpu · commit 5cfa5796

TL;DR: The "Run CPP tests" stage hit its Jenkins wall-clock timeout at test 159/203 because the ucx_threadpool/TestTransfer.* cases ran ~2.5–3x slower than the equivalent ucx/ucx_no_pt cases; PR #2057's change to sendAm() regresses the UCX threadpool send path. Fix by correcting the memory-type argument passed to sendAm() on the threadpool worker path so it doesn't force a slow AM send route.

Full analysis

Summary: Jenkins stage "Run CPP tests" (node 176) was aborted by a stage timeout ("Cancelling nested steps due to timeout" → exit code 143) while gtest-parallel was still on test 159/203.

Root cause: Not a hang — the log shows continuous per-test progress with no multi-minute gaps. It is a genuine slowdown: the ucx_threadpool variants of TestTransfer.* took 134–178s each versus 48–80s for the identical ucx/ucx_no_pt variants (a ~2.5–3x regression). This differential appears only on the UCX thread-pool worker path — exactly the code path touched by PR #2057, which alters how the memory type is specified in sendAm(). The regression inflated the gtest suite runtime beyond the ~49-minute stage timeout so it never reached tests 160–203.

Implicated commit: PR #2057 (PLUGINS/UCX: Improve the performance by specifying memory type in sendAm().), head commit [REDACTED:Hex High Entropy String]. (Author not shown in fetched data.)

File: src/plugins/ucx/ucx_backend.cpp — the sendAm() call/notification send path (the memory-type argument changed by PR #2057); the regression is exercised by ucx_threadpool/TestTransfer.* in test/gtest/.

Suggested fix: Review the sendAm() memory-type change on the thread-pool path. Most likely the memory type being passed forces UCX onto a slower AM protocol (e.g. tagging host notification/control buffers as a non-HOST/device type, disabling inlining/short path). Correct the memory-type parameter so notification/control sends still use the fast path, and re-run the ucx_threadpool/TestTransfer.* gtests to confirm per-test times return to the ~50–80s baseline. Do NOT raise the stage timeout — the slowdown is a code regression, not legitimate work growth.

Related: #2057

🛡️ This comment had 1 potential secret(s) redacted (Hex High Entropy String). See request_id f7935c77-5b26-426a-be4f-e9dac5b32c12 in the triage console for the audit trail.

[svc-nixl]

🤖 CI Triage Agent — nixl-ci-gpu · commit 5cfa5796

TL;DR: The "Run CPP tests" stage was killed (exit 143) at 15:20 because the gtest suite didn't finish in time — the ucx_threadpool/TestTransfer.* variants run ~2.5–3× slower than their non-threadpool counterparts (130–147 s vs ~50 s each), blowing the stage time budget with only 158/203 tests done. This is a slow-but-progressing suite, not a single hang; the fix is to shard/parallelize the gtest run (or split the ucx_threadpool cases into their own job) rather than raise the limit blindly.

Full analysis

Summary: Jenkins nixl-ci-gpu #3099 aborted in the "Run CPP tests" stage; the process was SIGTERM'd mid-gtest at test 159/203 after ~53 min.

Root cause: The gtest suite ran with gtest-parallel --workers=1 (serial), and the ucx_threadpool transfer/error-handling test variants are dramatically slower than the plain ucx variants — e.g. ucx/TestTransfer.RandomSizes = 57 s vs ucx_threadpool/TestTransfer.RandomSizes = 147 s; NotificationOnly 50 s vs 137 s; SelfNotification 50 s vs 136 s. The threadpool block (tests 154–158) alone consumed ~13 min and tests 159–203 never ran. Timestamps show continuous per-test completions right up to the kill (test 158 finished at 15:18:16, kill at 15:20:06 while test 159 ucx_threadpool/TestTransfer.ListenerCommSize — sibling takes ~79 s — was still running), so this is genuine slowness/progress exhausting the wall clock, not a stuck/hung process. The added per-test cost of the threadpool engine variant is the driver.

Implicated commit: Most likely 209abbe9b98e (Ilia Yastrebov, "Refactoring: derive nixlUcxThreadPoolEngine from nixlUcxThreadEngine", #1906) which shapes the threadpool engine behavior exercised by the slow variants; the ucx_threadpool gtest matrix itself amplifies total runtime. Exact regression point cannot be pinned from a single build's log.

File: test/gtest gtest-parallel invocation (gtest-parallel --workers=1 --serialize_test_cases ./bin/gtest) as seen in the stage log at 14:32:47; UCX threadpool engine under src/plugins/ucx.

Suggested fix: Speed up / parallelize the CPP test run instead of raising the timeout:

1. Increase gtest-parallel --workers=1 to multiple workers (GPU permitting) so the ~200 tests aren't fully serialized, OR split the ucx_threadpool/* parameterized cases into a separate parallel Jenkins stage.
2. Investigate why the threadpool engine variant is ~3× slower per transfer test (e.g. progress-thread startup/teardown or connection setup overhead per test in the derived nixlUcxThreadPoolEngine); reducing per-test iteration counts or reusing agents across the threadpool cases would cut total time.
3. Only if the above is insufficient and the log shows continuous progress (as it does here) consider a modest time-limit bump — but the real problem is test-suite duration, so prefer parallelization.

Related: #1906 (threadpool engine refactor), #2040 (single UCX connection per handle), #2044/#2045 (UCX in-flight transfer release changes); no existing issue specifically tracks the ucx_threadpool gtest slowness.