M1 MaxでINT8ConvRotを実行しMetal・MPSMatrix・MLXを比べる

前編では、Anima-Turbo INT8ConvRot量子化版をM1 MaxのComfyUIへ読み込み、通常経路、CPUフォールバック、逆量子化回避策を試した。通常経路はMPSにtorch._int_mmがなく停止し、CPUフォールバックは実用にならないほど遅く、逆量子化回避策はint8演算ではなくfp16相当の計算だった。

今回はComfyUIを使わずに、INT8ConvRotに近い行列積をPyTorch、Metal、MPSMatrix、MLXから直接実行し、量子化テンソルをApple GPUへ渡せるかだけでなく、通常のfp16行列積より速くなるかまで測ってみる。

検証環境

項目	内容
マシン	Mac (M1 Max, 64GB RAM)
OS	macOS 26.5
PyTorch	2.13.0
MLX	0.32.0（リポジトリ外の一時Python 3.12環境）

MPSのweight-only INT8経路を単体で試す

前編で確認したのは、INT8ConvRotが通常使うint8×int8経路がないことまで。PyTorch 2.13.0にはMPS実装を持つ別の内部演算torch._weight_int8pack_mmがある。これはfp16などの活性テンソルとint8の重み、スケールを受け取り、量子化重みを使った行列積をMPS上で実行するweight-only INT8演算だ。

ただしtorch._int_mmと同一でもないし完全な代替なわけでもない。torch._int_mmがint8の活性とint8の重みを掛けてint32で積算するのに対し、_weight_int8pack_mmへ渡す活性はfp16のままになる。INT8ConvRotがCUDAで使うint8×int8 GEMMを再現する経路ではないが、重み全体を毎回fp16へ戻す現行パッチよりは量子化テンソルを直接使う範囲が広い。

計算をどう変えたか

通常の線形層は、活性 $X \in \mathbb{R}^{M \times K}$ と重み $W \in \mathbb{R}^{N \times K}$ から $Y = XW^\top$ を計算する。ConvRotでは256要素ごとの正規化Hadamard行列 $H$ を使い、モデル保存前に重みを $W_{\text{rot}} = WH^\top$ へ回転してある。推論時には活性側を $X_{\text{rot}} = XH$ へ回転する。この実装の $H$ は対称かつ $HH = I$ なので、量子化誤差を無視すれば次の関係が成立する。

X_{\text{rot}} W_{\text{rot}}^\top = (XH)(WH^\top)^\top = XHHW^\top = XW^\top

QuaRotなどの回転量子化と同じく、回転の目的は答えを変えることではなく、行の一部だけに大きな外れ値が集中するのを散らし、INT8へ丸めたときの誤差を抑えることにある。

$H$ は4×4のHadamard行列のKronecker積を256×256まで積み上げて $\sqrt{256} = 16$ で割って正規化したもので、回転は最終次元を256要素ずつのグループに区切って掛ける。量子化はどちらも行単位で、丸めた値は $[-128, 127]$ へクランプする。数式の記号の読み方はAIの記事でよく出るベクトルと行列の入門、この手の最大絶対値/127をスケールに使うint8量子化はHugging Faceの8bit行列積の解説を読んどくとなんとなくわかる。

元のINT8ConvRotが行う処理を数式で書くとこんな感じ。まずモデル作成時に一度だけ、重みを回転して量子化する。スケールは出力チャンネルごとに持つ。

W_{\text{rot}} = WH^\top,\qquad s_W[n] = \frac{\max_k |W_{\text{rot}}[n,k]|}{127},\qquad Q_W = \mathrm{round}\left(\frac{W_{\text{rot}}}{s_W}\right)

線形層を呼ぶたびに、活性を回転・量子化し、int8同士の行列積をint32で積算してから2つのスケールで実数へ戻す。

X_{\text{rot}} = XH,\qquad s_X[m] = \frac{\max_k |X_{\text{rot}}[m,k]|}{127},\qquad Q_X = \mathrm{round}\left(\frac{X_{\text{rot}}}{s_X}\right)

A = Q_X Q_W^\top \;(\text{int32}),\qquad Y = A \odot s_X s_W^\top

前編の逆量子化パッチは、 $Q_W \odot s_W$ をfp16へ戻し、さらに逆回転してから通常の $XW^\top$ を実行していた。つまり重みを毎回すべて展開する。一方、今回のweight-only経路は活性だけを回転し、回転済みint8重みをそのまま内部演算へ渡すように変えた。

x_rot = rotate(x, h, group_size=256)     # 上の X_rot = XH
out = torch._weight_int8pack_mm(
    x_rot.contiguous(),              # fp16 [M, K]
    qweight.contiguous(),            # int8  [N, K]、上の Q_W
    weight_scale.reshape(-1).half(), # fp16  [N]、上の s_W
)

この変更では重みの逆量子化を明示的に行わないが、活性のquantize_rowwiseとint32積算は消えている。したがって計算上は「元のW8A8 ConvRot」ではなく、「ConvRotを保ったW8A16相当のweight-only経路」である。

とりあえずConvRotのHadamard回転を含む単体スクリプトを作ってテストしてみることにした。通常のfp16行列積と、現行パッチ相当の呼び出しごとにConvRot重みを逆量子化してから行うfp16行列積、そして活性だけをオンラインでHadamard回転してint8重みを_weight_int8pack_mmへ直接渡す経路、それぞれを試して比較してみた。

量子化済みUNetのsafetensorsヘッダーを確認すると、int8の重みは448個あり、形状は6種類だった。全形状でMPS演算が通ることを確認し、活性の行数をM=1024、ConvRotのグループサイズを256に固定して計測した。各経路は5回ウォームアップした後、20回実行を7セット測り、1回あたりの中央値を採用した。モデルロード時に一度だけ行う重みの量子化時間は含めていない。

重み形状（N×K）	該当レイヤー数	fp16	毎回逆量子化	weight-only INT8
2048×2048	168	1.065ms	1.902ms	1.506ms
256×2048	84	0.200ms	0.348ms	0.424ms
6144×256	84	0.474ms	0.807ms	0.909ms
2048×1024	56	0.572ms	1.013ms	0.847ms
8192×2048	28	4.130ms	8.046ms	5.296ms
2048×8192	28	4.283ms	8.275ms	5.383ms

448レイヤーをすべて同じM=1024で1回ずつ呼ぶと仮定し、レイヤー数で単純加重すると、fp16が約503ms、毎回逆量子化が約930ms、weight-only INT8が約711msになる。weight-only INT8は現行パッチ相当より約24%短いが、通常のfp16より約41%長い。細長い256×2048と6144×256の2形状では、毎回逆量子化する経路よりも遅かった。

計算結果も確認した。同じ量子化重みを使った毎回逆量子化経路とweight-only INT8経路の差は、6形状で最大絶対誤差0.00049〜0.00391、平均相対誤差0.20〜0.24%だった。少なくとも単体演算としては、ConvRotの回転方向やスケールの適用を間違ったりしたまま速度だけ測っている状態ではない。

この結果から、Mac GPU上で量子化重みを使う経路そのものは作れたと言える。ただし今回速くなったのは、既存の毎回逆量子化する回避策との比較。素のfp16行列積にはどれも全く速度が追いつかず、モデルカードが示すint8×int8の高速化を検証したことにもなってない。また_weight_int8pack_mmは先頭にアンダースコアが付くPyTorchの内部APIなので、将来の互換性も保証されていない。

int8×int8をMetalカーネルで実行する

とはいえ、weight-onlyで止めず、元のINT8ConvRotと同じ計算段階まで近づけた。次にPyTorch 2.13.0のtorch.mps.compile_shaderを使い、Metal Shading Languageでchar×charをintへ積算する行列積を書いた。16×16の出力タイルを1つのthreadgroupへ割り当て、K方向は32要素ずつthreadgroupメモリへ読み込む。

Metalカーネルの積算部分は次の処理である。実際のコードでは境界判定とタイルの読み込みを前後に加えている。

int acc = 0;
for (uint k0 = 0; k0 < k_size; k0 += 32) {
    // qxとqweightから32要素ずつthreadgroupメモリへ読み込む
    threadgroup_barrier(mem_flags::mem_threadgroup);
    for (uint kk = 0; kk < 32; ++kk) {
        acc += int(tile_a[lid.y][kk]) * int(tile_b[lid.x][kk]);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
}
out[row * n_size + col] = acc;

Python側ではMetalソースを実行時コンパイルし、出力用のint32テンソルと量子化済みの2テンソルを渡す。threadsは16×16のタイルに合わせてMとNを16の倍数へ切り上げ、タイルからはみ出た分はカーネル内の境界判定で余ったスレッドを捨てる。

library = torch.mps.compile_shader(METAL_SOURCE)
kernel = library.int8_gemm_tiled

m_size, k_size = qx.shape
n_size = qweight.shape[0]
out_i32 = torch.empty((m_size, n_size), device="mps", dtype=torch.int32)

dispatch_m = ((m_size + 15) // 16) * 16
dispatch_n = ((n_size + 15) // 16) * 16
kernel(
    out_i32, qx, qweight, m_size, n_size, k_size,
    threads=(dispatch_n, dispatch_m, 1),
    group_size=(16, 16, 1),
)
torch.mps.synchronize()

このカーネルの前後に、先ほど示した $XH$ の回転、活性の行単位INT8量子化、 $s_X$ と $s_W$ の外積による復元を足した。これで処理の並びは元のcomfy_kitchen.int8_linearと同じになる。違うのは、CUDAのtorch._int_mmが内部で使う最適化済みcuBLASLtの代わりに、自作の基本的なMetalカーネルを呼ぶ点と、回転・量子化・スケール復元を別々のPyTorch演算として実行している点である。

まずM=37, N=80, K=256で、MPS側の出力をCPU上のtorch._int_mmと整数のまま比較した。全2,960要素で不一致は0、最大差も0だった。少なくとも「int8へ見せかけて内部ではfp16行列積を呼ぶ」コードではなく、MPS上でint8×int8をint32へ積算している。

次にAnimaで一番数が多いN=2048, K=2048の層（168層）をM=256で測った。例によって5回ウォームアップし、10回実行を7セット測った中央値である。

経路	計算に含めたもの	1回あたり
fp16	$XW^\top$	0.359ms
weight-only	活性回転 + fp16×int8内部演算	0.490ms
自作Metal GEMMのみ	量子化済みint8×int8→int32	5.260ms
自作Metal版の全体	活性回転 + 行単位量子化 + Metal GEMM + スケール復元	5.619ms

計算経路としては元のConvRotへかなり近づいたが、速度はfp16の約15.6倍、weight-only経路の約11.5倍遅かった。ボトルネックは5.260msを占める自作GEMMそのもので、活性の回転や量子化ではない。Apple GPU向けにint8行列積を実行する経路は作れたが、基本的なタイル実装だけでは、フレームワーク内の最適化済みfp16行列積を上回れなかった。

int8×int8演算そのものはMPS上で実行でき、整数出力もCPU参照値と一致した。ただし、今回書けたのは基本的なタイル版であり、Apple GPU向けに高度に最適化されたGEMMではない。将来、最適化済みint8行列積が提供され、ConvRot、活性量子化、スケール復元まで融合できれば速くなる可能性は残る。しかし今回利用できた実装とAPIの範囲では、その仮定を実測で裏付けられなかった。

この段階では、自作Metal版をComfyUIへ組み込む意味はない。組み込めば元の量子化計算をMPSだけで完走できる可能性はあるが、線形層が約11倍以上遅くなるため、CPUフォールバックとは別の形で推論全体を遅くする。Metal側のint8 GEMMを大幅に最適化するか、INT32積算まで含めて利用できる高速な標準演算に置き換えないと実用にならない。weight-only経路をcomfy_kitchenへ組み込む実験はその次に残るが、こちらも単体ではfp16を上回っていないため、生成全体を速くする見込みは薄い。

MPSGraphとMPSMatrixを直接呼ぶ

自作Metalより最適化されたApple純正経路を見落としていないか確認した。検証時のOSはmacOS 26.5で、PyTorchやMLXを介さずObjective-CからMPSGraphへMPSDataTypeInt8の2行列を直接渡した。

MPSGraphはINT8テンソル自体を作成できる。しかしmatrixMultiplicationをコンパイルすると、次のエラーで停止した。

'mps.matmul' op operand #0 must be tensor of floating point values
or tensor of complex values, but got 'tensor<2x4xsi8>'

PyTorchのtorch._int_mmだけが未対応なのではなく、macOS 26.5のMPSGraph行列積もINT8入力を受け付けなかった。

次に、MPSGraphより古いMPSMatrixMultiplicationへ直接MPSMatrixを渡した。こちらはINT8入力を受け付ける。ただしINT32出力を指定すると、Only outputs of MPSDataTypeFloat16 and MPSDataTypeFloat32 are supported for this input typeとして停止する。元のtorch._int_mmと同じINT32出力ではないが、INT8の2行列からfp16またはfp32を直接出す経路は存在した。

呼び出しの中心部分はこんな感じ。チェックポイントの重みは[N,K]だが、この測定ではロード時に一度だけ[K,N]へ並べ替える前提にした。

NSUInteger aRowBytes =
    [MPSMatrixDescriptor rowBytesForColumns:k dataType:MPSDataTypeInt8];
NSUInteger bRowBytes =
    [MPSMatrixDescriptor rowBytesForColumns:n dataType:MPSDataTypeInt8];
NSUInteger cRowBytes =
    [MPSMatrixDescriptor rowBytesForColumns:n dataType:MPSDataTypeFloat32];

MPSMatrixDescriptor *aDesc =
    [MPSMatrixDescriptor matrixDescriptorWithRows:m
                                           columns:k
                                          rowBytes:aRowBytes
                                          dataType:MPSDataTypeInt8];
MPSMatrixDescriptor *bDesc =
    [MPSMatrixDescriptor matrixDescriptorWithRows:k
                                           columns:n
                                          rowBytes:bRowBytes
                                          dataType:MPSDataTypeInt8];
MPSMatrixDescriptor *cDesc =
    [MPSMatrixDescriptor matrixDescriptorWithRows:m
                                           columns:n
                                          rowBytes:cRowBytes
                                          dataType:MPSDataTypeFloat32];

MPSMatrix *a = [[MPSMatrix alloc] initWithBuffer:aBuffer descriptor:aDesc];
MPSMatrix *b = [[MPSMatrix alloc] initWithBuffer:bBuffer descriptor:bDesc];
MPSMatrix *c = [[MPSMatrix alloc] initWithBuffer:cBuffer descriptor:cDesc];

MPSMatrixMultiplication *mm =
    [[MPSMatrixMultiplication alloc] initWithDevice:device
                                      transposeLeft:NO
                                     transposeRight:NO
                                         resultRows:m
                                      resultColumns:n
                                    interiorColumns:k
                                              alpha:1.0
                                               beta:0.0];

[mm encodeToCommandBuffer:commandBuffer
                leftMatrix:a
               rightMatrix:b
              resultMatrix:c];

小さい2×4と4×3の行列では、出力70, 80, 90, 26, 28, 30がCPUで計算した内積と一致した。ただしAPIが公開する出力はfp16またはfp32なので、内部の積算型がINT32であるとは断定できない。

Metal版と同じM=256, N=2048, K=2048の形状を、MPSMatrixMultiplicationのfp16経路と比較した。INT8→fp16では、K=8192でも最大積算値がfp16範囲へ収まるようalpha=1/4096を指定した。後段で行・列スケールを掛ける際に4096を戻せるが、INT32のままスケール復元する元実装とは丸め位置が異なる。

MPSMatrix経路	1回あたり
fp16×fp16→fp16	0.321ms
int8×int8→fp16	0.967ms
int8×int8→fp32	1.075ms
自作MLX Metalのint8×int8→int32	5.166ms

Apple純正経路は自作Metalより約4.8倍短くなった。それでもINT8→fp32はfp16の約3.3倍、最良のINT8→fp16でも約3.0倍遅い。

Animaの6形状をM=1024で測った結果も同じ傾向だった。

重み形状（N×K）	fp16→fp16	INT8→fp16	INT8→fp32
2048×2048	1.201ms	3.321ms	3.507ms
256×2048	0.211ms	0.494ms	0.576ms
6144×256	0.484ms	1.278ms	1.382ms
2048×1024	0.620ms	1.671ms	1.753ms
8192×2048	4.641ms	13.118ms	13.663ms
2048×8192	4.808ms	13.344ms	14.074ms

448レイヤー分を単純加重すると、fp16が約559ms、INT8→fp16が約1,541ms、INT8→fp32が約1,628msになる。最良ケースでもfp16の約2.8倍であり、ここには活性のConvRot、量子化、行・列スケールの復元をまだ含めていない。したがって、このMPSMatrix経路をComfyUIへ組み込んでも生成速度は逆転しない。

2026年に追加されたMetal Performance Primitivesの量子化tensor経路も使えるかと思ったが、AppleのMetal Performance Primitives Programming Guideは、GPU Neural Acceleratorを使う対象をM5としている。AppleのMetal 4案内も量子化形式によるNeural Accelerator利用をM5 ProとM5 Max向けとしており、M1 Maxで今回使える高速経路ではない。

MPSGraphとPyTorchには今回使える経路がなく、MPSMatrixにはINT8入力から浮動小数点を出す経路があった。しかしApple純正実装を直接使ってもfp16より遅く、元のINT8ConvRotが必要とするINT32出力にも一致しなかった。

ComfyUIとPyTorchを外してMLXで測る

ここまでの遅さがComfyUIやPyTorch固有のものか確かめるため、リポジトリ外の一時Python 3.12環境へMLX 0.32.0を入れ、M1 MaxのGPUを直接使う単体ベンチを作った。ComfyUIサーバー、HTTP API、PyTorchは一切使っていない。

MLXのネイティブ8bit経路は、回転済み重みをgroup_size=128のアフィン形式へ量子化し、4個の8bit値をuint32へパックしてquantized_matmulへ渡した。活性はfp16のままなので、これはMLX形式のweight-only経路であり、Anima-Turbo INT8ConvRotの保存形式とは異なる。

qweight, scales, biases = mx.quantize(
    weight_rot,
    group_size=128,
    bits=8,
    mode="affine",
)
out = mx.quantized_matmul(
    x_rot,
    qweight,
    scales,
    biases,
    group_size=128,
    bits=8,
    mode="affine",
)
mx.eval(out)

活性も量子化するMLX標準のQQLinearは、現在mxfp8とnvfp4だけに対応している。mxfp8を同じM=256, N=2048, K=2048の形状で呼ぶと、M1 Maxでは[QQMatmul] NYI for the general caseになり実行できなかった。そのため、符号付きint8の活性と重みをint32へ積算する経路は、MLXのmetal_kernelで前段と同じ16×16タイルのカーネルを書いて比較した。

まずM=37, N=80, K=256で、MLXカスタムMetalのint32出力をNumPyのint32行列積と比較した。全2,960要素で不一致は0、最大差も0だった。PyTorch版と同様、fp16へ逃がさずint8×int8を実行できている。

M=256, N=2048, K=2048の結果は次の通り。同じように5回ウォームアップし、10回実行を7セット測った中央値である。

MLX経路	1回あたり
fp16行列積	0.550ms
fp16行列積 + ConvRot	0.589ms
ネイティブアフィン8bit GEMMのみ	0.533ms
ネイティブアフィン8bit + ConvRot	0.597ms
カスタム符号付きINT8 GEMMのみ	5.166ms
カスタム符号付きINT8ConvRot全体	5.410ms

この形状だけを見ると、MLXネイティブ8bit GEMM単体はfp16より約3%短い。ただし活性のConvRotまで含めるとfp16より約8%長くなった。正確な符号付きINT8ConvRotはfp16の約9.8倍遅く、PyTorchから呼んだMetal版と同じ傾向だった。

さらに量子化済みUNetの6形状をM=1024で測った。

重み形状（N×K）	fp16	MLXアフィン8bit GEMM	ConvRot込み
2048×2048	1.234ms	1.322ms	1.512ms
256×2048	0.405ms	0.479ms	0.582ms
6144×256	0.642ms	0.687ms	0.739ms
2048×1024	0.752ms	0.796ms	0.879ms
8192×2048	4.158ms	4.583ms	4.784ms
2048×8192	4.251ms	4.656ms	5.182ms

448レイヤー分をレイヤー数で単純加重すると、fp16が約573ms、アフィン8bit GEMMが約623ms、ConvRot込みが約693msになった。ネイティブ8bit GEMMだけでもfp16より約9%長く、ConvRot込みでは約21%長い。

したがって「ComfyUI APIを通したから遅かった」という理由ではないことが確定。ComfyUIとPyTorchを完全に外しても、今回の形状ではMLXの8bit weight-onlyはfp16を上回らず、元のINT8ConvRotと同じ符号付きint8×int8を自作するとさらに遅かった。ただしこれはM1 MaxとMLX 0.32.0での結果で、MLX形式のint4や、対応ハードウェア上のmxfp8を試したらまた違う結果になるかもしれない。

なぜ画像生成まで行わなかったか

自作の符号付きINT8ConvRotについては、生成まで回さなくても判断できる条件が揃ったので生成を試さなかった。一番数が多い2048×2048の層だけでfp16の約9.8倍遅く、時間の大半をint8 GEMMが占めているため、この実装をモデルへ組み込んでも、どう転んでも速度が速くならない。

一方、MLXネイティブのweight-only 8bitは、GEMM単体だけなら2048×2048の層で約3%短かった。この程度の差は実際の活性形状、モデルロード、非線形層、VAEデコード、メモリ使用量などで変わる。こちらは、同じモデル、解像度、ステップ数、シードで画像生成全体を比較するまで確定できないと言えばできない。

ただし、公開中のPure MLX版Animaに含まれるのは、bf16/int4のTransformer、テキストエンコーダー、LLM adapter、VAEの重みと設定ファイルである。公開リポジトリ内には、それらを読み込んでサンプリングからVAEデコードまで実行するPythonコードが含まれていない。モデルカードの利用手順もダウンロードまでで、今回のAnima-Turbo INT8ConvRotとは重み形式も異なる。

oMLXが公開している対応範囲はLLM、VLM、埋め込み、リランカーで、画像生成のAPIは含まれない。そのため、ComfyUIをoMLXのAPIへ置き換えるだけでは今回の生成比較はできない。ここから生成全体を測るには、Cosmos Transformer、テキスト経路、サンプラー、VAEを含むAnimaのMLX推論パイプラインを別途実装し、TurboのINT8ConvRot重みを読めるようにするところから始まる。

要するにやってみても時間はかかるが得るものがなさそうなので、生成しなかった。