B200 Slow GEMM (16384, 1024, 8192) Only when run in a particular order

Describe the bug
GEMM Shape (16384, 1024, 8192) is slow on TE.Linear but only when it runs as after a second kernel even though I set a 20 second cooldown between the first & second GEMM.

If I try just 1 GEMM (16384, 1024, 8192) then it gives a more reasonable TFLOP/s.

If i try multiple GEMMs without (16384, 1024, 8192) then it gives reasonable TFLOP/s for all the GEMMs too.

At first i thought it was thermal throttling but with a 20 second cooldown between GEMMs, it doesn't seem like the case.

I have ran the 2 GEMM reprod script multiple times too & can confirm it gives ~550 TFLOP/s every time on my environment.

(16384, 1024, 8192) is a real world training shape for llama2 70B TP=8

It could be possible that this is just a bug on my HGX B200 environment but I do not currently have access to another B200 environment to test it.

Run 2 GEMMs with 20 second cooldown between each GEMM

Benchmark results for Realistic GEMM shapes with warmup=30 and repeats=200
+---------------------+----------------------------------------+
| Shape (M, N, K)     | FP8 TE.Linear (autocast, bias=False)   |
+=====================+========================================+
| (8192, 8192, 8192)  | 2377.1 TFLOPS                          |
+---------------------+----------------------------------------+
| (16384, 1024, 8192) | 551.0 TFLOPS                           |
+---------------------+----------------------------------------+

Run Just 1 GEMM

Benchmark results for Realistic GEMM shapes with warmup=30 and repeats=200
+---------------------+----------------------------------------+
| Shape (M, N, K)     | FP8 TE.Linear (autocast, bias=False)   |
+=====================+========================================+
| (16384, 1024, 8192) | 1226.9 TFLOPS                          |
+---------------------+----------------------------------------+

Run Multiple GEMMs but doesn't include (16384, 1024, 8192)

Benchmark results for Realistic GEMM shapes with warmup=30 and repeats=200
+---------------------+----------------------------------------+
| Shape (M, N, K)     | FP8 TE.Linear (autocast, bias=False)   |
+=====================+========================================+
| (16384, 8192, 1280) | 1705.6 TFLOPS                          |
+---------------------+----------------------------------------+
| (16384, 8192, 7168) | 2513.5 TFLOPS                          |
+---------------------+----------------------------------------+
| (16384, 3584, 8192) | 1779.8 TFLOPS                          |
+---------------------+----------------------------------------+

Steps/Code to reproduce bug

Command

 docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -it --rm nvcr.io/nvidia/pytorch:25.04-py3

Reprod Script

import time
import torch
import tabulate
from triton.testing import do_bench
try:
    import transformer_engine.pytorch as te
    from transformer_engine.common import recipe
    import_te = True
except:
    import_te = False

torch.manual_seed(0)
repeats = 200
warmup = 30
timeout = 20
is_nvidia = "nvidia" in torch.cuda.get_device_name(0).lower()

device = 'cuda'
dtype_bf16 = torch.bfloat16

# GEMM Shapes
shapes = [
    (8192,8192,8192),
    (16384, 1024, 8192),
]

results = []

# FP8 Recipe (for scaling)
if import_te:
    fp8_format = recipe.Format.HYBRID
    fp8_recipe = recipe.DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo="max")

for (m, n, k) in shapes:
    # FLOPS
    nFLOPS = 2 * m * n * k

    # TE Linear (with FP8 autocast) benchmark
    if import_te:
        input_tensor = torch.randn(m, k, device=device)
        linear_layer = te.Linear(k, n, bias=False).to(device)
        with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
            ms_te_linear = do_bench(lambda: linear_layer(input_tensor), warmup=warmup, rep=repeats)
        tflops_te_linear = nFLOPS / ms_te_linear * 1e-9
        time.sleep(timeout)
    else:
        tflops_te_linear = 0.0

    results.append([
        f"({m}, {n}, {k})",
        f"{tflops_te_linear:.1f} TFLOPS",
    ])

# Print results
headers = [
    "Shape (M, N, K)",
    "FP8 TE.Linear (autocast, bias=False)",
]
print(f"Benchmark results for Realistic GEMM shapes with {warmup=} and {repeats=}")
print(tabulate.tabulate(results, headers=headers, tablefmt="grid"))

Expected behavior

for TFLOP/s to be performant on this shape regardless of the order of GEMMs

Environment overview (please complete the following information)

• Environment location: NGC Container inside DataCrunch VM
• Method of Transformer Engine install: NGC docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -it --rm nvcr.io/nvidia/pytorch:25.04-py3

Environment details
Driver Version: 570.133.20
Device details

• HGX NVIDIA 8xB200 SXM 1000W
• vboost=default

root@ad1fa1f4d57a:/workspace# nvidia-smi boost-slider -l
+-------------------------------------------------+
| GPU Boost Slider                                |
| GPU     Slider       Max Value    Current Value |
|=================================================|
|   0     vboost           4              0       |
+-------------------------------------------------+
|   1     vboost           4              0       |
+-------------------------------------------------+
|   2     vboost           4              0       |
+-------------------------------------------------+
|   3     vboost           4              0       |
+-------------------------------------------------+
|   4     vboost           4              0       |
+-------------------------------------------------+
|   5     vboost           4              0       |
+-------------------------------------------------+
|   6     vboost           4              0       |
+-------------------------------------------------+
|   7     vboost           4              0       |
+-------------------------------------------------+

Run the supplied Python reproduction in the NGC PyTorch 25.04 container on an HGX B200, starting with the two-shape benchmark around te.Linear and FP8 autocast. Compare the reported TFLOP/s when the (16384, 1024, 8192) GEMM runs first versus after another GEMM, while retaining the stated warmup, repeats, and cooldown. Done means the ordering-dependent slowdown is reproduced and its Transformer Engine entry point is identified.