nccl-ep: fix HT LSA barrier allocation for runtime CTA counts - #14
Open
EvanZhang118 wants to merge 1 commit into
Open
nccl-ep: fix HT LSA barrier allocation for runtime CTA counts#14EvanZhang118 wants to merge 1 commit into
EvanZhang118 wants to merge 1 commit into
Conversation
HT dispatch accepts a runtime communication grid size through max_num_sms and NCCL_EP_COMM_SMS. Each dispatch CTA uses blockIdx.x as its LSA barrier session ID, but DevComm creation allocates only the fixed default count plus one. Grids larger than the 16-CTA default can therefore access sessions outside the allocated range and hang with the synchronization guard enabled. Allocate barrier sessions from the resolved communication SM count in both DevComm creation paths. Use the JIT-specialized block count for the combine tail so its session remains disjoint from the dispatch CTA sessions. Signed-off-by: Evan <91120837+EvanZhang118@users.noreply.github.com>
EvanZhang118
changed the title
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
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.
Problem
HT dispatch resolves its communication grid size at runtime from
ncclEpGroupConfig_t::max_num_smsorNCCL_EP_COMM_SMS. Each dispatch CTA usesblockIdx.xas its LSA barrier session ID when the HT synchronizationguard is enabled.
The DevComm requirements instead allocate a fixed
NCCL_EP_HT_DISPATCH_BLOCKS + 1sessions, which is 17 with the default configuration. A runtime grid larger than 16 CTAs therefore addresses LSA barrier sessions outside the allocated range and can hang. The combine-tail barrier also uses the default dispatch block count as its session ID rather than the JIT-specialized grid size.Fix
lsaBarrierCountfrom the resolved communication SM count in both the LSA-only and cross-LSA DevComm creation paths.NBLOCKSvalue for the combine-tail barrier session, keeping it disjoint from dispatch sessions[0, NBLOCKS).No public API or default launch configuration changes.