TL;DR: SM120 の壊れた MoE GEMM タイルを修正するためのカスタム CUTLASS カーネルを構築しました。55 tok/s (WSL2) → 119 (ネイティブ Linux) → 142 (ドライバ/設定最適化) → 282 tok/s (カスタム K=64 カーネル)。FlashInfer へ PR を提出し、事前構築済みの Docker イメージを利用可能です。
問題点
NVFP4 MoE モデル(Qwen3.5-397B、DeepSeek など)を RTX PRO 6000、RTX 5090、または DGX Spark 上で実行している場合 — 基本的にはどの SM120 Blackwell ワークステーション GPU でも — おそらく次のようなメッセージを見たことがあるでしょう:
Failed to initialize cutlass TMA WS grouped gemm 自動チューナーは、GPU の 99KB の共有メモリを超えるため、すべての SM120 GEMM タイルをスキップします。データセンター Blackwell (B200) は 228KB の SMEM を持っており、タイルはそれ用に設計されています。ワークステーションの GPU は遅いフォールバックカーネルで行き詰まります。
結果: スループットの 50% 以上を活用できていません。
解決策
問題は K=128 のタイルが SM120 の SMEM より多くを必要とすることです。K=64 のタイルは収まりますが、CUTLASS にバグがありました。TMA のスケールファクターのレイアウトは K≥128 を想定しており、K=64 のときレイアウトの不一致を生じさせます(Blk_SF=4 だが K=64 には K 軸にスケールファクターが 2 個しかありません)。
CUTLASS の sm120_blockscaled_mma_builder.inl をパッチして、以下を行いました:
-
EffBlk_SF = min(K/SFVectorSize, Blk_SF)を計算して K<128 を処理 - MMA の要件を超える場合にスケールファクターを基本ブロックに折り畳む
これにより K=64 タイルが SM120 の 99KB SMEM 上で正しくコンパイルされ、実行可能になります。
結果
ハードウェア: 4x NVIDIA RTX PRO 6000 Blackwell (各 96GB GDDR7、SM 12.0) モデル: Qwen3.5-397B-A17B-NVFP4、TP=4、MTP=5 環境: CUDA 13.2、Driver 595.45.04、vLLM 0.17.1rc1、FlashInfer 0.6.6 Sehyo バージョンの Qwen3.5-397-a17b-NVFP4。
| ユーザー | 前 (tok/s) | 後 (tok/s) | 改善 |
|---|---|---|---|
| 1 | 142 | 283 | +99% |
| 4 | 250 | 850 | +240% |
| 8 | 510 | 1,283 | +151% |
WSL2 からの全経緯:
| 設定 | 1 ユーザー tok/s |
|---|---|
| WSL2 ベースライン | 55 |
| ネイティブ Linux | 119 |
| + MTP=5 + 設定チューニング | 134 |
| + ドライバ 595 + CUDA 13.2 + iommu=pt | 142 |
| + Cu |
