diff --git a/docs/npu.md b/docs/npu.md index fee09c40..23ac8e6f 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -556,7 +556,7 @@ FACTS §7.13 established that layout by known-answer write and §7.15 against sh ### The `.fconst` MLP section -The MLP section fills the rest of the layer block after the second norm, in three parts: +The MLP section fills the rest of the layer block after that block's **post-attention norm**, in three parts: | part | shape | layout | |---|---|---| @@ -588,6 +588,19 @@ different kernel shapes"; that was this reader's error, not the packs': the sect the wrong boundary, so the geometry check (`total == stride`) failed and was read as a shape difference. With §7.24's boundaries the gap closes to the byte on all four (FACTS §7.24.1). +### Which header opens the post-attention norm + +A layer block opens with a 64-byte epsilon header, then `input_layernorm`, then the attention section, +then **another epsilon header** and the `post_attention_layernorm` weight. Which header that second one is +depends on how many the block carries, and that is measured per pack, not assumed: **2-header blocks** +(Qwen2.5, Llama) carry exactly one extra, so it is the block's second header; **4-header blocks** (Qwen3) +carry three more, because the first two open the per-head `k_norm` and `q_norm` sections, so the +post-attention header is the block's **last** one. Reading the second header unconditionally is correct +only on the 2-header packs — on a Qwen3 block it imports the `k_norm` gains and the padding behind them as +`post_attention_layernorm` (1,791 of 2,048 values zero and max 68.0 on Qwen3-1.7B, against a real norm +that is zero-free and 0.9919 correlated with the upstream checkpoint). Both readers — the norm and the MLP +section's first byte — now go through `FconstGeom::post_attn_hdr()` (engine#1003, FACTS §7.25). + Which tile belongs to which projection is **measured, not assumed**: on the shipped 0.5B pack the tiles' bias subvolumes matched the pack's own q/k/v bias vectors, which identifies the runs as k (2 tiles), q (14), v (2), o (14) — 32 tiles, the payload's whole size. A pack whose payload is diff --git a/npu/pack_import.cpp b/npu/pack_import.cpp index d4a52fd5..bd45fc33 100644 --- a/npu/pack_import.cpp +++ b/npu/pack_import.cpp @@ -1020,7 +1020,7 @@ Plan plan(const std::string& pack_dir) { "or 4 (q/k/v, two per-head QK norms, o_proj) -- recorded 2026-11-04 on Qwen3"); } // the MLP section starts right after the block's post-attention norm - const int64_t mlp_off = g.hdr.back() + kNormOff + H * 2; + const int64_t mlp_off = g.post_attn_hdr() + kNormOff + H * 2; const MlpGeom mg = mlp_geom(H, d.intermediate, mlp_off, g.stride); p.mlp_note = mg.why; if (mg.ok()) { @@ -1359,9 +1359,9 @@ Report import_pack(const std::string& pack_dir, const std::string& out_dir, cons const int64_t base = int64_t(L) * g.stride; r.norm_anchors_total += 2; r.norm_anchors_ok += int(eps_at(fc, base, p.dims.eps_bits)) + - int(eps_at(fc, base + g.h1, p.dims.eps_bits)); + int(eps_at(fc, base + g.post_attn_hdr(), p.dims.eps_bits)); add_norm(base + kNormOff, "model.layers." + std::to_string(L) + ".input_layernorm.weight"); - add_norm(base + g.h1 + kNormOff, + add_norm(base + g.post_attn_hdr() + kNormOff, "model.layers." + std::to_string(L) + ".post_attention_layernorm.weight"); } // the final norm is refused in the plan: the 1,856-byte tail is not an RMSNorm weight diff --git a/npu/pack_import.h b/npu/pack_import.h index 2235d717..5b4b975c 100644 --- a/npu/pack_import.h +++ b/npu/pack_import.h @@ -210,7 +210,9 @@ int64_t group_bytes(int64_t k, int rows); // `bf16_bits_of()`; the default is bf16 1e-6, which is what the Qwen family declares. struct FconstGeom { int64_t stride = 0; // bytes per layer block - int64_t h1 = 0; // offset of the second (post-attention) header within a block + int64_t h1 = 0; // offset of the block's SECOND header (the post-attention one only when a block + // carries two; on a 4-header block it opens the k_norm section -- use + // post_attn_hdr() for the post-attention norm and the MLP section) int period = 0; // headers per block int blocks = 0; // block starts seen // Every in-block header offset after the block's first (so hdr[0] == h1), relative to the block @@ -218,6 +220,12 @@ struct FconstGeom { // one payload) and 3 on the Qwen3 packs (q/k/v, then one per-head QK norm each, then o_proj), so // which layout a pack uses is measured from the pack rather than guessed from a version. std::vector hdr; + // Where the block's post-attention section starts: `h1` on a 2-header block (Qwen2.5/Llama), the LAST + // header on a 4-header block (Qwen3, where the second and third open the per-head QK norms). Both the + // post_attention_layernorm weight and the first byte of the MLP section are read from here, so this is + // the one rule the two readers share -- reading `h1` unconditionally imported the k_norm section's + // padding as post_attention_layernorm on every Qwen3 pack (engine#1003). + int64_t post_attn_hdr() const { return hdr.empty() ? h1 : hdr.back(); } std::string why; // non-empty when the geometry could not be established }; unsigned short bf16_bits_of(float v); diff --git a/tests/npu_pack_import_test.cpp b/tests/npu_pack_import_test.cpp index afc19ba1..54321a46 100644 --- a/tests/npu_pack_import_test.cpp +++ b/tests/npu_pack_import_test.cpp @@ -968,6 +968,111 @@ void test_qwen3_attn_sections() { } } +// engine#1003: the post_attention_layernorm weight is the bytes after the block's post-attention header. +// On a 2-header block that is the second header; on a Qwen3 4-header block it is the FOURTH, because the +// second and third open the per-head QK norms. Reading the second one on every pack imported the k_norm +// section's padding as post_attention_layernorm -- measured on Qwen3-1.7B: 1,920 of 2,048 values zero, +// max 6.97, against the pack's real post-attention norm (all nonzero, 0.9919 correlated with the upstream +// checkpoint). The crafted pack below puts a different recognisable value behind each header, so the check +// is a known answer rather than a shape: the imported weight must be the last header's ramp and must NOT +// be the second header's marker. +void test_post_attn_norm_header(const std::string& root) { + const fs::path dir = fs::path(root) / "post-attn-hdr"; + fs::remove_all(dir); + fs::create_directories(dir / "cache"); + const int H = 128, L = 8; // locate_fconst wants >= 6 headers in the first GiB + const int64_t stride = 1 << 20; // the locator ignores periods below 1e6 + const int64_t h2 = 262144, h3 = 524288, h4 = 786432; // the three in-block headers after the first + std::string fc(size_t(int64_t(L) * stride + 64 + H * 2), '\0'); + auto put_bf16 = [&](int64_t off, float v) { + const uint16_t bits = pi::bf16_bits_of(v); + fc[size_t(off)] = char(bits & 0xFF); + fc[size_t(off) + 1] = char(bits >> 8); + }; + for (int l = 0; l < L; ++l) { + const int64_t base = int64_t(l) * stride; + for (int64_t h : {int64_t(0), h2, h3, h4}) put_bf16(base + h, 1e-6f); // bf16 0x3586, the Qwen eps + for (int i = 0; i < H; ++i) { + put_bf16(base + h2 + 64 + 2 * i, -3.0f); // the k_norm section: NOT the post-attention norm + put_bf16(base + h4 + 64 + 2 * i, float(i + 1)); // the post-attention norm + } + } + put_bf16(int64_t(L) * stride, 1e-6f); // the tail envelope's header + put_file((dir / "config.json").string(), + R"({"model_type":"qwen3","hidden_size":128,"intermediate_size":256,"num_hidden_layers":8,)" + R"("num_attention_heads":2,"num_key_value_heads":1,"head_dim":64,"vocab_size":128,)" + R"("rope_theta":1000000.0,"tie_word_embeddings":true,"rms_norm_eps":1e-06})"); + put_file((dir / "model.onnx.data").string(), std::string(size_t(H) * size_t(H) * 4, '\x11')); + put_file((dir / "model.onnx").string(), + pb::model({pb::tensor("model.embed_tokens.weight", {H, H}, "model.onnx.data", 0, + int64_t(H) * H * 4)}, {})); + const std::string tok = (dir / "dd_metastate_Llm_Token_Token_crafted.fconst").string(); + put_file(tok, fc); + + const pi::FconstGeom g = pi::locate_fconst(tok, 0x3586); + check(g.stride == stride && g.period == 4, + "post-attn header: the crafted 4-header .fconst is locatable (stride " + + std::to_string(g.stride) + ", period " + std::to_string(g.period) + + (g.why.empty() ? "" : ", why: " + g.why) + ")"); + check(g.h1 == h2, "post-attn header: h1 is the block's second header (" + std::to_string(g.h1) + ")"); + check(g.post_attn_hdr() == h4, + "post-attn header: the post-attention header on a 4-header block is the LAST one (" + + std::to_string(g.post_attn_hdr()) + " vs " + std::to_string(h4) + ")"); + check(g.post_attn_hdr() != g.h1, + "post-attn header: it is not h1, which is where the k_norm section opens"); + { // the 2-header case must be unchanged by the rule + std::string two(size_t(int64_t(L) * stride + 64 + H * 2), '\0'); + for (int l = 0; l < L; ++l) { + const int64_t base = int64_t(l) * stride; + two[size_t(base)] = char(0x86); + two[size_t(base) + 1] = char(0x35); + two[size_t(base + h2)] = char(0x86); + two[size_t(base + h2) + 1] = char(0x35); + } + two[size_t(int64_t(L) * stride)] = char(0x86); + const std::string p2 = (dir / "dd_metastate_Llm_Token_Token_two.fconst").string(); + put_file(p2, two); + const pi::FconstGeom g2 = pi::locate_fconst(p2, 0x3586); + check(g2.period == 2 && g2.post_attn_hdr() == g2.h1, + "post-attn header: on a 2-header block it is still h1 (" + std::to_string(g2.h1) + ")"); + fs::remove(p2); + } + + pi::Options allow; + allow.allow_incomplete = true; // the crafted block has no tiles, so the payload legs are refused + const std::string out = (fs::temp_directory_path() / + ("oga_pack_import_pa." + std::to_string(::getpid()))).string(); + fs::remove_all(out); + try { + const pi::Report r = pi::import_pack(dir.string(), out, allow); + check(r.stride == stride, "post-attn header: the import read the crafted stride"); + check(fs::exists(fs::path(out) / "model.q4nx"), + "post-attn header: the import wrote a container to inspect"); + if (fs::exists(fs::path(out) / "model.q4nx")) { + onebit::npu::Model m(out); + for (int l : {0, L - 1}) { + const std::string n = "model.layers." + std::to_string(l) + + ".post_attention_layernorm.weight"; + const auto& t = m.tensor(n); + bool ramp = t.bytes >= size_t(H) * 2, not_marker = true; + for (int i = 0; ramp && i < H; ++i) { + const uint16_t want = pi::bf16_bits_of(float(i + 1)); + ramp = t.data[2 * i] == uint8_t(want & 0xFF) && t.data[2 * i + 1] == uint8_t(want >> 8); + const uint16_t marker = pi::bf16_bits_of(-3.0f); + not_marker = not_marker && !(t.data[2 * i] == uint8_t(marker & 0xFF) && + t.data[2 * i + 1] == uint8_t(marker >> 8)); + } + check(ramp, n + " is the bytes after the block's LAST header"); + check(not_marker, n + " is not the k_norm section's marker (the engine#1003 bug)"); + } + } + } catch (const std::exception& e) { + check(false, std::string("post-attn header: the crafted pack imported: ") + e.what()); + } + fs::remove_all(out); + fs::remove_all(dir); +} + void test_fconst_eps(const std::string& root) { const fs::path dir = fs::path(root) / "fconst-eps"; fs::remove_all(dir); fs::create_directories(dir); @@ -1038,6 +1143,7 @@ int main(int argc, char** argv) { test_fconst_eps(golden_dir.empty() ? "." : golden_dir); test_attn_payload_off(); test_qwen3_attn_sections(); + test_post_attn_norm_header(golden_dir.empty() ? "." : golden_dir); if (pack_ran) test_pack(pack, golden_dir); else std::printf("skip pack: no pack directory given (the crafted and golden legs only)\n"); } catch (const std::exception& e) {