Tech14 min read

INT8ConvRot GEMM vs fp16 on M1 Max with Metal, MPSMatrix, and MLX (fp16 Wins)

IkesanContents

In part 1, I loaded the Anima-Turbo INT8ConvRot quantized build into ComfyUI on an M1 Max and tried the normal path, the CPU fallback, and a dequantize workaround. The normal path stops because MPS has no torch._int_mm, the CPU fallback is too slow to be usable, and the dequantize workaround computes the fp16 equivalent instead of int8 math.

This time I drop ComfyUI and run matrix multiplications close to INT8ConvRot directly from PyTorch, Metal, MPSMatrix, and MLX, measuring not just whether quantized tensors can be handed to the Apple GPU but whether they beat a plain fp16 matmul. The short version: int8×int8 does execute on the M1 Max GPU, and none of the paths beat fp16.

Test environment

ItemDetails
MachineMac (M1 Max, 64GB RAM)
OSmacOS 26.5
PyTorch2.13.0
MLX0.32.0 (temporary Python 3.12 environment outside the repo)

Testing the MPS weight-only INT8 path in isolation

Part 1 only established that the int8×int8 path INT8ConvRot normally uses is missing. PyTorch 2.13.0 has another internal op with an MPS implementation, torch._weight_int8pack_mm. It takes activations in fp16 or similar, int8 weights, and scales, and runs a matmul using the quantized weights on MPS, a weight-only INT8 op.

It is neither identical to torch._int_mm nor a full substitute. torch._int_mm multiplies int8 activations by int8 weights and accumulates in int32, while the activations passed to _weight_int8pack_mm stay in fp16. It does not reproduce the int8×int8 GEMM that INT8ConvRot uses on CUDA, but it uses quantized tensors directly over a wider range than the current patch, which dequantizes the whole weight back to fp16 on every call.

What changed in the computation

A normal linear layer computes Y=XWY = XW^\top from activations XRM×KX \in \mathbb{R}^{M \times K} and weights WRN×KW \in \mathbb{R}^{N \times K}. ConvRot uses a normalized Hadamard matrix HH applied per 256 elements; the weights are rotated to Wrot=WHW_{\text{rot}} = WH^\top before the model is saved, and at inference time the activations are rotated to Xrot=XHX_{\text{rot}} = XH. This implementation’s HH is symmetric with HH=IHH = I, so ignoring quantization error the following holds.

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

As with rotation quantization schemes like QuaRot, the point of the rotation is not to change the answer. It spreads out the large outliers that concentrate in a few rows, cutting the error when rounding to INT8.

HH is built by stacking Kronecker products of 4×4 Hadamard matrices up to 256×256 and dividing by 256=16\sqrt{256} = 16 to normalize, and the rotation is applied to the last dimension in groups of 256 elements. Quantization is row-wise on both sides, and rounded values are clamped to [128,127][-128, 127]. For how to read the notation, see the intro to vectors and matrices in AI articles; for this max-abs/127-scale style of int8 quantization, Hugging Face’s 8-bit matmul explainer gives a rough idea.

Written out, the original INT8ConvRot does the following. Once at model creation, the weights are rotated and quantized, with one scale per output channel.

Wrot=WH,sW[n]=maxkWrot[n,k]127,QW=round(WrotsW)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)

Every time the linear layer is called, the activations are rotated and quantized, the int8×int8 matmul accumulates in int32, and the two scales bring the result back to real values.

Xrot=XH,sX[m]=maxkXrot[m,k]127,QX=round(XrotsX)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=QXQW  (int32),Y=AsXsWA = Q_X Q_W^\top \;(\text{int32}),\qquad Y = A \odot s_X s_W^\top

Part 1’s dequantize patch turned QWsWQ_W \odot s_W back into fp16, un-rotated it, and then ran the normal XWXW^\top. In other words, it expands the full weight on every call. The weight-only path here instead rotates only the activations and passes the rotated int8 weights straight to the internal op.

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

This change never explicitly dequantizes the weights, but the activation-side quantize_rowwise and the int32 accumulation are gone. Computationally it is therefore not the original W8A8 ConvRot but a W8A16-equivalent weight-only path that keeps ConvRot.

I put together a standalone script with the ConvRot Hadamard rotation and compared three paths: a plain fp16 matmul, the current-patch equivalent that dequantizes the ConvRot weights on every call before an fp16 matmul, and the path that rotates only the activations online and hands the int8 weights directly to _weight_int8pack_mm.

Checking the safetensors header of the quantized UNet, there are 448 int8 weights in 6 shapes. I confirmed the MPS op runs for every shape, then measured with the activation row count fixed at M=1024 and the ConvRot group size at 256. Each path got 5 warmup runs, then 7 sets of 20 runs, taking the median per-call time. The one-time weight quantization at model load is not included.

Weight shape (N×K)Layersfp16Dequantize every callweight-only INT8
2048×20481681.065ms1.902ms1.506ms
256×2048840.200ms0.348ms0.424ms
6144×256840.474ms0.807ms0.909ms
2048×1024560.572ms1.013ms0.847ms
8192×2048284.130ms8.046ms5.296ms
2048×8192284.283ms8.275ms5.383ms

Assuming all 448 layers are each called once at the same M=1024 and weighting simply by layer count, fp16 comes to about 503ms, dequantize-every-call to about 930ms, and weight-only INT8 to about 711ms. Weight-only INT8 is about 24% shorter than the current-patch equivalent but about 41% longer than plain fp16. For the two skinny shapes, 256×2048 and 6144×256, it was even slower than the dequantize-every-call path.

I also checked the outputs. With the same quantized weights, the difference between the dequantize-every-call path and the weight-only INT8 path was a max absolute error of 0.00049-0.00391 and a mean relative error of 0.20-0.24% across the 6 shapes. At least as a standalone op, this is not a case of measuring speed while getting the ConvRot rotation direction or the scale application wrong.

A path that uses quantized weights on the Mac GPU can be built, then. But what got faster here is only relative to the existing dequantize-every-call workaround. None of these come anywhere near the plain fp16 matmul, and this does not verify the int8×int8 speedup the model card claims either. _weight_int8pack_mm is also a PyTorch internal API with a leading underscore, so future compatibility is not guaranteed.

Running int8×int8 in a Metal kernel

Still, rather than stopping at weight-only, I pushed to the same computation stage as the original INT8ConvRot. Using PyTorch 2.13.0’s torch.mps.compile_shader, I wrote a matmul in Metal Shading Language that accumulates char×char into int. Each 16×16 output tile is assigned to one threadgroup, and the K dimension is loaded into threadgroup memory 32 elements at a time.

The accumulation part of the Metal kernel is below. The actual code adds bounds checks and tile loads around it.

int acc = 0;
for (uint k0 = 0; k0 < k_size; k0 += 32) {
    // load 32 elements each from qx and qweight into threadgroup memory
    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;

On the Python side, the Metal source is compiled at runtime and given an int32 output tensor plus the two quantized tensors. threads rounds M and N up to multiples of 16 to match the 16×16 tile, and threads that fall outside the tile are discarded by the bounds check inside the kernel.

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()

Before and after this kernel I added the XHXH rotation shown earlier, row-wise INT8 quantization of the activations, and the restore via the outer product of sXs_X and sWs_W. The sequence of operations now matches the original comfy_kitchen.int8_linear. The differences are that instead of the optimized cuBLASLt that CUDA’s torch._int_mm uses internally, it calls my basic hand-written Metal kernel, and that the rotation, quantization, and scale restore run as separate PyTorch ops.

First, at M=37, N=80, K=256, I compared the MPS output against torch._int_mm on CPU, staying in integers. Across all 2,960 elements there were 0 mismatches and the max difference was 0. At minimum, this is not code that pretends to be int8 while calling an fp16 matmul internally; it accumulates int8×int8 into int32 on MPS.

Next I measured the most common layer in Anima, N=2048, K=2048 (168 layers), at M=256. As before, 5 warmup runs, then the median of 7 sets of 10 runs.

PathWhat is includedPer call
fp16XWXW^\top0.359ms
weight-onlyactivation rotation + fp16×int8 internal op0.490ms
custom Metal GEMM onlyquantized int8×int8→int325.260ms
custom Metal full pipelineactivation rotation + row-wise quantization + Metal GEMM + scale restore5.619ms

The computation path is now quite close to the original ConvRot, but it ran about 15.6x slower than fp16 and about 11.5x slower than the weight-only path. The bottleneck is the custom GEMM itself at 5.260ms, not the activation rotation or the quantization. A path that runs int8 matmuls on the Apple GPU exists now, but a basic tiled implementation alone cannot beat the framework’s optimized fp16 matmul.

The int8×int8 op itself runs on MPS and the integer output matches the CPU reference. What I could write, though, is a basic tiled version, not a GEMM highly optimized for Apple GPUs. If an optimized int8 matmul ships someday and ConvRot, activation quantization, and scale restore can be fused into it, the possibility of a speedup remains. Within the implementations and APIs available this time, that assumption could not be backed by measurement.

At this stage there is no point integrating the custom Metal version into ComfyUI. It might let the original quantized computation run entirely on MPS, but the linear layers get 11x+ slower, which slows down overall inference in a different way than the CPU fallback. It will not be practical without either heavily optimizing the Metal int8 GEMM or replacing it with a fast standard op that covers the INT32 accumulation. Integrating the weight-only path into comfy_kitchen remains as the next experiment, but since that path does not beat fp16 on its own either, the odds of it speeding up generation overall are slim.

Calling MPSGraph and MPSMatrix directly

To make sure I was not overlooking an Apple-native path more optimized than my Metal kernel, I passed two MPSDataTypeInt8 matrices to MPSGraph directly from Objective-C, without PyTorch or MLX, on macOS 26.5.

MPSGraph can create INT8 tensors. But compiling matrixMultiplication stops with this error.

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

It is not just PyTorch’s torch._int_mm that lacks support: the MPSGraph matmul on macOS 26.5 does not accept INT8 inputs either.

Next I passed MPSMatrix objects directly to the older MPSMatrixMultiplication. This one accepts INT8 inputs. Specifying INT32 output, however, stops with Only outputs of MPSDataTypeFloat16 and MPSDataTypeFloat32 are supported for this input type. It is not the same INT32 output as the original torch._int_mm, but a path that takes two INT8 matrices and directly produces fp16 or fp32 does exist.

The core of the call looks like this. The checkpoint weights are [N,K], but for this measurement I assumed a one-time rearrangement to [K,N] at load.

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];

For small 2×4 and 4×3 matrices, the output 70, 80, 90, 26, 28, 30 matched dot products computed on CPU. Since the API only exposes fp16 or fp32 outputs, though, I cannot conclude that the internal accumulation type is INT32.

I compared the same M=256, N=2048, K=2048 shape as the Metal version against MPSMatrixMultiplication’s fp16 path. For INT8→fp16, I specified alpha=1/4096 so the maximum accumulated value stays within fp16 range even at K=8192. The 4096 can be reapplied later when the row and column scales are applied, but the rounding position differs from the original implementation, which restores scales while still in INT32.

MPSMatrix pathPer call
fp16×fp16→fp160.321ms
int8×int8→fp160.967ms
int8×int8→fp321.075ms
custom MLX Metal int8×int8→int325.166ms

The Apple-native path came out about 4.8x shorter than my custom Metal kernel. Even so, INT8→fp32 is about 3.3x slower than fp16, and the best case, INT8→fp16, is still about 3.0x slower.

Measuring Anima’s 6 shapes at M=1024 showed the same trend.

Weight shape (N×K)fp16→fp16INT8→fp16INT8→fp32
2048×20481.201ms3.321ms3.507ms
256×20480.211ms0.494ms0.576ms
6144×2560.484ms1.278ms1.382ms
2048×10240.620ms1.671ms1.753ms
8192×20484.641ms13.118ms13.663ms
2048×81924.808ms13.344ms14.074ms

Weighting the 448 layers simply gives about 559ms for fp16, about 1,541ms for INT8→fp16, and about 1,628ms for INT8→fp32. Even the best case is about 2.8x fp16, and that does not yet include the activation ConvRot, the quantization, or the row/column scale restore. Plugging this MPSMatrix path into ComfyUI will not flip generation speed either.

I also wondered whether the quantized tensor path in Metal Performance Primitives, added in 2026, could help, but Apple’s Metal Performance Primitives Programming Guide targets the GPU Neural Accelerator at M5, and Apple’s Metal 4 page also lists Neural Accelerator use with quantized formats for M5 Pro and M5 Max. It is not a fast path M1 Max can use this time.

MPSGraph and PyTorch have no usable path here, while MPSMatrix has a path from INT8 inputs to floating-point outputs. But even calling Apple’s own implementation directly is slower than fp16, and it does not match the INT32 output the original INT8ConvRot requires.

Measuring with MLX, without ComfyUI or PyTorch

To check whether the slowness so far is specific to ComfyUI or PyTorch, I installed MLX 0.32.0 into a temporary Python 3.12 environment outside the repo and built a standalone benchmark that uses the M1 Max GPU directly. No ComfyUI server, no HTTP API, no PyTorch.

For MLX’s native 8-bit path, I quantized the rotated weights into the affine format with group_size=128, packing four 8-bit values into a uint32, and passed them to quantized_matmul. The activations stay in fp16, so this is an MLX-style weight-only path, different from the Anima-Turbo INT8ConvRot storage format.

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’s standard QQLinear, which quantizes the activations too, currently supports only mxfp8 and nvfp4. Calling mxfp8 at the same M=256, N=2048, K=2048 shape fails on M1 Max with [QQMatmul] NYI for the general case. For the path that accumulates signed int8 activations and weights into int32, I wrote the same 16×16 tiled kernel as before with MLX’s metal_kernel and compared.

First, at M=37, N=80, K=256, I compared the MLX custom Metal int32 output against a NumPy int32 matmul. All 2,960 elements matched with 0 mismatches and a max difference of 0. Like the PyTorch version, it runs int8×int8 without escaping to fp16.

The results at M=256, N=2048, K=2048 are below, again with 5 warmup runs and the median of 7 sets of 10 runs.

MLX pathPer call
fp16 matmul0.550ms
fp16 matmul + ConvRot0.589ms
native affine 8-bit GEMM only0.533ms
native affine 8-bit + ConvRot0.597ms
custom signed INT8 GEMM only5.166ms
custom signed INT8ConvRot full pipeline5.410ms

Looking at this shape alone, the native MLX 8-bit GEMM by itself is about 3% shorter than fp16. Including the activation ConvRot, it turns about 8% longer than fp16. The exact signed INT8ConvRot is about 9.8x slower than fp16, the same trend as the Metal version called from PyTorch.

I then measured the 6 shapes of the quantized UNet at M=1024.

Weight shape (N×K)fp16MLX affine 8-bit GEMMWith ConvRot
2048×20481.234ms1.322ms1.512ms
256×20480.405ms0.479ms0.582ms
6144×2560.642ms0.687ms0.739ms
2048×10240.752ms0.796ms0.879ms
8192×20484.158ms4.583ms4.784ms
2048×81924.251ms4.656ms5.182ms

Weighting the 448 layers by layer count gives about 573ms for fp16, about 623ms for the affine 8-bit GEMM, and about 693ms with ConvRot. The native 8-bit GEMM alone is about 9% longer than fp16, and about 21% longer with ConvRot.

That settles one question: “it was slow because of the ComfyUI API” is not the reason. Even with ComfyUI and PyTorch completely removed, MLX’s 8-bit weight-only does not beat fp16 at these shapes, and hand-building the same signed int8×int8 as the original INT8ConvRot is slower still. This is a result for M1 Max and MLX 0.32.0, though; MLX-format int4, or mxfp8 on supported hardware, might tell a different story.

Why I did not run image generation

For the custom signed INT8ConvRot, the conditions to judge without generating were already in place, so I did not try. The most common 2048×2048 layer alone is about 9.8x slower than fp16 and the int8 GEMM takes most of that time, so no matter how it plays out, wiring this implementation into the model cannot make it faster.

MLX’s native weight-only 8-bit, on the other hand, was about 3% shorter for the GEMM alone at the 2048×2048 layer. A gap that small shifts with the actual activation shapes, model load, nonlinear layers, VAE decode, and memory use. Strictly speaking, this one cannot be settled without comparing full image generation on the same model, resolution, step count, and seed.

However, the published Pure MLX Anima contains the bf16/int4 Transformer, text encoder, LLM adapter, and VAE weights plus config files. The public repo includes no Python code that loads them and runs sampling through VAE decode. The model card’s usage steps end at download, and the weight format also differs from this Anima-Turbo INT8ConvRot.

oMLX’s published coverage is LLMs, VLMs, embeddings, and rerankers, with no image generation API. Swapping ComfyUI for the oMLX API alone would not enable this generation comparison. Measuring full generation from here starts with implementing a separate MLX inference pipeline for Anima, covering the Cosmos Transformer, the text path, the sampler, and the VAE, and making it read the Turbo INT8ConvRot weights.

In short, it would take time and likely give nothing back, so I did not generate.