AI Navigate

55 → 282 トークン/秒: 4x RTX PRO 6000 Blackwell で Qwen3.5-397B を高速に動かす方法

Reddit r/LocalLLaMA / 2026/3/15

📰 ニュースTools & Practical Usage

要点

  • 著者は CUTLASS をパッチして K=64 タイルをサポートさせることで、SM120 の MoE GEMM タイルの問題を修正し、SM120 GPU で推論を高速化しました。
  • 連続的な最適化を経て、Qwen3.5-397B モデルを 4x RTX PRO 6000 Blackwell で、WSL2 の 55 tok/s からネイティブ Linux の 282 tok/s へスループットを向上させました。
  • 根本原因は SM120 の 99KB SMEM 制限と、K < 128 の場合の CUTLASS TMA レイアウトのバグでした。パッチは EffBlk_SF の計算調整と、ハードウェア制約に合わせたスケールファクターの折り畳みを導入します。
  • 結果には、構成ごとの詳細な性能指標、ハードウェアおよびソフトウェア環境、FlashInfer への PR 提出と、事前構築済みの Docker イメージの利用可能性が含まれます。

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 をパッチして、以下を行いました:

  1. EffBlk_SF = min(K/SFVectorSize, Blk_SF) を計算して K<128 を処理
  2. 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