diff --git a/CMakeLists.txt b/CMakeLists.txt index 9262b7018..657ac6f1c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2192,6 +2192,12 @@ if(ONEBIT_NPU_HOST AND NOT ONEBIT_NPU AND NOT WIN32 AND Python3_Interpreter_FOUN COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tests/batch_check_state_ci.py $ qwen2vl) set_tests_properties(batch_check_state_qwen2vl PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 120) + # Hy-Memory (hy_v3, engine#1786): sigmoid routing with a selection bias, router [H][E], a leading dense layer, + # [N][K] expert slices under FLM names, expert_hidden_dim, against a numpy hy_v3 reference. + add_test(NAME batch_check_state_hy_v3 + COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tests/batch_check_state_ci.py + $ hy_v3) + set_tests_properties(batch_check_state_hy_v3 PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 120) # GPT-OSS (engine#1786): a container written in FastFlowLM's D2560/MX2560 form, sinks, o bias, yarn, the top-k # MoE with clamped SwiGLU, against a numpy GPT-OSS forward on the same bytes. add_test(NAME gptoss_forward diff --git a/docs/npu.md b/docs/npu.md index 08148b335..4cb90d1df 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -2112,6 +2112,46 @@ the SwiGLU clamp (0.92), the YaRN attention factor (0.994), the YaRN frequency r (0.99967). A batched prefill of gpt_oss takes the per-token path. No kernel set for its shapes was built or run (the NPU is off). +**Hy-Memory-30B-A3B-NPU2: hy_v3 runs on the existing MoE path (engine#1786).** Tencent's hy_v3 routes as DeepSeek-V3 +does (our llama.cpp fork's `src/models/hy-v3.cpp`): sigmoid router probabilities, the top 8 of 128 chosen on probability +plus an expert bias, weighted by the unbiased probabilities normalised to sum 1 and scaled, one ungated shared expert, +and a dense first layer. The forward already ran that routing. What it did not do was read the pack's way of saying +it. Each of these, alone, made the reply ` of of of`: + +- **The config keys.** `moe_router_use_sigmoid` is the sigmoid (it was read as softmax), `router_scaling_factor` is the + weight scale, `first_k_dense_replace` is the leading dense count with the dense width from `intermediate_size`, and + `expert_hidden_dim` (768) is the expert width. `moe_intermediate_size` says 1536, but the experts and the shared + expert are `[768][2048]` in `tencent/Hy-Memory-30B-A3B` itself. +- **The router** is stored `[H][E]`, the transpose of upstream's `[E][H]`, checked by a ranged read of the checkpoint. + A router of that shape is transposed at load. **The expert bias** is the F32 `moe_expert_bias`, equal to upstream's + `mlp.expert_bias`. +- **The expert slices are `[N][K]`.** FastFlowLM's Qwen MoE packs transpose each expert, and the forward read every + `*_exps_proj` stack that way. This pack does not transpose: its gate stack is `[3072][8][4736]`, 24 x 8 tiles per + expert for 768 x 2048. The grid is now decided from the first MoE layer's own gate stack. Against upstream expert + 0 of layer 1, `[N][K]` gives gate and down corr 0.9974, and the old reading gives 0.008. The same rule also fires + on `MiniCPM-V-4.7-35B-A3B-NPU2` (the only other container on the box it fires on). That reading is right there too: + its layer-0 expert 0 against `openbmb/MiniCPM-V-4.7-35B-A3B`'s own `gate_up_proj` gives corr 0.997 as `[N][K]` and + 0.024 as `[K][N]`. That pack does not run on this forward either way ("MoE designs missing", engine#1788), so no + answer changes. +- **A leading dense layer without MLA** gets the DG/DU/DD designs and `dense_mlp` the MLA path already had. A batched + prefill of such a model takes the per-token path. + +On the host the known-answer prompt (no-think template) answers `Paris` and ``. Against a numpy hy_v3 forward +of the same weights, decoded through `q4nx_dequant.c`, the last position's logits match at corr 1.000000, max |d| +3e-5 of 26.6, with the same top 5. `batch_check_state_hy_v3` fabricates a two-layer pack written this way and holds +the forward to that reference: corr 1.0000000. Each of these mutants fails it: + +| mutant | corr | +|---|---| +| `[K][N]` slices | 0.59 | +| no router transpose | 0.78 | +| softmax gating | 0.9965 | +| no scale | 0.94 | +| bias ignored | 0.87 | +| expert width from `moe_intermediate_size` | refused at the expert gate | + +No kernel set for its shapes was built or run (the NPU is off). + **Per-head q/k normalisation is read from the container (engine#1788).** Whether the forward normalises q and k per head used to be decided by `config.json`'s `model_type` alone: on for `qwen3*` and `gemma*`. A pack of the same architecture under another name got it wrong in silence. FastFlowLM's MiniCPM-V 4.7 diff --git a/npu/forward/include/q4nx_forward.h b/npu/forward/include/q4nx_forward.h index 65e40b8b6..4c38db9aa 100644 --- a/npu/forward/include/q4nx_forward.h +++ b/npu/forward/include/q4nx_forward.h @@ -545,6 +545,7 @@ class Q4nxNpuForward { int mla_ = 0; int q_lora_ = 0, kv_lora_ = 0, rope_dim_ = 0, qk_nope_ = 0, v_head_ = 0; int n_expert_ = 0, top_k_ = 0, expert_im_ = 0, dense_im_ = 0; + bool moe_slice_nk_ = false; // per-expert slices on an [N][K] grid despite FLM names (hy_v3) // Which MoE container convention this container uses, decided from its own tensors in // init_impl (moe_container.h): the engine converter's `mlp.experts.*` names with a plain // [N][K] per-expert tile grid, or the FLM-style `*_exps_proj` names whose writer transposes diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index 687a0a188..6fc4e52a9 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -1315,11 +1315,48 @@ bool Q4nxNpuForward::init_impl(xrt::device* dev, const char* model_path, q4nx_config_int(cfg_path.c_str(), "num_experts_per_tok", 0)); const int moe_im = q4nx_config_int(cfg_path.c_str(), "moe_intermediate_size", 0); expert_im_ = moe_im > 0 ? moe_im : c_.IM; + // hy_v3 states both: moe_intermediate_size 1536 and expert_hidden_dim 768, and its experts (and the shared + // expert) are [768][2048] in tencent/Hy-Memory-30B-A3B's own checkpoint, so expert_hidden_dim is the width. + const int ehd = q4nx_config_int(cfg_path.c_str(), "expert_hidden_dim", 0); + if (ehd > 0) expert_im_ = ehd; dense_im_ = q4nx_config_int(cfg_path.c_str(), "dense_intermediate_size", 0); leading_dense_ = q4nx_config_int(cfg_path.c_str(), "leading_dense_block_count", 0); expert_gating_sigmoid_ = q4nx_config_int(cfg_path.c_str(), "expert_gating_func", 1) == 2; // 2 = sigmoid expert_w_scale_ = read_config_float(cfg_path, "\"expert_weights_scale\"", 1.0f); + // Hy-Memory (hy_v3, engine#1786) states the same DeepSeek-V3 routing under its own keys: the + // sigmoid as moe_router_use_sigmoid, the scale as router_scaling_factor, the leading dense block + // as first_k_dense_replace, and the dense FFN's width as intermediate_size. Our llama.cpp fork's + // hy-v3.cpp routes it as deepseek2 does (sigmoid probs, top-k on probs + expert bias, the + // unbiased probs sum-normalised and scaled), which is the path moe_mlp runs. + { + std::ifstream cf(cfg_path); + const std::string ct((std::istreambuf_iterator(cf)), std::istreambuf_iterator()); + if (ct.find("\"moe_router_use_sigmoid\": true") != std::string::npos || + ct.find("\"moe_router_use_sigmoid\":true") != std::string::npos) + expert_gating_sigmoid_ = true; + if (ct.find("\"expert_weights_scale\"") == std::string::npos) + expert_w_scale_ = read_config_float(cfg_path, "\"router_scaling_factor\"", expert_w_scale_); + if (leading_dense_ == 0) + leading_dense_ = q4nx_config_int(cfg_path.c_str(), "first_k_dense_replace", 0); + if (leading_dense_ > 0 && dense_im_ == 0) dense_im_ = c_.IM; + } + // The per-expert tile grid, from the first MoE layer's own gate stack rather than the name convention: + // FLM-named Qwen MoE containers transpose each expert ([K][N] slices), Hy-Memory-30B-A3B's do not -- its + // gate stack is [E * N/32][K/256][chunk] = [3072][8][4736] for N = 768, K = 2048. Read the wrong way, every + // expert is a permutation and the model answers " of of of". + { + char gn[160]; + snprintf(gn, sizeof gn, "model.layers.%d.mlp.gate_exps_proj.weight", leading_dense_ > 0 ? leading_dense_ : 0); + const TensorDesc* gt = model_tensor_by_name(mw_, gn); + if (gt && gt->ndim == 3 && n_expert_ > 0 && gt->shape[0] % n_expert_ == 0 && + gt->shape[0] / n_expert_ * Q4NX_TILE_ROWS == expert_im_ && + gt->shape[1] * Q4NX_TILE_COLS == q4nx_pad_k(c_.H) && expert_im_ != q4nx_pad_k(c_.H)) { + moe_slice_nk_ = true; + fprintf(stderr, " [moe] per-expert slices are [N][K] (%s is [%lld][%lld][...])\n", gn, + (long long)gt->shape[0], (long long)gt->shape[1]); + } + } if (top_k_ <= 0 || top_k_ > n_expert_ || expert_im_ <= 0) { fprintf(stderr, " [moe] config incomplete (n_expert=%d top_k=%d im=%d)\n", n_expert_, top_k_, expert_im_); @@ -1382,8 +1419,28 @@ bool Q4nxNpuForward::init_impl(xrt::device* dev, const char* model_path, q4nx_pad_n(dims.IM), tag); dims.fused_gu = file_exists(gu) || !file_exists(g) ? 1 : 0; } + // A MoE with a leading dense layer and no MLA (hy_v3): the routed experts' G/U/D are at the expert width, and the + // dense layer gets its own DG/DU/DD at intermediate_size -- the arrangement the MLA path already uses. + const bool moe_lead_dense = n_expert_ > 0 && !mla_ && leading_dense_ > 0 && dense_im_ > 0; + if (moe_lead_dense) { dims.IM = expert_im_; dims.fused_gu = 0; } ndesigns_ = q4nx_derive_designs(&dims, designs_); if (ndesigns_ <= 0) { err_ = "design derivation failed"; return false; } + if (moe_lead_dense && ndesigns_ + 3 <= Q4NX_MAX_DESIGNS) { + const struct { Q4nxDesign role; int K, N; const char* tag; } dd[3] = { + {Q4NX_DESIGN_DG, c_.H, dense_im_, "DG"}, {Q4NX_DESIGN_DU, c_.H, dense_im_, "DU"}, + {Q4NX_DESIGN_DD, dense_im_, c_.H, "DD"}}; + for (const auto& x : dd) { + memset(&designs_[ndesigns_], 0, sizeof(designs_[ndesigns_])); + designs_[ndesigns_].M = dims.M; designs_[ndesigns_].K = q4nx_pad_k(x.K); designs_[ndesigns_].N = q4nx_pad_n(x.N); + designs_[ndesigns_].real_k = x.K; designs_[ndesigns_].real_n = x.N; + designs_[ndesigns_].role = x.role; + snprintf(designs_[ndesigns_].elf_name, Q4NX_ELF_NAME_MAX, "full_i8_%s_K%d_N%d.elf", x.tag, + q4nx_pad_k(x.K), q4nx_pad_n(x.N)); + ndesigns_++; + } + fprintf(stderr, " [forward] MoE with %d leading dense layer(s): experts at IM %d, dense FFN at %d\n", + leading_dense_, expert_im_, dense_im_); + } // Batched expert gate/up (MoE): one GEMM per projection over ALL top-k experts // at once, instead of one per expert. Derived here rather than in // q4nx_derive_designs because it depends on top_k, which the dims struct does @@ -2433,7 +2490,26 @@ bool Q4nxNpuForward::load_layer(int l) { } if (!deq(&lw->o_proj_weight, wo_.data(), H, q4nx_pad_k(NH * hd))) return false; } - if (n_expert_ > 0) { + if (n_expert_ > 0 && !mla_ && l < leading_dense_) { + // the leading dense layer of a non-MLA MoE (hy_v3): dense_mlp's DG/DU/DD weights + char dn[160]; + auto load_dense = [&](const char* fmt, std::vector& v, int rows, int cols) -> bool { + snprintf(dn, sizeof dn, fmt, l); + TensorDesc* d = model_tensor_by_name(mw_, dn); + if (!d) { err_ = std::string("leading dense layer: model.q4nx has no ") + dn; return false; } + v.resize((size_t)rows * (size_t)cols); + const uint8_t* raw = (const uint8_t*)model_tensor_data(mw_, d); + if (!raw || dequant_any(raw, (size_t)d->data_size, rows, cols, v.data(), d) != 0) { + err_ = std::string("leading dense layer: cannot read ") + dn; + return false; + } + return true; + }; + return load_dense("model.layers.%d.mlp.gate_proj.weight", dense_g_, dense_im_, H) && + load_dense("model.layers.%d.mlp.up_proj.weight", dense_u_, dense_im_, H) && + load_dense("model.layers.%d.mlp.down_proj.weight", dense_d_, H, dense_im_); + } + if (n_expert_ > 0 && l >= leading_dense_) { // a leading dense layer (hy_v3's first_k_dense_replace) runs the dense FFN // Qwen MoE layer (qwen3moe / qwen35moe): a router + top-k routed experts // instead of a dense gate/up/down. The router is bf16; its bias is OPTIONAL // (Qwen MoE has none, deepseek2 does). The experts stay pooled -- one @@ -2459,7 +2535,31 @@ bool Q4nxNpuForward::load_layer(int l) { return true; }; if (!load_bf16_named(mn.router, moe_router_, n_expert_ * H)) { err_ = "moe router"; return false; } - if (!load_bf16_named(mn.router_b, moe_router_b_, n_expert_)) moe_router_b_.clear(); + { + // A router stored [H][E] (Hy-Memory-30B-A3B's moe_router.weight is upstream's [E][H] transposed, + // checked against tencent/Hy-Memory-30B-A3B) is read flat as [E][H] by moe_mlp: transpose it. + const TensorDesc* rd = td_named(mn.router); + if (rd && rd->ndim == 2 && rd->shape[0] == H && rd->shape[1] == n_expert_ && H != n_expert_) { + std::vector t((size_t)n_expert_ * H); + for (int k = 0; k < H; k++) + for (int e = 0; e < n_expert_; e++) t[(size_t)e * H + k] = moe_router_[(size_t)k * n_expert_ + e]; + moe_router_.swap(t); + } + } + if (!load_bf16_named(mn.router_b, moe_router_b_, n_expert_)) { + moe_router_b_.clear(); + // hy_v3: the expert-selection bias is `moe_expert_bias`, F32 in FastFlowLM's pack (bf16 from our packer) + const TensorDesc* bd = td_named("moe_expert_bias"); + const void* bp = bd ? model_tensor_data(mw_, const_cast(bd)) : nullptr; + if (bp && !strcmp(bd->dtype, "F32") && (size_t)bd->data_size >= sizeof(float) * (size_t)n_expert_) { + const float* fp = (const float*)bp; + moe_router_b_.assign(fp, fp + n_expert_); + } else if (bp && !strcmp(bd->dtype, "BF16") && (size_t)bd->data_size >= 2 * (size_t)n_expert_) { + const uint16_t* hp = (const uint16_t*)bp; + moe_router_b_.resize((size_t)n_expert_); + for (int e = 0; e < n_expert_; e++) moe_router_b_[(size_t)e] = bf16_to_f32(hp[e]); + } + } auto pool = [&](const char* suffix, const uint8_t** p, size_t* bytes) -> bool { TensorDesc* d = td_named(suffix); if (!d) { fprintf(stderr, " [moe] missing %s\n", nm2); return false; } @@ -2616,7 +2716,31 @@ bool Q4nxNpuForward::load_gdn_layer(int l) { return true; }; if (!load_bf16_named(mn.router, moe_router_, n_expert_ * H)) { err_ = "moe router"; return false; } - if (!load_bf16_named(mn.router_b, moe_router_b_, n_expert_)) moe_router_b_.clear(); + { + // A router stored [H][E] (Hy-Memory-30B-A3B's moe_router.weight is upstream's [E][H] transposed, + // checked against tencent/Hy-Memory-30B-A3B) is read flat as [E][H] by moe_mlp: transpose it. + const TensorDesc* rd = td_named(mn.router); + if (rd && rd->ndim == 2 && rd->shape[0] == H && rd->shape[1] == n_expert_ && H != n_expert_) { + std::vector t((size_t)n_expert_ * H); + for (int k = 0; k < H; k++) + for (int e = 0; e < n_expert_; e++) t[(size_t)e * H + k] = moe_router_[(size_t)k * n_expert_ + e]; + moe_router_.swap(t); + } + } + if (!load_bf16_named(mn.router_b, moe_router_b_, n_expert_)) { + moe_router_b_.clear(); + // hy_v3: the expert-selection bias is `moe_expert_bias`, F32 in FastFlowLM's pack (bf16 from our packer) + const TensorDesc* bd = td_named("moe_expert_bias"); + const void* bp = bd ? model_tensor_data(mw_, const_cast(bd)) : nullptr; + if (bp && !strcmp(bd->dtype, "F32") && (size_t)bd->data_size >= sizeof(float) * (size_t)n_expert_) { + const float* fp = (const float*)bp; + moe_router_b_.assign(fp, fp + n_expert_); + } else if (bp && !strcmp(bd->dtype, "BF16") && (size_t)bd->data_size >= 2 * (size_t)n_expert_) { + const uint16_t* hp = (const uint16_t*)bp; + moe_router_b_.resize((size_t)n_expert_); + for (int e = 0; e < n_expert_; e++) moe_router_b_[(size_t)e] = bf16_to_f32(hp[e]); + } + } auto pool = [&](const char* suffix, const uint8_t** p, size_t* bytes) -> bool { TensorDesc* d = td_named(suffix); if (!d) { fprintf(stderr, " [moe] missing %s\n", nm2); return false; } @@ -4259,7 +4383,8 @@ bool Q4nxNpuForward::step_batch_impl(const std::vector& tokens, int start_p const bool sandwich = gemma4_layer_ && ple_dim_ == 0 && sem_.extra_layer_norms == 2; // gpt_oss runs its MoE, sinks and o bias in step() only, so a batched prefill takes the per-token path const bool simple_dense = - !mla_ && ple_dim_ == 0 && (!gemma4_layer_ || sandwich) && kv_self && !sem_.gpt_oss; + !mla_ && ple_dim_ == 0 && (!gemma4_layer_ || sandwich) && kv_self && !sem_.gpt_oss && + !(n_expert_ > 0 && leading_dense_ > 0); // a leading dense layer of a non-MLA MoE runs in step() only if (!simple_dense || getenv("NPU_NO_BATCH")) { if (inject_embeds_) { err_ = "step_batch_embeds: this model's batched body falls back to per-token step(), which needs " @@ -4655,7 +4780,7 @@ bool Q4nxNpuForward::step_batch_impl(const std::vector& tokens, int start_p for (int m = 0; m < M; m++) rmsnorm(X.data() + (size_t)m * H, (sandwich ? pre_ff_norm_ : post_norm_)[l].data(), H, XN.data() + (size_t)m * H); - if (n_expert_ > 0) { + if (n_expert_ > 0 && l >= leading_dense_) { // a leading dense layer (hy_v3's first_k_dense_replace) runs the dense FFN // MoE FFN: the router/top-k pick different experts per row, so this // stays per-token (the attention GEMMs above are already batched). // Grouping rows by expert so each expert's weight read serves the @@ -4984,6 +5109,7 @@ bool Q4nxNpuForward::mla_expert(const uint8_t* pool, size_t per_expert, int e, // 0.00855 on [K][N]. int rows = K, cols = N; moe_slice_geometry(moe_format_, N, K, &rows, &cols); + if (moe_slice_nk_) { rows = N; cols = K; } // the container's own grid says [N][K] (init_impl) if (dequant_any(pool + (size_t)e * per_expert, per_expert, rows, cols, out.data()) != 0) return false; return true; @@ -5844,7 +5970,14 @@ bool Q4nxNpuForward::step_impl(int token, int pos, std::vector& logits, i if (!nancheck("moe")) return false; continue; } - if (n_expert_ > 0) { + if (n_expert_ > 0 && !mla_ && l < leading_dense_) { + if (!dense_mlp(l, xn, d)) return false; + { LhScope _tlh(LH_RES); + for (int k = 0; k < H; k++) x[k] += d[k] * residual_scale_; } + if (!nancheck("dense")) return false; + continue; + } + if (n_expert_ > 0 && l >= leading_dense_) { // a leading dense layer (hy_v3's first_k_dense_replace) runs the dense FFN // Qwen MoE layer (qwen3moe / qwen35moe): the FFN is router -> top-k // experts, not a dense gate/up/down. moe_mlp writes the whole FFN // delta into d. (deepseek2/GLM MoE layers never reach here -- the MLA diff --git a/tests/batch_check_state_ci.py b/tests/batch_check_state_ci.py index eb6f7ff0f..1edbd2146 100755 --- a/tests/batch_check_state_ci.py +++ b/tests/batch_check_state_ci.py @@ -15,7 +15,7 @@ # limitations under the License. """The batch-check state pin (tests/batch_check_state_test.cpp, engine#1905/#1946) on a container made here. -usage: tests/batch_check_state_ci.py [lfm2|lfm3|qwen2|qwen2vl|qwen3_5] (needs numpy) +usage: tests/batch_check_state_ci.py [lfm2|lfm3|qwen2|qwen2vl|qwen3_5|hy_v3] (needs numpy) batch_check_state_test needs a Q4NX container, and CI has none, so on its own it skips (77) there and a regression of #1909 -- the NPU_BATCH_CHECK oracle resetting the prefix it replays from -- lands green. This writes a two-layer qwen3 @@ -52,6 +52,14 @@ {"type": "mrope", "mrope_section": [16, 24, 24]}: the forward refused that at load, and on a text-only prompt (all three position planes equal) it has to give exactly the plain-RoPE reference. +With `hy_v3` it is a two-layer Hy-Memory-shaped MoE (engine#1786), written the way FastFlowLM's Hy-Memory-30B-A3B-NPU2 +is: `model.layer.N.` names, a dense layer 0 (`first_k_dense_replace`), a sigmoid router stored [H][E] with an F32 +`moe_expert_bias`, `router_scaling_factor`, one ungated shared expert, and expert stacks `*_exps_proj` whose +per-expert slices are on an [N][K] grid declared 3-D ([E*N/32][K/256][5120]). The expert width is +`expert_hidden_dim` (512) while `moe_intermediate_size` says 1024, as the real pack does (its experts are [768][2048] +upstream with moe_intermediate_size 1536). Held to hyv3_reference() on the container's own weights; the forward +answered " of of of" on the real pack while any one of these was wrong, and batched == sequential cannot see it. + With `qwen3_5` it is a qwen3_5 hybrid (engine#1889): a linear-attention (GDN) layer, then a full-attention layer with q/k norms, written under `language_model.*` so the packer takes its qwen3_5 path (renames, the conv1d transpose, F32 `ssm_a` / `ssm_dt.bias`). The same rule, scoped to "not a conv layer", still refused every qwen3_5 hybrid: a GDN @@ -348,6 +356,133 @@ def rope(x): return rms(x, W("model.norm.weight"))[-1] @ W("lm_head.weight").T +HY = dict(H=256, NH=2, NKV=1, HD=128, E=4, TOPK=2, EIM=512, DIM=768, SCALE=2.5) + + +def make_hyv3(d, np): + """A two-layer hy_v3 in FastFlowLM's naming; returns (weights for the reference, names to re-declare 3-D).""" + rng = np.random.default_rng(1786) + c = HY + H, E, EIM = c["H"], c["E"], c["EIM"] + w = lambda *sh, sc=0.05: (rng.standard_normal(sh) * sc).astype(np.float32) + nrm = lambda n: (1.0 + 0.1 * rng.standard_normal(n)).astype(np.float32) + W = {"model.embed_tokens.weight": w(VOCAB, H, sc=0.5), "model.norm.weight": nrm(H), "lm_head.weight": w(VOCAB, H, sc=0.2)} + for l in range(2): + p = f"model.layer.{l}." + W.update({p + "input_layernorm.weight": nrm(H), p + "post_attention_layernorm.weight": nrm(H), + p + "self_attn.q_proj.weight": w(c["NH"] * c["HD"], H), p + "self_attn.k_proj.weight": w(c["NKV"] * c["HD"], H), + p + "self_attn.v_proj.weight": w(c["NKV"] * c["HD"], H), p + "self_attn.o_proj.weight": w(H, c["NH"] * c["HD"]), + p + "self_attn.q_norm.weight": nrm(c["HD"]), p + "self_attn.k_norm.weight": nrm(c["HD"])}) + if l == 0: + W.update({p + "mlp.gate_proj.weight": w(c["DIM"], H), p + "mlp.up_proj.weight": w(c["DIM"], H), + p + "mlp.down_proj.weight": w(H, c["DIM"])}) + else: + W.update({p + "moe_router.weight": w(H, E, sc=0.5), # [H][E], as the pack + # large enough to change which experts are picked: the bias is a selection rule + p + "moe_expert_bias": np.array([2.0, -2.0, 1.0, -1.0], np.float32)[:E], + p + "mlp.gate_exps_proj.weight": w(E * EIM, H), p + "mlp.up_exps_proj.weight": w(E * EIM, H), + p + "mlp.down_exps_proj.weight": w(E * H, EIM), # stacked [N][K] per expert + p + "mlp.share_gate_exps_proj.weight": w(EIM, H), p + "mlp.share_up_exps_proj.weight": w(EIM, H), + p + "mlp.share_down_exps_proj.weight": w(H, EIM)}) + os.makedirs(d) + write_safetensors(os.path.join(d, "model.safetensors"), W) + cfg = {"architectures": ["HYV3ForCausalLM"], "model_type": "hy_v3", "hidden_size": H, "num_hidden_layers": 2, + "num_attention_heads": c["NH"], "num_key_value_heads": c["NKV"], "head_dim": c["HD"], + "intermediate_size": c["DIM"], "vocab_size": VOCAB, "rms_norm_eps": 1e-5, "max_position_embeddings": 64, + "rope_parameters": {"rope_theta": 1000000.0, "rope_type": "default"}, "tie_word_embeddings": False, + "num_experts": E, "num_experts_per_tok": c["TOPK"], "moe_intermediate_size": 2 * EIM, + "expert_hidden_dim": EIM, "first_k_dense_replace": 1, "moe_router_use_sigmoid": True, + "moe_router_enable_expert_bias": True, "router_scaling_factor": c["SCALE"], "route_norm": True, + "num_shared_experts": 1, "qk_norm": True, "hidden_act": "silu", "torch_dtype": "float32"} + with open(os.path.join(d, "config.json"), "w") as f: + json.dump(cfg, f) + with open(os.path.join(d, "tokenizer.json"), "w") as f: + f.write("{}") + with open(os.path.join(d, "tokenizer_config.json"), "w") as f: + json.dump({"add_bos_token": False}, f) + return W + + +def redeclare_3d(q4nx_dir, names): + """The packer writes [tiles, 5120] row-major over the 32 x 256 grid; FastFlowLM declares an expert stack as + [rows/32][cols/256][5120]. Same bytes, so only the header changes (its length is kept by space padding).""" + path = os.path.join(q4nx_dir, "model.q4nx") + with open(path, "r+b") as f: + n = struct.unpack("hts", q, k) / np.sqrt(HD), -np.inf) + sc = np.exp(sc - sc.max(-1, keepdims=True)) + x = x + np.einsum("hts,shd->thd", sc / sc.sum(-1, keepdims=True), v).reshape(T, NH * HD) @ W(p + "self_attn.o_proj.weight").T + h = rms(x, W(p + "post_attention_layernorm.weight")) + if l == 0: + x = x + (silu(h @ W(p + "mlp.gate_proj.weight").T) * (h @ W(p + "mlp.up_proj.weight").T)) @ W(p + "mlp.down_proj.weight").T + continue + probs = 1.0 / (1.0 + np.exp(-(h @ W(p + "moe_router.weight")))) + bias = W(p + "moe_expert_bias") + G, U, D = W(p + "mlp.gate_exps_proj.weight"), W(p + "mlp.up_exps_proj.weight"), W(p + "mlp.down_exps_proj.weight") + out = np.zeros_like(x) + for t in range(T): + sel = np.argsort(-(probs[t] + bias))[:c["TOPK"]] + wts = probs[t, sel] / probs[t, sel].sum() * c["SCALE"] + for e, wt in zip(sel, wts): + g, u, dn = G[e * EIM:(e + 1) * EIM], U[e * EIM:(e + 1) * EIM], D[e * H:(e + 1) * H] + out[t] += wt * ((silu(h[t] @ g.T) * (h[t] @ u.T)) @ dn.T) + out += (silu(h @ W(p + "mlp.share_gate_exps_proj.weight").T) * (h @ W(p + "mlp.share_up_exps_proj.weight").T)) @ \ + W(p + "mlp.share_down_exps_proj.weight").T + x = x + out + return rms(x, W("model.norm.weight"))[-1] @ W("lm_head.weight").T + + def tmp_root(): """tests/tmp_root.h's rule (engine#2029): scratch goes on disk, not in tmpfs. ONEBIT_TEST_TMP if set, else ~/.cache/scratch-tests/tmp, else the system temp directory.""" @@ -366,7 +501,7 @@ def tmp_root(): def main(): - if len(sys.argv) not in (2, 3) or (len(sys.argv) == 3 and sys.argv[2] not in ("lfm2", "lfm3", "qwen2", "qwen2vl", "qwen3_5")): + if len(sys.argv) not in (2, 3) or (len(sys.argv) == 3 and sys.argv[2] not in ("lfm2", "lfm3", "qwen2", "qwen2vl", "qwen3_5", "hy_v3")): print(__doc__.strip().splitlines()[2]) return 2 family = sys.argv[2] if len(sys.argv) == 3 else "qwen3" @@ -377,14 +512,23 @@ def main(): return 77 with tempfile.TemporaryDirectory(prefix="batch_check_state_ci.", dir=tmp_root()) as tmp: src, out = os.path.join(tmp, "hf"), os.path.join(tmp, "q4nx") - shapes = make_model(src, np, family) + if family == "hy_v3": + Whyv = make_hyv3(src, np) + shapes = {k: v.shape for k, v in Whyv.items()} + else: + shapes = make_model(src, np, family) r = subprocess.run([sys.executable, TOOL, "pack", src, out], capture_output=True, text=True) if r.returncode != 0: print(f"FAIL packing the fabricated container: {r.stderr[-600:]}") return 1 + if family == "hy_v3": + c = HY + redeclare_3d(out, {"model.layer.1.mlp.gate_exps_proj.weight": (c["E"] * c["EIM"] // 32, c["H"] // 256), + "model.layer.1.mlp.up_exps_proj.weight": (c["E"] * c["EIM"] // 32, c["H"] // 256), + "model.layer.1.mlp.down_exps_proj.weight": (c["E"] * c["H"] // 32, c["EIM"] // 256)}) env = dict(os.environ, ONEBIT_TEST_Q4NX=out) logits_out = os.path.join(tmp, "sequential.f32") - if family in ("lfm3", "qwen2", "qwen2vl"): + if family in ("lfm3", "qwen2", "qwen2vl", "hy_v3"): env["ONEBIT_TEST_LOGITS_OUT"] = logits_out env.pop("NPU_INFER_HOST_GEMM", None) # init_host is host-GEMM already; nothing else should steer it r = subprocess.run([sys.argv[1]], env=env, capture_output=True, text=True, timeout=110) @@ -397,9 +541,9 @@ def main(): # arm, not a reason to skip it. print("FAIL batch_check_state_test skipped on the fabricated container (init_host refused it)") return 1 - if family in ("lfm3", "qwen2", "qwen2vl") and r.returncode == 0: + if family in ("lfm3", "qwen2", "qwen2vl", "hy_v3") and r.returncode == 0: ours = np.fromfile(logits_out, dtype=np.float32)[:VOCAB] - ref = (lfm3_reference if family == "lfm3" else qwen2_reference)(out, np, [1, 2, 3, 4, 5, 6], shapes) + ref = {"lfm3": lfm3_reference, "hy_v3": hyv3_reference}.get(family, qwen2_reference)(out, np, [1, 2, 3, 4, 5, 6], shapes) corr = float(np.corrcoef(ours, ref)[0, 1]) ok = (corr >= LFM3_MIN_CORR and float(np.abs(ours - ref).max()) <= LFM3_MAX_DIFF * float(np.abs(ref).max()) and int(np.argmax(ours)) == int(np.argmax(ref)))