[JAX] Synchronize tutorial speedometer measurements

docs/examples/jax/attention.out

@@ -1,9 +1,9 @@
 # SINGLE_GPU_OUTPUT_START
 Native JAX bf16 GQA + SWA:
-Mean time: 5.109810829162598 ms
+Mean time: 51.31755525362678 ms
 
 TE DotProductAttention GQA + SWA:
-Mean time: 0.09856224060058594 ms
+Mean time: 1.1168592551257461 ms
 # SINGLE_GPU_OUTPUT_END
 
 # MLA_OUTPUT_START

docs/examples/jax/attention_context_parallel.out

@@ -1,16 +1,16 @@
 # SINGLE_GPU_OUTPUT_START
 Single-GPU THD GQA + SWA:
-Mean time: 126.68747901916504 ms
+Mean time: 124.06466079410166 ms
 # SINGLE_GPU_OUTPUT_END
 
 # RING_OUTPUT_START
 THD CP Ring stripe_size=1:
-Mean time: 57.16729164123535 ms
-Speedup vs single GPU: 2.22x
+Mean time: 52.09881339687854 ms
+Speedup vs single GPU: 2.38x
 # RING_OUTPUT_END
 
 # AG_OUTPUT_START
 THD CP AllGather stripe_size=512:
-Mean time: 53.79219055175781 ms
-Speedup vs single GPU: 2.36x
+Mean time: 53.20774579886347 ms
+Speedup vs single GPU: 2.33x
 # AG_OUTPUT_END

docs/examples/jax/attention_context_parallel.py

@@ -346,11 +346,11 @@ def run_single_gpu_bench():
         result = grad_fn((q, k, v), sequence_descriptor, dout)
     jax.block_until_ready(result)
 
-    start = time.time()
+    start = time.perf_counter()
     for _ in range(timing_iters):
         result = grad_fn((q, k, v), sequence_descriptor, dout)
     jax.block_until_ready(result)
-    mean_ms = (time.time() - start) * 1000 / timing_iters
+    mean_ms = (time.perf_counter() - start) * 1000 / timing_iters
     print(f"Mean time: {mean_ms} ms")
     return mean_ms
 
@@ -374,15 +374,15 @@ def run_context_parallel_bench(
             )
         jax.block_until_ready(result)
 
-        start = time.time()
+        start = time.perf_counter()
         for _ in range(timing_iters):
             result = grad_fn(
                 sharded["qkv"],
                 sharded["sequence_descriptor"],
                 sharded["dout"],
             )
         jax.block_until_ready(result)
-        end = time.time()
+        end = time.perf_counter()
 
     mean_ms = (end - start) * 1000 / timing_iters
     print(f"Mean time: {mean_ms} ms")

docs/examples/jax/attention_context_parallel.rst

@@ -162,7 +162,7 @@ while keeping those sharding controls visible.
       :start-after: # AG_OUTPUT_START
       :end-before: # AG_OUTPUT_END
 
-On four GB200s, Ring is roughly **2.22x faster** and AllGather roughly **2.36x
+On four GB200s, Ring is roughly **2.38x faster** and AllGather roughly **2.33x
 faster** than the equivalent single-GPU fused-attention forward+backward pass.
 These results are specific to this workload and system. Applications with long
 segments and wide attention windows generally have more attention computation

docs/examples/jax/attention_single_gpu.rst

@@ -80,7 +80,7 @@ implementations.
       :start-after: # SINGLE_GPU_OUTPUT_START
       :end-before: # SINGLE_GPU_OUTPUT_END
 
-On a single GB200, this run is roughly **52x faster** for the fwd+bwd of this
+On a single GB200, this run is roughly **46x faster** for the fwd+bwd of this
 BSHD GQA + SWA example. This compares TE ``DotProductAttention`` against the
 native JAX baseline above, which materializes attention scores with XLA ops; it
 is not a comparison against ``jax.nn.dot_product_attention(...,

docs/examples/jax/dense.out

@@ -6,16 +6,16 @@ Variable collections: ['params']
 {'params': {'Dense_0': {'kernel': ((8192, 32768), dtype('float32'))}}}
 
 bf16 baseline:
-Mean time: 18.056 ms
+Mean time: 18.265029820613563 ms
 
 TE MXFP8BlockScaling:
-Mean time: 11.260 ms
+Mean time: 10.855670000892133 ms
 # SINGLE_GPU_OUTPUT_END
 
 # MULTI_GPU_OUTPUT_START
 bf16 DP=2/TP=2:
-Mean time: 5.516 ms
+Mean time: 5.609957461711019 ms
 
 TE MXFP8BlockScaling DP=2/TP=2:
-Mean time: 3.712 ms
+Mean time: 3.6635108408518136 ms
 # MULTI_GPU_OUTPUT_END

docs/examples/jax/dense.rst

@@ -75,7 +75,7 @@ same input for both models.
       :start-after: # SINGLE_GPU_OUTPUT_START
       :end-before: # SINGLE_GPU_OUTPUT_END
 
-On a single GB200, that's roughly **1.6× faster** for the fwd+bwd of one large
+On a single GB200, that's roughly **1.7× faster** for the fwd+bwd of one large
 Dense — and the only code change was passing ``dot_general=te_dot_general_cls()``
 into ``nn.Dense``.
 

docs/examples/jax/quickstart_jax_utils.py

@@ -23,8 +23,10 @@ def speedometer(
     warmup_iters: int = 50,
     rngs: Dict[str, jax.random.PRNGKey] = None,
 ) -> None:
-    """Measure average runtime for a JAX module
-    Perform forward and backward passes .
+    """Measure average forward/backward runtime for a JAX module.
+
+    Wait for warmup and timed results so asynchronous dispatch does not skew
+    the elapsed wall time.
     """
     if autocast_kwargs is None:
         autocast_kwargs = {"enabled": False}
@@ -36,16 +38,19 @@ def speedometer(
     train_step_fn = create_train_step_fn(model_apply_fn, autocast_kwargs, forward_kwargs)
 
     # Warm up runs
+    result = None
     for _ in range(warmup_iters):
         rngs, step_rngs = _split_step_rngs(rngs)
-        loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_rngs)
+        result = train_step_fn(variables, input, output_grad, step_rngs)
+    jax.block_until_ready(result)
 
     # Timing runs
-    start = time.time()
+    start = time.perf_counter()
     for _ in range(timing_iters):
         rngs, step_rngs = _split_step_rngs(rngs)
-        loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_rngs)
-    end = time.time()
+        result = train_step_fn(variables, input, output_grad, step_rngs)
+    jax.block_until_ready(result)
+    end = time.perf_counter()
 
     print(f"Mean time: {(end - start) * 1000 / timing_iters} ms")