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

オープン
  • 何か問題がありますか？
#5,042 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 3 日以内に返信

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
22/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
python
領域
compilers, performance

調査の方向性

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.

索引モデルが issue の本文から書いたものです。

説明

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
主要言語
Python
スター
3k
フォーク
421
平均マージ
4日 7時間
マージ済み PR（30日）
62

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

pytorch/TensorRT のほかの issue

pytorch/TensorRT の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。