@@ -32,18 +32,16 @@ def sample():
3232 return [sample () for _ in range (iters )]
3333
3434
35- def time_cuda_graph_phase_us (torch , fn , warmup : int , iters : int , interval , align ) -> list [float ]:
35+ def time_cuda_graph_phase_us (torch , fn , warmup : int , iters : int , interval ) -> list [float ]:
3636 """Time one event-record interval captured inside graph replay.
3737
38- `align()` enqueues a device-side rank barrier before each replay, so replays start together
39- rather than ~75us apart (b200 EP16), which the cross-rank MAX would report as latency.
38+ The graph contains an all-reduce before the timing events to reduce rank entry skew.
4039 """
4140 for _ in range (max (0 , warmup )):
4241 fn ()
4342 torch .cuda .synchronize ()
4443 samples = []
4544 for _ in range (iters ):
46- align ()
4745 fn ()
4846 torch .cuda .synchronize ()
4947 samples .append (interval [0 ].elapsed_time (interval [1 ]) * 1000.0 )
@@ -223,53 +221,24 @@ def poison(tensor):
223221
224222
225223class GraphTiming :
226- """Captured dispatch -> combine pairs timed by event nodes inside the graph, each timed replay
227- started behind a device-side rank barrier so the cross-rank MAX is the operation, not launch
228- skew."""
224+ """Measure dispatch/combine latency with CUDA events inside a replayed CUDA graph.
225+
226+ Each replay first runs an all-reduce to bring ranks closer together before timed work
227+ begins. The all-reduce is part of the graph but precedes the timing events, so its cost
228+ is excluded from the reported latency.
229+ """
229230
230231 graph = True
231232 # Handle attributes dispatch writes; the replay value check poisons them.
232233 DISPATCH_OUTPUT_FIELDS = ("recv_x" , "recv_scales" , "dispatch_output" )
233234
234235 def __init__ (self , backend ):
235236 self .backend = backend
236- self .align_cycles = None
237237 self ._align_token = None
238238
239239 def components (self ):
240240 return ["roundtrip" , "dispatch" , "combine" ]
241241
242- def calibrate_align_spin (self , spin_us = 100.0 ):
243- """Size the post-barrier spin to `spin_us` of wall time on this GPU.
244-
245- The spin lets every host enqueue its replay before the stream reaches it. `_sleep` counts
246- SM cycles, so a fixed count left ~15us of skew between differently clocked gb200 ranks.
247- """
248- import torch
249-
250- probe = 200_000
251- start = torch .cuda .Event (enable_timing = True )
252- end = torch .cuda .Event (enable_timing = True )
253- torch .cuda ._sleep (probe // 10 ) # ramp clocks before the measured spin
254- start .record ()
255- torch .cuda ._sleep (probe )
256- end .record ()
257- torch .cuda .synchronize ()
258- elapsed_us = max (start .elapsed_time (end ) * 1000.0 , 1e-3 )
259- self .align_cycles = max (1 , int (probe * spin_us / elapsed_us ))
260-
261- def align (self ):
262- """Enqueue a device-side rank barrier on the current stream, without a host sync."""
263- import torch
264- import torch .distributed as dist
265-
266- if self ._align_token is None :
267- self ._align_token = torch .zeros (1 , device = self .backend .device )
268- if self .align_cycles is None :
269- self .calibrate_align_spin ()
270- dist .all_reduce (self ._align_token )
271- torch .cuda ._sleep (self .align_cycles )
272-
273242 def capture_pairs (self , problem , staged , pairs , marks ):
274243 """Capture `pairs` back-to-back pairs into one graph; `marks` picks the windows that get
275244 event nodes ("pair", "dispatch", "combine"). Returns (graph, {mark: (starts, ends)},
@@ -278,6 +247,9 @@ def capture_pairs(self, problem, staged, pairs, marks):
278247 import torch .distributed as dist
279248
280249 b = self .backend
250+ if self ._align_token is None :
251+ self ._align_token = torch .zeros (1 , device = b .device )
252+ dist .all_reduce (self ._align_token )
281253 stamps = {mark : ([graph_event () for _ in range (pairs )], [graph_event () for _ in range (pairs )])
282254 for mark in marks }
283255
@@ -290,6 +262,14 @@ def record(mark, edge, i):
290262 graph = torch .cuda .CUDAGraph ()
291263 combined = handle = None
292264 with torch .cuda .graph (graph , capture_error_mode = "relaxed" ):
265+ # A separate all-reduce followed by a host-launched graph can let CPU launch
266+ # delays spread the ranks apart again, inflating the measured communication time.
267+ # Capture the all-reduce and its stream dependencies in this graph so each replay
268+ # reaches the timed work without another host launch after alignment. Place it
269+ # before the first timing event to exclude alignment cost, and only once per graph
270+ # to preserve back-to-back dispatch/combine pairs. This reduces entry skew; it
271+ # does not guarantee that every rank starts the timed work simultaneously.
272+ dist .all_reduce (self ._align_token )
293273 for i in range (pairs ):
294274 record ("pair" , 0 , i )
295275 record ("dispatch" , 0 , i )
@@ -313,10 +293,9 @@ def component(self, name, problem, warmup, iters):
313293 staged = self .backend .warm_and_hoist_stage (problem , warmup )
314294 mark = "pair" if name == "roundtrip" else name
315295 graph , stamps , combined , _ = self .capture_pairs (problem , staged , 1 , (mark ,))
316- self .calibrate_align_spin () # clocks move with load and temperature: per timed series
317296 starts , ends = stamps [mark ]
318297 samples = time_cuda_graph_phase_us (
319- torch , graph .replay , warmup , iters , (starts [0 ], ends [0 ]), align = self . align
298+ torch , graph .replay , warmup , iters , (starts [0 ], ends [0 ])
320299 )
321300 combined .fill_ (float ("nan" ))
322301 torch .cuda .synchronize ()
@@ -329,16 +308,14 @@ def component(self, name, problem, warmup, iters):
329308
330309 def chain (self , problem , staged , iters , drop ):
331310 """Each sibling is ONE graph of `iters` unrolled pairs, the shape of a decode graph, replayed
332- once untimed (upload) and once aligned and timed ."""
311+ once untimed (upload) and once timed, with alignment captured before the first event ."""
333312 import torch
334313
335314 floors , floor_stamps , _ , _ = self .capture_pairs (problem , staged , iters , ("dispatch" , "combine" ))
336315 period , period_stamps , combined , _ = self .capture_pairs (problem , staged , iters , ("pair" ,))
337- self .calibrate_align_spin ()
338316 for graph in (floors , period ):
339317 graph .replay ()
340318 torch .cuda .synchronize ()
341- self .align ()
342319 graph .replay ()
343320 torch .cuda .synchronize ()
344321 return _chain_result (period_stamps ["pair" ], floor_stamps , drop , combined )
0 commit comments