gemm: asymmetric INT8xINT4 (4-bit weights) on AIE2P 4x16x16 mmul — bit-exact + W4A8 llama decode

.gitignore

@@ -11,6 +11,9 @@ build/*
 **/_build/**
 **/build/**
 **/build_elf/**
+**/build_elf_*/
+**/build_int8_gemm/
+**/build_pb_*/
 *.exe
 *.csv
 secret_github_token

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.

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

aie_kernels/aie2p/mm.cc

@@ -79,12 +79,22 @@ template <typename T_in,
           unsigned s,
           unsigned t,
           bool b_row_maj = true,
-          bool c_row_maj = true>
+          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<r, s, t, T_in, T_in, accauto>;
+    using MMUL = aie::mmul<r, s, t, T_in, T_inB, accauto>;
+
+    // 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<T_inB, int4> ? 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<T_in, MMUL::size_A> A0;
                     aie::vector<T_in, MMUL::size_A> A1;
-                    aie::vector<T_in, MMUL::size_B> B0;
-                    aie::vector<T_in, MMUL::size_B> B1;
+                    aie::vector<T_inB, MMUL::size_B> B0;
+                    aie::vector<T_inB, MMUL::size_B> 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<MMUL::size_A>(pA1);
+                    pA1 += MMUL::size_A;
+                    A1 = aie::load_v<MMUL::size_A>(pA2);
+                    pA2 += MMUL::size_A;
+                    if constexpr (b_row_maj) {
+                        B0 = aie::load_v<MMUL::size_B>(pB1);
+                        pB1 += B_ADV * colB;
+                        B1 = aie::load_v<MMUL::size_B>(pB2);
+                        pB2 += B_ADV * colB;
+                    } else {
+                        B0 = aie::transpose(aie::load_v<MMUL::size_B>(pB1), t, s);
+                        pB1 += B_ADV;
+                        B1 = aie::transpose(aie::load_v<MMUL::size_B>(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<MMUL::size_A>(pA1);
+                            aie::vector<T_in, MMUL::size_A> A0n =
+                                aie::load_v<MMUL::size_A>(pA1);
                             pA1 += MMUL::size_A;
-                            A1 = aie::load_v<MMUL::size_A>(pA2);
+                            aie::vector<T_in, MMUL::size_A> A1n =
+                                aie::load_v<MMUL::size_A>(pA2);
                             pA2 += MMUL::size_A;
                             if constexpr (b_row_maj) {
-                                B0 = aie::load_v<MMUL::size_B>(pB1);
-                                pB1 += MMUL::size_B * colB;
-                                B1 = aie::load_v<MMUL::size_B>(pB2);
-                                pB2 += MMUL::size_B * colB;
+                                aie::vector<T_inB, MMUL::size_B> B0n =
+                                    aie::load_v<MMUL::size_B>(pB1);
+                                pB1 += B_ADV * colB;
+                                aie::vector<T_inB, MMUL::size_B> B1n =
+                                    aie::load_v<MMUL::size_B>(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<MMUL::size_B>(pB1), t, s);
-                                pB1 += MMUL::size_B;
-                                B1 = aie::transpose(aie::load_v<MMUL::size_B>(pB2), t, s);
-                                pB2 += MMUL::size_B;
+                                aie::vector<T_inB, MMUL::size_B> B0n = aie::transpose(
+                                    aie::load_v<MMUL::size_B>(pB1), t, s);
+                                pB1 += B_ADV;
+                                aie::vector<T_inB, MMUL::size_B> B1n = aie::transpose(
+                                    aie::load_v<MMUL::size_B>(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 <unsigned m, unsigned k, unsigned n>
+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<int8, int32, (m / r), (k / s), (n / t), r, s, t, is_b_row_maj,
+                                      is_c_row_maj, int4>(pA, reinterpret_cast<const int4 *>(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<ctype_out, DIM_M, DIM_N>(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<DIM_M, DIM_K, DIM_N>(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<ctype_out, DIM_M, DIM_N>(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"

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;

aie_kernels/generic/add.cc

@@ -23,6 +23,12 @@ template <typename T_in, typename T_out> 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;

aie_kernels/generic/mul.cc

@@ -20,6 +20,10 @@ template <typename T_in, typename T_out> 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);

aie_kernels/generic/rope.cc

@@ -10,6 +10,7 @@
 template <typename T, int N>
 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 <typename T, int N>
 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;