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

Open
  • Is something wrong?
#5,042 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
22/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python
Domain
compilers, performance

Research direction

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.

Written by the indexing model from the issue text.

Description

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
Dominant language
Python
Stars
3k
Forks
421
Avg merge
4d 7h
Merged PRs (30d)
62

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from pytorch/TensorRT

All issues in pytorch/TensorRT

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.