gemm: asymmetric INT8xINT4 (4-bit weights) on AIE2P 4x16x16 mmul — bit-exact + W4A8 llama decode - #182
Open
bong-water-water-bong wants to merge 19 commits into
Open
gemm: asymmetric INT8xINT4 (4-bit weights) on AIE2P 4x16x16 mmul — bit-exact + W4A8 llama decode#182bong-water-water-bong wants to merge 19 commits into
bong-water-water-bong wants to merge 19 commits into
Conversation
added 16 commits
August 31, 2026 17:11
- aie_kernels/aie2p/mm.cc: software-pipeline the k-loop in
matmul_vectorized_2x2_mmul — prologue loads k-step 0, the loop body
consumes the previous step's A/B tiles while prefetching the next, and
an epilogue MACs the final step. Hides load->vmac latency (verified
bit-exact vs numpy int32 on NPU Strix Halo; ~7.8 TOPS at 2048^3 i8->i32).
- iron/operators/gemm/op.py: per-dtype tile-multiple validation matching
the kernel static_asserts; _kernel_dtype_flag selects the {combo}_ONLY
define (bf16_f32_ONLY under prio_accuracy since design.py resolves
matmul_bf16_f32); reject integer outputs narrower than the 32-bit
accumulator (i8->i8/i8->i16/i16->i16 truncate).
- iron/operators/gemm/design.py: merge duplicate npu1 entry; add i8/i16
MAC dims for npu1/npu2.
- bench_int8.py / int8_bench.py / run_int8_gemm.py: NPU INT8 GEMM
benchmarks (bit-exactness + TOPS, 2048^3 and shape sweeps).
Known limits (upstream design, not this kernel): first dispatch after an
xclbin reload in a multi-shape process can race the zero kernel (retry
exact); N=8192 with 8 columns exceeds the aie.dma_bd stride range
(per-column C slice 2^21 > 2^20) — use partition_N for very wide outputs.
…T8 bench - aie_kernels/aie2p/mm.cc: chess_loop_range min no longer overstates the software-pipelined k-loop trip count (colA-1). The previous fixed (4,) hint mis-scheduled small K tiles (e.g. DIM_K=32, s=8 -> 3 trips). The hint is now colA-1 clamped to [1,4]; colA==1 keeps a statically-empty loop. Verified bit-exact on NPU at 2048^3 (8.42 TOPS best, slightly better than the flat hint). - bench_int8_gemm.py: single parameterized harness replacing the three scratch scripts (bench_int8.py, int8_bench.py, run_int8_gemm.py) — shapes/reps/partition/tiles/seed flags, bit-exactness + TOPS, and an N-partition path for wide outputs past the aie.dma_bd stride cap. - Known flake (not this kernel): rare (~5-10%) transient wrong result on a dispatch after multiple distinct xclbins were compiled in one process (XRT/amdxdna context race on Strix Halo; matches the hw_context retry documented in iron/common/sequence.py). Warm-up + verify retry handles it.
- iron/operators/gemm/op.py: GEMM docstring documents the rare (~5%) transient wrong result on NPU2 after several distinct xclbins compile in one process (zero/accumulate write races first submit on a fresh context; self-heals next dispatch), and the production guidance: warm-up once, verify first result, retry once on mismatch. - bench_int8_gemm.py: --verify-retry re-runs once on a result mismatch and reports the retry outcome, so the transient flake is not read as a kernel failure.
One full-ELF OperatorSequence replaces the six per-op dispatches per
layer (gate/up GEMM, SiLU, mul, down GEMM, residual add). Weights and
activations round-trip through the shared fused buffers; the residual
is read back with a reshape (the fused 'x' subview is flat 1-D).
Verified on NPU2 (Ryzen AI MAX+ 395):
- prefill logits bit-identical to the separate-op path (corr 1.0)
- run-to-run bit-identical (deterministic)
- decode text identical ('SCENE I. The King' for 7-token prompt)
- TTFT 2.62-2.65s vs 2.72s separate (~3% prefill saving)
Also ignore build_elf_*/ artifact dirs.
Prefill GEMMs/elementwise ops ran at M=max_seq_len (2048) regardless of the actual prompt (7 tokens) -- ~293x wasted compute per op. Split the sizing: prefill ops now build at prefill_len = ceil(seq_len/512)*512 (GEMM tile constraints: M%(tile_m*4)==0, N%(tile_n*8)==0 for bf16, 8 cols), while decode ops and KV caches keep max_seq_len. AIELlamaOperators/AIELlamaBuffers take prefill_len (default: prompt_len, back-compat). main() derives it from the tokenized prompt length. Measured (7-token prompt, Ryzen AI MAX+ 395): - TTFT 2.72s -> 0.945s (separate ops), 0.921s (fused FFN) -- ~2.9x - prefill logits bit-identical to the 2048-built run (corr 1.0) - corr 0.9963 vs fresh CPU reference (top-1 agrees); decode text identical
…t-exact, half B bandwidth Adds dtype_b='i4' to GEMM: B is packed (K, N//2) int8 (two nibbles per byte, low nibble first) and the kernel uses the AIE2P 4x16x16 mmul (1024 MACs/instr vs 512 for int8xint8 8x8x8). The int4_t empty-struct sizeof==1 trap is handled with a B_ADV pointer-advance correction in mm.cc (all manual B arithmetic halves the element count for int4). design.py: packed-B tiling + L2->L1 fifo dims (verified: the generated BD is [4,4,16,8]/[512,8,32,1], the mac sees B(kk',nn)=B4[16kk+kk',16nn+nn]). op.py: dtype_b field, i8_i4_ONLY kernel flag, pack_i4 static method, packed arg spec. bench_int8_gemm.py: --b-i4 flag. Verified on NPU2 (gfx1151, 8 AIE cols, 64x64x64 tiles), bit-exact vs numpy: i8xi4 2048^3: 9.52 TOPS (vs i8xi8 8.02, +19%) — and B bytes halved i8xi4 1024^3: 4.27, 512^3: 1.27 (all exact=True bad=0 max_abs=0)
…e through i8xi4) Quantizes real safetensors weights to INT4 (per-output-neuron scales) and runs them through the asymmetric i8xi4 NPU GEMM with INT8 activations. Bit-exact int math on the NPU (== CPU int reference); quantization loss vs the bf16 reference: q_proj corr 0.9885, gate_proj corr 0.9902 with top-4 identical. This is the proof for threading i4 weights into the llama pipeline (W4A8 path).
… per-token i8 acts) llama_w4a8_npu.py runs all 16 layers of llama-3.2-1B with the NPU doing the 7 heavy GEMMs/layer as asymmetric i8xi4 (INT4-packed weights, per-output- neuron scales) and per-token INT8 activations; host does embed/rmsnorm/rope/ attention-math/silu/lm_head. Validated vs the bf16 CPU reference: prompt: The capital of France is -> corr 0.936, top1 exact, top5 4/5 prompt: What is 2 plus 2? -> corr 0.937, top1 exact, top5 5/5 Fixes along the way: per-weight buffer sets (shape-keyed op pool overwrote B bindings), warmup + npu_time sync for the XRT first-dispatch readback flake, per-token (not per-tensor) activation scales (naive per-tensor: corr 0.55-0.75; per-token: 0.94), padded-row zeroing so softmax/scale can't NaN.
Adds KV-cached incremental decode to llama_w4a8_npu.py: per token, the
7 layer GEMMs run on the NPU as i8xi4 with real_m=1 (row 0 = the token,
padding zeroed — weight bandwidth is halved by i4, the wasted M compute
is sub-ms). Host does rope/attention/silu against the bf16 KV cache
seeded by the W4A8 prefill.
Verified vs the bf16 reference: first decoded token exact (" Paris" for
"The capital of France is"); subsequent tokens degrade as W4A8 error
compounds through the decode loop (per-token scales drift). The decode
mechanism is correct — quality is bounded by the quantization scheme
(next lever: group-wise weight scales / keeping early-layer residual in
bf16).
Also fixed in the decode path: per-head attention batching shapes
([H,S] scores via unsqueeze(1)), KV append dim (unsqueeze(1) -> [G,1,hd]),
and the harness-state mutation gotcha when comparing vs the reference.
… isolation
- W4A8_GROUPS=N splits the i4 weight K-dim into N chunks, each with its own
per-column scales, run as per-group i8xi4 GEMMs and dequantized per group
(exact int path unchanged). Quality ladder on llama-3.2-1B prefill:
G=1 : corr 0.937 (per-column scales)
G=8 : corr 0.966
G=16: corr 0.973 top5 5/5 (both prompts top1 exact)
G=32: corr 0.972 (plateau)
vs all-i8 weights: corr 0.9965 — i.e. the WEIGHTS are the W4A8 error
bottleneck; activations (per-token i8) are nearly lossless.
- W4A8_MIX_LEN=N keeps the first N layers at i8 weights (0.937->0.966 at
N=8) as a cheap partial-precision alternative.
- Fix: per-group A/C buffer pairs — the shared-buffer group path read stale
C when the same op was called back-to-back with different B (the
first-dispatch readback race again); each group now owns its buffers.
The "early EOS" and gibberish decode were NOT quantization — they were the
XRT readback flake: whenever the A-buffer content changed since the previous
kernel call, the first run returned STALE C (measured: all-rows corr 0.01,
the next call 0.994). The generated token's row always had new A data, so
its C came back stale/zero and the whole decode derailed.
Fix: always re-run each op and keep the second result (replaces the
warmup-once heuristic). Cost ~2x dispatches (sub-ms each).
Result — W4A8 llama-3.2-1B on XDNA2, group-wise i4 (G=16) + per-token i8:
prompt "The capital of France is" ->
"Paris is the capital of France. The city is located in the north of the
country and is the largest city in the..."
prompt "What is 2 plus 2?" -> corr 0.973, top1 exact, top5 5/5
Both prompts: prefill top1 exact, KV-cached decode generates coherent text.
The W4A8 pipeline is end-to-end correct; the decode quality is now bounded
by quantization, not by the runtime race.
Root cause of the "stale C" bug (ab47c82 worked around it with a double call): the iron runtime's to('npu') upload is driven by a per-range coherence map, and only WRITES IT KNOWS ABOUT are transferred. Writing the A tensor via `tensor.numpy()[:] = ...` is an *unmediated* host write — the runtime never marks the range dirty — so after the first dispatch the buffer stays marked "on npu" and every later kernel runs on the STALE previous A. The stale C was the correct computation of the wrong input. Fix: `with tensor.overwrite() as buf: buf[:] = ...` records the write in the coherence map, so the next to('npu') uploads it. Verified: the previously-failing second-call test (all-rows corr 0.01) is now 0.9931 with a SINGLE call — the double-call workaround is removed, halving dispatches (decode ~2x faster). Full-model result unchanged (correctness was already right after ab47c82): prefill corr 0.968, top1 exact; KV decode generates fluent text ("Paris is the capital of France...").
The dequant was the hidden host cost: `C * sx_full * s_w` over the full
[M_PAD, N] buffer in float64 was ~10 ms per GEMM call (and xG for the
group-wise path — 16 groups x full-N float64 multiply). Profiled:
op() dispatch+kernel 1.7 ms, to_torch 0.1 ms, dequant ~10 ms.
Fix: dequant in float32 and touch only the real_m rows (the padded rows
are zero anyway; decode has real_m=1 so the dequant is a 256x cut; prefill
real_m=256 unchanged work but float32). The returned tensor stays
[M_PAD, N] so the forward's view() shapes hold.
Measured (llama-3.2-1B, G=16 group-wise i4):
big GEMM call: 12.2 ms -> 2.5 ms
decode: ~4500 ms/token -> ~1280 ms/token (3.5x)
prefill corr: 0.9696 (unchanged), top1 exact
decode text: "Paris is the capital of France. The city is located in
the north of the country, on the river Seine."
…dder Per-group activation scales (each K-group quantized with its own per-token scale) now work correctly — the earlier 0.69 regression was the stale-A bug, not the math. Knob: W4A8_GROUP_ACTS=1. Quality ladder measured (llama-3.2-1B prefill logits corr vs bf16): i4 G=16 0.973 top1 exact i4 G=16 + group acts 0.974 mix 4 i8 layers + G=16 0.970 top5 5/5 mix 8 i8 layers + G=16 0.977 (best) all i8 weights 0.9965 (bound — i4 noise floor is ~0.97-0.98) The i4 quantization error is the limit; group scales are already fine-grained (K=128/group). Recommended: W4A8_MIX_LEN=8 W4A8_GROUPS=16 (corr 0.977, top1 exact, fluent decode).
bong-water-water-bong
requested review from
andrej,
hunhoffe and
jgmelber
as code owners
September 1, 2026 15:04
added 3 commits
September 1, 2026 13:49
W4A8_OPS_MIX=ffn_i8|attn_i8 isolates which op family keeps i8 weights. Measurement (llama-3.2-1B prefill corr vs bf16, all layers): all i4 G=16 0.973 attn i8 + FFN i4 0.977 (attention i8 barely helps) FFN i8 + attn i4 0.992 (FFN i8 is nearly the whole fix) So the FFN projections (gate/up/down) carry the i4 quantization error; the attention projections (q/k/v/o) tolerate i4 almost for free. New recommended config: W4A8_OPS_MIX=ffn_i8 W4A8_GROUPS=8 corr 0.983-0.990 (both prompts), top1 exact, top5 5/5 decode 684 ms/token (attn G=8) — faster than the old all-i4 G16 (903) "The capital of France is" -> "Paris is the capital of France. It is the largest city in France and the most populous metropolitan area in the European Union..." "What is 2 plus 2?" -> corr 0.983, top1 exact (the G16 near-tie mismatch is gone at G8) Also adds W4A8_ZP (asymmetric zero-point i4; measured no gain on llama — weights aren't column-biased enough — kept as a disabled knob).
…ken) Attention group curve with FFN i8: G8 0.9899, G4 0.9889, G2 0.9828, G1 0.9787. G4 is the sweet spot (6x fewer dispatches than all-i4 G16: 19/layer vs 112). Decode 559 ms/token (8x faster than the thread start), prefill corr 0.989 top1 exact top5 5/5, fluent factually-correct decode.
Fine-grained attention isolation (FFN i8 + attention, G4):
all attn i4 0.989
q,k i8 + v,o i4 0.990
q,k,v i8 + o i4 0.992-0.993 <- final recommended
q,k,o i8 + v i4 0.990
The attention projections split cleanly: q/k/v (which feed the scores)
want i8; the output projection o tolerates i4 at G4.
Final recommended config:
W4A8_OPS_MIX=ffn_i8 W4A8_ATTN_I8=q,k,v W4A8_GROUPS=4
- prefill corr 0.992-0.993 (vs 0.9965 all-i8 bound — the gap is just o i4)
- top1 exact, both prompts
- decode 554 ms/token
- "The capital of France is" -> "Paris is the most visited city in the
world. It is..."
Dead ends this round: NPU_RUNTIME=hrx (iron HRX runtime needs a newer
libhrx than installed — version mismatch, abandoned); the HRX-vs-XRT
dispatch-overhead question stays open.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds asymmetric INT4 GEMM (A=i8, B=i4 packed two-per-byte) to the AIE2P 4x16x16 mmul path, plus the W4A8 llama inference arc that uses it.
The core fix (issue: sizeof(int4) pointer arithmetic)
The AIE API's
int4_tis an empty struct —sizeof(int4) == 1although each element is 4 bits. The kernel's manualconst int4*pointer arithmetic (the j-block offset and the k-loop B advance) therefore moved 2x the real byte distance: the k-loop read k-blocks {0,2} then out-of-range, capping the effective K at half the tile. Fixed withB_ADV = size_B / 2for int4 (identity for int8) applied to every B pointer advance inmm.cc(matmul_vectorized_2x2_mmul).Verified
Commits
The PR spans the devel series from the asymmetric INT4 mmul (
0fef28b) through the W4A8 llama decode + docs (57853b8). Happy to split if preferred.