[Common] Comm+GEMM overlap API updated to support cuBlasMp backend (incl. framework API) #2443
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
…rk extensions Signed-off-by: Alp Dener <[email protected]>
denera
force-pushed
the
common/tp-overlap-cublasmp
branch
from
9c8fe95 to
334e8c4
Compare
…entirely Signed-off-by: Alp Dener <[email protected]>
denera
force-pushed
the
common/tp-overlap-cublasmp
branch
from
908bbc2 to
69cf235
Compare
for more information, see https://pre-commit.ci
ptrendx
reviewed
Dec 2, 2025
|
|
||
| CommOverlapCore::CommOverlapCore(int64_t nccl_comm_ptr, int tp_rank, int tp_size, int num_comm_sm, | ||
| bool is_p2p, bool atomic_gemm) { | ||
| _with_cublasmp = true; |
Member
There was a problem hiding this comment.
This should probably check if TE was built with cublasmp and error out otherwise.
ptrendx
reviewed
Dec 2, 2025
|
|
||
| #define NVTE_COMM_OVERLAP_MAX_STREAMS 3 | ||
|
|
||
| /* \brief Check if TE is built with cuBlasMp. |
ptrendx
reviewed
Dec 2, 2025
| void ub_barrier(ExtComm comm); | ||
|
|
||
| int64_t get_nccl_comm_ptr(std::string comm_name) { | ||
| NVTE_CHECK(backend_is_nccl, "Cannot get nccComm_t ptr if backend is not NCCL."); |
Member
There was a problem hiding this comment.
This error message could be more descriptive - e.g. something like "chosen backend for the communication-computation overlap (cuBLASMp) requires NCCL communicator, but the passed ProcessGroup uses a different backend."
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Description
This PR adds support for the NVTE cuBlasMp bindings in the Comm+GEMM overlap API.
Type of change
Checklist: