Prime Intellectは、DeepSeekの融合GPUカーネル「Mega MoE」の学習拡張版となるPrime Mega MoEを公開した。DeepSeekのカーネルがフォワードパスのみを対象としていたのに対し、Prime Intellectは高速なバックワードパスを追加し、他の計算とオーバーラップさせることで、推論だけでなく学習にも使用可能なカーネルとした。同社によると、学習フレームワークprime-rlにおいて、Qwen3 MoEモデルの完全な学習ステップは最大34%高速化され、TorchTitanとの比較では最大55%高速化された。
MoE（mixture of experts、専門家混合）モデルでは、各トークンは異なるGPUに分散された専門家の小さなサブセットにのみ送られる。したがって、各レイヤーは、学習に使用されるバックプロパゲーションのフェーズであるバックワードパスと同様に、フォワードパスにおいてトークンを送信し、専門家に計算を行わせ、結果を返す必要がある。標準的な設定では、これらの操作は分離されている：GPUはトークンをソートし、通信を待ち、計算し、再び待機する。これらの待機時間は、計算がアイドル状態の時間となる。
学習に対応した融合カーネル
2026年4月にDeepGEMMライブラリに追加されたDeepSeekのMega MoEは、この一連の処理を単一のカーネル呼び出しに置き換える。カーネル（kernel）はGPU上で直接実行される最適化されたプログラムである。これは、専門家並列ディスパッチ、専門家の2つの線形レイヤー、SwiGLU、および結合ステップを融合する。あるトークングループが転送中である間に、カーネルは既に到着した別のグループの計算を行い、GPU上のみに配置されたワークバッファにより中間テンソーを回避する。ただし、このカーネルにはフォワードパスの実装のみが含まれていた。
Prime Intellectは、不足していたバックワードパスを、同様にオーバーラップする形で実装した。サポートされるデータ型は、BF16、Blackwell GPU専用のMXFP8、およびHopper GPU専用のブロックFP8である。prime-rlの設定では、fp8_bwdモードも提供される：フォワードパスはBF16、バックワードパスはFP8で行う。
自主発表された数値
全ての結果はPrime Intellect自身によって公開されている。学習ステップは8基のH200ノードを使用し、FP8精度、16,384トークンのシーケンス長、完全なアクティベーション・チェックポインティングで実施された。32基のH200に分散されたQwen3-30B-A3Bでは、完全な1ステップの所要時間は、prime-rlベースラインで2.78秒、TorchTitan（FP8）で3.07秒、Prime Mega MoEで2.08秒となり、34%および48%の高速化となった。128基のH200に分散されたQwen3-235B-A22Bでは、所要時間は12.72秒、14.98秒、9.65秒となり、32%および55%の高速化となった。スループットは、30BモデルではTorchTitanに対してGPUあたり毎秒5,337トークンから7,886トークンに、235Bモデルでは1,094トークンから1,697トークンに向上した。
単一ノードでレイヤーごとに測定すると、フォワードパスでの高速化がバックワードパスよりも顕著である。Qwen3-30B-A3Bモデルのフォワードパスは12.0 msから3.9 msへ、3.11倍の高速化となった。Qwen3-235B-A22Bモデルのバックワードパスは30.5 msから24.6 msへ、1.24倍の高速化となった。したがって、「最大3.11倍」という数値はフォワードパスのみを指す。一方、TorchTitanは最大41%の高速化を主張しているが、これは異なる設定（B200 GPU、6710億パラメータモデル、BF16ベースライン）でのものであり、情報源は直接比較可能ではないと判断している。
Prime Mega MoEはApache 2.0ライセンスの下で公開されている。prime-rlでは、ディスパッチのバックエンドとして有効化される：[trainer.model.moe.dispatch]の下でtype = 「mega_moe」と設定し、bf16、mxfp8、fp8、fp8_bwdの中から精度を選択するだけでよい。学習カーネルは、prime-kernelsリポジトリで入手可能な事前コンパイル済みDeepGEMMホイールに含まれており、標準版のDeepGEMMには含まれていない。
