This post covers getting Intel AMX to do a matrix multiply, and why the memory layout of one of the input matrices is the detail that causes the most confusion.
A note on the title: the layout is called VNNI, after the AVX-512 dot-product instructions it was designed for. AMX inherited it.
What AMX actually is
One sentence: a normal register holds a list, a tile register holds a grid.
AVX-512 register (1D) AMX tile register (2D)
┌────────────────────┐ ┌────────────────────┐
│ 3 7 1 4 9 2 …│ │ 3 7 1 4 … │
└────────────────────┘ │ 9 2 5 8 … │
one row of 16 numbers │ 6 0 3 1 … │
│ ⋮ │
└────────────────────┘
16 rows × 64 bytes
Once the silicon can hold a grid, it can multiply two grids in one
instruction. Eight tiles, tmm0 through tmm7, and a unit called the TMUL
that takes two tiles and accumulates into a third.
That is the whole idea. Everything else is plumbing.
Matmul, by hand, with small numbers
Before touching intrinsics, do it on paper. A is 2×3, B is 3×2:
A B
┌ ┐ ┌ ┐
│ 1 2 3 │ │ 1 0 │
│ 4 5 6 │ │ 0 1 │
└ ┘ │ 2 2 │
└ ┘
C[0][0] walks row 0 of A across column 0 of B:
row0 of A: 1 2 3
× × ×
col0 of B: 1 0 2
↓ ↓ ↓
1 + 0 + 6 = 7
All four cells:
C
┌ ┐
│ 7 8 │
│ 16 17 │
└ ┘
Twelve multiplies. Three ways to execute them:
- Scalar CPU: twelve instructions.
- AVX-512: load
[1,2,3], load[1,0,2], multiply, horizontal-sum. One dot product at a time, four passes. - AMX: load A into one tile, B into another, one instruction. All twelve multiplies happen inside the TMUL grid; four output cells come out at once.
┌──────────┐ ┌──────────┐
│ A │ │ B │
│ (tmm1) │ │ (tmm2) │
└────┬─────┘ └────┬─────┘
│ │
└──────► TMUL ◄──────┘ one instruction
│
▼
┌──────────┐
│ C │
│ (tmm0) │
└──────────┘
The TMUL is physically shaped like the output. Every cell computes its own dot product in parallel. The instruction takes fixed latency regardless of how many cells are active — unused cells waste cycles silently.
The real numbers
Swap in the actual hardware shape. For BF16:
A B C
16 × 32 × 32 × 16 = 16 × 16
(bf16) (bf16) (fp32)
256 output cells, each a dot product of length 32. That’s 8,192 multiply-adds, in one instruction, in about 16 cycles. 1,024 FLOPs per cycle per core.
An AVX-512 FMA does 32 FLOPs per cycle. The 32× ratio is the entire reason AMX exists.
A digression on the FLOP count
Count the toy by hand: per output cell you get K multiplies and K−1 adds, so 2K−1 operations, not 2K. Every FLOP-counting convention says 2K. Why?
Because the instruction is tmm0 += tmm1 × tmm2 — it always accumulates. C is a
running total across the K-loop, never a fresh assignment:
C[0][0] = 0 ← the accumulator exists before you start
C[0][0] += 1·1 ← mult, add
C[0][0] += 2·0 ← mult, add
C[0][0] += 3·2 ← mult, add
3 mults + 3 adds = 2K
The extra add is the +=. It’s a real FMA the hardware really performs. So
FLOPs = 2·M·N·K, and for one tile op that’s 2 · 16 · 16 · 32 = 16,384.
Ceiling number one: the length of K
Back to the toy. The 3 in “A is 2×3” is the dot-product length.
On AMX that dimension is fixed at 32 for BF16 and 64 for INT8.
K = 32 → every multiply unit busy 100% of peak
K = 16 → half the grid computes nothing 50%
K = 4 → most of the grid idles 12.5%
The instruction takes 16 cycles regardless. Short K means you paid for the whole grid and used a corner of it. LIBXSMM documents this floor, and it is the most common reason a hand-rolled AMX kernel underperforms.
For a real K = 4096 you fire the instruction 128 times into the same C tile:
C_tile ← 0
for k_chunk in 0..127:
load A_chunk, B_chunk
C_tile += A_chunk × B_chunk
store C_tile ← once, at the very end
C never leaves the register file. Only A and B stream. That’s the amortization,
and it’s why += is the natural semantics.
The layout, which is the actual point of this post
The part that is not obvious from the ISA reference.
The TMUL does not fetch one element at a time. It fetches a 4-byte dword — for BF16 that is exactly two values. It multiplies both pairs and accumulates. Two K-steps per fetch, always.
one fetch from A: ( a[k], a[k+1] )
one fetch from B: ( b[k], b[k+1] )
↓
a[k]·b[k] + a[k+1]·b[k+1]
Dword is a 4-byte chunk of memory. The hardware does not know what your matrix means.
For A this is free. Consecutive K values live next to each other in a row.
For B, consecutive K values are in different rows. That’s the problem. So you pre-shuffle B until the pairs sit together.
Watch it happen
Toy with an even K. A is 2×4, B is 4×2:
A B
┌ ┐ ┌ ┐
│ 1 2 3 4 │ │ 1 0 │ k=0
│ 5 6 7 8 │ │ 0 1 │ k=1
└ ┘ │ 2 2 │ k=2
│ 3 1 │ k=3
└ ┘
n=0 n=1
By hand, C[0][0] = 1·1 + 2·0 + 3·2 + 4·3 = 19.
The transformation:
B_packed[ k/2 ][ 2n + (k mod 2) ] = B[k][n]
Walk every element:
B[0][0]=1 → row 0, col 0 B[0][1]=0 → row 0, col 2
B[1][0]=0 → row 0, col 1 B[1][1]=1 → row 0, col 3
B[2][0]=2 → row 1, col 0 B[2][1]=2 → row 1, col 2
B[3][0]=3 → row 1, col 1 B[3][1]=1 → row 1, col 3
A (unchanged) B (on paper) B_packed (in memory)
┌ ┐ ┌ ┐ ┌─────────┬─────────┐
│ 1 2 3 4 │ │ 1 0 │ k=0 │ 1 0 │ 0 1 │ ← k=0,1
│ 5 6 7 8 │ │ 0 1 │ k=1 ├─────────┼─────────┤
└ ┘ │ 2 2 │ k=2 │ 2 3 │ 2 1 │ ← k=2,3
dword0 dword1 │ 3 1 │ k=3 └─────────┴─────────┘
└ ┘ n=0 n=1
n=0 n=1 each cell = one dword = one K-pair
Now the multiply, two fetches:
Fetch 1:
A row 0, dword 0 = (1, 2) ← k = 0,1
B_packed row 0, n=0 dword = (1, 0) ← k = 0,1
1·1 + 2·0 = 1
Fetch 2:
A row 0, dword 1 = (3, 4) ← k = 2,3
B_packed row 1, n=0 dword = (2, 3) ← k = 2,3
3·2 + 4·3 = 18
C[0][0] = 19 ✓
The packing exists so fetch 2 lands on (2, 3) — the k=2 and k=3 values of
column 0, adjacent.
What happens if you skip it
Feed B row-major, reinterpreted as 2 rows of 4:
row 0: [ 1 0 0 1 ]
row 1: [ 2 2 3 1 ]
Right shape. Loads fine. No fault. And then:
Fetch 2:
A row 0, dword 1 = (3, 4)
row 1, n=0 dword = (2, 2) ← B[2][0] and B[2][1] !
two different COLUMNS, same k
3·2 + 4·2 = 14
C[0][0] = 15 ✗ (should be 19)
Wrong answer. Silent. The hardware multiplied a column-0 value against a column-1 value because it has no idea what your matrix means. It just fetches dwords.
This is the bug that cost me two days. And if you test with a matrix of all ones, everything passes — the layout error is invisible.
Toy #1: doing it by hand
The smallest AMX program that computes something.
// Build: gcc -O2 -mamx-tile -mamx-bf16 amx_toy.c -o amx_toy
// Needs: Sapphire Rapids or newer, Linux 5.16+
#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <immintrin.h>
#include <sys/syscall.h>
#include <unistd.h>
#define ARCH_REQ_XCOMP_PERM 0x1023
#define XFEATURE_XTILEDATA 18
typedef struct {
uint8_t palette_id;
uint8_t start_row;
uint8_t reserved[14];
uint16_t colsb[16]; // bytes per row, per tile
uint8_t rows[16]; // rows, per tile
} __attribute__((packed)) tileconfig;
#define BF16_ONE 0x3F80 // top 16 bits of the fp32 pattern for 1.0f
int main(void) {
// Ask the kernel for permission. Skip this and the first tile
// instruction dies with SIGILL. Nobody warns you about this.
if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) {
fprintf(stderr, "no AMX here\n");
return 1;
}
// tile0 = C: 16x16 fp32. tile1 = A: 16x32 bf16. tile2 = B: 16x32 bf16.
// All three are 16 rows x 64 bytes. The hardware counts bytes, not elements.
tileconfig cfg;
memset(&cfg, 0, sizeof(cfg));
cfg.palette_id = 1;
for (int t = 0; t < 3; t++) { cfg.rows[t] = 16; cfg.colsb[t] = 64; }
_tile_loadconfig(&cfg);
uint16_t A[16][32], B[16][32];
float C[16][16];
for (int m = 0; m < 16; m++)
for (int k = 0; k < 32; k++) A[m][k] = BF16_ONE;
for (int r = 0; r < 16; r++)
for (int c = 0; c < 32; c++) B[r][c] = BF16_ONE;
memset(C, 0, sizeof(C));
_tile_loadd(0, C, 16 * sizeof(float)); // strides are in BYTES
_tile_loadd(1, A, 32 * sizeof(uint16_t));
_tile_loadd(2, B, 32 * sizeof(uint16_t));
_tile_dpbf16ps(0, 1, 2); // C += A * B
_tile_stored(0, C, 16 * sizeof(float));
_tile_release(); // let the core leave the AMX
// power license
printf("C[0][0] = %.1f (expected 32.0)\n", C[0][0]);
return 0;
}
Three things that will trip you up:
- The
arch_prctlcall is mandatory. Linux gates tile state behind per-process permission. Skip it and the first tile instruction dies with SIGILL. - B must be VNNI-packed. All-ones hides the bug — use distinct values.
- The tile shape is fixed. One
_tile_dpbf16psis always 16×32 × 32×16. A real GEMM is a loop nest around this, holding C tiles intmm0..3, A intmm4..5, B intmm6..7.
Toy #2: doing it with LIBXSMM
LIBXSMM is an open-source JIT library out of Intel Labs (Heinecke et al., SC'16). On BF16 it outperforms oneDNN by around 1.45×, largely because oneDNN does not keep B in blocked layout and suffers cache-conflict misses at leading dimension 4096.
Same toy, LIBXSMM style. Note that it’s column-major, BLAS convention.
/* C(2x2) = A(2x4) * B(4x2), bf16 in, fp32 out.
* LIBXSMM is COLUMN-MAJOR: A[m][k] lives at a[m + k*lda]. */
#include <libxsmm.h>
#include <stdio.h>
int main(void) {
const libxsmm_blasint m = 2, n = 2, k = 4;
float a_f32[8] = { 1,5, 2,6, 3,7, 4,8 }; /* [1 2 3 4; 5 6 7 8] */
float b_f32[8] = { 1,0,2,3, 0,1,2,1 }; /* [1 0; 0 1; 2 2; 3 1] */
libxsmm_bfloat16 A[8], B[8], B_vnni[8];
float C[4] = { 0 };
libxsmm_rne_convert_fp32_bf16(a_f32, A, 8);
libxsmm_rne_convert_fp32_bf16(b_f32, B, 8);
/* --- pack B. The ugly detail, as a library call. --- */
libxsmm_meltw_unary_shape ushape = libxsmm_create_meltw_unary_shape(
n, k, k, k,
LIBXSMM_DATATYPE_BF16, LIBXSMM_DATATYPE_BF16, LIBXSMM_DATATYPE_BF16);
libxsmm_meltwfunction_unary vnni = libxsmm_dispatch_meltw_unary(
LIBXSMM_MELTW_TYPE_UNARY_TRANSFORM_NORM_TO_VNNI2,
ushape, LIBXSMM_MELTW_FLAG_UNARY_NONE);
libxsmm_meltw_unary_param uparam;
uparam.in.primary = B;
uparam.out.primary = B_vnni;
vnni(&uparam);
/* --- dispatch: JIT a kernel for exactly these sizes --- */
libxsmm_gemm_shape shape = libxsmm_create_gemm_shape(
m, n, k,
m, k, m, /* lda, ldb, ldc */
LIBXSMM_DATATYPE_BF16, LIBXSMM_DATATYPE_BF16,
LIBXSMM_DATATYPE_F32, LIBXSMM_DATATYPE_F32);
libxsmm_bitfield flags = LIBXSMM_GEMM_FLAGS('N','N')
| LIBXSMM_GEMM_FLAG_VNNI_B /* B is pre-packed */
| LIBXSMM_GEMM_FLAG_BETA_0; /* C = A*B, not += */
libxsmm_gemmfunction kernel =
libxsmm_dispatch_gemm(shape, flags, LIBXSMM_GEMM_PREFETCH_NONE);
if (!kernel) { fprintf(stderr, "dispatch failed\n"); return 1; }
/* --- run --- */
libxsmm_gemm_param p;
p.a.primary = A;
p.b.primary = B_vnni;
p.c.primary = C;
kernel(&p);
/* [19 12 ; 43 28], column-major */
printf("C = [ %.0f %.0f ; %.0f %.0f ]\n", C[0], C[2], C[1], C[3]);
return 0;
}
No tile intrinsics anywhere. arch_prctl, tileconfig, _tile_release — all
handled internally. The same source compiles for AVX-512, AVX2, and ARM SVE
without #ifdef or -march.
LIBXSMM_GEMM_FLAG_VNNI_B is a promise, not a request. It tells the kernel
“B is already packed” but does not pack it for you. Forget the transform, keep
the flag, and the same silent-wrong-answer bug reappears.
What “dispatch” means
The two-phase API looks odd at first but the idea is simple.
An ordinary GEMM implementation does bookkeeping in the inner loop — is this the last iteration? is K a multiple of 32? do I need a cleanup path? Those branches cost cycles, and the same binary must handle k=4 and k=4096.
LIBXSMM turns this around: what if m, n, k were known before generating the code?
You: "m=2, n=2, k=4, bf16 in, fp32 out, B pre-packed."
│
[ D I S P A T C H ]
│
emits x86 bytes into an executable page:
tileloadd tmm1, [rdi]
tileloadd tmm2, [rsi]
tdpbf16ps tmm0, tmm1, tmm2
tilestored [rdx], tmm0
ret
│
function pointer
▼
You: kernel(¶m);
No loop counters, no branches. The loop was unrolled at code-generation time because the trip count was a constant — just not at your compile time.
| cost | how often | |
|---|---|---|
| dispatch | microseconds | once, at startup |
| run | nanoseconds | millions of times |
Calling libxsmm_dispatch_gemm inside your hot loop is the classic mistake. It’s
cached by shape, so it isn’t fatal, but the hash lookup costs more than the
kernel.
The VNNI transform uses the same shape → dispatch → run pattern, which
is why that block of code looks heavy for what it does — same three steps.
meltw is Intel-speak for “elementwise”, unary means one input tensor
one output tensor, and NORM_TO_VNNI2 is one entry in a menu that also
includes ReLU, transpose, etc.
In practice you rarely do the transform at runtime. B is the weight matrix — weights do not change between tokens. Pack once at model load and store them packed. A is the activations, new every token, but A needs no packing. Only B requires the layout shuffle.
BF16 vs INT8 vs INT4
A common assumption is that lower precision always means faster math. Partially true.
AMX has no INT4 instruction. The ISA has tdpbf16ps for BF16, the
tdpb{ss,su,us,uu}d family for INT8, and tdpfp16ps for FP16 on Granite Rapids.
That is the complete set. Any “AMXINT4” path means 4-bit storage,
unpacked to INT8 before the tile op.
| elem/row | K per op | ops/cycle/core | |
|---|---|---|---|
| BF16 | 32 | 32 | 1,024 FLOP |
| INT8 | 64 | 64 | 2,048 OP |
| INT4 | — | — | 2,048 OP (via INT8) |
INT8 doubles BF16 purely because twice as many elements fit in a 64-byte row. INT4 buys zero extra compute, and costs you the unpack.
So why does anyone use it? Because prefill and decode are different animals.
PREFILL — compute-bound. Weights reused across many tokens.
BF16 ████████████████ ~288 TFLOPS
INT8 ████████████████████████████████ ~576 TOPS (2x)
INT4 ██████████████████████████████ ~576 minus unpack
DECODE — bandwidth-bound. Each weight read once, used once.
bytes/param → arithmetic intensity → ceiling
BF16 2.0 B 1 OP/B ████ ~0.85 TFLOPS
INT8 1.0 B 2 OP/B ████████ ~1.7 TOPS (2x)
INT4 0.5 B 4 OP/B ████████████████ ~3.4 TOPS (4x)
During decode you pay for bytes moved, not ops issued. INT4 halves the bytes. The TMUL is idle either way — waiting on DRAM at something like 0.2% utilization — so the unpack is effectively free. During prefill you pay for ops, INT4 gives the hardware nothing extra, and the shift-and-mask sequence eats into your gains.
There is also the capacity angle. For a 671B-parameter model (e.g. DeepSeek-V3):
| format | footprint |
|---|---|
| BF16 | ~1.34 TB |
| INT8 | ~671 GB |
| INT4 | ~336 GB |
INT4 is what makes single-node inference at this scale feasible at all. It is not a speed trick — it is a fitting trick.
How you store a number and how you multiply it are two separate choices.
Portions of this post were drafted with the assistance of AI and subsequently reviewed and edited for accuracy.