Feature/mnnvl p2p by bjoo · Pull Request #1641 · lattice/quda

bjoo · 2026-07-08T08:17:35Z

Summary

This work adds several modifications, the chief of which is stream-gated transport for MNNVL.

Description

Previous intra node P2P mechanism used CUDA/HIP IPC mechanism. The sender copied data to a shared buffer, and once the copy was completed created a remote IPC event. An MPI doorbell message was sent to the receiver. The receiver waiter first on the MPI doorbell and then began waiting on the remote IPC event to complete to know that the transfer was done.

Multi-node NVLink use CUDA Fabric handles for remote memory operations which must be allocated using CUDA Virtual Memory Management (VMM) allocators. Remote memory allocated using VMM does not support remote IPC events. The solution to this is to use an approach we will refer to as stream-gating. In this case once the transfer of the shared memory is complete the sender will issue cuStreamWrite64 operation to a piece of remote memory which serves as the doorbell. The receiver sits in a polling loop over the various directions and can check the remote doorbell location and finally unblock with a cuStreamWait64 operation. Follow on kernels can be queued on the stream behind this operation so once the cuStreamWait64 operation fires subsequent kernels can dispatch without further host involvement, giving this method a very low latency.

Implementation notes

We implented the following:

MNNVL build selected with CMake option -DQUDA_MNNVL=ON.
Added stream gating as a user choice transport option selectable with the environment variable QUDA_P2P_TRANSPORT=stream_gated|event. The choice is only meaningful in CUDA builds on the non MNNVL path. For MNNVL builds, event is not available. On Non-MNNVL builds event is default.
MNNVL comm buffer allocations proceed via VMM and allocate a single handle for their memory which they return as Fabric handles, both for P2P and even for regular MPI. MPI and P2P no longer share buffers with NVSHMEM when NVSHMEM is enabled.
In MNNVL stream gating, P2P Remote Write operations rely on being able to flush streams using cuStreamWait64 using a specific FLUSH flag. Whether the system in question provides the FLUSH capability is tested for at startup. On systems where the FLUSH capability is not available, P2P remote write operations are explicitly disabled and use Copy Engines only.
in order to prevent proliferation of #ifdef lines in the main include/lib dirs the signaling API was abstracted and implementation was pushed down into the appropriate include/targets and lib/targets directory.
The implementation should (in our basic testing did) can work accross MNNVL-Clique boundaries, gracefully falling back to MPI communications in directions where MNNVL-P2P is not available. In these use cases it may be worth setting variable UCX_CUDA_IPC_ENABLE_MNNVL=n so that 'cross-clique' MPI traffic doesn't attempt to use UCX's own MNNVL
In order for fabric handles to work, the nodes must have an IMEX daemon running
Testing was carried out using Wilson and Staggered Dslash tests, some of the hisq force tests and using short MILC trajectories. Chroma testing (especially MG) is currently still TBD.

The implementation was carried out in large part by Claude Code (Opus 4.7 and Opus 4.8). The expecation is that 'stream-gate' works best in the case where we have low latency requirements (e.g. Multi-Grid or strong scaling). For larger message sizes stream gating and the event based approach should perform similarly with each other as in that case higher latencies can be amortized (we chacked this on non-MNNVL builds where both are available).

…abstraction Introduce the default-off QUDA_MNNVL CMake option and refactor the existing P2P halo doorbell into a comm_p2p_* abstraction. Pure refactor: no functional change when QUDA_MNNVL=OFF.

Add a tag-dispatched comm-buffer allocator (DeviceCommBuffer kind) backed by cuMemAlloc with tracker tags, and remove the dead pinned-malloc / MPI-kind machinery it replaces.

Add the QudaP2PSignal enum and stream-mem-op primitives, migrate all call sites onto the unified comm_p2p_signal_* API, and give stream-gated signalling a dedicated flag-buffer lifecycle. Signalling is now orthogonal to the data-movement policy; event mode remains the default.

Expose stream-gated as a P2P sub-policy on tp.aux.y and wire it through DslashFusedExterior/FusedGDR and GaugeField::exchangeGhost (4-phase halo). Default the P2P signalling transport to stream_gated.

Under QUDA_MNNVL, back P2P and MPI comm buffers with fabric VMM allocations (contiguous, RDMA-capable, separate from the NVSHMEM heap), pass fabric handles to UCX, and establish cross-clique reachability via a symmetric fabric import-probe. All gated by QUDA_MNNVL (off by default).

… MNNVL Enable the stream-gated zero-copy-pack send/recv and GDR dslash policies under MNNVL. Offer the REMOTE_WRITE P2P sub-policy only when the device can flush remote writes (capability-gated; QUDA_P2P_FORCE_FLUSH overrides), else fall back to the copy engine. Adds the QUDA_P2P_FAKE_CLIQUE study hook.

Fix multi-node mixed P2P/non-P2P halo signalling, split-grid bufferIndex and signal-counter resets on communicator switch, extended gauge-ghost multi-dim corruption, HISQ force fermion-halo signalling, the dslash run-ahead race, and NVSHMEM build wiring. NB: Remote-Write policy had issues on MNNVL (receive side visibility - needs receiver cuStreamWait64(..., CU_STREAM_WAIT_VALUE_FLUSH)) For now: check for the availability of this at P2P init time - if not available, turn off remote-write. (Remote write attempts can still be forced with an env. var.) Add staggered_dslash options and stream-gated ctests. Tested and validated using MNNVL=ON/OFF, NVSHMEM=ON/OFF.

… code) Move all MNNVL #ifdefs and <cuda.h> out of the target-agnostic headers/sources into include/targets/cuda + lib/targets/cuda: - comm_target.h fabric primitives (fabric_handle_size / open_fabric_probe / try_import_fabric_handle / close_fabric_probe) declared unconditionally; impls in comm_target.cpp are real under QUDA_MNNVL, cheap stubs otherwise. HIP header gets inline stubs. Generic code calls them under if constexpr(comm_build_is_mnnvl()). - Move get_p2p_fabric_handle / get_p2p_buffer_size / get_p2p_buffer_generation (return CUmemFabricHandle) to new include/targets/cuda/malloc_target.h; drop <cuda.h> and the fabric-accessor block from malloc_quda.h. - Remove dead comm_gather_clique_id + get_fabric_clique_id (old NVML-cliqueId reachability, superseded by the import probe); un-guard the live comm_gather_fabric_handle member; delete the dead free-function forwarders. - communicator_quda.h P2P reachability probe (#ifdef/#else) and lattice_field.cpp neighbour-event exchange (#ifndef) -> if constexpr(comm_build_is_mnnvl()). Keeps add_compile_definitions(QUDA_MNNVL) global: comm_build_is_mnnvl() is a constexpr in a header included by generic TUs, so the define must agree across all TUs (no per-target split).

Ghost2() merely delegated to Ghost2P2P(); the split was arbitrary. The underlying buffer (ghost_recv_buffer_p2p_d) is the shared non-shmem recv buffer used by P2P, GDR and MPI-host-staged transports alike, not P2P-specific -- so keep the transport-neutral name Ghost2() and drop Ghost2P2P(). Ghost2Shmem() remains for the symmetric/NVSHMEM recv base. Behavior-identical (Ghost2() always returned Ghost2P2P()); the sole Ghost2P2P() call site (dslash.h) now calls Ghost2(). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…pers, trim comments - Drop unused #include <device.h> from communicator_{mpi,qmp,single}.cpp (branch-added, no device:: symbol used; not present in develop). - Rename file-local ghost_comm_buffer_malloc_/_free_ -> ghost_comm_buffer_malloc/_free in lattice_field.cpp (no macro sibling / provenance threading; all callers in-file). - comm_quda.h: trim verbose doc comments on the P2P/NVSHMEM/stream-gated APIs. - communicator_qmp.cpp: add FIXME note on the byte-wise Allgather emulation. - tests/CMakeLists.txt: stream-gated-p2p staggered tests use --test Dslash. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

bjoo · 2026-07-08T13:51:03Z

All but CSCS test runs finished. CSCS build passed but test job didn't execute - for what may be non QUDA reasons

hummingtree · 2026-07-09T17:27:28Z

+      memcpy(&fh, peer_handle, sizeof(CUmemFabricHandle));
+      CUmemGenericAllocationHandle h;
+      CUresult err = cuMemImportFromShareableHandle(&h, &fh, CU_MEM_HANDLE_TYPE_FABRIC);
+      if (err != CUDA_SUCCESS) return false;


Should we clear the CU error after this?

And should we check to see what is the error type to make sure we are seeing the expected error type when the remote is not accessible?

This is a reachability probe, so a failure is an expected outcome (peer not reachable over fabric → we symmetrically downgrade that link to MPI), not an error condition. On the driver API the status comes back in the return value rather than as a sticky per-thread error (unlike the runtime's cudaGetLastError), so there's nothing latched to clear here — but happy to add an explicit note/clear if you'd prefer defense-in-depth. We could refine the error checking though to see if there is something bad there.

Make sense. So when the remote is not accessible, just make sure that the error type is expected? From the API doc it looks like only CUDA_ERROR_NOT_SUPPORTED and CUDA_ERROR_NOT_PERMITTED are expected?

Claude is on it :)

hummingtree · 2026-07-09T17:51:01Z

+    unsigned int wait_flags = CU_STREAM_WAIT_VALUE_GEQ;
+    if (stream_gated_remote_flush_supported) wait_flags |= CU_STREAM_WAIT_VALUE_FLUSH;
+    CUresult err = cuStreamWaitValue64(target::cuda::get_stream(stream), local, expected, wait_flags);


Why do we use GEQ here instead of EQ? I know by design they are the same but just wondering which one is better.

This is the monotonic counter protocol: the receiver waits for expected = ++recv_counter while the sender writes an ever-increasing ++send_counter into the slot. By the time a WaitValue64 is actually evaluated the peer may already have advanced the slot past expected (the next generation's write can land first). EQ would miss that exact value and block forever; GEQ is satisfied as soon as the slot reaches at-least the expected generation.

By the time a WaitValue64 is actually evaluated the peer may already have advanced the slot past expected (the next generation's write can land first).

I think the double buffered bufferIndex is targeting to avoid this kind of over-by-one-step (I am inventing some words here). I am pretty sure it would work but maybe worth a documentation here.

hummingtree · 2026-07-09T17:55:02Z

+    access.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+    access.location.id = device_id;
+    access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+    err = cuMemSetAccess(ptr_d, padded, &access, 1);


Maybe it's time to add a macro/etc for these driver calls so we don't need to repeat the if (err != CUDA_SUCCESS) errorQuda( every time.

Claude suggests that yes a macro is maybe worth it? It offered to generalize target::cuda::set_driver_error(err, name, func, file, line)

Up to you :)

hummingtree · 2026-07-09T18:08:35Z

+        = static_cast<char *>(ghost_remote_send_buffer_p2p_d[bufferIndex][dim][dir]) + ghost_offset[dim][(dir + 1) % 2];
+      qudaMemcpyP2PAsync(ghost_dst, my_face_dim_dir_p2p_d[bufferIndex][dim][dir], ghost_face_bytes[dim], stream);
+    }
+    comm_p2p_signal_send_done(FieldKind::COLOR_SPINOR, bufferIndex, dim, dir, stream, QudaP2PSignal::STREAM_GATED);


I am not sure but maybe we should check to see if STREAM_GATED is enabled/supported when we use the enum QudaP2PSignal::STREAM_GATED directly? (to be able to detect more errors from the decision trees up-stream)

Or maybe just rename this function to include GatedStream to make it clear.

hummingtree · 2026-07-09T18:12:57Z

+                  // qudaStreamWaitEvent(packEnd) on that stream, so the P2P copy serializes
+                  // with the pack output but runs in parallel with the other directions.
+                  PROFILE(if (dslash_comms)
+                            halo.sendStartStream(2 * i + dir, device::get_stream(2 * i + dir), dslashParam.remote_write),


Yeah I think renaming this to sendStartStreamGated (or something like that) would make it clear that this function is not general purpose.

hummingtree

Thanks @bjoo for the great work - it looks nice and neat to me. I have left a few comments out of my curiosity.

- try_import_fabric_handle: keep the fabric-import probe non-fatal (failed import -> downgrade link to MPI) but warn on an *unexpected* CUresult and clear the latched error, so a real misconfiguration is not masked as a benign fallback. - Add CHECK_CU_DRIVER macro (mirrors CHECK_CUDA_ERROR via set_driver_error) and use it for the six context-free driver calls; per-(dim,dir) sites keep their explicit errorQuda. - assert_stream_gated_resolved() guard in the three STREAM_GATED enum dispatchers: error if STREAM_GATED is requested but is not the resolved P2P transport. - Rename sendStartStream -> sendStartStreamGated to reflect that it is the stream-gated send path.

bjoo · 2026-07-12T14:31:58Z

Pushed some cleanup changes suggested by @hummingtree .

bjoo and others added 13 commits July 7, 2026 03:58

MNNVL P2P (1/7): add QUDA_MNNVL build flag and comm_p2p_* signalling …

44a1a51

…abstraction Introduce the default-off QUDA_MNNVL CMake option and refactor the existing P2P halo doorbell into a comm_p2p_* abstraction. Pure refactor: no functional change when QUDA_MNNVL=OFF.

MNNVL P2P (2/7): per-kind device comm-buffer allocator

2b8d8c3

Add a tag-dispatched comm-buffer allocator (DeviceCommBuffer kind) backed by cuMemAlloc with tracker tags, and remove the dead pinned-malloc / MPI-kind machinery it replaces.

MNNVL P2P (4/7): wire stream-gated P2P as a dslash sub-policy

53baf93

Expose stream-gated as a P2P sub-policy on tp.aux.y and wire it through DslashFusedExterior/FusedGDR and GaugeField::exchangeGhost (4-phase halo). Default the P2P signalling transport to stream_gated.

MNNVL P2P: apply clang-format

e01924a

CI Cleanups recommended by Evan

c098b48

applied clang format

08af2d3

bjoo requested review from a team as code owners July 8, 2026 08:17

bjoo added 3 commits July 8, 2026 04:45

Satisfy CI 'value may be used unitialized' error

5e43f98

Fixed CI unused variable issues

571c6df

Fix missing ROCm stub picked up by CI

e4acc54

hummingtree reviewed Jul 9, 2026

View reviewed changes

Comment thread lib/targets/cuda/malloc.cpp

hummingtree reviewed Jul 9, 2026

View reviewed changes

Comment thread lib/color_spinor_field.cpp

hummingtree reviewed Jul 9, 2026

View reviewed changes

Comment thread lib/gauge_field.cpp

hummingtree approved these changes Jul 9, 2026

View reviewed changes

bjoo added 2 commits July 12, 2026 10:28

Apply clang-format (clang-format 11)

a2572a0

Uh oh!

Conversation

bjoo commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Description

Implementation notes

Uh oh!

bjoo commented Jul 8, 2026

Uh oh!

hummingtree Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

hummingtree Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

hummingtree left a comment

Choose a reason for hiding this comment

Uh oh!

bjoo commented Jul 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

bjoo commented Jul 8, 2026 •

edited

Loading

hummingtree Jul 9, 2026 •

edited

Loading

hummingtree Jul 9, 2026 •

edited

Loading