A MoE combine fused with the next RMSNorm gets a slow one-block-per-token kernel (bf16)

Aperta
  • Qualcosa non va?
#5,042 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 3 giorni

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
22/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python
Ambito
compilers, performance

Direzione di ricerca

Start by running python repro_combine_norm.py to reproduce the three engine timings and kernel lines: combine 273 us, norm 44 us, fused 392 us. The fusion decision lives in the TensorRT/Torch-TensorRT lowering and kernel selection, not in the Python model code, so the reported kernel names (__myl_...) and launch configs are the main clues. Done means the fused engine costs no more than the two engines run apart, or the root cause is shown to be outside this repository.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

What is wrong

In a mixture-of-experts (MoE) transformer, each token adds up its top-k expert outputs (the "combine"). The next block then starts with a residual add and an RMSNorm. TensorRT fuses all of this into one kernel.

Alone, the combine gets a fast kernel with one thread per element. Once the RMSNorm (a per-token sum) is fused in, the kernel switches to one 128-thread block per token, at about 44% occupancy. The fused kernel is slower than the two kernels run apart.

Who hits it: MoE models that have a norm right after the expert sum.

Related: #5008 (bf16 reduce over the top-k axis is slow, the other way this combine is often written).

What this blocks

No policy fails because of this. It costs speed.
A mixture-of-experts model runs this step in every MoE layer on every call, so the slowdown below adds up across layers.
A fix makes such models faster with no change to the model.

Results

A100 80 GB, driver 580.126.09, torch 2.15.0, torch_tensorrt 2.15.0, TensorRT 11.4.0.106, CUDA 13.2. 3467 tokens, top 8, hidden size 3072, bf16.

Engine Main kernel Grid x block Registers Est. occupancy
combine 273 us 41604 x 256 32 100%
norm 44 us 3467 x 128 40 75%
fused 392 us 3467 x 128 72 44%

Fused: 392 us. Apart: 273 + 44 = 317 us. Fusing makes this step 24% slower.

Expected

The fused graph should cost no more than the two parts apart. For example, keep one thread per element for the combine and run the norm as its own kernel, or issue the gather loads for all k slots before using them in the per-token schedule.

Repro

python repro_combine_norm.py. Needs torch, torch_tensorrt and a CUDA GPU. It builds 3 engines from the same inputs and prints each engine's time and kernel list.

repro_combine_norm.py 
import json, os, statistics, tempfile

import torch
import torch_tensorrt
from torch.profiler import ProfilerActivity, profile

T, K, C = 3467, 8, 3072  # tokens, experts per token, hidden size
torch.manual_seed(0)
dev = "cuda"
order = torch.randperm(T * K, device=dev)
inverse = torch.empty_like(order)
inverse[order] = torch.arange(T * K, device=dev)
inverse = inverse.view(T, K)  # expert-sorted row of token t's slot k
rows = torch.randn(T * K, C, device=dev, dtype=torch.bfloat16)  # expert outputs, in expert order
gate = torch.rand(T, K, device=dev)
gate = gate / gate.sum(-1, keepdim=True)
x = torch.randn(T, C, device=dev, dtype=torch.bfloat16)  # residual stream
weight = torch.randn(C, device=dev) * 0.1


def combine(rows, inverse, gate):
    acc = torch.zeros(T, C, device=dev, dtype=torch.float32)
    for s in range(K):
        acc = acc + rows[inverse[:, s]].float() * gate[:, s].unsqueeze(-1)
    return acc.to(torch.bfloat16)


def rmsnorm(y):
    y = y.float()
    y = y * torch.rsqrt(y.pow(2).mean(-1, keepdim=True) + 1e-6) * (1 + weight)
    return y.to(torch.bfloat16)


class Combine(torch.nn.Module):
    def forward(self, rows, inverse, gate):
        return combine(rows, inverse, gate)


class Norm(torch.nn.Module):
    def forward(self, x, h):
        return rmsnorm(x + h)


class Fused(torch.nn.Module):
    def forward(self, x, rows, inverse, gate):
        return rmsnorm(x + combine(rows, inverse, gate))


def bench(fn, iters=50):
    for _ in range(5):
        fn()
    torch.cuda.synchronize()
    ts = []
    for _ in range(iters):
        s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        s.record(); fn(); e.record(); torch.cuda.synchronize()
        ts.append(s.elapsed_time(e))
    return statistics.median(ts) * 1000


def kernels(fn):
    with profile(activities=[ProfilerActivity.CUDA]) as prof:
        fn(); torch.cuda.synchronize()
    path = os.path.join(tempfile.mkdtemp(), "trace.json")
    prof.export_chrome_trace(path)
    out = []
    for e in json.load(open(path))["traceEvents"]:
        if e.get("cat") == "kernel":
            a = e.get("args", {})
            out.append((round(e["dur"], 1), e["name"][:60], a.get("grid"), a.get("block"),
                        a.get("registers per thread"), a.get("est. achieved occupancy %")))
    return out


h = combine(rows, inverse, gate)
cases = {"combine": (Combine(), (rows, inverse, gate)), "norm": (Norm(), (x, h)),
         "fused": (Fused(), (x, rows, inverse, gate))}
print(torch.cuda.get_device_name(), "torch", torch.__version__, "torch_tensorrt", torch_tensorrt.__version__)
with torch.no_grad():
    for name, (mod, args) in cases.items():
        ref = mod(*args)
        ep = torch.export.export(mod, args)
        trt = torch_tensorrt.dynamo.compile(ep, arg_inputs=list(args), min_block_size=1, disable_tf32=True)
        diff = (trt(*args).float() - ref.float()).abs().max().item()
        print(f"{name:8s} engine {bench(lambda: trt(*args)):7.1f} us, max abs vs eager {diff:.3g}")
        for us, k, grid, block, regs, occ in kernels(lambda: trt(*args)):
            print(f"           {us:7.1f} us  grid {grid} block {block} regs {regs} est. occupancy {occ}%  {k}")

Actual

Output of the script (A100 80 GB, driver 580.126.09; warnings removed). The "engine" times include the call overhead of each Torch-TensorRT module; compare the kernel lines.

NVIDIA PG509-210 torch 2.15.0+cu132 torch_tensorrt 2.15.0
combine  engine   448.9 us, max abs vs eager 0.00781
             273.0 us  grid [41604, 1, 1] block [256, 1, 1] regs 32 est. occupancy 100%  __myl_SlicSlicReshGathCastMulSlicSlicReshGathCastMulSlicSlic
norm     engine   239.6 us, max abs vs eager 0.0156
              44.2 us  grid [3467, 1, 1] block [128, 1, 1] regs 40 est. occupancy 75%  __myl_AddCastMulMeanAddSqrtDivMulMulCast_0x8b3364a1de2f0c92f
fused    engine   568.7 us, max abs vs eager 0.0156
             392.8 us  grid [3467, 1, 1] block [128, 1, 1] regs 72 est. occupancy 44%  __myl_SlicSlicReshGathCastMulSlicSlicReshGathCastMulSlicSlic
Lingua principale
Python
Stelle
3k
Fork
421
Merge medio
4g 7h
PR unite (30g)
62

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di pytorch/TensorRT

Tutte le issue di pytorch/TensorRT

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.