perf(ep): optimize InterNodeV1LL small-token latency on EP16 - #613
Draft
isytwu wants to merge 2 commits into
Draft
perf(ep): optimize InterNodeV1LL small-token latency on EP16#613isytwu wants to merge 2 commits into
isytwu wants to merge 2 commits into
Conversation
torch.cuda.current_stream() re-resolves the device and builds a Stream object every call (~4.9us measured vs ~0.16us for the raw binding it wraps). _launch_multi and _resolve_launch_params each ran an import statement per call; both moved to module level. At small token counts the host submission path is what paces the GPU, so this is latency, not bookkeeping.
Gathering a token reads from up to numExpertPerToken peer GPUs over xGMI, where latency rather than bandwidth is the cap. WarpAccumLF issues AccumNum*Unroll of those reads before accumulating any, so they overlap; WarpAccum keeps only AccumNum in flight and moves 4B/lane. The intra-node combine path has used the 16B load-first form for a while, the v1 internode path had not. Falls back to the 4B path when the staging stride is not 16B-aligned (CombineVecAligned), and sizes each warp's slice to a whole vector step instead of the old fixed warpsPerToken=4, since a slice shorter than one step is slower than not vectorizing at all. Measured: EpCombineInterNodeV1KernelLowLatency mean 56.9 -> 46.2us (-19%) on EP16 at 4 tokens, hidden 6144.
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
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.
Motivation
Reduce host-side overhead and xGMI gather latency for small-token InterNodeV1LL workloads on EP16/MI300X, while improving tuning reliability.
Technical Details
Test Plan
Test Result
Submission Checklist