Skip to content

perf(ep): optimize InterNodeV1LL small-token latency on EP16 - #613

Draft
isytwu wants to merge 2 commits into
mainfrom
perf/ep16-kernel-opt-final
Draft

perf(ep): optimize InterNodeV1LL small-token latency on EP16#613
isytwu wants to merge 2 commits into
mainfrom
perf/ep16-kernel-opt-final

Conversation

@isytwu

@isytwu isytwu commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Reduce host-side overhead and xGMI gather latency for small-token InterNodeV1LL workloads on EP16/MI300X, while improving tuning reliability.

Technical Details

  • Cache persistent dispatch/combine tensor views.
  • Reduce stream lookup and repeated Python import overhead.
  • Use a configurable relative tuning margin, median-based scoring, and deterministic tie-breaking.
  • Add configurable low-latency benchmark shapes and tuned EP16 MI300X launch configurations.
  • Use a 16-byte load-first gather to overlap xGMI reads.
  • Fall back to the existing 4-byte path when alignment requirements are not met.
  • Partition hidden dimensions in whole vector steps to avoid expensive scalar tails.

Test Plan

  • Run EP16 correctness tests on 2×8 MI300X GPUs with 4, 8, and 16 tokens.
  • Test hidden dimensions 1024, 4096, 6144, and 7168.
  • Exercise multiple block, warp, and RDMA-block configurations.
  • Verify eager and HIP graph replay results.
  • Compare host enqueue, end-to-end, and Kineto kernel latency.

Test Result

  • All tested ranks completed with zero correctness errors.
  • HIP graph replay was bit-exact with eager execution.
  • Cached tensor views reduced 4-token end-to-end latency from 214 µs to 157 µs.
  • Raw current-stream lookup measured approximately 0.16 µs versus 4.9 µs previously.
  • The InterNodeV1LL combine kernel improved from 56.9 µs to 46.2 µs at 4 tokens and hidden size 6144, a 19% reduction.

Submission Checklist

isytwu added 2 commits August 28, 2026 10:14
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant