前回、C++で実装したTensorRTプラグインを使ったMoEの推論処理の推論速度を比較する。
比較対象は、
- Dense MoE
- Sparse MoE (自前実装CUDAカーネル)
- Sparse MoE (CUTLASS Grouped GEMM)
の3パターンとする。
比較条件
- SwinTransformerのStage 0/1をMoE化
- Stage 0の解像度 8x8、State 1の解像度 4x4
- Stage 0のhidden_features 256、State 1のhidden_features 512
- Expert数4
- top_k=2
- バッチサイズ128
- FP16
- RTX 4090で測定
前回実装した推論コマンドを呼び出す、以下のベンチマークスクリプトを使用する。
train_cifar10/scripts/benchmark_moe_trt.py at feature/moe · TadaoYamaoka/train_cifar10 · GitHub
測定結果
CUDA Graph 有効
batch=128 warmup=5 iters=10 cuda_graph=True
| target | avg_latency_ms | throughput_images_per_s |
|---|---|---|
| Dense MoE | 0.582870 | 219603.000 |
| Sparse Plugin MoE | 1.790130 | 71503.000 |
| Sparse Plugin MoE (CUTLASS) | 0.947411 | 135105.000 |
CUDA Graph 無効
batch=128 warmup=5 iters=10 cuda_graph=False
| target | avg_latency_ms | throughput_images_per_s |
|---|---|---|
| Dense MoE | 1.028980 | 124395.000 |
| Sparse Plugin MoE | 2.067490 | 61910.700 |
| Sparse Plugin MoE (CUTLASS) | 1.332480 | 96061.200 |
考察
Dense MoEが一番早く、Sparse MoEでは、CUTLASSを使った方が速いという結果になった。
以下に結果を考察する。
1. 概要
主因は、今回の条件では MoE を sparse にするメリットが小さく、sparse 実行のオーバーヘッドが相対的に大きいためであると考えられる。 結果を比較すると以下の通りである。
| 手法 | Latency | 備考 |
|---|---|---|
| Dense MoE | 0.583 ms | 最速 |
| Sparse Plugin MoE | 1.790 ms | Dense比 約3.07倍低速 |
| Sparse Plugin MoE (CUTLASS) | 0.947 ms | Dense比 約1.63倍低速 |
CUTLASS 版は非 CUTLASS 版より約 1.89 倍高速であり、CUTLASS grouped GEMM 自体は有効に機能している。しかし、依然として Dense に及ばないという点が極めて重要である。
2. 理論メリットの限定性
今回の MoE は E=4, top_k=2 であるため、sparse 化による理論上の計算量削減幅が小さい。
- Dense MoE: 全 expert (E=4) を計算。
- Sparse MoE: top-k expert (K=2) のみを計算。
expert MLP の計算量のみに着目すれば、Sparse は Dense の約半分となる。しかし、モデル全体には Patch embedding、LayerNorm、Window attention、TensorRT の入出力処理などが含まれる。 モデル全体のレイテンシに対する MoE expert 計算の支配率が 100% でない限り、全体の短縮幅は 2 倍未満に留まる。この時点で Sparse が勝利するには、routing や packing などの追加コストを極小化する必要がある。
3. Dense MoE における TensorRT の最適化
Dense 版は全 expert を規則的な tensor 演算（torch.einsum 等）として実行する。これは算術量こそ増えるものの、GPU および TensorRT にとっては極めて扱いやすい。
- contiguous な tensor 構造
- 静的かつ規則的な shape
- gather / scatter の抑制
- TensorRT の標準オペレータとしての最適化
- cuBLAS / Tensor Core tactic への適合
特に E=4 のような少数 expert 設定では、「計算の無駄」よりも「規則的な dense matmul による高速化」のメリットが上回る。
4. Sparse Plugin におけるオーバーヘッド
Sparse MoE は計算量を削減する代償として、不規則なメモリアクセスや多数のカーネル実行を必要とする。具体的には、routing、prefix sum、token assignment の packing、scatter-add、atomic add といった処理が Dense 版にはほぼ不要な追加コストとして発生する。 E=4, K=2 という設定では、削減できる計算量よりも、これら不規則処理のオーバーヘッドが支配的になりやすい。
5. Plugin 境界による最適化の制限
plugin モードでは、MoE 全体が trt.plugins::CustomMoE という単一の custom node となり、TensorRT から見て内部が不透明（opaque）になる。
- TensorRT による layer fusion が適用されない。
- 内部のカーネルに対して TensorRT の tactic selection が効かない。
- カーネル起動やメモリアクセスの効率が、独自実装の品質に完全に依存する。
対して Dense 版は標準 ONNX オペレータのグラフであるため、TensorRT が広範囲にわたってグラフ最適化を適用できる。この差は極めて大きい。
6. Batch Size 128 における実行効率
今回のバッチサイズ 128 では、特に stage 0 等でトークン数が多くなる（例：128 * 64 = 8192 tokens）。 Dense MoE はこのまとまったトークン数を高 occupancy・高 arithmetic intensity で処理できる。一方、Sparse MoE では expert ごとにトークンが分割されるため、個々の GEMM サイズが小さくなり、metadata のセットアップやスケジューリングのオーバーヘッドによる効率低下の影響を受けやすくなる。
7. CUDA Graph の限界
cuda_graph=True の設定により CPU 側の起動オーバーヘッドは削減されているが、GPU 上での routing 計算、packing copy、atomicAdd といったデバイス側のコストは消失しない。今回の結果は、CPU launch ではなくデバイス側のオーバーヘッドが支配的であることを示唆している。
8. 結論：Dense が勝る理由と今後の展望
GPU 最適化において「無駄な計算をしても規則的な演算の方が速い」という典型的な事例である。
優先度の高い仮説
- E=4, K=2 では計算削減の恩恵がオーバーヘッドを相殺できない。
- Dense ONNX は TensorRT の標準最適化を最大限に享受している。
- Sparse plugin 内の routing や atomic 処理が依然として重い。
改善に向けたアプローチ：
今後 Sparse plugin で Dense を上回るには、少数 expert への特化（kernel fusion の徹底）、atomicAdd の排除、あるいはトークン数や設定に応じて Dense と Sparse を動的に切り替えるハイブリッド戦略の検討が有効である。
まとめ
TensorRTプラグインで実装したMoEの推論速度を比較した。
結果、Dense MoEが一番速いという結果になった。Sparse MoEではCUTLASSを使用することで約1.89 倍速くなることが分かった。
今回測定したSwinTransformerはモデルサイズが小さく、CIFAR-10の画像解像度も小さいため、Sparse化のオーバーヘッドが大きくかえって遅くなるという結果になった。
また、実装したCUDAカーネルはkernel fusionなどの最適化を行っていないため、改良の余地がある。
LLMの推論ライブラリとして人気のvLLMは、デバイスごとの手書きCUDAカーネルで最適化を行っている。
そこで使われている手法を参考すれば、まだまだ最適化できると思っている。
SwinTransformerでの実験はこれまでにして、次は、dlshogiのモデルにMoEを組み込むことを試したい。
