diff --git a/npu/forward/include/q4nx_forward.h b/npu/forward/include/q4nx_forward.h index c64be7992..33c53eb49 100644 --- a/npu/forward/include/q4nx_forward.h +++ b/npu/forward/include/q4nx_forward.h @@ -243,7 +243,15 @@ class Q4nxNpuForward { // quantisation, a wrong 128-K group boundary, an int8 intermediate or a saturated accumulator. bool host_gemm_i8_ = false; // Weight codes + per-128-K multipliers for the integer host arm, quantised once per tensor (engine#1809). - struct HostI8Tensor { int groups = 0; int N = 0; std::vector codes; std::vector d; }; + struct HostI8Tensor { + int groups = 0; + int N = 0; + // engine#1994: which layout `codes` is in. A blocked image is ceil(N/128)*ceil(K/64)*8192 bytes; K*N only fits + // when both are aligned, and indexing a K*N buffer with blocked_index() walks past it otherwise. + bool blocked = false; + std::vector codes; + std::vector d; + }; // Keyed on (LAYER, wkey, K, N): `wkey` alone is a per-role slot that repeats in every layer and `Bmat` is the one // scratch buffer each layer is decoded into, so keying on it served layer 0's weights to every later layer // (engine#1830 review). Routed experts pass wkey < 0 and are never cached: their weights change per token. diff --git a/npu/forward/include/q4nx_weight_layout.h b/npu/forward/include/q4nx_weight_layout.h index 722407646..18586aa63 100644 --- a/npu/forward/include/q4nx_weight_layout.h +++ b/npu/forward/include/q4nx_weight_layout.h @@ -62,6 +62,15 @@ inline size_t rowmajor_index(size_t i, size_t j, int KD, int ND) { return i * (size_t)ND + j; } +// Bytes a blocked image of KD x ND occupies: ceil(ND/128) tiles across, ceil(KD/64) rows of tiles, 8192 B per tile. +// engine#1994: a caller that allocates KD*ND and indexes with blocked_index() reads and writes past the buffer whenever +// the geometry is not 128/64-aligned -- the two expressions are each correct on their own, which is why it survived review: +// `150144 % 8192 = 2688 = 21 * 128` and `11008 = 86 * 128`, so every measured shape was aligned. +inline size_t blocked_bytes(int KD, int ND) { + const size_t nt = ((size_t)ND + 127) / 128, kt = ((size_t)KD + 63) / 64; + return nt * kt * 8192; +} + inline bool blocked_fits(int KD, int ND) { return KD > 0 && ND > 0 && KD % 64 == 0 && ND % 128 == 0; } diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index 09cba1ef7..79eb19ae2 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -3142,7 +3142,13 @@ void Q4nxNpuForward::host_gemm_i8(const float* A, int a_k, int M, int K, int N, if (!cacheable || it == host_i8_cache_.end()) { HostI8Tensor t; t.N = N; - t.codes.assign((size_t)K * (size_t)N, 0); + /* engine#1994: a blocked image is ceil(N/128)*ceil(K/64)*8192 bytes, which equals K*N only on aligned geometry. + * Allocating K*N and indexing with blocked_index() runs past the vector for e.g. N=8512 (66.5 tiles) or an LM head + * whose last vocabulary block is not 128-aligned (phi-3-mini: 32064 % 8192 = 7488) -- the packer writes past the + * buffer and the MAC reads the same out-of-range offsets back, so the reference returns plausible wrong numbers. + * The row-major layout is safe for any geometry and is only slower; correctness first. */ + t.blocked = q4nx_layout::blocked_fits(K, N); + t.codes.assign(t.blocked ? q4nx_layout::blocked_bytes(K, N) : (size_t)K * (size_t)N, 0); // RESOLVED toward the int8 packer (this branch) rather than main's corrected int4 rule (engine#1839). // Both are self-consistent; the question is what the DEVICE consumes. The device's body image comes from // packB_blocked() -- one multiplier per column, int8 codes -- and this branch measured the consequence: @@ -3160,7 +3166,7 @@ void Q4nxNpuForward::host_gemm_i8(const float* A, int a_k, int M, int K, int N, // blocked=1: the image is written in the blocked microtile order, which is the fast path AND the layout the // device actually reads, so the MAC below walks it with q4nx_layout::blocked_index(). (blocked=0 was measured // at 174.5 s of warm-up on this container versus 32.2 s blocked -- the row-major scatter is the slow one.) - q4nx_pack_B_int8_blocked(t.codes.data(), Bmat, K, N, K, N, /*blocked=*/1, t.d.data(), &ssum); + q4nx_pack_B_int8_blocked(t.codes.data(), Bmat, K, N, K, N, /*blocked=*/t.blocked ? 1 : 0, t.d.data(), &ssum); if (!cacheable) { scratch = std::move(t); cached = &scratch; @@ -3199,9 +3205,11 @@ void Q4nxNpuForward::host_gemm_i8(const float* A, int a_k, int M, int K, int N, float* dst = out + (size_t)m * (size_t)N; for (int n = 0; n < N; n++) { int32_t acc = 0; - for (int k = 0; k < K; k++) - acc += (int32_t)((int)t.codes[q4nx_layout::blocked_index((size_t)k, (size_t)n, K, N)] * - (int)a8[(size_t)k]); + for (int k = 0; k < K; k++) { + const size_t wi = t.blocked ? q4nx_layout::blocked_index((size_t)k, (size_t)n, K, N) + : q4nx_layout::rowmajor_index((size_t)k, (size_t)n, K, N); + acc += (int32_t)((int)t.codes[wi] * (int)a8[(size_t)k]); + } dst[(size_t)n] = t.d[(size_t)n] * (ascale * (float)acc); } } @@ -5359,16 +5367,24 @@ bool Q4nxNpuForward::run_lm_head(const std::vector& h, std::vector for (int n = n0; n < n1; n++) for (int k = 0; k < H; k++) bt[(size_t)k * Tn + (n - n0)] = lm_at(n, k); c.N = Tn; - c.codes.assign((size_t)H * (size_t)Tn, 0); + /* engine#1994: same sizing rule as the body GEMM above -- a blocked image is ceil(Tn/128)*ceil(H/64)*8192 + * bytes, and the LAST vocabulary block is where this bites: phi-3-mini's 32064 vocabulary ends in a 7488-wide + * block (58.5 tiles of 128), so K*Tn under-allocates by 196,608 B and both the packer and the MAC walk off + * the end of the vector. Row-major for unaligned geometry, blocked where it fits. */ + c.blocked = q4nx_layout::blocked_fits(H, Tn); + c.codes.assign(c.blocked ? q4nx_layout::blocked_bytes(H, Tn) : (size_t)H * (size_t)Tn, 0); c.d.assign((size_t)Tn, 0.0f); double ssum = 0.0; - q4nx_pack_B_int8_blocked(c.codes.data(), bt.data(), H, Tn, H, Tn, /*blocked=*/1, c.d.data(), &ssum); + q4nx_pack_B_int8_blocked(c.codes.data(), bt.data(), H, Tn, H, Tn, /*blocked=*/c.blocked ? 1 : 0, + c.d.data(), &ssum); } for (int n = 0; n < Tn; n++) { int32_t acc = 0; - for (int k = 0; k < H; k++) - acc += (int32_t)((int)c.codes[q4nx_layout::blocked_index((size_t)k, (size_t)n, H, Tn)] * - (int)a8i[(size_t)k]); + for (int k = 0; k < H; k++) { + const size_t wi = c.blocked ? q4nx_layout::blocked_index((size_t)k, (size_t)n, H, Tn) + : q4nx_layout::rowmajor_index((size_t)k, (size_t)n, H, Tn); + acc += (int32_t)((int)c.codes[wi] * (int)a8i[(size_t)k]); + } float v = c.d[(size_t)n] * (ascale * (float)acc); if (sem_.logit_softcap > 0) v = sem_.logit_softcap * std::tanh(v / sem_.logit_softcap); logits[(size_t)(n0 + n)] = v; diff --git a/tests/npu_pack_test.cpp b/tests/npu_pack_test.cpp index e8f354a95..aa3db1117 100644 --- a/tests/npu_pack_test.cpp +++ b/tests/npu_pack_test.cpp @@ -23,6 +23,7 @@ #include "full_elf.h" #include "pack.h" #include "q4nx_pack.h" +#include "q4nx_weight_layout.h" #include "batch_width.h" #include "golden/host_gemm_i8_fixture.h" @@ -112,7 +113,36 @@ static void batch_width_arms() { std::printf(" batch width: 7 arms (min not first, empty slots, fallback)\n"); } +/* engine#1994: `host_gemm_i8`'s weight image is sized K*N but indexed with the BLOCKED tile grid, so any N that is not + * 128-aligned reads and writes past the vector. The two expressions are each correct alone, which is why it survived + * review: `150144 % 8192 = 2688 = 21x128` and `11008 = 86x128`, so every measured shape was aligned. The failure needs a + * vocabulary whose last block is not (`32064 % 8192 = 7488`, phi-3-mini) or a body width that is not (`8512 = 66.5x128`). + * + * The arms pin the sizing rule, because the size is the defect: `blocked_bytes()` must equal K*N exactly when the geometry + * is aligned (so the fast path is unchanged) and must EXCEED K*N when it is not (so a blocked image cannot be allocated + * as K*N), and `blocked_fits()` must be false for exactly those geometries. */ +static void blocked_image_size_arms() { + using namespace q4nx_layout; + // aligned: the two agree, which is why nothing measured ever saw the bug + check(blocked_fits(4096, 11008) && blocked_bytes(4096, 11008) == size_t(4096) * 11008, + "blocked size: an aligned geometry (4096x11008, the measured family) is K*N exactly"); + check(blocked_fits(4096, 2688) && blocked_bytes(4096, 2688) == size_t(4096) * 2688, + "blocked size: the measured LM-head block (2688 = 21x128) is K*N exactly"); + // unaligned: the blocked image is BIGGER than K*N, which is the out-of-bounds write + check(!blocked_fits(3072, 32064), + "blocked size: 32064 is 250.5 tiles of 128, so blocked_fits() is false (phi-3-mini's vocabulary)"); + check(blocked_bytes(2048, 8512) > size_t(2048) * 8512, + "blocked size: 8512 (66.5 tiles) needs MORE than K*N -- allocating K*N and indexing blocked is the overflow"); + check(!blocked_fits(2048, 8512), "blocked size: and blocked_fits() refuses it"); + check(blocked_bytes(2048, 8512) - size_t(2048) * 8512 == 131072, + "blocked size: the overrun is 131,072 B for a 2048x8512 body, the figure in the issue"); + // the coarse predicate: K%64 and N%128 + check(!blocked_fits(2047, 8192), "blocked size: a K that is not 64-aligned is refused too"); + std::printf(" blocked size: 8 arms (aligned == K*N, unaligned > K*N with the issue's 131,072 B)\n"); +} + int main(int argc, char** argv) { + blocked_image_size_arms(); batch_width_arms(); quant_int4_arms(); for (int G : {8, 9, 16}) {