04. ハードウェア対応の畳み込みカーネル

03で「フィルタをグループ化すると畳み込みを GEMM にまとめられる」と述べました。このページでは、その GEMM を GPU のテンソルコアで実際にどう計算するか ——論文の核心である 2段階ブロックアルゴリズム を掘り下げます。

1. 畳み込みは Toeplitz 行列の積

まず出発点。離散畳み込みは Toeplitz 行列との行列積と数学的に等価 です。長さ $\ell_h$ の因果 FIR フィルタ $h$ を、長さ $\ell$ の入力 $x$ （通常 $\ell \gg \ell_h$ ）に適用した出力は、

y_t = \sum_{k=0}^{t} h_{t-k}\, x_k \qquad (h_k = 0 \text{ for } k < 0 \text{ or } k \ge \ell_h)

これを行列の形に書くと、フィルタ係数が斜めに並ぶ 下三角 Toeplitz 行列 $T$ との積になります。

\mathbf{y} = T\mathbf{x}, \qquad T = \begin{pmatrix} h_0 & 0 & 0 & \cdots \\ h_1 & h_0 & 0 & \cdots \\ h_2 & h_1 & h_0 & \cdots \\ \vdots & \ddots & \ddots & \ddots \end{pmatrix}

StripedHyena 2 の演算子（SE/MR/LI）はいずれも grouped depthwise convolution ですが、従来 CNN 向けの実装（im2col・Winograd）はこれに最適化されていません。長いフィルタは FFT ベースになりますが、ハードウェア利用率が低い。そこで論文は、filter grouping を活かした direct な multi-pass ブロックアルゴリズム を採用します。

2. ブロック畳み込み：長い系列を分割する

古典的なデジタル信号処理の ブロック畳み込み（block convolution） を使います。入力と出力を サイズ $\ell_b$ のチャンク に分割し、 $\ell_b \times \ell_b$ の部分ブロックごとに掛けます。

\hat{y}_n = \sum_{k} H_{n-k}\, \hat{x}_k

ここで $\hat{x}_k, \hat{y}_n$ は $k, n$ 番目の入力・出力チャンク、 $H_k$ は Toeplitz $T$ を $\ell_b \times \ell_b$ に分割した部分行列です。ポイントは、フィルタのサポート（長さ $\ell_h$ ）を超えた部分ブロックはすべてゼロになり、スキップできる こと。 $\ell_h \ll \ell$ のとき、これが大きな効率になります。

3. 2段階ブロックアルゴリズム

SE・MR では $\ell_h \ll \ell$ です。さらに $\ell_h \le 2\ell_b$ （フィルタ長がブロックの2倍以内）なら、特に効率的な 2段階（two-stage）アルゴリズム が使えます。このとき Toeplitz $T$ は、対角ブロック部（ $H_0$ ） と オフ対角部（ $H_1$ ） の2つだけに分解できます。

2段階分解。

T

は対角の

H_0

（第1段）とオフ対角の

H_1

（第2段）だけになる（論文式8）

具体例として $\ell = 6$ （系列長）、 $\ell_h = 4$ （フィルタ長）、 $\ell_b = 3$ （ブロックサイズ）を取ると、

H_0 = \begin{pmatrix} h_0 & 0 & 0 \\ h_1 & h_0 & 0 \\ h_2 & h_1 & h_0 \end{pmatrix}, \qquad H_1 = \begin{pmatrix} h_3 & h_2 & h_1 \\ 0 & h_3 & h_2 \\ 0 & 0 & h_3 \end{pmatrix}

$H_0$ は「現在のチャンク $\hat{X}_n$ に整列するフィルタ点」を担当。
$H_1$ は「前のチャンク $\hat{X}_{n-1}$ からはみ出してくる（spillover）点」、つまりチャンク境界をまたぐ係数を担当。

各出力チャンクは、たった2つの行列積で計算できます。

\hat{Y}_n = H_0\, \hat{X}_n + H_1\, \hat{X}_{n-1} \qquad (\hat{X}_{-1} = 0)

$\ell_h \le 2\ell_b$ という条件が、「 $H_1$ より先のオフ対角ブロックは出てこない」ことを保証しています。

各出力チャンクは「現在チャンク×H₀」と「前チャンク×H₁」の和。境界の spillover を H₁ が引き受ける

4. テンソルコアへ：GEMM 化と再利用

この分解が GPU で速い理由は3つあります。

$H_0, H_1$ は全チャンク・全チャネルで不変 — フィルタが決まれば $H_0, H_1$ は固定。一度オンチップメモリ（共有メモリ）にロードすれば、すべてのチャンクで再利用できます。
03 のグループ化と組み合わさる — グループ内の全チャネルが同じ $H_0, H_1$ を共有するので、小さな GEMV ではなく 大きな GEMM にまとめられます。テンソルコアのサイズ $d_g$ と $\ell_b = d_g$ を揃えれば、 $\hat{Y}_n = H_0\hat{X}_n + H_1\hat{X}_{n-1}$ は 2つのフル GEMM として実行できます。
「現在チャンク」と「前チャンク」を並列／パイプライン化 — 第1段（ $H_0$ ）と第2段（ $H_1$ ）は独立に走らせられます。

論文の前進カーネル（Algorithm 1）を擬似コードで示すと、こうなります。

入力 v, q, k を ℓ_b × d_g のブロックに分割
for ブロック i = 0 .. ⌈ℓ/ℓ_b⌉ − 1:
    v_i, q_i, k_i, H_0, H_1 をオンチップメモリにロード
    y_i = H_0 · v_i              # 第1 GEMM（対角・block-diagonal）
    if i > 0:
        y_i += H_1 · v_{i-1}     # 第2 GEMM（オフ対角・spillover）
    y_i = q_i ⊙ y_i              # ゲーティング
return y

なぜ GEMM 化がそんなに効くのか

GPU のテンソルコアは 密な行列積（GEMM） に特化したユニットで、桁違いのスループットを出します。depthwise 畳み込みを素朴に書くと小さな行列ベクトル積（GEMV）の山になり、テンソルコアを活かせません。「グループ化 → GEMM 化 → テンソルコア」 という一連の co-design が、StripedHyena 2 の速さの正体です。これは 03 の filter grouping と一体の工夫です。

5. 実測：演算子レベルのスループット

この2段階アプローチの効果は実測でも明確です。Hyena-MR（フィルタ長 128）を、素朴な PyTorch 畳み込み（F.conv1d）と比べると、2段階ブロックカーネルは レイテンシ・スループットともに大幅に改善 します。

さらに演算子レベルで他の主要演算子と比較すると、Hyena-SE・Hyena-MR は、最適化された MHA（FlashAttention2・SDPA）や Mamba2・xLSTM・DeltaNet を上回るスループット を、バッチサイズ1・幅 4096 の設定で達成します。

演算子別スループットの概念図（H100・幅4096）。Hyena-SE/MR が attention・SSM 系を上回る（論文 Fig. 3.2・B.4 の趣旨）

6. 計算コストと chunk 方向の並列化

03 のグループ化は「チャネル方向」に GEMM をまとめる方法でした。論文の付録では、グループ化なしでも GEMM 化する別アプローチ——チャンク方向に並列化 する方法も示されています。

各出力チャンクは $\hat{Y}_n = H_0\hat{X}_n + H_1\hat{X}_{n-1}$ で、 $H_0, H_1$ は全チャンクで不変です。そこで GPU カーネルは次の手順を取ります。

フィルタ preload： $H_0, H_1$ をオンチップメモリ（共有メモリ）に一度だけロードし、全チャンクで再利用。
チャンク読み込み：現在チャンク $\hat{X}_n$ （と前チャンク $\hat{X}_{n-1}$ ）をレジスタ／共有メモリへ。
テンソルコア GEMM： $H_0\hat{X}_n$ と $H_1\hat{X}_{n-1}$ を計算し、 $\hat{Y}_n$ に累積。
書き戻し： $\hat{Y}_n$ をグローバルメモリへ。

計算コスト（cost model） も明快です。各チャンクは $(\ell_b\times\ell_b)\times(\ell_b\times d)$ の GEMM を2回、すなわち $2\ell_b^2 d$ FLOPS。系列全体（ $\lceil\ell/\ell_b\rceil$ チャンク）では、

2\,\ell_b^2\, d\, \left\lceil \frac{\ell}{\ell_b} \right\rceil \ \text{FLOPS}

フィルタが短い（ $\ell_h \le 2\ell_b$ ）ため $H_2$ 以降のオフ対角ブロックが現れず、この見積もりで収まります。系列長 $\ell$ に対して線形（チャンク数に比例）であり、Attention の $O(\ell^2)$ と対照的です。

7. 逆伝播：two-pass backward カーネル

学習には逆伝播が必要ですが、フィルタの勾配計算は「全体での累積（global accumulation）」 を要します（同じフィルタが全位置で共有されるため）。これを1つのカーネルに詰め込むと非効率なので、論文は 2段（two-pass）の back-to-back カーネル で実装します。

1段目：forward と同じブロック構造を保ったまま、ブロックごとに部分的なフィルタ勾配を累積。
2段目：部分勾配を reduction（集約） して最終的なフィルタ勾配を得る。

ポイントは、1段目で部分勾配を coalesced（連続した）形式 で書き出すこと。これにより2段目を単純なベクトル化 reduction にできます。

また、Toeplitz 因子 $H_0, H_1$ は その場で効率的に生成（materialize） します。前述の2段階分解で見た Toeplitz 行列を、メモリ上に陽に展開せず、行・列インデックスの差からフィルタ係数を引くだけ で構築できます（Triton カーネルの例）。

import triton
import triton.language as tl

@triton.jit
def load_toeplitz(h_ptr, FILTER_LEN: tl.constexpr, CHUNK_SIZE: tl.constexpr):
    # 行 r と列 c のインデックス差 (r - c) が、そのセルが参照するフィルタ係数の添字
    r = tl.arange(FILTER_LEN - 1, CHUNK_SIZE + FILTER_LEN - 1)[None, :]
    c = tl.arange(0, CHUNK_SIZE)[:, None]
    idx = r - c
    mask = (idx >= 0) & (idx < FILTER_LEN)         # フィルタ範囲外は 0 埋め
    return tl.load(h_ptr + idx, mask=mask, other=0.0)

「インデックス差がフィルタ係数の添字になる」のは、まさに Toeplitz 行列が 対角方向に同じ値が並ぶ 性質そのものです。陽に行列を作らずインデックス計算で済むので、メモリと帯域を節約できます。

8. まとめ

畳み込みは Toeplitz 行列の積。SE/MR は $\ell_h \le 2\ell_b$ なら 2段階（ $H_0$ ＋ $H_1$ ） に分解できる。
$\hat{Y}_n = H_0\hat{X}_n + H_1\hat{X}_{n-1}$ という 2つの GEMM に落ち、 $H_0/H_1$ の再利用とグループ化で テンソルコアをフル活用。
結果、Hyena-SE/MR は MHA・Mamba2・xLSTM・DeltaNet を上回るスループット。

ここまでは1台の GPU 内の話でした。100 万トークンを 複数 GPU に分散 して学習するには、また別の工夫が要ります。最後の 05. 長系列の分散学習で、context parallelism と FFT 畳み込みを見ていきます。

1. 畳み込みは Toeplitz 行列の積​

2. ブロック畳み込み：長い系列を分割する​

3. 2段階ブロックアルゴリズム​

4. テンソルコアへ：GEMM 化と再利用​

5. 実測：演算子レベルのスループット​

6. 計算コストと chunk 方向の並列化​

7. 逆伝播：two-pass backward カーネル​

8. まとめ​