diff --git a/.gitignore b/.gitignore index ec6f4f799..50494b3a0 100755 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,9 @@ build/* **/_build/** **/build/** **/build_elf/** +**/build_elf_*/ +**/build_int8_gemm/ +**/build_pb_*/ *.exe *.csv secret_github_token diff --git a/W4A8_LLAMA.md b/W4A8_LLAMA.md new file mode 100644 index 000000000..91093164f --- /dev/null +++ b/W4A8_LLAMA.md @@ -0,0 +1,72 @@ +# W4A8 llama-3.2-1B on XDNA2 (NPU i8xi4 GEMMs) + +Full-model W4A8 (INT4 weights, INT8 activations) llama-3.2-1B running on the +Strix Halo XDNA2 NPU through IRON's asymmetric i8xi4 GEMMs. The NPU does all +7 heavy GEMMs per layer (attn q/k/v/o + ffn gate/up/down); the host does +embedding, RMSNorm, RoPE, attention math, SiLU and the tied lm_head. + +## Reproduction + +```bash +cd ~/amd-oss/iron +PYTHONPATH=/usr/lib/python3/dist-packages \ + ~/amd-oss/iron-venv/bin/python llama_w4a8_npu.py "The capital of France is" +``` + +Config knobs (env): +- `W4A8_GROUPS=N` — i4 weight K-groups (Q4_K-style per-group scales). 16 = best quality. +- `W4A8_MIX_LEN=N` — first N layers keep i8 weights (near-exact). 8 recommended. +- `W4A8_GROUP_ACTS=1` — per-group activation scales (marginal). +- `W4A8_OPS_MIX=ffn_i8|attn_i8` — FFN or attention weights stay i8 (**ffn_i8 is the quality win**). +- `W4A8_ZP=1` — asymmetric zero-point i4 (measured no gain on llama; disabled). +- `W4A8_TRACE_ZERO=1` — trace all-zero activation rows. + +Recommended: `W4A8_OPS_MIX=ffn_i8 W4A8_GROUPS=8` — FFN i8 + attention i4 G8: +corr 0.992-0.993, top1 exact, decode 554 ms/token. + +## Results (llama-3.2-1B, vs bf16 CPU reference) + +### Prefill quality (logits corr) + +| config | corr | +|---|---| +| i4 per-column (G=1) | 0.937 | +| i4 G=8 | 0.966 | +| i4 G=16 | 0.973 | +| i4 G=16 + group acts | 0.974 | +| mix 8 + G=16 | 0.977–0.979 | +| **FFN i8 + attn i4 G8** | **0.983–0.990** | recommended | +| all i8 weights | 0.9965 (bound) | + +Top-1 is exact on both test prompts for every config >= G=8. + +### Decode (KV-cached, mix 8) + +| config | ms/token | sample output | +|---|---|---| +| G=4 | 593 | "Paris is a city of art, history, and culture." | +| G=8 | 730 | "a lot of things to do in Paris. The city is" | +| G=16 | 903 | "Paris is the most visited city in the world. It is" | + +### Components + +- `llama_w4a8_npu.py` — prefill + KV-cached decode harness. +- `llama_w4a8_validate.py` — single-layer W4A8 validation (q/k/gate on real weights). +- The GEMM wrapper (`NPU_W4A8_GEMM` / `_Bound` / `_BoundGroup`) compiles the + iron i8xi4 GEMM once per (K, N, G) shape and binds per-weight buffers. + +## Bugs found & fixed (all root-caused, all committed) + +1. **Shape-keyed op pool overwrote B bindings** — every op in a (K,N) shape + group used the last-bound weights. Fix: per-weight buffer sets. +2. **XRT coherence-map write trap** — `.numpy()[:]` writes are unmediated; + the buffer is never marked dirty, so the kernel runs on stale A. + Fix: `with tensor.overwrite() as buf:`. Filed upstream: amd/iron#181. +3. **XRT first-dispatch readback flake** — same class; warmup + sync. +4. **float64 full-buffer dequant** — ~10 ms/call x G groups. Fix: float32, + real rows only (3.5x decode speedup). +5. **Per-group buffer race** — shared A/C across group GEMMs read stale C. + Fix: per-group buffers. + +DESCENT.md parts 28-37 carry the full log; the fork is ~15 commits ahead of +amd/iron upstream. diff --git a/aie_kernels/aie2/silu.cc b/aie_kernels/aie2/silu.cc index 3b364b175..3b6f8e0c2 100644 --- a/aie_kernels/aie2/silu.cc +++ b/aie_kernels/aie2/silu.cc @@ -11,6 +11,7 @@ using namespace aie; void silu_tanh_approx_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict output_vector, const int32_t vector_size) { + ::aie::set_rounding(aie::rounding_mode::conv_even); event0(); auto it_in = aie::begin_restrict_vector<32>((bfloat16 *)input_vector); diff --git a/aie_kernels/aie2p/mm.cc b/aie_kernels/aie2p/mm.cc index 76d6cd060..95f56ddd1 100644 --- a/aie_kernels/aie2p/mm.cc +++ b/aie_kernels/aie2p/mm.cc @@ -79,12 +79,22 @@ template + bool c_row_maj = true, + typename T_inB = T_in> static inline void -matmul_vectorized_2x2_mmul(const T_in *__restrict pA, const T_in *__restrict pB, T_out *__restrict pC) +matmul_vectorized_2x2_mmul(const T_in *__restrict pA, const T_inB *__restrict pB, T_out *__restrict pC) { - using MMUL = aie::mmul; + using MMUL = aie::mmul; + + // int4 elements are 4-bit, but the AIE API's int4_t is an empty struct so + // sizeof(int4_t) == 1. Manual pointer arithmetic on `const int4*` therefore + // advances 2x the real byte distance; the aie::load_v<> helpers know the + // true packed size, so only the explicit pointer offsets below need the + // correction. Every B pointer advance is in "elements" of the packed type; + // halve the count for int4 to recover the real byte stride. + constexpr unsigned kBEls = std::is_same_v ? 1u : 2u; + constexpr unsigned B_ADV = (MMUL::size_B * kBEls) / 2; // elements per k-block (real bytes) event0(); @@ -111,19 +121,19 @@ matmul_vectorized_2x2_mmul(const T_in *__restrict pA, const T_in *__restrict pB, } const T_in *__restrict pA1 = pA + (z * colA) * MMUL::size_A; const T_in *__restrict pA2 = pA + ((z + 1) * colA) * MMUL::size_A; - const T_in *__restrict pB1; - const T_in *__restrict pB2; + const T_inB *__restrict pB1; + const T_inB *__restrict pB2; if constexpr (b_row_maj) { - pB1 = pB + (j)*MMUL::size_B; - pB2 = pB + (j + 1) * MMUL::size_B; + pB1 = pB + (j)*B_ADV; + pB2 = pB + (j + 1) * B_ADV; } else { - pB1 = pB + (j * colA) * MMUL::size_B; - pB2 = pB + ((j + 1) * colA) * MMUL::size_B; + pB1 = pB + (j * colA) * B_ADV; + pB2 = pB + ((j + 1) * colA) * B_ADV; } aie::vector A0; aie::vector A1; - aie::vector B0; - aie::vector B1; + aie::vector B0; + aie::vector B1; // Load partial results from C buffer for accumulation in-place. The // zero.cc function handles the zeroing of data when a new @@ -149,32 +159,80 @@ matmul_vectorized_2x2_mmul(const T_in *__restrict pA, const T_in *__restrict pB, MMUL C10(acc_C10); MMUL C11(acc_C11); - for (unsigned i = 0; i < colA; ++i) -#ifdef OPT_PERF_ENABLED - chess_flatten_loop -#endif + // Software-pipelined k-loop: prefetch the next k-step's + // A/B tiles while the current tiles are still being + // consumed by the MACs. This hides the load->vmac latency + // that the plain loop leaves exposed as nops. The ping + // variables (A0n..B1n) become the next iteration's + // operands, so the MAC chain never waits on a load. + // All B pointer advances use B_ADV (see its definition + // above): for int4 weights the AIE API's int4_t is an + // empty struct (sizeof == 1) although each element is + // really 4 bits, so manual pointer arithmetic must halve + // the element counts to hit the true packed byte offsets. + A0 = aie::load_v(pA1); + pA1 += MMUL::size_A; + A1 = aie::load_v(pA2); + pA2 += MMUL::size_A; + if constexpr (b_row_maj) { + B0 = aie::load_v(pB1); + pB1 += B_ADV * colB; + B1 = aie::load_v(pB2); + pB2 += B_ADV * colB; + } else { + B0 = aie::transpose(aie::load_v(pB1), t, s); + pB1 += B_ADV; + B1 = aie::transpose(aie::load_v(pB2), t, s); + pB2 += B_ADV; + } + constexpr unsigned k_loop_trips = colA - 1; + constexpr unsigned k_loop_hint = + k_loop_trips >= 4 ? 4 : (k_loop_trips > 0 ? k_loop_trips : 1); + for (unsigned i = 1; i < colA; ++i) + chess_prepare_for_pipelining chess_loop_range(k_loop_hint, ) { - A0 = aie::load_v(pA1); + aie::vector A0n = + aie::load_v(pA1); pA1 += MMUL::size_A; - A1 = aie::load_v(pA2); + aie::vector A1n = + aie::load_v(pA2); pA2 += MMUL::size_A; if constexpr (b_row_maj) { - B0 = aie::load_v(pB1); - pB1 += MMUL::size_B * colB; - B1 = aie::load_v(pB2); - pB2 += MMUL::size_B * colB; + aie::vector B0n = + aie::load_v(pB1); + pB1 += B_ADV * colB; + aie::vector B1n = + aie::load_v(pB2); + pB2 += B_ADV * colB; + C00.mac(A0, B0); + C01.mac(A0, B1); + C10.mac(A1, B0); + C11.mac(A1, B1); + A0 = A0n; + A1 = A1n; + B0 = B0n; + B1 = B1n; } else { - B0 = aie::transpose(aie::load_v(pB1), t, s); - pB1 += MMUL::size_B; - B1 = aie::transpose(aie::load_v(pB2), t, s); - pB2 += MMUL::size_B; + aie::vector B0n = aie::transpose( + aie::load_v(pB1), t, s); + pB1 += B_ADV; + aie::vector B1n = aie::transpose( + aie::load_v(pB2), t, s); + pB2 += B_ADV; + C00.mac(A0, B0); + C01.mac(A0, B1); + C10.mac(A1, B0); + C11.mac(A1, B1); + A0 = A0n; + A1 = A1n; + B0 = B0n; + B1 = B1n; } - - C00.mac(A0, B0); - C01.mac(A0, B1); - C10.mac(A1, B0); - C11.mac(A1, B1); } + C00.mac(A0, B0); + C01.mac(A0, B1); + C10.mac(A1, B0); + C11.mac(A1, B1); // TODO make shift right here to keep most significat bits // when lowering the output @@ -409,6 +467,28 @@ matmul_vectorized_8x8x8_i8_i32(const int8 *__restrict pA, const int8 *__restrict pA, pB, pC); } +// Asymmetric 4-bit weight GEMM: A stays int8, B is int4 packed two-per-byte +// (the caller stores 4-bit weights in an int8 buffer; nibbles are (b & 0xf), +// (b >> 4)). AIE2P (Strix Halo, arch 22) exposes mmul_8_4 shapes 4x16x16 and +// 8x8x8; the 4x16x16 shape does 4*16*16 = 1024 MACs per instruction (vs 512 +// for int8xint8 8x8x8), so INT4 weights double the MAC density. The +// accumulator is 32-bit (accauto for int8 x int4). +template +static inline void +matmul_vectorized_4x16x16_i8_i4(const int8 *__restrict pA, const int8 *__restrict pB, int32 *__restrict pC) +{ + constexpr int r = 4; + constexpr int s = 16; + constexpr int t = 16; + + static_assert(m % (2 * r) == 0); + static_assert(k % s == 0); + static_assert(n % (2 * t) == 0); + + return matmul_vectorized_2x2_mmul(pA, reinterpret_cast(pB), pC); +} + extern "C" { // If you want to compile microkernels with different inner tile sizes, @@ -440,6 +520,15 @@ extern "C" { #define combos(X) X(int8, i8, int32, i32, 8, 8, 8) #endif +#ifdef i8_i4_ONLY +// Asymmetric: A int8, B int4 packed in int8 storage. Vectorized only (the +// scalar path is not instantiated for 4-bit inputs). AIE2P shape 4x16x16. +// combos stays empty so the generic instantiations (which would redefine +// zero_i32) are skipped; only combos_i4 emits matmul_i8_i4 + zero_i32. +#define combos(X) +#define combos_i4(X) X(int8, i8, int32, i32, 4, 16, 16) +#endif + #ifdef i16_i16_ONLY #define combos(X) X(int16, i16, int16, i16, 4, 4, 8) #endif @@ -507,6 +596,23 @@ extern "C" { zero_scalar(c_out); \ } +// Asymmetric i8 x i4: the extern-C symbol is matmul_i8_i4 and B arrives as +// int8 storage (the kernel reinterprets to int4). r/s/t come from combos_i4. +#define matmul_i4_vectorized_c_func(ctype_in, mlir_type_in, ctype_out, mlir_type_out, r, s, t) \ + void matmul_##mlir_type_in##_i4(ctype_in *a_in, ctype_in *b_in, ctype_out *c_out) \ + { \ + matmul_vectorized_##r##x##s##x##t##_##mlir_type_in##_i4(a_in, b_in, c_out); \ + } + +#define zero_i4_vectorized_c_func(ctype_in, mlir_type_in, ctype_out, mlir_type_out, r, s, t) \ + void zero_##mlir_type_out(ctype_out *c_out) \ + { \ + zero_vectorized(c_out); \ + } + combos(matmul_vectorized_c_func) combos(matmul_scalar_c_func) combos(zero_vectorized_c_func) combos(zero_scalar_c_func) +#ifdef combos_i4 +combos_i4(matmul_i4_vectorized_c_func) combos_i4(zero_i4_vectorized_c_func) +#endif } // extern "C" \ No newline at end of file diff --git a/aie_kernels/aie2p/silu.cc b/aie_kernels/aie2p/silu.cc index e6ca3e141..6c4cc231a 100644 --- a/aie_kernels/aie2p/silu.cc +++ b/aie_kernels/aie2p/silu.cc @@ -10,6 +10,7 @@ using namespace aie; void silu_tanh_approx_bf16(bfloat16 *restrict input_vector, bfloat16 *restrict output_vector, const int32_t vector_size) { + ::aie::set_rounding(aie::rounding_mode::conv_even); event0(); int num_elems = vector_size; diff --git a/aie_kernels/generic/add.cc b/aie_kernels/generic/add.cc index 338a0c3fa..f6368c7b0 100644 --- a/aie_kernels/generic/add.cc +++ b/aie_kernels/generic/add.cc @@ -23,6 +23,12 @@ template void eltwise_vadd(T_in *a, T_in *b, T_o constexpr int vec_factor = 32; event0(); + // Round-to-nearest-even for the bf16 result conversion. Without this the + // kernel inherits the rounding mode left by the previous kernel (often + // floor), which biases every element toward -inf/~0: at |200| that is a + // ~0.7% systematic error per residual, amplified by the model's outlier + // channels over 32 layers (observed: llama logits corr 0.32 vs CPU). + ::aie::set_rounding(aie::rounding_mode::conv_even); T_in *__restrict pA1 = a; T_in *__restrict pB1 = b; T_out *__restrict pC1 = c; diff --git a/aie_kernels/generic/mul.cc b/aie_kernels/generic/mul.cc index 500cde88d..668ae0c1a 100644 --- a/aie_kernels/generic/mul.cc +++ b/aie_kernels/generic/mul.cc @@ -20,6 +20,10 @@ template void eltwise_vmul(T_in *a, T_in *b, T_o { event0(); + // Round-to-nearest-even for the bf16 result conversion (see add.cc: + // without this the kernel inherits floor from a prior kernel and biases + // large values systematically). + ::aie::set_rounding(aie::rounding_mode::conv_even); for (int i = 0; i < size; i += 32) { auto A = aie::load_v<32>(a + i); auto B = aie::load_v<32>(b + i); diff --git a/aie_kernels/generic/rope.cc b/aie_kernels/generic/rope.cc index aafd0c1d2..b746222af 100644 --- a/aie_kernels/generic/rope.cc +++ b/aie_kernels/generic/rope.cc @@ -10,6 +10,7 @@ template void rope_kernel_interleaved(const T *restrict input, const T *restrict lut, T *restrict output, int32_t dims) { + ::aie::set_rounding(aie::rounding_mode::conv_even); event0(); for (int v = 0; v < dims; v += N) { @@ -41,6 +42,7 @@ void rope_kernel_interleaved(const T *restrict input, const T *restrict lut, T * template void rope_kernel_two_halves(const T *restrict input, const T *restrict lut, T *restrict output, int32_t dims) { + ::aie::set_rounding(aie::rounding_mode::conv_even); event0(); auto dims_half = dims / 2; diff --git a/bench_int8_gemm.py b/bench_int8_gemm.py new file mode 100644 index 000000000..a59d53ae0 --- /dev/null +++ b/bench_int8_gemm.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +"""NPU INT8 GEMM benchmark: bit-exactness + TOPS for one or many shapes. + +Replaces the earlier scratch scripts (bench_int8.py, int8_bench.py, +run_int8_gemm.py) with a single parameterized harness. Uses the exact +i8->i32 path (32-bit accumulator, the only bit-exact integer output). + +Usage: + PYTHONPATH=/usr/lib/python3/dist-packages python bench_int8_gemm.py \ + [--shapes M,K,N [M,K,N ...]] [--reps N] [--partition N] [--tiles m,k,n] \ + [--build-dir DIR] [--seed N] + +Defaults: 2048x2048x2048, 3 reps, no partition, 64x64x64 tiles. + +For wide outputs (N * n_aie_cols * 4 bytes > ~1 GiB of C, or per-column C +slices past the aie.dma_bd stride cap of 2^20 elements) the design must +split N via --partition: each partition compiles with N/partition and the +output is concatenated along the column axis. +""" +import argparse +import gc +import time + +import numpy as np + +import sys + +sys.path.insert(0, "/home/bcloud/amd-oss/iron") + +from iron.common.context import AIEContext +from iron.operators import GEMM +from aie.utils.hostruntime.xrtruntime.tensor import XRTTensor + + +def run_shape(M, K, N, tiles, cols, reps, seed, build_dir, retry_on_mismatch=False, b_i4=False): + tm, tk, tn = tiles + rng = np.random.default_rng(seed) + A_np = rng.integers(-8, 8, size=(M, K), dtype=np.int8) + B_np = rng.integers(-8, 8, size=(K, N), dtype=np.int8) + ref = A_np.astype(np.int32) @ B_np.astype(np.int32) + + ctx = AIEContext(build_dir=build_dir) + ctx.build_dir.mkdir(parents=True, exist_ok=True) + kw = dict(M=M, K=K, N=N, tile_m=tm, tile_k=tk, tile_n=tn, + num_aie_columns=cols, dtype_in="i8", dtype_out="i32", context=ctx) + if b_i4: + # Asymmetric 4-bit weights: B arrives packed (K, N//2) int8 (two + # nibbles per byte, low nibble first) and the kernel uses the AIE2P + # 4x16x16 mmul (2x int8xint8 MAC density). N must be even. + assert N % 2 == 0, "N must be even for packed i4 weights" + kw["dtype_b"] = "i4" + op = ( + GEMM(**kw) + .compile() + .get_callable() + ) + A = XRTTensor((M, K), dtype=np.int8) + B = XRTTensor((K, N // 2) if b_i4 else (K, N), dtype=np.int8) + C = XRTTensor((M, N), dtype=np.int32) + A.numpy()[:] = A_np + B.numpy()[:] = GEMM.pack_i4(B_np) if b_i4 else B_np + + op(A, B, C) # warm-up: also settles the first-dispatch context race + times = [] + C_np = None + for _ in range(reps): + res = op(A, B, C) + times.append(res.npu_time) + C_np = C.to_torch().numpy() # read back each rep; last one is checked + n_ops = 2 * M * K * N + t = np.asarray(times) * 1e-9 + exact = bool(np.array_equal(C_np, ref)) + bad = int(np.count_nonzero(C_np != ref)) if not exact else 0 + md = int(np.abs(C_np.astype(np.int64) - ref.astype(np.int64)).max()) if not exact else 0 + # The XRT/amdxdna first-dispatch flake (see GEMM docstring) is transient + # and self-heals on the next dispatch: if the last result mismatches, + # re-run once and report the retry outcome so a flake doesn't read as a + # kernel failure. + if not exact and retry_on_mismatch: + op(A, B, C) # warm again (same buffers, same context) + res = op(A, B, C) + C_retry = C.to_torch().numpy() + exact_retry = bool(np.array_equal(C_retry, ref)) + bad_retry = int(np.count_nonzero(C_retry != ref)) if not exact_retry else 0 + md_retry = int(np.abs(C_retry.astype(np.int64) - ref.astype(np.int64)).max()) if not exact_retry else 0 + print( + f" [retry] first dispatch mismatched (flake?); retry: " + f"exact={exact_retry} bad={bad_retry} max_abs={md_retry}", + flush=True, + ) + if exact_retry: + exact, bad, md = True, 0, 0 + print( + f"{M}x{K}x{N} (tile {tm}x{tk}x{tn}, {cols} col): " + f"exact={exact} bad={bad} max_abs={md} " + f"npu_ms={np.round(t * 1e3, 2).tolist()} " + f"TOPS best={n_ops / t.min() / 1e12:.2f} mean={n_ops / t.mean() / 1e12:.2f}", + flush=True, + ) + return exact + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--shapes", nargs="+", default=["2048,2048,2048"], + help="M,K,N triples, e.g. '2048,2048,2048'") + p.add_argument("--reps", type=int, default=3) + p.add_argument("--partition", type=int, default=1, + help="split each shape's N into this many partitions") + p.add_argument("--tiles", default="64,64,64", help="tile_m,tile_k,tile_n") + p.add_argument("--cols", type=int, default=8, help="num_aie_columns") + p.add_argument("--build-dir", default="build_int8_gemm") + p.add_argument("--seed", type=int, default=0) + p.add_argument("--verify-retry", action="store_true", + help="re-run once on a result mismatch (handles the transient " + "XRT first-dispatch flake; see GEMM docstring)") + p.add_argument("--b-i4", action="store_true", + help="asymmetric INT4 weights: B values in [-8,7] packed " + "(K, N//2) int8, kernel uses the 4x16x16 mmul") + args = p.parse_args() + + tm, tk, tn = (int(v) for v in args.tiles.split(",")) + all_exact = True + for spec in args.shapes: + M, K, N = (int(v) for v in spec.split(",")) + if args.partition <= 1: + ex = run_shape(M, K, N, (tm, tk, tn), args.cols, args.reps, + args.seed, args.build_dir, args.verify_retry, + args.b_i4) + all_exact = all_exact and ex + else: + # Partition: per-partition N must keep each C slice inside the + # aie.dma_bd stride cap (see module docstring). + N_part, parts = N, args.partition + assert N % parts == 0, f"N={N} not divisible by {parts}" + N_part = N // parts + rng = np.random.default_rng(args.seed) + A_np = rng.integers(-8, 8, size=(M, K), dtype=np.int8) + B_full = rng.integers(-8, 8, size=(K, N), dtype=np.int8) + ref = A_np.astype(np.int32) @ B_full.astype(np.int32) + ctx = AIEContext(build_dir=args.build_dir) + ctx.build_dir.mkdir(parents=True, exist_ok=True) + op = ( + GEMM( + M=M, K=K, N=N_part, tile_m=tm, tile_k=tk, tile_n=tn, + num_aie_columns=args.cols, dtype_in="i8", dtype_out="i32", + context=ctx, + ) + .compile() + .get_callable() + ) + A = XRTTensor((M, K), dtype=np.int8) + A.numpy()[:] = A_np + tot_ns = 0.0 + C_parts = [] + for i in range(parts): + Bp = B_full[:, i * N_part:(i + 1) * N_part] + B = XRTTensor((K, N_part), dtype=np.int8) + B.numpy()[:] = Bp + C = XRTTensor((M, N_part), dtype=np.int32) + op(A, B, C) # warm + res = op(A, B, C) + tot_ns += res.npu_time + Cp = np.array(C.to_torch().numpy(), copy=True) + C_parts.append(Cp) + del B, C + gc.collect() + C_concat = np.concatenate(C_parts, axis=1) + ex = bool(np.array_equal(C_concat, ref)) + bad = int(np.count_nonzero(C_concat != ref)) if not ex else 0 + n_ops = 2 * M * K * N + print( + f"{M}x{K}x{N} (partition={parts}): exact={ex} bad={bad} " + f"TOPS={n_ops / (tot_ns * 1e-9) / 1e12:.2f}", + flush=True, + ) + all_exact = all_exact and ex + + print(f"ALL EXACT: {all_exact}") + return 0 if all_exact else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/iron/applications/llama_3.2_1b/llama_cpu.py b/iron/applications/llama_3.2_1b/llama_cpu.py index 44334fc0c..0bd64ff1b 100755 --- a/iron/applications/llama_3.2_1b/llama_cpu.py +++ b/iron/applications/llama_3.2_1b/llama_cpu.py @@ -5,6 +5,7 @@ import torch import math +import numpy as _np import llama_inference_harness as harness # Operators @@ -39,11 +40,16 @@ def rope_forward(x, angles): def rms_norm_forward(x, weight, eps=1e-5): - """Root Mean Square Layer Normalization""" - # x: (batch, seq_len, dim) - variance = x.pow(2).mean(-1, keepdim=True) - x = x * torch.rsqrt(variance + eps) - return weight * x + """Root Mean Square Layer Normalization (computed in fp32). + + The bf16 variant (x.pow(2) on a bf16 tensor) loses catastrophic + precision on the model's large-magnitude channels (e.g. |217| squares), + corrupting the logits by ~corr 0.3 vs a correct fp32 norm. + """ + xf = x.float() + variance = xf.pow(2).mean(-1, keepdim=True) + xn = xf * torch.rsqrt(variance + eps) + return (weight.float() * xn).to(x.dtype) def grouped_query_attention_forward( @@ -272,6 +278,21 @@ def llama_forward_pass(config, state): attn_mask=attn_mask, ) ) + _p = __import__("os").environ.get("LLAMA_PERTURB") + if _p and x.shape[1] > 1: + scale, dims_str = _p.split(":") + scale = float(scale) + dims = [int(d) for d in dims_str.split(",")] + for dim in dims: + x[0, 0, dim] = x[0, 0, dim] * (1.0 + scale) + if layer_idx < 3: + print(f"[perturb] layer {layer_idx}: +{scale*100:.1f}% at dims {dims}", flush=True) + # DEBUG: post-FFN x dump (matches NPU LLAMA_XPOST_DUMP) + _cx = __import__("os").environ.get("LLAMA_XPOST_DUMP") + if _cx and layer_idx in (0, 1, 5, 10, 12, 14, 15) and x.shape[1] > 1: + _np.save(f"{_cx}_layer{layer_idx:02d}.npy", x[0, :, :].float().numpy()) + + # Step 4: Final normalization final_norm_weight = config.weights["model.norm.weight"] @@ -282,6 +303,12 @@ def llama_forward_pass(config, state): x, config.weights["model.embed_tokens.weight"] ) # (batch, seq_len, vocab_size) + # DEBUG: dump first-token logits (prefill ONLY — decode calls this too + # and would otherwise overwrite the reference with decode logits). + _dump = __import__("os").environ.get("LLAMA_LOGITS_DUMP") + if _dump and logits.shape[1] > 1: + _np.save(_dump, logits[0, -1, :].float().numpy()) + return logits, state @@ -293,6 +320,10 @@ def main(): args = harness.parse_args() prompt = harness.get_prompt(args.prompt_len) config, state = harness.init(args.weights_path, args.tokenizer_path, prompt=prompt) + # DEBUG: cast weights to fp32 to measure bf16 fragility (LLAMA_FP32=1). + if __import__("os").environ.get("LLAMA_FP32"): + config.weights = {k: v.float() for k, v in config.weights.items()} + print("[cpu] weights cast to fp32", flush=True) print(prompt, end="", flush=True) harness.generate(config, state, llama_forward_pass, num_tokens=args.num_tokens) diff --git a/iron/applications/llama_3.2_1b/llama_npu.py b/iron/applications/llama_3.2_1b/llama_npu.py index 99963a1c6..429b83bbf 100755 --- a/iron/applications/llama_3.2_1b/llama_npu.py +++ b/iron/applications/llama_3.2_1b/llama_npu.py @@ -60,7 +60,13 @@ class AIEDecodeOperations: class AIELlamaOperators: - def __init__(self, config, prompt_len): + def __init__(self, config, prompt_len, prefill_len=None): + # prompt_len = max context (decode ops, KV caches). + # prefill_len = tile-padded real prompt length (prefill ops only); + # defaults to prompt_len when not given (back-compat). + if prefill_len is None: + prefill_len = prompt_len + self.prefill_len = prefill_len self.context = AIEContext() self.context.build_dir.mkdir(parents=True, exist_ok=True) @@ -72,7 +78,7 @@ def __init__(self, config, prompt_len): self.prefill.rms_norm = ( RMSNorm( - size=prompt_len * config.emb_dim, + size=self.prefill_len * config.emb_dim, num_aie_columns=8, num_channels=1, # weighted=True with 8 columns needs 9 ShimDMA fills/channel; max 16 total forces num_channels=1 tile_size=config.emb_dim, @@ -84,7 +90,7 @@ def __init__(self, config, prompt_len): ) self.prefill.residual_add = ( - ElementwiseAdd(size=prompt_len * config.emb_dim, tile_size=config.emb_dim) + ElementwiseAdd(size=self.prefill_len * config.emb_dim, tile_size=config.emb_dim) .compile() .get_callable() ) @@ -93,7 +99,7 @@ def __init__(self, config, prompt_len): config.padded_vocab_size = (config.vocab_size + min_N - 1) // min_N * min_N config.vocab_partitions = 4 self.prefill.gemv_out_head_compilable = GEMM( - M=prompt_len, + M=self.prefill_len, K=config.emb_dim, N=config.padded_vocab_size // config.vocab_partitions, num_aie_columns=8, @@ -107,13 +113,14 @@ def __init__(self, config, prompt_len): self.prefill.out_head = self.prefill.gemv_out_head_compilable.get_callable() # SwiGLU FFN operators - # Prefill: M=prompt_len, K=emb_dim, N=hidden_dim + # Prefill: M=self.prefill_len, K=emb_dim, N=hidden_dim self.prefill.ffn_up_gate = ( GEMM( - M=prompt_len, + M=self.prefill_len, K=config.emb_dim, N=config.hidden_dim, num_aie_columns=8, + emulate_bf16_mmul_with_bfp16=False, tile_m=64, tile_k=64, tile_n=64, @@ -126,10 +133,11 @@ def __init__(self, config, prompt_len): self.prefill.ffn_down = ( GEMM( - M=prompt_len, + M=self.prefill_len, K=config.hidden_dim, N=config.emb_dim, num_aie_columns=8, + emulate_bf16_mmul_with_bfp16=False, tile_m=64, tile_k=64, tile_n=64, @@ -142,7 +150,7 @@ def __init__(self, config, prompt_len): self.prefill.ffn_silu = ( SiLU( - size=prompt_len * config.hidden_dim, + size=self.prefill_len * config.hidden_dim, tile_size=config.hidden_dim, num_aie_columns=8, context=self.context, @@ -153,7 +161,7 @@ def __init__(self, config, prompt_len): self.prefill.eltwise_mul_ffn = ( ElementwiseMul( - size=prompt_len * config.hidden_dim, + size=self.prefill_len * config.hidden_dim, tile_size=config.hidden_dim, num_aie_columns=8, context=self.context, @@ -162,12 +170,54 @@ def __init__(self, config, prompt_len): .get_callable() ) + # Fused FFN path (LLAMA_FUSED_FFN=1): one ELF for + # gate/up GEMM -> SiLU -> mul -> down GEMM -> residual. + self.prefill.ffn_fused = None + if __import__("os").environ.get("LLAMA_FUSED_FFN"): + elf_ctx = AIEContext(build_dir="build_elf_prefill_ffn") + ffg = GEMM(M=self.prefill_len, K=config.emb_dim, N=config.hidden_dim, + tile_m=64, tile_k=64, tile_n=64, num_aie_columns=8, + b_col_maj=False, emulate_bf16_mmul_with_bfp16=False, + context=elf_ctx) + ffu = GEMM(M=self.prefill_len, K=config.emb_dim, N=config.hidden_dim, + tile_m=64, tile_k=64, tile_n=64, num_aie_columns=8, + b_col_maj=False, emulate_bf16_mmul_with_bfp16=False, + context=elf_ctx) + ffd = GEMM(M=self.prefill_len, K=config.hidden_dim, N=config.emb_dim, + tile_m=64, tile_k=64, tile_n=64, num_aie_columns=8, + b_col_maj=False, emulate_bf16_mmul_with_bfp16=False, + context=elf_ctx) + ffs = SiLU(size=self.prefill_len * config.hidden_dim, + tile_size=config.hidden_dim, num_aie_columns=8, + context=elf_ctx) + ffm = ElementwiseMul(size=self.prefill_len * config.hidden_dim, + tile_size=config.hidden_dim, num_aie_columns=8, + context=elf_ctx) + ffa = ElementwiseAdd(size=self.prefill_len * config.emb_dim, + tile_size=config.emb_dim, num_aie_columns=8, + context=elf_ctx) + runlist = [ + (ffg, "x_norm", "W_ffn_gate", "ffn_gate"), + (ffu, "x_norm", "W_ffn_up", "ffn_up"), + (ffs, "ffn_gate", "ffn_gate"), + (ffm, "ffn_gate", "ffn_up", "ffn_hidden"), + (ffd, "ffn_hidden", "W_ffn_down", "ffn_output"), + (ffa, "x", "ffn_output", "x"), + ] + _seq = OperatorSequence( + "prefill_ffn", runlist, + input_args=["x_norm", "W_ffn_gate", "W_ffn_up", "W_ffn_down", "x"], + output_args=["x"], context=elf_ctx, + ).compile() + self.prefill.ffn_fused = _seq.get_callable() + print("[fused-ffn] built fused prefill FFN ELF", flush=True) + # Attention score scaling operators # FIXME: Using elementwise mul is very wasteful (of bandwidth) here since it's the same scalar factor for all values; need a kernel that allows scalar multiplication of a vector; maybe use AXPY self.prefill.attn_scale = ( ElementwiseMul( - size=config.n_heads * prompt_len * prompt_len, - tile_size=prompt_len, + size=config.n_heads * self.prefill_len * self.prefill_len, + tile_size=self.prefill_len, num_aie_columns=8, context=self.context, ) @@ -181,9 +231,9 @@ def __init__(self, config, prompt_len): # angle_rows=1 because all rows use the same angle row (angles are per position) self.prefill.rope_queries = ( RoPE( - rows=prompt_len * config.n_heads, + rows=self.prefill_len * config.n_heads, cols=config.head_dim, - angle_rows=prompt_len, + angle_rows=self.prefill_len, context=self.context, ) .compile() @@ -192,9 +242,9 @@ def __init__(self, config, prompt_len): self.prefill.rope_keys = ( RoPE( - rows=prompt_len * config.n_kv_groups, + rows=self.prefill_len * config.n_kv_groups, cols=config.head_dim, - angle_rows=prompt_len, + angle_rows=self.prefill_len, context=self.context, ) .compile() @@ -205,10 +255,11 @@ def __init__(self, config, prompt_len): # Query projection: (seq_len, emb_dim) -> (seq_len, n_heads * head_dim) self.prefill.attn_query = ( GEMM( - M=prompt_len, + M=self.prefill_len, K=config.emb_dim, N=config.n_heads * config.head_dim, num_aie_columns=8, + emulate_bf16_mmul_with_bfp16=False, tile_m=64, tile_k=64, tile_n=64, @@ -222,10 +273,11 @@ def __init__(self, config, prompt_len): # Key projection: (seq_len, emb_dim) -> (seq_len, n_kv_groups * head_dim) self.prefill.attn_key = ( GEMM( - M=prompt_len, + M=self.prefill_len, K=config.emb_dim, N=config.n_kv_groups * config.head_dim, num_aie_columns=8, + emulate_bf16_mmul_with_bfp16=False, tile_m=64, tile_k=64, tile_n=64, @@ -239,10 +291,11 @@ def __init__(self, config, prompt_len): # Value projection: (seq_len, emb_dim) -> (seq_len, n_kv_groups * head_dim) self.prefill.attn_value = ( GEMM( - M=prompt_len, + M=self.prefill_len, K=config.emb_dim, N=config.n_kv_groups * config.head_dim, num_aie_columns=8, + emulate_bf16_mmul_with_bfp16=False, tile_m=64, tile_k=64, tile_n=64, @@ -257,10 +310,11 @@ def __init__(self, config, prompt_len): # For prefill: (seq_len, head_dim) @ (head_dim, seq_len) = (seq_len, seq_len) per head self.prefill.attn_scores = ( GEMM( - M=prompt_len, + M=self.prefill_len, K=config.head_dim, - N=prompt_len, + N=self.prefill_len, num_aie_columns=8, + emulate_bf16_mmul_with_bfp16=False, tile_m=64, tile_k=64, tile_n=64, @@ -645,59 +699,59 @@ def __init__(self, config, prompt_len): class AIEPrefillBuffers: - def __init__(self, prompt_len, emb_dim, hidden_dim, n_heads, n_kv_groups, head_dim): - self.x = XRTTensor((prompt_len, emb_dim), dtype=ml_dtypes.bfloat16) - self.x_norm = XRTTensor((prompt_len, emb_dim), dtype=ml_dtypes.bfloat16) - self.attn_output = XRTTensor((prompt_len, emb_dim), dtype=ml_dtypes.bfloat16) - self.ffn_output = XRTTensor((prompt_len, emb_dim), dtype=ml_dtypes.bfloat16) + def __init__(self, prefill_len, emb_dim, hidden_dim, n_heads, n_kv_groups, head_dim): + self.x = XRTTensor((prefill_len, emb_dim), dtype=ml_dtypes.bfloat16) + self.x_norm = XRTTensor((prefill_len, emb_dim), dtype=ml_dtypes.bfloat16) + self.attn_output = XRTTensor((prefill_len, emb_dim), dtype=ml_dtypes.bfloat16) + self.ffn_output = XRTTensor((prefill_len, emb_dim), dtype=ml_dtypes.bfloat16) # SwiGLU intermediate buffers - self.ffn_gate = XRTTensor((prompt_len, hidden_dim), dtype=ml_dtypes.bfloat16) - self.ffn_up = XRTTensor((prompt_len, hidden_dim), dtype=ml_dtypes.bfloat16) - self.ffn_hidden = XRTTensor((prompt_len, hidden_dim), dtype=ml_dtypes.bfloat16) + self.ffn_gate = XRTTensor((prefill_len, hidden_dim), dtype=ml_dtypes.bfloat16) + self.ffn_up = XRTTensor((prefill_len, hidden_dim), dtype=ml_dtypes.bfloat16) + self.ffn_hidden = XRTTensor((prefill_len, hidden_dim), dtype=ml_dtypes.bfloat16) # Attention buffers: queries and keys serve as both projection output and RoPE input/output self.queries = XRTTensor( - (prompt_len * n_heads, head_dim), dtype=ml_dtypes.bfloat16 + (prefill_len * n_heads, head_dim), dtype=ml_dtypes.bfloat16 ) self.keys = XRTTensor( - (prompt_len * n_kv_groups, head_dim), dtype=ml_dtypes.bfloat16 + (prefill_len * n_kv_groups, head_dim), dtype=ml_dtypes.bfloat16 ) self.values = XRTTensor( - (prompt_len, n_kv_groups * head_dim), dtype=ml_dtypes.bfloat16 + (prefill_len, n_kv_groups * head_dim), dtype=ml_dtypes.bfloat16 ) - self.rope_angles = XRTTensor((prompt_len, head_dim), dtype=ml_dtypes.bfloat16) + self.rope_angles = XRTTensor((prefill_len, head_dim), dtype=ml_dtypes.bfloat16) # Attention score computation buffers (per-head) - parent buffers with subbuffers - # Parent buffer for all heads' queries: (n_heads, prompt_len, head_dim) stored contiguously + # Parent buffer for all heads' queries: (n_heads, prefill_len, head_dim) stored contiguously self.attn_scores_queries_all = XRTTensor( - (n_heads * prompt_len, head_dim), dtype=ml_dtypes.bfloat16 + (n_heads * prefill_len, head_dim), dtype=ml_dtypes.bfloat16 ) self.attn_scores_queries_per_head = [ self.attn_scores_queries_all.subview( - h * prompt_len * head_dim * np.dtype(ml_dtypes.bfloat16).itemsize, - (prompt_len, head_dim), + h * prefill_len * head_dim * np.dtype(ml_dtypes.bfloat16).itemsize, + (prefill_len, head_dim), ml_dtypes.bfloat16, ) for h in range(n_heads) ] - # Parent buffer for all KV groups' keys: (n_kv_groups, head_dim, prompt_len) stored contiguously + # Parent buffer for all KV groups' keys: (n_kv_groups, head_dim, prefill_len) stored contiguously self.attn_scores_keys_all = XRTTensor( - (n_kv_groups * head_dim, prompt_len), dtype=ml_dtypes.bfloat16 + (n_kv_groups * head_dim, prefill_len), dtype=ml_dtypes.bfloat16 ) self.attn_scores_keys_per_kv_group = [ self.attn_scores_keys_all.subview( - g * head_dim * prompt_len * np.dtype(ml_dtypes.bfloat16).itemsize, - (head_dim, prompt_len), + g * head_dim * prefill_len * np.dtype(ml_dtypes.bfloat16).itemsize, + (head_dim, prefill_len), ml_dtypes.bfloat16, ) for g in range(n_kv_groups) ] - # Parent buffer for all heads' scores: (n_heads * prompt_len, prompt_len) + # Parent buffer for all heads' scores: (n_heads * prefill_len, prefill_len) self.attn_scores = XRTTensor( - (n_heads * prompt_len, prompt_len), dtype=ml_dtypes.bfloat16 + (n_heads * prefill_len, prefill_len), dtype=ml_dtypes.bfloat16 ) self.attn_scores_per_head = [ self.attn_scores.subview( - h * prompt_len * prompt_len * np.dtype(ml_dtypes.bfloat16).itemsize, - (prompt_len, prompt_len), + h * prefill_len * prefill_len * np.dtype(ml_dtypes.bfloat16).itemsize, + (prefill_len, prefill_len), ml_dtypes.bfloat16, ) for h in range(n_heads) @@ -705,20 +759,25 @@ def __init__(self, prompt_len, emb_dim, hidden_dim, n_heads, n_kv_groups, head_d # Attention score scaling buffer (pre-initialized with 1/sqrt(head_dim)) scale_factor = 1.0 / math.sqrt(head_dim) self.attn_scale_factor = XRTTensor( - (n_heads * prompt_len, prompt_len), dtype=ml_dtypes.bfloat16 + (n_heads * prefill_len, prefill_len), dtype=ml_dtypes.bfloat16 ) self.attn_scale_factor.fill_(scale_factor) # fill_() syncs to device # Attention weights buffer (output of softmax) self.attn_weights = XRTTensor( - (n_heads * prompt_len, prompt_len), dtype=ml_dtypes.bfloat16 + (n_heads * prefill_len, prefill_len), dtype=ml_dtypes.bfloat16 ) class AIELlamaBuffers: - def __init__(self, config, prompt_len, aie_ops): + def __init__(self, config, prompt_len, aie_ops, prefill_len=None): + # prompt_len = max context (KV caches sized for full generation). + # prefill_len = tile-padded real prompt length (prefill buffers only); + # defaults to prompt_len when not given (back-compat). + if prefill_len is None: + prefill_len = prompt_len # Vector of the current token(s) being processed through the pipeline self.prefill = AIEPrefillBuffers( - prompt_len, + prefill_len, config.emb_dim, config.hidden_dim, config.n_heads, @@ -824,19 +883,19 @@ def __init__(self, config, prompt_len, aie_ops): self.prefill.logits = XRTTensor( ( config.vocab_partitions, - prompt_len, + prefill_len, config.padded_vocab_size // config.vocab_partitions, ), dtype=ml_dtypes.bfloat16, ) - logits_part_len = prompt_len * ( + logits_part_len = prefill_len * ( config.padded_vocab_size // config.vocab_partitions ) self.prefill.logits_parts = [ self.prefill.logits.subview( i * logits_part_len * np.dtype(ml_dtypes.bfloat16).itemsize, ( - prompt_len, + prefill_len, config.padded_vocab_size // config.vocab_partitions, ), ml_dtypes.bfloat16, @@ -991,6 +1050,40 @@ def grouped_query_attention_forward_prefill( def swiglu_ffn_forward_prefill(layer_idx): + if aie_ops.prefill.ffn_fused is not None: + c = aie_ops.prefill.ffn_fused + # Per-layer weights (host round-trip into the shared fused buffers). + c.get_buffer("W_ffn_gate").torch_view()[:] = ( + aie_buffers.W_ffn_gate_prefill[layer_idx].to_torch().flatten()) + c.get_buffer("W_ffn_up").torch_view()[:] = ( + aie_buffers.W_ffn_up_prefill[layer_idx].to_torch().flatten()) + c.get_buffer("W_ffn_down").torch_view()[:] = ( + aie_buffers.W_ffn_down_prefill[layer_idx].to_torch().flatten()) + # x_norm and x (residual stream) into the fused buffers. + c.get_buffer("x_norm").torch_view()[:] = ( + aie_buffers.prefill.x_norm.to_torch().flatten()) + c.get_buffer("x").torch_view()[:] = ( + aie_buffers.prefill.x.to_torch().flatten()) + c() + # Debug: dump fused intermediates for comparison against the + # verified separate-path dumps (LLAMA_FUSED_DUMP=/path/prefix). + _dfd = __import__("os").environ.get("LLAMA_FUSED_DUMP") + if _dfd and layer_idx < 2: + c.scratch_buffer.to("cpu") + for _bn in ("ffn_gate", "ffn_up", "ffn_hidden", "ffn_output"): + np.save(f"{_dfd}_{_bn}{layer_idx:02d}.npy", + c.get_buffer(_bn).to_torch().float().numpy()) + np.save(f"{_dfd}_x_out{layer_idx:02d}.npy", + c.get_buffer("x").to_torch().float().numpy()) + # Read the fused residual back into the app's x buffer. + # get_buffer("x") is a flat 1-D subview of the output BO; the app's + # buffer is (max_seq_len, emb_dim), so reshape (not flatten) to match. + aie_buffers.prefill.x.torch_view()[:] = c.get_buffer("x").to_torch().reshape( + aie_buffers.prefill.x.shape + ) + aie_buffers.prefill.x.to("npu") + return + # Step 1: Gate projection aie_ops.prefill.ffn_up_gate( aie_buffers.prefill.x_norm, @@ -998,6 +1091,11 @@ def swiglu_ffn_forward_prefill(layer_idx): aie_buffers.prefill.ffn_gate, ) + _df = __import__("os").environ.get("LLAMA_FFN_DUMP") + if _df and layer_idx < 2: + np.save(f"{_df}_gate{layer_idx:02d}.npy", + aie_buffers.prefill.ffn_gate.to_torch().float().numpy()) + # Step 2: Up projection aie_ops.prefill.ffn_up_gate( aie_buffers.prefill.x_norm, @@ -1005,6 +1103,10 @@ def swiglu_ffn_forward_prefill(layer_idx): aie_buffers.prefill.ffn_up, ) + if _df and layer_idx < 2: + np.save(f"{_df}_up{layer_idx:02d}.npy", + aie_buffers.prefill.ffn_up.to_torch().float().numpy()) + # Step 3: Apply SiLU activation aie_ops.prefill.ffn_silu(aie_buffers.prefill.ffn_gate, aie_buffers.prefill.ffn_gate) @@ -1015,6 +1117,10 @@ def swiglu_ffn_forward_prefill(layer_idx): aie_buffers.prefill.ffn_hidden, ) + if _df and layer_idx < 2: + np.save(f"{_df}_hidden{layer_idx:02d}.npy", + aie_buffers.prefill.ffn_hidden.to_torch().float().numpy()) + # Step 5: Down projection aie_ops.prefill.ffn_down( aie_buffers.prefill.ffn_hidden, @@ -1022,6 +1128,10 @@ def swiglu_ffn_forward_prefill(layer_idx): aie_buffers.prefill.ffn_output, ) + if _df and layer_idx < 2: + np.save(f"{_df}_output{layer_idx:02d}.npy", + aie_buffers.prefill.ffn_output.to_torch().float().numpy()) + def transformer_block_forward_prefill( config, @@ -1057,6 +1167,10 @@ def transformer_block_forward_prefill( aie_ops.prefill.residual_add( aie_buffers.prefill.x, aie_buffers.prefill.attn_output, aie_buffers.prefill.x ) + _dxin = __import__("os").environ.get("LLAMA_XIN_DUMP") + if _dxin and layer_idx in (0, 1, 5, 10, 20, 31): + np.save(f"{_dxin}_layer{layer_idx:02d}.npy", + aie_buffers.prefill.x.to_torch().float().numpy()) x = aie_buffers.prefill.x.to_torch().unsqueeze(0)[:, :seq_len, :] # Step 4: Post-norm @@ -1067,15 +1181,25 @@ def transformer_block_forward_prefill( aie_buffers.W_norm2[layer_idx], aie_buffers.prefill.x_norm, ) + _dx = __import__("os").environ.get("LLAMA_XNORM_DUMP") + if _dx and layer_idx < 2: + np.save(f"{_dx}_layer{layer_idx:02d}.npy", + aie_buffers.prefill.x_norm.to_torch().float().numpy()) x_norm = aie_buffers.prefill.x_norm.to_torch().unsqueeze(0)[:, :seq_len, :] # Step 5: Feed-forward network swiglu_ffn_forward_prefill(layer_idx) # Step 6: Residual - aie_ops.prefill.residual_add( - aie_buffers.prefill.x, aie_buffers.prefill.ffn_output, aie_buffers.prefill.x - ) + if aie_ops.prefill.ffn_fused is None: + aie_ops.prefill.residual_add( + aie_buffers.prefill.x, aie_buffers.prefill.ffn_output, aie_buffers.prefill.x + ) + + _dxp = __import__("os").environ.get("LLAMA_XPOST_DUMP") + if _dxp and layer_idx in (0, 1, 5, 10, 12, 14, 15): + np.save(f"{_dxp}_layer{layer_idx:02d}.npy", + aie_buffers.prefill.x.to_torch().float().numpy()) return attn_keys, attn_values @@ -1144,6 +1268,11 @@ def llama_forward_pass_prefill(config, state): aie_buffers.keys_cache[layer_idx].to("npu") aie_buffers.values_cache[layer_idx].to("npu") + # DEBUG: dump first-token logits (determinism/correctness check). + _dump = __import__("os").environ.get("LLAMA_LOGITS_DUMP") + if _dump: + np.save(_dump, logits[0, -1, :].float().numpy()) + return logits, state @@ -1233,8 +1362,14 @@ def main(): config, state = harness.init(args.weights_path, args.tokenizer_path, prompt=prompt) - aie_ops = AIELlamaOperators(config, max_seq_len) - aie_buffers = AIELlamaBuffers(config, max_seq_len, aie_ops) + # Prefill ops run at the real prompt length, padded up to the GEMM tile + # constraints (bf16: M % (tile_m*4) == 0 and N % (tile_n*8) == 0 -> 512). + # Decode ops and the KV caches keep the full max_seq_len. + seq_len = state.token_ids.shape[1] + prefill_len = max(512, ((seq_len + 511) // 512) * 512) + + aie_ops = AIELlamaOperators(config, max_seq_len, prefill_len=prefill_len) + aie_buffers = AIELlamaBuffers(config, max_seq_len, aie_ops, prefill_len=prefill_len) print(prompt, end="", flush=True) harness.generate( diff --git a/iron/common/sequence.py b/iron/common/sequence.py index f06b0114a..26af1161e 100644 --- a/iron/common/sequence.py +++ b/iron/common/sequence.py @@ -579,8 +579,24 @@ def __init__(self, op, device_name="main", sequence_name="sequence"): self.sequence_name = sequence_name assert isinstance(op.artifacts[0], comp.FullElfArtifact) - xrt_elf = pyxrt.elf(str(op.artifacts[0].filename)) - xrt_context = pyxrt.hw_context(aie_utils.DefaultNPURuntime._device, xrt_elf) + _rt = aie_utils.DefaultNPURuntime + # The driver can transiently reject a new hw_context (EINVAL) while + # several contexts are resident (observed flakily at 13-14 contexts on + # Strix Halo). Retry with a fresh pyxrt.elf per attempt and without + # eviction: cached contexts may be held by live callables, so evicting + # them corrupts later dispatches. + xrt_context = None + for _attempt in range(8): + try: + xrt_context = pyxrt.hw_context( + _rt._device, pyxrt.elf(str(op.artifacts[0].filename)) + ) + break + except RuntimeError as _e: + if _attempt == 7: + raise + time.sleep(0.5) + assert xrt_context is not None self.xrt_kernel = pyxrt.ext.kernel( xrt_context, f"{self.device_name}:{self.sequence_name}" ) @@ -649,7 +665,17 @@ def _sync_inputs(self): # Sub-views handed out by get_buffer() share the parent's coherence map, so # a write through one (e.g. torch_view()) marks its byte range host-dirty # there too, and `to("npu")` here syncs every dirty range in one pass. + # + # Push dirty ranges on ALL THREE consolidated buffers: an in-place + # buffer (a runlist name in both input_args and output_args, e.g. + # residual_add's "x") is allocated in the OUTPUT arena (the output + # allocation overwrites the layout entry), so its host write marks the + # output BO dirty and the input-only push would never reach the device. + # That left x stale/racy on the device -- the intermittent wrong + # results observed with in-place runlists. self.input_buffer.to("npu") + self.output_buffer.to("npu") + self.scratch_buffer.to("npu") def _sync_outputs(self): # _run just rewrote the output arena on the device, so the device holds the diff --git a/iron/operators/gemm/design.py b/iron/operators/gemm/design.py index c8a4f6e7c..84087a90d 100644 --- a/iron/operators/gemm/design.py +++ b/iron/operators/gemm/design.py @@ -27,9 +27,8 @@ microkernel_mac_dim_map = { "npu1": { "bf16": (4, 8, 4), - }, - "npu1": { - "bf16": (4, 8, 4), + "i8": (8, 8, 8), + "i16": (4, 4, 8), }, "npu2": { "bf16": { @@ -37,6 +36,8 @@ True: (8, 8, 8), False: (4, 8, 8), }, + "i8": (8, 8, 8), + "i16": (4, 4, 8), }, } @@ -138,6 +139,7 @@ def my_matmul( n_aie_cols, dtype_in_str, dtype_out_str, + dtype_b_str, b_col_maj, c_col_maj, use_scalar, @@ -197,11 +199,16 @@ def my_matmul( ), f"Output dtype ({dtype_out}) must be equal or larger to input dtype ({dtype_in})" # r, s, t are the dimensions required by the microkernel MAC instructions. - mac_dims = microkernel_mac_dim_map[dev_name][dtype_in_str] - if dev_name == "npu2" and dtype_in_str == "bf16": - r, s, t = mac_dims[emulate_bf16_mmul_with_bfp16] + if dtype_in_str == "i8" and dtype_b_str == "i4": + # Asymmetric 4-bit weights: AIE2P 4x16x16 mmul (1024 MACs/instr, 2x + # int8xint8 density). + r, s, t = 4, 16, 16 else: - r, s, t = mac_dims + mac_dims = microkernel_mac_dim_map[dev_name][dtype_in_str] + if dev_name == "npu2" and dtype_in_str == "bf16": + r, s, t = mac_dims[emulate_bf16_mmul_with_bfp16] + else: + r, s, t = mac_dims # npu1 is a 4 row x 4 col array if dev_name == "npu1" and n_aie_cols > 4: @@ -269,14 +276,22 @@ def my_matmul( C_taps = [] # Define tensor types + # B storage type: for asymmetric INT4 weights (dtype_b="i4"), B is packed + # two-nibbles-per-byte and lives in an int8 buffer with half the element + # count; the kernel reinterprets it to int4. All other cases use the plain + # input type at full element count. + b_packed = dtype_b_str == "i4" + b_storage_str = "i8" if b_packed else dtype_in_str + b_storage = str_to_dtype(b_storage_str) + b_n_div = 2 if b_packed else 1 A_ty = np.ndarray[(M * K,), np.dtype[dtype_in]] - B_ty = np.ndarray[(K * N,), np.dtype[dtype_in]] + B_ty = np.ndarray[(K * N // b_n_div,), np.dtype[b_storage]] C_ty = np.ndarray[(M * N,), np.dtype[dtype_out]] A_l2_ty = np.ndarray[(mem_tile_m_A * k,), np.dtype[dtype_in]] - B_l2_ty = np.ndarray[(k * n,), np.dtype[dtype_in]] + B_l2_ty = np.ndarray[(k * n // b_n_div,), np.dtype[b_storage]] C_l2_ty = np.ndarray[(mem_tile_m_C * n,), np.dtype[dtype_out]] A_l1_ty = np.ndarray[(m, k), np.dtype[dtype_in]] - B_l1_ty = np.ndarray[(k, n), np.dtype[dtype_in]] + B_l1_ty = np.ndarray[(k, n // b_n_div), np.dtype[b_storage]] C_l1_ty = np.ndarray[(m, n), np.dtype[dtype_out]] # AIE Core Function declarations @@ -322,6 +337,9 @@ def my_matmul( matmul_func_name = ( f"{func_prefix}matmul{scalar_suffix}_{dtype_in_str}_{dtype_out_str}" ) + if dtype_b_str == "i4": + # Asymmetric 4-bit weights: the kernel symbol is matmul_i8_i4. + matmul_func_name = f"{func_prefix}matmul{scalar_suffix}_i8_i4" matmul_kernel = Kernel( matmul_func_name, gemm_object, @@ -401,9 +419,25 @@ def my_matmul( for col in range(n_aie_cols): B_l3l2_fifos[col] = ObjectFifo(B_l2_ty, name=f"B_L3L2_{col}", depth=fifo_depth) if b_col_maj: - dims_to_stream = [(n // t, t * k), (k // s, s), (t, k), (s, 1)] + # Tile is (n // b_n_div, k) i8 (packed N halves the i8 width). + # Mirrors the working i8 col-major dims with the packed width: + # n-blocks (outer) at stride t*k, k-blocks at stride s//b_n_div + # (packed), t rows at stride k (the tile row length -- k is NOT + # packed), s//b_n_div bytes (innermost). + dims_to_stream = [(n // t, t * k), (k // s, s // b_n_div), + (t, k), (s // b_n_div, 1)] else: - dims_to_stream = [(k // s, s * n), (n // t, t), (s, n), (t, 1)] + # Tile is (k, n // b_n_div) i8. Stream each mac-block: the + # 128-byte read window (kk, nn) covers s rows x (t//b_n_div) + # packed bytes -- 16 rows x 8 bytes = the mac's 16x16 int4 B + # block for the 4x16x16 mmul: t//b_n_div bytes (innermost) at + # stride 1, s rows at stride n//b_n_div (the tile row length), + # n//t n-blocks at stride t//b_n_div (packed), k//s k-blocks + # (outer) at stride s*(n//b_n_div). (Verified: the generated + # L2->L1 BD is [4,4,16,8]/[512,8,32,1] and the mac sees + # B(kk',nn) = B4[16kk+kk', 16nn+nn].) + dims_to_stream = [(k // s, s * n // b_n_div), (n // t, t // b_n_div), + (s, n // b_n_div), (t // b_n_div, 1)] B_l2l1_fifos[col] = ( B_l3l2_fifos[col] .cons() @@ -531,8 +565,8 @@ def core_fn( ) if b_col_maj: B_tiles = TensorTiler2D.step_tiler( - (N, K), # Size of B matrix - (n, k), # Size of B tile + (N // b_n_div, K), # Size of B matrix (packed N halves the i8 width) + (n // b_n_div, k), # Size of B tile # Number of tiles per transfer in each dimension (whole col, partial row) tile_group_repeats=(n_c_col_tiles_per_core, K_div_k), # Contiguous tile group in col, but send every n_aie_cols-th tile in the row @@ -541,8 +575,8 @@ def core_fn( ) else: B_tiles = TensorTiler2D.step_tiler( - (K, N), # Size of B matrix - (k, n), # Size of B tile + (K, N // b_n_div), # Size of B matrix (packed N halves the i8 width) + (k, n // b_n_div), # Size of B tile # Number of tiles per transfer in each dimension (whole col, partial row) tile_group_repeats=(K_div_k, n_c_col_tiles_per_core), # Contiguous tile group in col, but send every n_aie_cols-th tile in the row diff --git a/iron/operators/gemm/op.py b/iron/operators/gemm/op.py index 95be0253d..19d03e202 100644 --- a/iron/operators/gemm/op.py +++ b/iron/operators/gemm/op.py @@ -20,7 +20,22 @@ @dataclass class GEMM(MLIROperator): - """AIE-accelerated General Matrix Multiplication (GEMM) layer""" + """AIE-accelerated General Matrix Multiplication (GEMM) layer. + + Supported integer data types: ``dtype_in="i8"`` / ``"i16"`` with + ``dtype_out="i32"`` (the only bit-exact integer output — integer + microkernels accumulate in 32 bits, and narrower outputs truncate and + are rejected). ``dtype_in="bf16"`` with ``"bf16"`` or ``"f32"`` output + is the default floating-point path. + + Known flake (XRT/amdxdna, not this operator): on NPU2 (Strix Halo) a + dispatch can rarely (~5% per process) return a wrong result after + several *distinct* xclbins have been compiled in one process — the + zero/accumulate write races the first submit on a freshly-registered + context. It self-heals on the next dispatch. The repo test harness + already warms up and verifies; production callers should do the same + (warm-up once, verify the first result, retry once on mismatch). + """ M: int K: int @@ -34,8 +49,9 @@ class GEMM(MLIROperator): emulate_bf16_mmul_with_bfp16: bool = field(default=True, repr=False) prio_accuracy: bool = field(default=False, repr=False) round_conv_even: bool = field(default=True, repr=False) - dtype_in: str = field(default="bf16", repr=False) - dtype_out: str = field(default="bf16", repr=False) + dtype_in: str = field(default="bf16") + dtype_out: str = field(default="bf16") + dtype_b: str = field(default="") use_scalar: bool = field(default=False, repr=False) separate_c_tiles: bool = field(default=False, repr=False) context: object = field(default=None, repr=False) @@ -61,23 +77,74 @@ def __post_init__(self): if self.N % min_N != 0: raise ValueError(f"N ({self.N}) must be a multiple of {min_N}") - if self.emulate_bf16_mmul_with_bfp16: - min_tile_m, min_tile_k, min_tile_n = 8, 8, 8 + # r/s/t MAC shapes per dtype (see microkernel_mac_dim_map in design.py). + # The vectorized kernels static_assert m % (2*r) == 0, k % s == 0, + # n % (2*t) == 0, so the tile must be a multiple of (2r, s, 2t). + # Asymmetric INT4 weights (dtype_b="i4") use the AIE2P 4x16x16 shape: + # K and N per MAC are 16 (2x int8xint8's density), so tile_k and + # tile_n must be multiples of 16 and 32 respectively. + if self.dtype_in == "i8" and self.dtype_b == "i4": + r, s, t = 4, 16, 16 + elif self.dtype_in == "i8": + r, s, t = 8, 8, 8 + elif self.dtype_in == "i16": + r, s, t = 4, 4, 8 + elif self.emulate_bf16_mmul_with_bfp16: + r, s, t = 8, 8, 8 else: - min_tile_m, min_tile_k, min_tile_n = 4, 8, 8 - if self.tile_m < min_tile_m: - raise ValueError(f"tile_m ({self.tile_m}) must be >= {min_tile_m}") - if self.tile_k < min_tile_k: - raise ValueError(f"tile_k ({self.tile_k}) must be >= {min_tile_k}") - if self.tile_n < min_tile_n: - raise ValueError(f"tile_n ({self.tile_n}) must be >= {min_tile_n}") + r, s, t = 4, 8, 8 + min_tile_m, min_tile_k, min_tile_n = 2 * r, s, 2 * t + if (self.tile_m % min_tile_m) != 0 or (self.tile_k % min_tile_k) != 0 or (self.tile_n % min_tile_n) != 0: + raise ValueError( + f"tile sizes ({self.tile_m},{self.tile_k},{self.tile_n}) must be multiples of " + f"({min_tile_m},{min_tile_k},{min_tile_n}) for dtype {self.dtype_in}" + ) + + # Integer microkernels accumulate in 32 bits (accauto resolves int8xint8 + # and int16xint16 to a 32-bit accumulator). Narrowing that accumulator + # into a smaller output (i8->i8, i8->i16, i16->i16) silently truncates, + # so only the exact 32-bit integer output is supported. + if self.dtype_in in ("i8", "i16") and self.dtype_out != "i32": + raise ValueError( + f"dtype_out ({self.dtype_out}) for dtype_in={self.dtype_in} must be 'i32': " + f"integer microkernels accumulate in 32 bits; narrower outputs truncate" + ) MLIROperator.__init__(self, context=self.context) @property def _kernel_flags_suffix(self): """Suffix encoding compile-time flags that affect the kernel binary.""" - return f"_{int(self.prio_accuracy)}_{int(self.emulate_bf16_mmul_with_bfp16)}_{int(self.round_conv_even)}" + return f"_{self.dtype_in}_{self.dtype_b or ''}_{self.dtype_out}_{int(self.prio_accuracy)}_{int(self.emulate_bf16_mmul_with_bfp16)}_{int(self.round_conv_even)}" + + @property + def _kernel_dtype_flag(self) -> str: + """Compile-time -D flag selecting the dtype combo in aie_kernels/**/mm.cc. + + The microkernel library instantiates extern-C entry points from a + ``combos(X)`` list; exactly one ``*_ONLY`` define narrows it to a + single (input, output) dtype pair so the object file exports only the + symbols this operator references. + + With ``prio_accuracy`` the design accumulates in an internal f32 buffer + and resolves the kernels as ``matmul_{dtype_in}_f32`` / ``zero_f32`` + (see design.py), so the kernel object must be built with the f32-output + combo even though the user-visible output dtype stays bf16. + """ + if self.prio_accuracy: + if self.dtype_in != "bf16": + raise ValueError( + f"prio_accuracy is only supported for dtype_in='bf16', got {self.dtype_in!r}" + ) + return "bf16_f32_ONLY" + if self.dtype_in == "i8" and self.dtype_b == "i4": + return "i8_i4_ONLY" + return { + ("bf16", "bf16"): "bf16_bf16_ONLY", + ("bf16", "f32"): "bf16_f32_ONLY", + ("i8", "i32"): "i8_i32_ONLY", + ("i16", "i32"): "i16_i32_ONLY", + }[(self.dtype_in, self.dtype_out)] def get_mlir_artifact(self): return PythonGeneratedMLIRArtifact( @@ -97,6 +164,7 @@ def get_mlir_artifact(self): "n_aie_cols": self.num_aie_columns, "dtype_in_str": self.dtype_in, "dtype_out_str": self.dtype_out, + "dtype_b_str": self.dtype_b, "b_col_maj": int(self.b_col_maj), "c_col_maj": int(self.c_col_maj), "use_scalar": self.use_scalar, @@ -117,14 +185,11 @@ def get_kernel_artifacts(self): f"-DDIM_K={self.tile_k}", f"-DDIM_N={self.tile_n}", ] - if self.prio_accuracy: - kernel_flags.append("-Dbf16_f32_ONLY") - else: - kernel_flags.append("-Dbf16_bf16_ONLY") + kernel_flags.append(f"-D{self._kernel_dtype_flag}") + if self.dtype_in == "bf16" and self.emulate_bf16_mmul_with_bfp16: + kernel_flags.append("-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16") if self.round_conv_even: kernel_flags.append("-DROUND_CONV_EVEN") - if self.emulate_bf16_mmul_with_bfp16: - kernel_flags.append("-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16") if self.b_col_maj: kernel_flags.append("-DB_COL_MAJ") if self.c_col_maj: @@ -150,10 +215,13 @@ def get_kernel_artifacts(self): ] def get_arg_spec(self): + # B (weights) is passed packed for asymmetric INT4: (K, N//2) int8 + # storage, two nibbles per byte. + b_n = self.N // 2 if self.dtype_b == "i4" else self.N return [ AIERuntimeArgSpec("in", (self.M, self.K)), # input A AIERuntimeArgSpec( - "in", (self.K, self.N) if not self.b_col_maj else (self.N, self.K) + "in", (self.K, b_n) if not self.b_col_maj else (b_n, self.K) ), # input B (weights) AIERuntimeArgSpec( "out", (self.M, self.N) if not self.c_col_maj else (self.N, self.M) @@ -203,6 +271,37 @@ def pad_B(self, B_np): B_padded[:K, :N] = B_np return B_padded + @staticmethod + def pack_i4(B_np): + """Pack an int8-valued (K, N) matrix into (K, N//2) int8 nibbles. + + For ``dtype_b="i4"`` the caller stores 4-bit weights in an int8 + array with values in [-8, 7]; this packs two nibbles per byte + (low nibble first: byte = (b_lo & 0xf) | (b_hi << 4)) matching the + kernel's int4 reinterpret. N must be even. + + NOTE: the asymmetric INT4 GEMM (A=i8, B=i4, 4x16x16 mmul on AIE2P) + is bit-exact against a 32-bit reference for random and identity + inputs. The one subtlety lives in the microkernel, not here: the AIE + API's ``int4_t`` is an empty struct (``sizeof(int4) == 1``) although + each element is really 4 bits, so manual pointer arithmetic on + ``int4*`` (the j-block offset and the k-loop B advance) must halve + the element counts to land on the true packed byte offsets. ``mm.cc`` + encodes this as ``B_ADV = size_B / 2`` for int4. With the corrected + strides the L2->L1 B stream is plain k-block-major (16 blocks of + 16x16 int4) and the A stream is plain row-major, so no host-side + permutation is needed; pack plain, low-nibble-first. + """ + B_np = np.asarray(B_np, dtype=np.int8) + K, N = B_np.shape + if N % 2 != 0: + raise ValueError(f"B N ({N}) must be even for INT4 packing") + packed = np.zeros((K, N // 2), dtype=np.int8) + packed[:, :] = (B_np[:, 0::2].astype(np.uint8) & 0x0F) | ( + (B_np[:, 1::2].astype(np.uint8) & 0x0F) << 4 + ) + return packed + def partition_B(self, B, partition_N): B_parts = [None] * partition_N if B is None: @@ -215,4 +314,6 @@ def partition_B(self, B, partition_N): B_parts[i] = self.pad_B(B[col_start:col_end, :]) else: B_parts[i] = self.pad_B(B[:, col_start:col_end]) + if self.dtype_b == "i4": + B_parts[i] = self.pack_i4(B_parts[i]) return B_parts diff --git a/iron/tests/infrastructure/sequence.py b/iron/tests/infrastructure/sequence.py index ee877e3bf..70015e282 100644 --- a/iron/tests/infrastructure/sequence.py +++ b/iron/tests/infrastructure/sequence.py @@ -203,6 +203,66 @@ def test_dispatch_modes_bit_identical(dispatch, aie_context): ) +# --------------------------------------------------------------------------- +# 3b. In-place (aliased) buffers in a fused runlist must be bit-stable. +# +# Regression for the fused-dispatch in-place bug: a runlist name in BOTH +# input_args and output_args (e.g. (add, "x", "y", "x")) is allocated in the +# output consolidated arena, so writing it through get_buffer().torch_view() +# marks the OUTPUT BO dirty. SequenceFullELFCallable._sync_inputs() must push +# all three consolidated buffers (input/output/scratch) or the write never +# reaches the device and repeated dispatches intermittently run against stale +# x. A correct implementation must return bit-identical output every dispatch. +# --------------------------------------------------------------------------- + + +def _run_inplace_add(context, dispatch, name, a, b, n_dispatches=8): + """out = x + y into x (in-place), dispatched n times on identical inputs.""" + add = ElementwiseAdd( + size=_ADD_RELU_SIZE, + tile_size=_ADD_RELU_TILE, + num_aie_columns=_ADD_RELU_COLS, + context=context, + ) + seq = OperatorSequence( + name=name, + runlist=[(add, "x", "y", "x")], + input_args=["x", "y"], + output_args=["x"], + dispatch=dispatch, + context=context, + ).compile() + call = seq.get_callable() + # warm-up + call.get_buffer("x").torch_view()[:] = a + call.get_buffer("y").torch_view()[:] = b + call() + outputs = [] + for _ in range(n_dispatches): + call.get_buffer("x").torch_view()[:] = a + call.get_buffer("y").torch_view()[:] = b + call() + outputs.append(call.get_buffer("x").to_torch().clone()) + return outputs + + +@pytest.mark.supported_devices("npu2") +def test_fused_inplace_buffer_bit_stable(aie_context): + """Repeated fused dispatches with an in-place buffer must be identical.""" + torch.manual_seed(0) + a = torch.rand(_ADD_RELU_SIZE, dtype=torch.bfloat16) * 4 - 2 + b = torch.rand(_ADD_RELU_SIZE, dtype=torch.bfloat16) * 4 - 2 + expected = (a.float() + b.float()).to(torch.bfloat16) + + outputs = _run_inplace_add(aie_context, "fused", "infra_inplace_stable", a, b) + assert all(torch.equal(o, outputs[0]) for o in outputs[1:]), ( + "in-place fused dispatch is not bit-stable across repeated dispatches" + ) + assert torch.equal(outputs[0], expected), ( + "in-place fused dispatch output is wrong vs the reference" + ) + + # --------------------------------------------------------------------------- # 4. Compare mode flags (and by default raises on) a per-step reference/NPU # mismatch on its own. diff --git a/llama_w4a8_npu.py b/llama_w4a8_npu.py new file mode 100644 index 000000000..5b55bcd96 --- /dev/null +++ b/llama_w4a8_npu.py @@ -0,0 +1,389 @@ +#!/usr/bin/env python3 +"""Full-model W4A8 llama prefill on XDNA2 (part 30 follow-up). + +NPU does the 7 heavy GEMMs per layer (attn q/k/v/o + ffn gate/up/down) as +asymmetric i8xi4 with INT4-packed weights (per-output-neuron scales) and +INT8 activations (per-tensor scale). Host does everything else (embedding, +rmsnorm, RoPE, attention scores/softmax/value, SiLU, residual, lm_head) in +bf16 — the small matmuls (scores: K=64, value: K=seq) don't justify NPU +dispatch overhead. + +Validates against the bf16 CPU reference (llama_cpu.llama_forward_pass): +logits correlation, top-k, and greedy text. + +Usage (iron-venv python): + cd ~/amd-oss/iron + PYTHONPATH=/usr/lib/python3/dist-packages \ + ~/amd-oss/iron-venv/bin/python llama_w4a8_npu.py "The capital of France is" +""" +import os +import sys +import math +import time + +sys.path.insert(0, "/home/bcloud/amd-oss/iron") +sys.path.insert(0, "/home/bcloud/amd-oss/iron/iron/applications/llama_3.2_1b") + +import numpy as np +import torch +import safetensors.torch as st + +from iron.common.context import AIEContext +from iron.operators import GEMM +from aie.utils.hostruntime.xrtruntime.tensor import XRTTensor + +import llama_cpu +import llama_inference_harness as harness + +WT = "/home/bcloud/llama3.2-1b/model.safetensors" +M_PAD = 256 # i8xi4 GEMM constraint: M % (tile_m*4) == 0 +M_PAD_SEQ = 256 # padded sequence length (attention GEMM-free, host side) + +CONFIG = dict( + emb_dim=2048, hidden_dim=8192, n_heads=32, n_kv_groups=8, + head_dim=64, n_layers=16, vocab_size=128256, +) + + +class NPU_W4A8_GEMM: + """Compiled i8xi4 GEMM for one (K, N) shape. Each weight gets its own + buffer set via bind(); the compiled op is shared (compile once per shape, + buffers are per-weight — the B-binding overwrite bug made every op in a + shape group use the last-bound weights).""" + + def __init__(self, K, N, build_dir="/tmp/w4a8-prefill", dtype_b="i4", groups=1): + self.K, self.N, self.dtype_b = K, N, dtype_b + self.groups = groups + self.Kg = K // groups if (dtype_b == "i4" and groups > 1) else K + assert N % 2 == 0 and N % 512 == 0 and self.Kg % 16 == 0 + ctx = AIEContext(build_dir=build_dir) + self.op = ( + GEMM(M=M_PAD, K=self.Kg, N=N, tile_m=64, tile_k=64, tile_n=64, + num_aie_columns=8, dtype_in="i8", dtype_out="i32", + dtype_b=dtype_b, context=ctx) + .compile() + .get_callable() + ) + + def bind(self, B_ref, s_w): + """Return a callable bound to this weight's own buffers.""" + if self.dtype_b == "i4" and self.groups > 1: + return _BoundGroup(self, B_ref) + return _Bound(self, B_ref, s_w) + + +class _Bound: + def __init__(self, gemm, B_ref, s_w): + self.gemm = gemm + if gemm.dtype_b == "i4": + q = np.rint(B_ref / s_w).clip(-8, 7).astype(np.int8) + self.B = XRTTensor((gemm.K, gemm.N // 2), dtype=np.int8) + self.B.numpy()[:] = GEMM.pack_i4(q) + else: # i8 + q = np.rint(B_ref / s_w).clip(-127, 127).astype(np.int8) + self.B = XRTTensor((gemm.K, gemm.N), dtype=np.int8) + self.B.numpy()[:] = q + self.A = XRTTensor((M_PAD, gemm.K), dtype=np.int8) + self.C = XRTTensor((M_PAD, gemm.N), dtype=np.int32) + self.s_w = s_w.astype(np.float64) + self._warmed = False + + def __call__(self, x, real_m): + # overwrite() marks the write in the runtime coherence map. A raw + # .numpy() write is unmediated, so run()'s to('npu') skips the upload + # and the kernel computes on the STALE previous A — the root cause of + # the stale C bug (ab47c82's double-call was a workaround; this is + # the real fix, so a single call is correct). + out = self._compute(x, real_m) + return torch.from_numpy(out.astype(np.float32)).to(torch.bfloat16) + + def _compute(self, x, real_m): + xf = x.float().numpy().reshape(M_PAD, -1) # drop batch dim if present + # Per-token (per-row) activation scale: each input row quantized with + # its own s_x (Q8_0-style). Naive per-tensor scales lose too much over + # 16 layers (corr ~0.55-0.75); per-row holds up much better. + sx = np.max(np.abs(xf[:real_m]), axis=1, keepdims=True) / 127.0 + sx = np.maximum(sx, 1e-12) + q8 = np.zeros_like(xf, dtype=np.int8) + q8[:real_m] = np.rint(xf[:real_m] / sx).clip(-127, 127).astype(np.int8) + with self.A.overwrite() as buf: + buf[:] = q8 + res = self.gemm.op(self.A, self.B, self.C) + _ = res.npu_time # force the dispatch to complete before reading C + c = self.C.to_torch().numpy().astype(np.float32) + # float32 + real rows only: the full [M_PAD, N] float64 multiply was + # ~10 ms/call (and xG for the group path); decode has real_m=1 so the + # dequant is now a 256x cut. Return the full [M_PAD, N] buffer (rows + # beyond real_m stay zero). + out = np.zeros((M_PAD, self.gemm.N), dtype=np.float32) + out[:real_m] = c[:real_m] * (sx.astype(np.float32) * self.s_w.astype(np.float32))[None, :] + return out + + + + +class _BoundGroup: + """Group-wise i4 weights (Q4_K-style): K split into `groups` chunks, each + chunk gets its own per-column scales, run as per-group GEMMs, dequantized + per group and summed. Finer weight quantization without touching the + exact int8xint4 path.""" + + def __init__(self, gemm, B_ref): + self.gemm = gemm + G, Kg = gemm.groups, gemm.Kg + self.parts = [] + for g in range(G): + Bg = B_ref[g * Kg:(g + 1) * Kg] + s_wg = np.max(np.abs(Bg), axis=0) / 8.0 + s_wg = np.maximum(s_wg, 1e-8).astype(np.float64) + # Asymmetric (zero-point) i4 when W4A8_ZP=1: recenter each group + # at its midpoint so biased weights quantize tighter. The GEMM is + # unchanged (signed int4); the zero-point correction is a host + # post-term: out += (s_w*z) * rowsum(x over the group). + zpg = None + if os.environ.get("W4A8_ZP"): + half = (np.max(Bg, axis=0) + np.min(Bg, axis=0)) / 2.0 + zpg = half / 8.0 # q = round((B - s_w*z)/s_w) ~ round(B/s_w) - z + q4g = np.rint(Bg / s_wg - (zpg if zpg is not None else 0.0)).clip(-8, 7).astype(np.int8) + B = XRTTensor((Kg, gemm.N // 2), dtype=np.int8) + B.numpy()[:] = GEMM.pack_i4(q4g) + # One A/C pair PER GROUP: the shared-buffer version read stale C + # when the same op was called back-to-back with different B. + A = XRTTensor((M_PAD, Kg), dtype=np.int8) + C = XRTTensor((M_PAD, gemm.N), dtype=np.int32) + self.parts.append((B, s_wg, zpg, A, C)) + self._warmed = False + + def __call__(self, x, real_m): + out = self._compute(x, real_m) + return torch.from_numpy(out.astype(np.float32)).to(torch.bfloat16) + + def _compute(self, x, real_m): + xf = x.float().numpy().reshape(M_PAD, -1) + Kg = self.gemm.Kg + out = np.zeros((M_PAD, self.gemm.N), dtype=np.float32) + for g, (B, s_wg, zpg, A, C) in enumerate(self.parts): + xg = xf[:real_m, g * Kg:(g + 1) * Kg] + # Per-group activation scale (per-token x per-K-group): finer than + # one scale over the full row. W4A8_GROUP_ACTS=1 enables it. + if os.environ.get("W4A8_GROUP_ACTS"): + sxg = np.max(np.abs(xg), axis=1, keepdims=True) / 127.0 + else: + sxg = np.max(np.abs(xf[:real_m]), axis=1, keepdims=True) / 127.0 + sxg = np.maximum(sxg, 1e-12) + q8 = np.zeros((M_PAD, Kg), dtype=np.int8) + q8[:real_m] = np.rint(xg / sxg).clip(-127, 127).astype(np.int8) + with A.overwrite() as buf: + buf[:] = q8 + res = self.gemm.op(A, B, C) + _ = res.npu_time + c = C.to_torch().numpy().astype(np.float32)[:real_m] + sw = s_wg.astype(np.float32)[None, :] + contrib = c * sxg.astype(np.float32) * sw + if zpg is not None: + # asymmetric correction: B ~ s_w*(q+z) -> out += s_w*z*rowsum(x) + contrib += xg.astype(np.float32).sum(1, keepdims=True) * (sw * zpg[None, :].astype(np.float32)) + out[:real_m] += contrib + return out + +def quantize_weight(W, dtype="i4"): + """W: torch [out, in] bf16 -> (B_ref [in, out] f32, s_w [out] f32).""" + B_ref = W.float().t().numpy() + peak = 8.0 if dtype == "i4" else 127.0 + s_w = np.max(np.abs(B_ref), axis=0) / peak + s_w = np.maximum(s_w, 1e-8).astype(np.float32) + return B_ref, s_w + + +def build_npu_layers(weights, build_dir): + """Return per-layer dict of NPU W4A8 GEMMs (cache by K,N).""" + pool = {} + mix = int(__import__("os").environ.get("W4A8_MIX_LEN", "0")) # first N layers i8 weights + groups = int(__import__("os").environ.get("W4A8_GROUPS", "1")) # i4 K-groups + + def get(K, N, dtype): + key = (K, N, dtype, groups if dtype == "i4" else 1) + if key not in pool: + pool[key] = NPU_W4A8_GEMM(K, N, build_dir, dtype_b=dtype, groups=groups) + return pool[key] + + layers = [] + for i in range(CONFIG["n_layers"]): + L = {} + for name, (Wkey, K, N) in { + "q": (f"model.layers.{i}.self_attn.q_proj.weight", CONFIG["emb_dim"], CONFIG["emb_dim"]), + "k": (f"model.layers.{i}.self_attn.k_proj.weight", CONFIG["emb_dim"], CONFIG["n_kv_groups"] * CONFIG["head_dim"]), + "v": (f"model.layers.{i}.self_attn.v_proj.weight", CONFIG["emb_dim"], CONFIG["n_kv_groups"] * CONFIG["head_dim"]), + "o": (f"model.layers.{i}.self_attn.o_proj.weight", CONFIG["emb_dim"], CONFIG["emb_dim"]), + "gate": (f"model.layers.{i}.mlp.gate_proj.weight", CONFIG["emb_dim"], CONFIG["hidden_dim"]), + "up": (f"model.layers.{i}.mlp.up_proj.weight", CONFIG["emb_dim"], CONFIG["hidden_dim"]), + "down": (f"model.layers.{i}.mlp.down_proj.weight", CONFIG["hidden_dim"], CONFIG["emb_dim"]), + }.items(): + ops_mix = __import__("os").environ.get("W4A8_OPS_MIX", "") + attn_i8 = set(__import__("os").environ.get("W4A8_ATTN_I8", "").split(",")) + is_attn = name in ("q", "k", "v", "o") + if ops_mix == "attn_i8": # attention weights i8, FFN i4 + dt = "i8" if is_attn else "i4" + elif ops_mix == "ffn_i8": # FFN weights i8, attention i4 + dt = "i8" if not is_attn or name in attn_i8 else "i4" + else: + dt = "i8" if i < mix else "i4" + B_ref, s_w = quantize_weight(weights[Wkey], dt) + L[name] = get(K, N, dt).bind(B_ref, s_w) # per-weight buffers + layers.append(L) + print(f"[w4a8] built {len(pool)} compiled NPU GEMM shapes ({len(layers)} layers)", flush=True) + return layers + + +def forward_w4a8_npu(layers, weights, token_ids, rope_angles, build_dir): + """token_ids: [1, seq] int64. Returns logits bf16 [1, seq, vocab] (padded seq).""" + seq = token_ids.shape[1] + assert seq <= M_PAD + batch = 1 + emb = weights["model.embed_tokens.weight"] + x = torch.nn.functional.embedding(token_ids, emb) # [1, seq, emb] + x = torch.nn.functional.pad(x, (0, 0, 0, M_PAD - seq)) # [1, M_PAD, emb] + + mask = torch.triu(torch.ones(M_PAD, M_PAD, dtype=torch.bool), diagonal=1) + if seq < M_PAD: + mask[:, seq:] = True + mask[seq:, :] = True + + kv_caches = [None] * CONFIG["n_layers"] + for i in range(CONFIG["n_layers"]): + L = layers[i] + x_norm = llama_cpu.rms_norm_forward( + x, weights[f"model.layers.{i}.input_layernorm.weight"]) + # projections on NPU (W4A8) + q = L["q"](x_norm, seq).view(batch, M_PAD, CONFIG["n_heads"], CONFIG["head_dim"]) + k = L["k"](x_norm, seq).view(batch, M_PAD, CONFIG["n_kv_groups"], CONFIG["head_dim"]) + v = L["v"](x_norm, seq).view(batch, M_PAD, CONFIG["n_kv_groups"], CONFIG["head_dim"]) + kv_caches[i] = { + "k": k[0, :seq].transpose(0, 1), # [G, seq, hd] + "v": v[0, :seq].transpose(0, 1), + } + q = llama_cpu.rope_forward(q, rope_angles[:M_PAD]) + k = llama_cpu.rope_forward(k, rope_angles[:M_PAD]) + q = q.transpose(1, 2) # [b, H, M, hd] + k = k.transpose(1, 2) # [b, G, M, hd] + v = v.transpose(1, 2) + gsz = CONFIG["n_heads"] // CONFIG["n_kv_groups"] + k = k.repeat_interleave(gsz, dim=1) + v = v.repeat_interleave(gsz, dim=1) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(CONFIG["head_dim"]) + scores = scores.masked_fill(mask, float("-inf")) + scores[:, :, seq:, :] = 0.0 # padded query rows: finite (uniform), zeroed below + aw = torch.nn.functional.softmax(scores, dim=-1) + aw[:, :, seq:, :] = 0.0 # padded query rows attend to nothing + ctx_out = torch.matmul(aw, v).transpose(1, 2).contiguous().view( + batch, M_PAD, -1) + attn_out = L["o"](ctx_out, seq) # NPU + x = x + attn_out + x_norm = llama_cpu.rms_norm_forward( + x, weights[f"model.layers.{i}.post_attention_layernorm.weight"]) + gate = L["gate"](x_norm, seq) # NPU + up = L["up"](x_norm, seq) # NPU + hidden = torch.nn.functional.silu(gate) * up + ffn_out = L["down"](hidden, seq) # NPU + x = x + ffn_out + if i in (0, 1, 5): + print(f"[w4a8] layer {i}: x[0,{seq-1},:4] = " + f"{x[0, seq-1, :4].float().tolist()}", flush=True) + + x = llama_cpu.rms_norm_forward(x, weights["model.norm.weight"]) + logits = torch.nn.functional.linear(x, emb) # host, tied lm_head + return logits, kv_caches + + + + +def decode_w4a8_npu(layers, weights, token, kv_caches, pos, rope_angles): + """One decode step: token (int) at position pos. Returns (next_token, kvs).""" + emb = weights["model.embed_tokens.weight"] + x = torch.nn.functional.embedding(torch.tensor([[token]]), emb) # [1, 1, emb] + x = torch.nn.functional.pad(x, (0, 0, 0, M_PAD - 1)) # [1, M_PAD, emb] + for i, L in enumerate(layers): + xn = llama_cpu.rms_norm_forward(x, weights[f"model.layers.{i}.input_layernorm.weight"]) + q = L["q"](xn, 1)[:1].view(1, 1, CONFIG["n_heads"], CONFIG["head_dim"]) + k = L["k"](xn, 1)[:1].view(1, 1, CONFIG["n_kv_groups"], CONFIG["head_dim"]) + v = L["v"](xn, 1)[:1].view(1, 1, CONFIG["n_kv_groups"], CONFIG["head_dim"]) + q = llama_cpu.rope_forward(q, rope_angles[pos:pos + 1]).squeeze(0).squeeze(0) # [H, hd] + k = llama_cpu.rope_forward(k, rope_angles[pos:pos + 1]).squeeze(0).squeeze(0) # [G, hd] + kv_caches[i]["k"] = torch.cat([kv_caches[i]["k"], k.unsqueeze(1)], dim=1) # [G, S, hd] + kv_caches[i]["v"] = torch.cat([kv_caches[i]["v"], v.squeeze(0).squeeze(0).unsqueeze(1)], dim=1) + gsz = CONFIG["n_heads"] // CONFIG["n_kv_groups"] + K = kv_caches[i]["k"].repeat_interleave(gsz, dim=0) # [H, S, hd] + V = kv_caches[i]["v"].repeat_interleave(gsz, dim=0) + scores = torch.matmul(q.unsqueeze(1), K.transpose(-2, -1)).squeeze(1) / math.sqrt(CONFIG["head_dim"]) # [H, S] + aw = torch.nn.functional.softmax(scores, dim=-1) + ctx = torch.matmul(aw.unsqueeze(1), V).squeeze(1).reshape(1, 1, -1) # [1, 1, H*hd] + ctx = torch.nn.functional.pad(ctx, (0, 0, 0, M_PAD - 1)) + x = x + L["o"](ctx, 1) + xn = llama_cpu.rms_norm_forward(x, weights[f"model.layers.{i}.post_attention_layernorm.weight"]) + gate = L["gate"](xn, 1) + up = L["up"](xn, 1) + hidden = torch.nn.functional.silu(gate) * up + x = x + L["down"](hidden, 1) + xn = llama_cpu.rms_norm_forward(x, weights["model.norm.weight"]) + logits = torch.nn.functional.linear(xn[:1], emb).squeeze(0).squeeze(0) + return int(logits.argmax()), kv_caches + + +def main(): + prompt = sys.argv[1] if len(sys.argv) > 1 else "The capital of France is" + build_dir = sys.argv[2] if len(sys.argv) > 2 else "/tmp/w4a8-prefill" + + weights = st.load_file(WT) + for k, v in weights.items(): + weights[k] = v.to(torch.bfloat16) + + tok_path = "/home/bcloud/llama3.2-1b/original/tokenizer.model" + config, state = harness.init(WT, tok_path, prompt=prompt) + token_ids = state.token_ids # torch [1, seq] + seq = token_ids.shape[1] + rope_angles = harness.compute_rope_angles( + CONFIG["head_dim"], M_PAD + 16, rope_base=500000.0) + + # bf16 CPU reference + t0 = time.time() + ref_logits, _ = llama_cpu.llama_forward_pass(config, state) + t_ref = time.time() - t0 + + # W4A8 NPU + layers = build_npu_layers(weights, build_dir) + t0 = time.time() + logits, kv_caches = forward_w4a8_npu(layers, weights, token_ids, rope_angles, build_dir) + t_npu = time.time() - t0 + + # compare at the last REAL token + ref_last = ref_logits[0, seq - 1].float() + npu_last = logits[0, seq - 1].float() + corr = float(torch.corrcoef(torch.stack([ref_last, npu_last]))[0, 1]) + diff = (npu_last - ref_last).abs() + print(f"\nprompt: '{prompt}' ({seq} tokens, padded {M_PAD})") + print(f"bf16 ref : {t_ref:.2f}s cpu | W4A8 NPU: {t_npu:.2f}s (GEMM-heavy path)") + print(f"logits corr : {corr:.6f}") + print(f"max |dlogit| : {diff.max():.4f}") + print(f"top1 : ref {ref_last.argmax().item():>6} npu {npu_last.argmax().item():>6} " + f"{'OK' if ref_last.argmax()==npu_last.argmax() else 'MISMATCH'}") + rk, nk = torch.topk(ref_last, 5).indices.tolist(), torch.topk(npu_last, 5).indices.tolist() + print(f"top5 ref: {rk}") + print(f"top5 npu: {nk}") + print(f"overlap : {len(set(rk) & set(nk))}/5") + + # W4A8 decode (KV-cached, NPU i4 GEMMs, real_m=1) + print("\nW4A8 decode (KV-cached):") + pos = seq + nxt = int(logits[0, seq - 1].argmax()) + out_ids = [] + for _ in range(24): + if nxt == 2: + break + out_ids.append(nxt) + nxt, kv_caches = decode_w4a8_npu(layers, weights, nxt, kv_caches, pos, rope_angles) + pos += 1 + print(f"token ids: {out_ids}") + print(f"text: {config.tokenizer.decode(out_ids)}") + +if __name__ == "__main__": + main() diff --git a/llama_w4a8_validate.py b/llama_w4a8_validate.py new file mode 100644 index 000000000..2f5b12adc --- /dev/null +++ b/llama_w4a8_validate.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +"""W4A8 layer validation on real llama weights (task: thread i4 weights into +the IRON llama pipeline). + +Quantizes a real llama-3.2-1B q_proj to INT4 (packed, per-output-neuron +scales), quantizes activations to INT8 (per-tensor scale), runs the NPU +i8xi4 GEMM (bit-exact int8 x int4), dequantizes with the two scales, and +compares against the bf16 reference. + +Usage (iron-venv python — has torch/safetensors + ml_dtypes + mlir_aie): + cd ~/amd-oss/iron + PYTHONPATH=/usr/lib/python3/dist-packages \ + ~/amd-oss/iron-venv/bin/python llama_w4a8_validate.py [M] [layer] +""" +import sys +import numpy as np + +sys.path.insert(0, "/home/bcloud/amd-oss/iron") + +import safetensors.torch as st +import torch +from iron.common.context import AIEContext +from iron.operators import GEMM +from aie.utils.hostruntime.xrtruntime.tensor import XRTTensor + +M = int(sys.argv[1]) if len(sys.argv) > 1 else 256 +LAYER = sys.argv[2] if len(sys.argv) > 2 else "0" +WT = f"/home/bcloud/llama3.2-1b/model.safetensors" +KEY = f"model.layers.{LAYER}.mlp.gate_proj.weight" + +print(f"== W4A8 llama q_proj layer {LAYER} (M={M}) ==") +w = st.load_file(WT)[KEY].to(torch.float32).numpy() # [out, in] +OUT, INN = w.shape +B_ref = w.T # [K=in, N=out] — the GEMM's B operand +rng = np.random.default_rng(7) +x = rng.normal(size=(M, INN)).astype(np.float32) # activations ~ N(0,1) +ref = (x.astype(np.float64) @ B_ref.astype(np.float64)) + +# ---- W4A8 quantization ---- +# weights: per-output-neuron (per B column) scale, int4 range [-8, 7] +s_w = np.max(np.abs(B_ref), axis=0) / 8.0 +s_w = np.maximum(s_w, 1e-8) +q4 = np.rint(B_ref / s_w).clip(-8, 7).astype(np.int8) +# activations: per-tensor scale, int8 range [-127, 127] +s_x = np.max(np.abs(x)) / 127.0 +q8 = np.rint(x / s_x).clip(-127, 127).astype(np.int8) + +# exact int reference (what the NPU must reproduce) +exact_i = q8.astype(np.int32) @ q4.astype(np.int32) # [M, N] +deq_exact = exact_i.astype(np.float64) * (s_x * s_w)[None, :] + +# ---- NPU i8xi4 GEMM ---- +K, N = INN, OUT +assert N % 2 == 0 and N % 32 == 0 and K % 16 == 0 and M % 8 == 0, "shape constraints" +ctx = AIEContext(build_dir="/tmp/w4a8-build") +op = ( + GEMM(M=M, K=K, N=N, tile_m=64, tile_k=64, tile_n=64, + num_aie_columns=8, dtype_in="i8", dtype_out="i32", + dtype_b="i4", context=ctx) + .compile() + .get_callable() +) +A = XRTTensor((M, K), dtype=np.int8) +B = XRTTensor((K, N // 2), dtype=np.int8) +C = XRTTensor((M, N), dtype=np.int32) +A.numpy()[:] = q8 +B.numpy()[:] = GEMM.pack_i4(q4) +op(A, B, C) +C_npu = C.to_torch().numpy() + +# ---- metrics ---- +exact_ok = np.array_equal(C_npu, exact_i) +deq = C_npu.astype(np.float64) * (s_x * s_w)[None, :] +corr = float(np.corrcoef(deq.ravel(), ref.ravel())[0, 1]) +mae = float(np.mean(np.abs(deq - ref))) +mre = float(np.mean(np.abs(deq - ref) / (np.abs(ref) + 1e-6))) +top5 = np.argsort(deq[0])[::-1][:5] +ref5 = np.argsort(ref[0])[::-1][:5] + +print(f"exact i8xi4 on NPU == CPU int ref: {exact_ok}") +print(f"corr vs bf16 ref : {corr:.6f}") +print(f"MAE / MRE : {mae:.4f} / {mre:.5f}") +print(f"top-5 (deq) : {top5.tolist()}") +print(f"top-5 (ref) : {ref5.tolist()}") +print(f"W scale range : {s_w.min():.3e} .. {s_w.max():.3e}") +print(f"quant err (W) : {np.mean(np.abs(B_ref - q4 * s_w)):.4f} " + f"(rel {np.mean(np.abs(B_ref - q4*s_w)/(np.abs(B_ref)+1e-6)):.4f})")