Mamba系状態空間モデルのカーネル実装と評価
← 技術
MCSoC 2026 原稿
射影演算とSSDのステップ1のPEの段構成(arXiv公開版)。 出典
Mamba-130Mのデコード時間の内訳(arXiv公開版)。モデル全体の同条件CPU/GPU高速化比較ではありません。 出典
Mamba系の推論処理を全結合の射影演算、SSD、再帰的な状態更新に分け、それぞれがCGLAの演算パイプラインと転送境界にどう対応するかを調べた研究です。
手法と評価条件の詳細は、 論文(arXiv公開版) にまとめています。
実装・評価した内容
- 射影演算のGEMM、SSDのステップ1、再帰状態更新のカーネルを実装・評価。
- カーネル演算だけでなく、DMAを含む呼出し時間とCPU参照との出力差を確認。
- Mamba-130Mのトークン単位の実行環境への組込みで、デコード時間の内訳を計測。
- 長い集約演算と短い集約演算の差、カーネル境界、射影演算のGEMVの課題を整理。
結果
論文では集約演算が長い射影演算と、短い集約演算・境界費用に制約されるSSDのステップ1を区別しています。
Mamba-130Mのデコード計測では射影演算のGEMVをボトルネックとして確認しています。