nixl_ep: Optimize low-latency NVLink path

[itayalroy]

This PR adds a set of optimizations for NIXL EP low-latency NVLink path.

Each optimization is in a separate commit, with the explanation included in the commit message.

On the TRT-LLM MoE communication benchmark (bench_moe_comm.py) over single-node NVLink, 8xH100, these changes bring NIXL EP dispatch to parity with DeepEP low latency and reduce the remaining combine gaps.

Full dispatch benchmark table

Batch	deepep low latency	nixl ep main	main vs deepep	nixl ep PR 1749	PR vs deepep
16	32.06 µs	38.34 µs	0.836x	32.55 µs	0.985x
32	37.61 µs	45.49 µs	0.827x	37.98 µs	0.990x
64	49.38 µs	61.12 µs	0.808x	49.69 µs	0.994x
128	75.21 µs	92.18 µs	0.816x	74.76 µs	1.006x
256	128.85 µs	147.32 µs	0.875x	128.20 µs	1.005x
512	229.28 µs	248.60 µs	0.922x	228.14 µs	1.005x

Combine active rank mask optimization

Batch	disable mask / speed-of-light	current: load acq	current vs speed-of-light	this PR: plain load	PR vs speed-of-light
16	25.03 µs	27.24 µs	0.919x	26.26 µs	0.953x
32	31.74 µs	33.79 µs	0.939x	33.06 µs	0.960x
64	45.65 µs	48.48 µs	0.942x	46.94 µs	0.973x
128	73.58 µs	76.75 µs	0.959x	74.37 µs	0.989x
256	137.18 µs	141.01 µs	0.973x	137.34 µs	0.999x
512	258.57 µs	266.23 µs	0.971x	258.68 µs	1.000x

Note that there is no visible performance impact on NIXL EP’s elastic.py.

Follow-up PRs will continue optimizing NIXL EP combine, with the goal of matching and then exceeding DeepEP low-latency performance.

Summary by CodeRabbit

• New Features

  • GPU-side P2P pointer caching to improve peer-to-peer GPU memory access performance.

• Chores

  • Reduced default index width from 64-bit to 32-bit.
  • Added a CUDA compile flag to alter PTX instruction selection for targeted memory operations.
  • Improved allocation, initialization, population, and cleanup of the GPU pointer cache to ensure consistent behavior.

[github-actions]

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

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

🚀

[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: 0317d4d0-9996-4759-8e42-0200167fa3c0

📥 Commits

Reviewing files that changed from the base of the PR and between a7a3bf3 and 2678857.

📒 Files selected for processing (3)

• examples/device/ep/csrc/kernels/api.cuh
• examples/device/ep/csrc/kernels/nixl_ep_ll.cu
• examples/device/ep/csrc/nixl_ep.cpp

---

📝 Walkthrough

Walkthrough

Adds a device-side per-rank P2P pointer cache (device array + kernel + host launcher), wires allocation/refresh/free in Buffer, updates kernels to use cached pointers and adjust mask-load semantics, and tweaks PTX opcode selection and Meson/config defaults (TOPK_IDX_BITS → 32, DISABLE_AGGRESSIVE_PTX_INSTRS).

Changes

P2P Pointer Cache System

Layer / File(s)	Summary
P2P pointer cache context structure and API examples/device/ep/csrc/kernels/api.cuh	gpu_nixl_ctx gains void** p2p_ptrs; declares cache_p2p_ptr(gpu_nixl_ctx*, int rank_id, cudaStream_t).
Cache population kernel and host wrapper examples/device/ep/csrc/kernels/nixl_ep_ll.cu	Adds cache_p2p_ptr_kernel that writes per-rank entries into nixl_ctx->p2p_ptrs[rank_id] and a host wrapper cache_p2p_ptr(...) that launches it on a stream.
Kernel pointer fetch and mask-check updates examples/device/ep/csrc/kernels/nixl_ep_ll.cu	p2p_ptr_get reads cached ctx.p2p_ptrs[dst_rank]; is_rank_masked gains use_acquire; combine loops use is_rank_masked<false, false>(...).
Buffer initialization, refresh, and destruction examples/device/ep/csrc/nixl_ep.cpp	Allocates/zero-inits device gpu_ctx.p2p_ptrs at init, calls ep_kernels::cache_p2p_ptr(gpu_ctx_ptr, remote_rank, comm_stream) after memory-view creation, clears entries on disconnect via cudaMemset, and frees/nulls the device buffer at destroy.

PTX Instructions and Build Configuration

Layer / File(s)	Summary
PTX instruction updates examples/device/ep/csrc/kernels/utils.cuh	Under DISABLE_AGGRESSIVE_PTX_INSTRS, LD_NC_FUNC uses ld.volatile.global.L2::256B; tma_store_wait<N>() uses cp.async.bulk.wait_group.
Meson and configuration build settings examples/device/ep/meson.build, examples/device/ep/csrc/kernels/configs.cuh	Meson adds -DDISABLE_AGGRESSIVE_PTX_INSTRS; topk_idx_bits default changed to 32 and propagated as -DTOPK_IDX_BITS=...; TOPK_IDX_BITS set to 32 in configs.cuh (so topk_idx_t becomes int32_t).

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Suggested reviewers

• ebarilanM
• rakhmets
• tomerg-nvidia
• brminich
• iyastreb

Poem

🐰 I cached P2P pointers in a hop and a beat,
Launched a kernel so lookups end neat,
Buffers allocated, freed with care,
PTX whispers tuned the GPU air,
Now ranks find each other fast and fleet.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'nixl_ep: Optimize low-latency NVLink path' directly summarizes the main objective of the PR, which is to introduce optimizations for the NIXL EP low-latency NVLink datapath across multiple commits.
Description check	✅ Passed	The PR description adequately covers the What (set of optimizations for NIXL EP low-latency NVLink), Why (benchmark improvements bringing dispatch to parity with DeepEP), and How (separate commits with detailed explanations and performance tables).
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/device/ep/csrc/kernels/utils.cuh`:
- Around line 179-183: The LD_NC_FUNC macro currently expands to
"ld.volatile.global.L2::256B" when DISABLE_AGGRESSIVE_PTX_INSTRS is defined,
which requires PTX ISA 7.4+/SM80+ assembler support; update the build
configuration to target PTX 7.4+ / compute capability SM80+ (e.g., set the
appropriate -arch / -gpu-architecture / CUDA_NVCC_FLAGS) so the assembler
accepts the ".L2::256B" cache-hint, or add a fallback macro value for older
PTX/SM targets; locate and adjust the LD_NC_FUNC definition and any build
scripts that influence PTX/SM targets to ensure compatibility.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: aac73b69-891d-42e2-a161-a8806cda6330

📥 Commits

Reviewing files that changed from the base of the PR and between 7af417b and 252b191.

📒 Files selected for processing (6)

• examples/device/ep/csrc/kernels/api.cuh
• examples/device/ep/csrc/kernels/nixl_ep_ll.cu
• examples/device/ep/csrc/kernels/utils.cuh
• examples/device/ep/csrc/nixl_ep.cpp
• examples/device/ep/csrc/nixl_ep.hpp
• examples/device/ep/meson.build

[itayalroy]

Changed P2P pointer caching to launch one small kernel per newly connected rank.

nixlGetPtr crashes on gap ranks, so this avoids calling it for unconnected ranks. Since this only runs on the connection path, and a kernel launch is ~2-5 us, this is fine for now and alignes with how we update other GPU state like the active mask.

The alternative is adding and populating a gpu_connected_ranks array, which is arguably worse, and also pretty messy.

I think we should consider making UCX return nullptr for gap ranks instead of crashing, but shouldn't block this PR

[itayalroy]

/build

[ebarilanM]

LGTM