Add host wrapper for AllToAllv collective #401
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Summary: Simplify the ChunkState API by: 1. Renaming methods from camelCase to snake_case for consistency with the rest of the pipes codebase (e.g., waitReadyToSend -> wait_ready_to_send) 2. Requiring ThreadGroup parameter for all public methods, removing the single-threaded overloads that were error-prone The ThreadGroup abstraction ensures proper synchronization: - For signals (ready_to_recv, ready_to_send): sync before leader writes - For waits (wait_ready_to_send, wait_ready_to_recv): all threads poll for better latency This is a breaking change to the ChunkState API but the only callers are within the pipes library (P2pNvlTransportDevice, P2pSyncBench). Differential Revision: D91412047
Summary: Add optional timeout support to ChunkState and SignalState wait methods that can block indefinitely on P2P NVLink I/O operations. When a timeout occurs, the kernel aborts via __trap(). Key changes: - Add TimeoutUtils.cuh with Timeout struct for efficient GPU-side timeout checking using clock64() and precomputed deadline cycles - Add TimeoutUtils.h with makeTimeout() host helper that queries GPU clock rate and converts milliseconds to cycles - Add optional Timeout parameter to ChunkState::wait_ready_to_send(), wait_ready_to_recv() and SignalState::wait_until() - Propagate Timeout parameters through P2pNvlTransportDevice methods: send(), recv(), send_one(), recv_one(), send_multiple(), recv_multiple() - Add TimeoutTrapTest with tests verifying timeout behavior Design decisions: - Timeout unit is milliseconds (intuitive for debugging) - Default timeout_ms=0 means infinite wait (backward compatible) - start() must be called once at kernel entry to capture reference time - check() is called in polling loops with leader-only optimization for ThreadGroup-based waits Differential Revision: D91412046
Differential Revision: D91439414
Summary: Integrate timeout support into the all_to_allv collective communication primitive. This allows the collective to detect and abort when peer communication hangs, rather than blocking indefinitely. Key changes: - Add required Timeout parameter to all_to_allv() function signature - Call timeout.start() at the beginning to initialize the deadline - Propagate timeout to P2pNvlTransportDevice send() and recv() calls - Update CollectiveBenchmark kernel to accept and pass timeout - Update AllToAllvBenchmark to create and pass Timeout to kernel This builds on the timeout infrastructure added in D91412046. Differential Revision: D91439413
Summary: Add a host-side wrapper function `launch_all_to_allv()` for the AllToAllv collective, similar to the dispatch host wrapper. This eliminates the need for a separate kernel wrapper in the benchmark. Key changes: - Add AllToAllv.h with host wrapper declaration - Add AllToAllv.cu with kernel and host wrapper implementation - Support optional cluster launch with spread scheduling for better load balancing across GPCs - Update AllToAllvBenchmark to use the new host wrapper instead of directly launching the kernel - Fix AllToAllvTest.cu to pass Timeout parameter to device function The host wrapper is named `launch_all_to_allv()` to distinguish it from the device function `all_to_allv()`. Differential Revision: D91440692
meta-cla
bot
added
the
CLA Signed
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
Add a host-side wrapper function
launch_all_to_allv()for the AllToAllvcollective, similar to the dispatch host wrapper. This eliminates the need
for a separate kernel wrapper in the benchmark.
Key changes:
load balancing across GPCs
directly launching the kernel
The host wrapper is named
launch_all_to_allv()to distinguish it fromthe device function
all_to_allv().Differential Revision: D91440692