動機​

Mamba 公式実装の本体は csrc/selective_scan/selective_scan_fwd_kernel.cuh にある CUDA カーネル。これが速度の核であり、Mamba を「論文上の理論」から「実機で動くアーキテクチャ」へ変えている部分。NVIDIA GPU 専用なので Apple Silicon では本来動かない。

参照実装の純 PyTorch 版（selective_scan_ref）も存在するが、for ループの素朴な漸化式評価で、長系列では非実用。Mamba の本質的な並列化（プレフィックススキャン）が抜けている。

そこで MSL で同等のカーネルを書くことにした。MLX の mx.fast.metal_kernel を介して JIT コンパイル・ディスパッチさせ、.metal ファイルを第一級資産として残す。

Selective scan の本質​

Mamba の隠れ状態は

という入力依存の係数を持つ漸化式。$\bar{A}_t = \exp(\Delta_t A)$ で、$\Delta_t, B_t, C_t$ は入力 $x_t$ から計算される（selective: 入力に応じてゲートが開閉する）。

漸化式 $h_t = a_t h_{t-1} + b_t$ は素直には逐次にしか解けないが、ペア $(a, b)$ に対する次の演算

は結合的 (associative)。よって prefix scan で $O(\log T)$ 段の並列ステップに落とせる（Blelloch 1990 / Martin & Cundy 2017）。Mamba カーネルがやっているのも本質的にこれ。

MSL での核心​

SIMD-group プリミティブ​

Metal は CUDA の warp に相当する SIMD-group（32 スレッド）を持ち、simd_prefix_inclusive_sum、simd_shuffle_up などの組み込み関数がある。ただしこれらは float スカラー専用。$(a, b)$ ペアの結合演算は自分で書く必要がある：

for (uint d = 1u; d < 32u; d <<= 1) {
    float a_prev = simd_shuffle_up(a, d);
    float b_prev = simd_shuffle_up(b, d);
    if (lane >= d) {
        b = a * b_prev + b;   // 順序重要: 先に b を更新（古い a を使う）
        a = a * a_prev;
    }
}

これで 32 レーンの SIMD-group 内で inclusive scan が完了。

Block-level scan（two-tier）​

1024 スレッド（= 32 SIMD-group × 32 lane）の threadgroup 全体で scan するために、SIMD-group の合計値を threadgroup memory に書き出し、1 つ目の SIMD-group がそれをさらに scan し、各スレッドが carry を加える、という二段構成にする。これは CUB の BlockScan + WARP_SCANS 戦略の MSL 版。

チャンク間 running prefix​

seqlen > 1024 の場合、smem_running_prefix 方式：各 SSM 状態 $s$ ごとに (carry_a[s], carry_b[s]) をチャンク間で持ち越す。新しいチャンクの先頭で前チャンクの累積を「左から」結合してから scan を実行：

これにより任意長の系列が単一カーネル呼び出しで処理できる。

観察：tg memory は単純な再利用には効かない​

愚直な「データを threadgroup memory に置いて K 回読み返す」パターンは、Apple Silicon の System Level Cache（CPU/GPU 共有）が黙って吸収してしまうため、global memory 直読みと差がほぼ出なかった。tg memory が真に必要なのはスレッド間通信（scan の中間値交換、running prefix の保管）であって、データキャッシュ代用ではない、というのが実測の結論。

カーネルから推論まで​

カーネルが組めたら、上に Python のモデル層を積む。

selective_scan (Metal kernel)
  ↓
MambaBlock           = in_proj → conv1d → SiLU → x_proj/dt_proj → SSM → out_proj
  ↓
MambaResidualBlock   = pre-norm RMSNorm + MambaBlock + residual
  ↓
MambaModel           = embeddings → N × ResidualBlock → norm_f → tied LM head
  ↓
generate / generate_fast

HF の state-spaces/mamba-*-hf 重みは：

1. backbone. 接頭辞を剥がす
2. conv1d.weight だけ PyTorch (out, in/g, k) → MLX (out, k, in/g) で transpose
3. それ以外（Linear, embeddings, A_log, D, norm）はそのまま mx.array に変換

の 2 ステップだけで MLX 側にロードできる。hidden_size / intermediate_size / num_hidden_layers という HF transformers 標準フィールドを優先するのがコツ（790m などで legacy d_model フィールドが壊れているため）。

O(L) インクリメンタルデコード​

Mamba の論文上の最大の魅力は「長文脈で一定速度」。これを実機で具現化するには、推論時に SSM の隠れ状態と conv1d の sliding window を呼び出し間で持ち越す必要がある。

conv_states, ssm_states = model.init_state(batch_size=1)
for token in prompt:
    logits, conv_states, ssm_states = model.step(token, conv_states, ssm_states)
# 以降は 1 トークンあたり O(1)

毎ステップは elementwise 演算のみ（SSM scan は不要、なぜなら状態を既に持っているから）：

これを z ゲートと out_proj で締める。

実測（M4 Max, mamba-130m, greedy decode）：

generate_fast は n=50 から n=2000 まで一貫して ~7 ms/token。これが Mamba の "linear-time decode" の正体。

モデルサイズ別の結果​

state-spaces/mamba-*-hf の全 5 サイズが load & generate 可能：

130m はサイズの限界で繰り返しに陥りやすいが、2.8b では「東京は最大の都市でもある」と付加的な事実までまとめて出してくる。greedy だけでこの差。

selective scan カーネル単体のピーク性能は seqlen=32k で ~187 GFLOPS、Unified Memory の実効帯域は vec4 ロードで ~290 GB/s（M4 Max の理論ピーク 410 GB/s の約 70%）。

残り課題​

• Prefill の高速化: 現状プロンプトを 1 トークンずつ step で流すため、長文脈プロンプトでは秒オーダー。selective_scan カーネルから最終 SSM 状態を抽出できれば、parallel scan で prefill して decode に O(1)/token で接続できる
• iPhone 上での Transformer vs Mamba ベンチマーク: 同じ規模で速度・精度を比較し、長文脈での Mamba 優位を可視化する
• 後方カーネル: 学習用の backward pass はまだ未実装

振り返り​

Mamba は論文の数式自体は短いが、「実機で線形時間」を実現する部分はカーネルにある。それを別ハードウェア向けに書き直してみると、初めて論文の主張の細部が手触りとして理解できる：

• 何を SRAM に閉じるべきで何を HBM に出すべきか（Apple Silicon ではキャッシュが吸収するので少し違う）
• なぜ A は対角でなければならないか（per-state の独立性で外側ループに置けるから）
• なぜ exp(ΔA) を exp2f + LOG2E で書くか（少しでも速い）
• 状態キャッシュがあれば本当に O(L) になるという主張の確認

論文を読むだけでは抜けていた解像度が、書いてみると一気に上がる。

参考文献​

• Albert Gu, Tri Dao. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" arXiv:2312.00752, 2023.
• Guy E. Blelloch. "Prefix Sums and Their Applications" CMU-CS-90-190, 1993.
• Eric Martin, Chris Cundy. "Parallelizing Linear Recurrent Neural Nets Over Sequence Length" arXiv:1709.04057, 2017.
• state-spaces/mamba — 公式実装
• createcentury/mamba-metal — 本記事のプロジェクト

作成日: 2026-05-16 / 最終更新日: 2026-05-16

参考文献​

• "Nonlinear Dynamic Inversion" Advanced Flight Control course notes, Aerostudents.

作成日: 2026-05-14 / 最終更新日: 2026-05-14

基礎代謝​

国立健康・栄養研究所の式（国立栄研式 / Ganpule の式）に当てはめる：

パラメータ $W = 81.0\text{ kg}$, $H = 176.5\text{ cm}$, $A = 27\text{ 歳}$ を代入：

1日あたりの基礎代謝量 ≈ 1,727 kcal

タンパク質​

知的・身体的パフォーマンス向上を狙うレンジは $1.2\text{ g} \sim 1.6\text{ g/kg}$。

脳の神経伝達物質や免疫細胞もタンパク質から作られるため、高負荷な生活を送るならこのあたりが「投資効率の良いゾーン」になる。

戦略​

• 吸収効率の最適化: 一度に吸収できる量には限界があるため、1日3〜4回に分散させる
• 「高密度」な選択: 脂質の少ない赤身肉、卵、必要に応じてプロテインを活用し、胃の容量をタンパク質に優先的に割り振る

具体的な食材​

食物繊維​

推奨量: 1日 21 g 以上

具体的な食材​

• ブロッコリー（約100g / 1/2房）: 食物繊維 約 4.4 g。タンパク質も含まれる、利回りの良い食材
• ニンジン（約100g / 1本弱）: 食物繊維 約 2.8 g。彩りと栄養バランスを整えるのに最適
• もち麦・玄米（1膳）: 食物繊維 約 2.0〜4.0 g。主食を白米から置き換えるだけでベースライン底上げ
• 納豆（1パック）: 食物繊維 約 3.3 g。タンパク質も同時確保

1日のモデル​

ビタミン​

ビタミンB群（B1, B2, B6, B12, ナイアシン等）​

• 役割: 脳や筋肉がエネルギーを産生する際の必須コンポーネント。不足すると「燃料はあるのにパワーが出ない（倦怠感）」状態に
• 推奨量（B1）: 1.4 mg/日

ビタミンD​

• 役割: 骨の健康だけでなく、免疫機能やメンタルの安定にも関与。室内での仕事や研究が多い場合は不足しがち
• 目安量: 8.5 µg/日
• 近年の研究で、ビタミンD濃度が高い人ほどテストステロン値が高い傾向。日光浴が不足しがちな現代人には必須の「ブーストパッチ」

ビタミンC​

• 役割: 抗酸化作用（身体の錆び取り）とコラーゲン合成。ハードな活動による酸化ストレスから細胞を保護
• 推奨量: 100 mg/日

その他のビタミン​

ミネラル​

鉄（Iron）​

• 役割: 全身に酸素を運ぶヘモグロビンの材料。不足すると脳への酸素供給が滞り集中力低下
• 推奨量: 7.5 mg/日

亜鉛（Zinc）​

• 役割: 数百種類の酵素の活性化に関わり、タンパク質合成や免疫応答、味覚の維持を支える
• 推奨量: 11 mg/日
• 「セックスミネラル」とも呼ばれる最重要項目。精子の生成やテストステロンの代謝に直接関わる

マグネシウム​

• 役割: エネルギー代謝、神経伝達、筋肉の弛緩。不足すると足がつりやすく、疲れが抜けにくくなる
• 推奨量: 340 mg/日

カルシウム​

• 役割: 骨の維持と神経の安定
• 推奨量: 750 mg/日〜（耐容上限量 2,500 mg）

微量ミネラル​

脂質​

n-3系脂肪酸（DHA、EPA、α-リノレン酸）​

• 役割: 脳の神経細胞の膜を柔軟に保ち、認知機能や炎症抑制に寄与。サバ缶や青魚に豊富
• 目安量: 2.0 g/日 以上

脂質の質についての注意​

• 飽和脂肪酸: 総エネルギーの 10% 相当以下 が目標。ステーキ（動物性脂質）に偏りすぎるとここが超過するため、魚や植物性オイルへ分散投資が推奨
• n-6系脂肪酸（リノール酸など）: 27歳男性の目安量は 11 g/日。主に植物油に含まれる
• コレステロール: かつては摂取制限があったが、現在は「脂質異常症の重症化予防」目的で、可能な限り控える、あるいは脂質の質（飽和脂肪酸）に気をつけるという方針

1日の食事プラン​

朝（または昼）​

オートミール (30 g) + 納豆 (1パック) + プロテイン + 粉末食物繊維 (5 g) + サバの味噌煮缶 (1缶) + もち麦ごはん (1膳)

→ 食物繊維 11 g とタンパク質 35 g を即時デプロイ。n-3系脂肪酸 (DHA/EPA) で脳を最適化。亜鉛・鉄も補給。炭水化物をインフラ化。

間食​

ギリシャヨーグルト + ゆで卵 (1個)

→ アミノ酸スコア100で、血中のアミノ酸濃度を一定に維持。

夕食​

赤身肉ステーキ (200〜300 g) + ブロッコリー & ニンジン (各100 g)

→ 亜鉛・鉄・ビタミンA・B群を大量投入。睡眠中のリカバリー準備。

コスト目安: 1日あたり 約 1,800 円 〜 2,300 円

参考文献​

• 厚生労働省「「日本人の食事摂取基準（2025年版）」策定検討会報告書」

作成日: 2026-05-14 / 最終更新日: 2026-05-14

動かす漸化式​

Mamba の隠れ状態 $h_t \in \mathbb{R}^{N}$ は、入力依存の係数で動く線形時不変ではない SSM：

$\Delta_t$ をステップサイズとして、連続系から離散化する典型は ZOH だが、Mamba 公式実装は近似的に

を使う（diagonal $A$ の各成分について scalar exp）。$\Delta_t, B, C$ は入力 $x_t$ から線形写像 + softplus などで都度作る — これが「selective」 (入力に応じてゲートが開閉する)。

なぜ「scan」か​

漸化式 $h_t = a_t h_{t-1} + b_t$ は 左畳み込みだが、ペア $(a, b)$ に対する次の演算

は結合的 (associative)。よって prefix scan (Blelloch 1990) で $O(\log T)$ 段の並列ステップに落とせる。Martin & Cundy (2017) と S5 (Smith et al. 2022) はこの観察を線形 RNN・SSM に持ち込んだ。Mamba のカーネルもこれを GPU の cub::BlockScan で具体化している。

実装上の thread_data[i] の中身がまさにこのペア：

// L221-222
thread_data[i] = make_float2(
    exp2f(delta_vals[r][i] * A_val[r]),                 // a_i = exp(Δ A)
    !kIsVariableB ? delta_u_vals[r][i] : B_vals[i] * delta_u_vals[r][i]  // b_i = ΔB · u
);

exp2f が使われているのは、$A$ を読み込む際に LOG2E を一度かけておく前処理 (L174-179) があるため。expf より高速。

カーネルの全体構造​

ファイルは大きく3レイヤ：

スレッド/ブロックの並び​

// L322
dim3 grid(params.batch, params.dim / kNRows);

1つの CUDA ブロック = (batch_id, dim_id)。各ブロックは

• 入力 $u, \Delta \in \mathbb{R}^{T}$ (1チャネル分)
• 重み $A \in \mathbb{R}^{N}$、入力依存 $B, C \in \mathbb{R}^{N \times T}$

を読んで、$y \in \mathbb{R}^{T}$ を返す。kNThreads が seqlen に応じて 32〜128 で切り替わる：

// L353-364
if (params.seqlen <= 128)  launch<32,  4>();
else if (seqlen <= 256)    launch<32,  8>();
else if (seqlen <= 512)    launch<32, 16>();
else if (seqlen <= 1024)   launch<64, 16>();
else                       launch<128, 16>();

短い系列で多くのスレッドを使うとオーバーヘッドが勝つのでチューニングされている。

チャンク化​

ブロック内 1 イテレーションで処理する系列長は

つまり最大でも 2048 トークン (128×16)。seqlen がこれを超える場合は

// L137
for (int chunk = 0; chunk < params.n_chunks; ++chunk) { ... }

でチャンクをループする。チャンク境界で状態を引き継ぐのが smem_running_prefix (L100, L244-247, L257-258)：scan の最後の prefix を共有メモリに保存し、次チャンクの初期 prefix として読む。これにより HBM への状態 readback を避ける（ハードウェア意識 の本体）。

1チャンク内の処理フロー​

1. load_input で u, delta を coalesced 読み込み
2. delta_softplus 適用 → delta_vals
3. delta_u_vals = delta * u, out_vals = D * u (skip connection)
4. for state_idx in [0, dstate):
     a. A_val を読み (LOG2E 倍済み)
     b. B_val, C_val を読み (selective なら BlockLoad、定数なら直接)
     c. thread_data = (exp2f(Δ A), ΔB u)  ← scan の入力タプル
     d. cub::BlockScan で InclusiveScan(SSMScanOp)
        → running_prefix を carry
     e. out_vals += scan_output.y * C
5. store_output で y を書き出し
6. (オプション) kHasZ: out *= z * sigmoid(z)  ← SwiGLU 風ゲート

state 次元 dstate (N) は外側の for ループになっていることに注意。並列 scan は時間方向で取り、状態次元は逐次。これは $A$ が対角行列だから各 state 成分が独立しているのを利用している（diagonal SSM の旨味）。

共有メモリ設計​

Selective_Scan_fwd_kernel_traits::kSmemSize (L63-69) は

• BlockLoad/Store の TempStorage (union 的に再利用)
• BlockScan の TempStorage

を合算したサイズ。さらにカーネル本体で kSmemSize の後ろに MAX_DSTATE * sizeof(scan_t) * kNRows を継ぎ足し、running prefix 用領域を確保する：

// L321
kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(scan_t);

48KB を超える場合は cudaFuncSetAttribute でダイナミック共有メモリの上限を引き上げる (L331-340)。

主要な最適化テクニック​

• exp2f + LOG2E 前処理: 浮動小数指数を expf でなく exp2f で。A 側に LOG2E を 1 回かけるだけで全 step に効く
• WARP_TRANSPOSE BlockLoad: ストライドアクセスをワープ単位で転置して coalesce
• WARP_SCANS BlockScan: warp-level 並列スキャンを採用 (RAKING より高速、コメント L60-61 に他の選択肢が残されている)
• kIsEvenLen 分岐: 系列長がチャンクで割り切れる場合は BLOCK_LOAD_DIRECT に切替 (L47-59)
• complex 数の自前 cexp2f: PyTorch の thrust::complex_exp が遅いので独自実装 (L229)
• kIsVariableB/C の compile-time 分岐: selective 性が無いケース (LTI) の不要な BlockLoad を消去 (L186-212)
• __launch_bounds__: kMinBlocks=3 or 5 で occupancy を明示 (L33, L73)

観察と疑問​

• kNRows == 1 しか実機で検証されていない (L312-314)。複数の dim を 1 ブロックで処理して reuse する余地が残っているが現状未開拓
• delta_softplus の境界 <= 20.f (L160): 浮動オーバーフロー対策のショートカット
• MAX_DSTATE の値は selective_scan.h 側にあるはず（読まないと不明）— state 次元の上限を決めている

Mamba のアーキテクチャ自体は SSM + selectivity のみだが、論文の主張する「線形時間で実用速度」は、このカーネルの チャンク化 × 結合的 scan × 状態は SRAM の 3 点で初めて成立している。

参考文献​

• Albert Gu, Tri Dao. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" arXiv:2312.00752, 2023.
• Guy E. Blelloch. "Prefix Sums and Their Applications" Technical Report CMU-CS-90-190, 1993.
• Eric Martin, Chris Cundy. "Parallelizing Linear Recurrent Neural Nets Over Sequence Length" arXiv:1709.04057, 2017.
• Jimmy T.H. Smith, Andrew Warrington, Scott W. Linderman. "Simplified State Space Layers for Sequence Modeling" arXiv:2208.04933, 2022.
• Wikipedia. "Leaky integrator"
• Wikipedia. "Zero-order hold"

作成日: 2026-05-11 / 最終更新日: 2026-05-14