Skip to content

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

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

perf(ep): optimize InterNodeV1LL small-token latency on EP16#613
isytwu wants to merge 9 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 6 commits August 31, 2026 11:40
@isytwu
perf(ep): trim per-call host overhead on the dispatch/combine path
8ee2580 
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.
@isytwu
perf(ep): use the 16B load-first gather in the v1_ll combine
e147684 
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.
@isytwu
fix(profiler): use the intranode_ll slot enum in intranode_ll.hpp
464dfde 
The generator groups slots by source file, so this file's slots live in
IntranodeLlSlot; initializing the intranode context made Slot:: resolve to
IntranodeSlot and every ENABLE_PROFILER=ON build failed to compile.
@isytwu
perf(profiler): decode trace buffers with numpy
37ea75a 
_parse_trace_events walked the full worst-case buffer (16384 events x 4096
warps = 134M int64) one .item() at a time, 141s per rank. Same output, 0.07s.
@isytwu
tools(profiler): aggregate EP kernel traces across ranks
ce03c19 
analyze_ep_kernel_trace.py shows one rank's timeline; this rolls all ranks up
into per-phase duration and share of the dispatch/combine window.
@isytwu
tune(ep): add mi308x InterNodeV1LL ep16 rules for hidden 6144
4e0577e 
Covers 4/8/16/32 tokens. Only mi300x carried this shape, so AUTO found no
rule on mi308x and fell back to the hard-coded (256, 128, 8).
@isytwu
isytwu force-pushed the perf/ep16-kernel-opt-final branch from 51cfbf7 to 4e0577e Compare August 31, 2026 03:47
@isytwu

isytwu commented Aug 31, 2026

Copy link
Copy Markdown
Collaborator Author
image

isytwu and others added 3 commits August 31, 2026 08:29
@isytwu
fix(ep): get the real stream pointer for kernel launch, not its packe…
8a1be90 
…d id

_current_stream() called torch._C._cuda_getCurrentStream(device)[0] to
avoid building a torch.cuda.Stream object (~4.8us) on every dispatch/combine
call. That call does not return a stream pointer despite what its result
was documented as here: it returns CUDAStream's packed stream_id (pool
index + per-pool index + priority folded into an int), not the
cudaStream_t/hipStream_t address that _launch's hipModuleLaunchKernel
needs.

This went unnoticed because the default stream's packed id happens to be 0,
which coincides with the null-stream sentinel HIP already treats as "the
current stream" -- so ordinary (non-graph) dispatch/combine calls kept
working by accident. Entering torch.cuda.graph() switches capture onto a
real non-default stream, whose packed id is a small nonzero int (e.g. 3).
Passed to hipModuleLaunchKernel as if it were a stream pointer, that value
addresses unmapped memory: the launch goes nowhere, capture_end() reports
"The CUDA Graph is empty", and later using the captured (empty) graph
corrupts the context (HIP error 709, hipErrorContextIsDestroyed).

Reproduced with:
  PYTHONPATH=$(pwd)/python:$(pwd) python3 \
    tests/python/ops/bench_dispatch_combine.py --world-size 8 --cmd bench
whose default path captures dispatch/combine into CUDA graphs
(_capture_split_graphs). Verified torch._C._cuda_getCurrentRawStream(device)
returns the same address as the previously-correct
torch.cuda.current_stream().cuda_stream in both the default-stream and
active-graph-capture cases, and is still ~8x cheaper (~0.6us) than
reconstructing the Stream wrapper -- so this keeps the intended host-latency
win while launching kernels on the stream that was actually asked for. The
same bench command now runs to completion (10 rounds, dispatch/combine/e2e
all report).
@isytwu
tune(ep): re-tune MI300X EP16 v1_ll at hidden 6144, tokens 4/8/16/32
a9116d0 
Full sweep on the current kernel (16B load-first combine gather + the
hierarchical dispatch grid barrier), same hardware/topology/config as the
rules being replaced: EP16, gfx942 MI300X, fp8_e4m3_fnuz->bf16, num_qp=1,
MORI_RDMA_TC=41.

  dispatch: 48.42 -> 40.70us (4 tok, -15.9%), 48.52 -> 42.17us (8 tok, -13.1%),
            49.84 -> 43.20us (16 tok, -13.3%), 53.54 -> 46.58us (32 tok, -13.0%)
  combine:  70.32 -> 57.54us (4 tok, -18.2%), 69.57 -> 58.04us (8 tok, -16.6%),
            69.93 -> 60.95us (16 tok, -12.8%), 75.33 -> 68.50us (32 tok, -9.1%)

Unaffected by the stream-pointer fix in the preceding commit: internode
--cmd tuning runs through run_bench_once, which does not use
torch.cuda.graph() (the only CUDA Graph capture in this file is under
--cmd stress), so these numbers were never exposed to that bug.
@jhchouuu @claude
feat(bench): socket bootstrap for shmem p2p benchmarks (no MPI)
2fb0f87 
PerfInit uses a socket bootstrap (SocketBootstrapNetwork + ShmemInit)
when MASTER_ADDR is set, mirroring the EP tests' RANK/WORLD_SIZE/
MASTER_ADDR launch. This lets p2p_{put,get}_{latency,bw} run across two
nodes with one process per node (no mpirun / MPI orchestration), which
is what internode IBGDA latency measurement needs. Falls back to
MPI_Init when MASTER_ADDR is unset; PerfFinalize guards MPI_Comm_free
for the socket path (local_comm == MPI_COMM_NULL).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@isytwu @jhchouuu