2 posts tagged with "SSM"

Mamba (state-spaces/mamba) の selective scan は CUDA カーネル前提で書かれており、Apple Silicon ではそのまま走らない。Metal Shading Language (MSL) で書き直し、HuggingFace の重みを直接ロードして推論まで通すプロジェクト mamba-metal を作った。本記事はその設計と検証結果の備忘録。

Mamba の中核は selective scan という線形時間の系列演算。論文だけ読むと「ハードウェア意識した実装」で済まされてしまう部分を、公式実装の forward カーネル csrc/selective_scan/selective_scan_fwd_kernel.cuh を読んで分解する。