NPU GGUF-on-NPU: converter preference, forward re-sync (MLA/MoE), streaming crash fix

CMakeLists.txt

@@ -388,6 +388,12 @@ if(ONEBIT_NPU)
         npu/forward/src/q4nx_pack.c
         npu/forward/src/q4nx_semantics.c
         npu/forward/src/q4nx_forward.cpp)
+    # NPU_INFER_DUMP_ATTN / _LN / _HEAD write debug dumps to a path taken from the environment;
+    # only debug builds carry them
+    option(ONEBIT_NPU_DEBUG_DUMPS "NPU forward: env-named debug dumps (developer builds)" OFF)
+    if(ONEBIT_NPU_DEBUG_DUMPS)
+        target_compile_definitions(onebit_npu_forward PRIVATE NPU_INFER_DEBUG_DUMPS)
+    endif()
     target_include_directories(onebit_npu_forward SYSTEM PUBLIC
         npu/forward/include npu/forward/engine-src ${ONEBIT_XRT_ROOT}/include)
     target_link_libraries(onebit_npu_forward PUBLIC

app/forward_serve.cpp

@@ -197,27 +197,41 @@ void handle_generate(ForwardEngine& e, const httplib::Request& req, httplib::Res
             });
             res.set_content(wrap(content, "stop"), "application/json");
         } else {
-            res.set_chunked_content_provider("text/event-stream", [&](size_t, httplib::DataSink& sink) {
-                run([&](const std::string& t) {
-                    json chunk;
-                    chunk["id"] = id;
-                    chunk["object"] = "chat.completion.chunk";
-                    chunk["created"] = created;
-                    chunk["model"] = e.id;
-                    if (chat)
-                        chunk["choices"] = json::array({{{"index", 0},
-                                                         {"delta", {{"content", t}}},
-                                                         {"finish_reason", nullptr}}});
-                    else
-                        chunk["choices"] = json::array({{{"index", 0}, {"text", t}, {"finish_reason", nullptr}}});
-                    const std::string d = "data: " + chunk.dump() + "\n\n";
-                    return sink.write(d.data(), d.size());
-                });
-                const std::string done = "data: [DONE]\n\n";
-                sink.write(done.data(), done.size());
-                sink.done();
+            // Build the SSE frames while the handler still owns its locals.  The
+            // content provider below runs AFTER this handler returns, so it must
+            // not capture anything by reference (id/created/object die here) --
+            // that was a dangling-reference crash (`std::system_error`).  Run the
+            // forward here too, so it is not invoked from httplib's write path.
+            json chunk_tpl;
+            chunk_tpl["id"] = id;
+            chunk_tpl["object"] = "chat.completion.chunk";
+            chunk_tpl["created"] = created;
+            chunk_tpl["model"] = e.id;
+            auto chunks = std::make_shared<std::vector<std::string>>();
+            run([&](const std::string& t) {
+                json chunk = chunk_tpl;
+                if (chat)
+                    chunk["choices"] = json::array({{{"index", 0},
+                                                     {"delta", {{"content", t}}},
+                                                     {"finish_reason", nullptr}}});
+                else
+                    chunk["choices"] = json::array({{{"index", 0}, {"text", t}, {"finish_reason", nullptr}}});
+                chunks->push_back("data: " + chunk.dump() + "\n\n");
                 return true;
             });
+            auto cursor = std::make_shared<size_t>(0);
+            res.set_chunked_content_provider("text/event-stream",
+                [chunks, cursor](size_t, httplib::DataSink& sink) {
+                    while (*cursor < chunks->size()) {
+                        const std::string& d = (*chunks)[*cursor];
+                        if (!sink.write(d.data(), d.size())) return false;
+                        ++*cursor;
+                    }
+                    const std::string done = "data: [DONE]\n\n";
+                    sink.write(done.data(), done.size());
+                    sink.done();
+                    return true;
+                });
         }
     } catch (const std::exception& ex) {
         res.status = 500;

npu/forward/engine-src/npu_engine_i8ctx_inc.h

@@ -349,6 +349,16 @@ struct I8Ctx {
     // weight BO per (layer, expert) at startup, pack+sync once, and pass the
     // BO handle directly at decode.
     std::unique_ptr<xrt::bo> make_weight_bo(xrt::device& d) {
+        // The xclbin path picks a group off the kernel object; the ELF path has
+        // no kernel object (`k` is null) and the amdxdna driver rejects a
+        // group-less xrt::bo, so use the same xrt::ext::bo allocation
+        // init_elf's mkbo uses.
+        if (elf_mode) {
+            xrt::ext::bo b{d, (size_t)KD * ND};
+            auto bo = std::make_unique<xrt::bo>(b);
+            if (void* m = bo->map()) memset(m, 0, (size_t)KD * ND);
+            return bo;
+        }
         int grp_w = k->group_id(4);
         // Weight BOs are written once (packB_into) and read every token by the
         // shim DMA. HOST_ONLY forces the device through the slow cache-coherent

npu/forward/include/q4nx_forward.h

@@ -40,10 +40,12 @@
 #include "model.h"
 #include "q4nx_pack.h"        // Q4nxDesignGeom / Q4nxModelDims (derived design table)
 #include "q4nx_semantics.h"   // Q4nxSemantics (config-derived architecture flags)
+#include <memory>
+#include <unordered_map>
 
 struct I8Ctx;
 
-namespace xrt { class device; }
+namespace xrt { class device; class bo; }
 
 struct Q4nxForwardConfig {
     int H = 1024, NH = 16, NKV = 8, HD = 128, IM = 3072, NL = 28, NV = 151936;
@@ -93,7 +95,14 @@ class Q4nxNpuForward {
     bool host_gemm_ = false;
     bool load_layer(int l);
     void npu_gemm(I8Ctx& ctx, const float* A, int K, int N,
-                  const float* Bmat, std::vector<float>& out);
+                  const float* Bmat, std::vector<float>& out, int wkey = -1);
+    // True when this layer's projection `wkey` is already packed in its own BO,
+    // so the caller can skip both the Q4NX dequant (load_layer) and the f32->B
+    // assembly for the rest of the run.
+    bool weight_cached(int key) const {
+        return weight_cache_ && cur_layer_ >= 0 && cur_layer_ < (int)wcache_.size() &&
+               wcache_[(size_t)cur_layer_].find(key) != wcache_[(size_t)cur_layer_].end();
+    }
     bool attention(const std::vector<float>& qkv, int pos, std::vector<float>& ctx,
                    int head_dim, int n_kv_heads, float theta, int rotary_dim,
                    bool proportional, int window, int kv_layer, bool append_kv);
@@ -144,16 +153,46 @@ class Q4nxNpuForward {
     std::vector<std::vector<float>> q_norm_, k_norm_;       // [NL][HD]
     std::vector<std::vector<float>> q_bias_, k_bias_, v_bias_;  // [NL][NH*HD]/[NL][NKV*HD] attention biases
     std::vector<float> rope_scale_;  // LongRoPE short_factor (per-half-dim), empty when none
-    // MiniCPM4 scales (config "embedding_scale" / "residual_scale"); 1.0 elsewhere.
-    // MiniCPM4 scales the embedding by 12.0 and every residual branch by
-    // scale_depth/sqrt(NL) (0.2475 for 32 layers).
+    // MiniCPM4 scales (config "embedding_scale" / "residual_scale"); 1.0 for
+    // every other architecture.  MiniCPM4 scales the embedding by 12 and each
+    // residual branch by scale_depth/sqrt(NL) (0.2475 for 32 layers); skipping
+    // them sends the hidden state to the wrong magnitude and the logits explode.
     float embed_scale_ = 1.0f;
     float residual_scale_ = 1.0f;
-    // MiniCPM4 logit_scale (config "logit_scale", 16.0).  The final logits are
-    // DIVIDED by it; 1.0 = no-op everywhere else.  A constant positive divisor
-    // cannot move an argmax, so this only matters for the probability
-    // distribution (sampling/temperature).
+    // MiniCPM4 logit_scale (config "logit_scale", 16.0 for GLM's sibling
+    // MiniCPM4).  The final logits are DIVIDED by it.  1.0 = no-op everywhere
+    // else.  A constant positive divisor cannot move an argmax, so this only
+    // matters for the probability distribution (sampling/temperature).
     float logit_scale_ = 1.0f;
+
+    // ---- MLA + MoE (deepseek2 / GLM-4.7-Flash) ----
+    // Set from config.json in init_impl; mla_ == 0 leaves every path below
+    // untouched.  Each layer's dequantized weights are held here and refilled by
+    // load_layer(), sized from the model's own dims.
+    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;
+    int leading_dense_ = 0;
+    int expert_gating_sigmoid_ = 1;   // llama.cpp deepseek2: 2 = sigmoid, 1 = softmax
+    float expert_w_scale_ = 1.0f;     // applied after the top-k weight normalisation
+    std::vector<float> w_qa_, w_qb_, w_kva_, w_o_;      // per layer, [out][in]
+    std::vector<float> mla_qan_, mla_kvan_;             // q_a / kv_a layernorms
+    std::vector<float> w_kb_, w_vb_;                    // [NH][N][K], N padded for k_b
+    int kb_rows_ = 0;                                   // k_b's padded row count (256)
+    std::vector<std::vector<float>> mla_kn_, mla_kr_, mla_v_;  // per-layer MLA caches
+    std::vector<float> moe_router_, moe_router_b_, share_g_, share_u_, share_d_;
+    std::vector<float> dense_g_, dense_u_, dense_d_;    // leading dense layer only
+    // Routed experts stay packed until selected: per layer the Q4NX pools them as
+    // [n_expert][tiles][5120].
+    const uint8_t* exp_g_ = nullptr; const uint8_t* exp_u_ = nullptr; const uint8_t* exp_d_ = nullptr;
+    size_t exp_g_bytes_ = 0, exp_u_bytes_ = 0, exp_d_bytes_ = 0;
+
+    bool mla_attn(int l, const std::vector<float>& x, std::vector<float>& o, int pos);
+    bool moe_mlp(int l, const std::vector<float>& xn, std::vector<float>& d);
+    bool dense_mlp(int l, const std::vector<float>& xn, std::vector<float>& d);
+    bool mla_rope(std::vector<float>& q, int NH, int qh, int nope, std::vector<float>& kr, int pos);
+    bool mla_expert(const uint8_t* pool, size_t per_expert, int e, int N, int K,
+                    std::vector<float>& out);
     std::vector<float> final_norm_;                         // [H]
 
     // Per-layer dequantized projections (reused each layer, chunked by design)
@@ -163,6 +202,18 @@ class Q4nxNpuForward {
     std::vector<std::vector<float>> kcache_, vcache_;
 
     std::vector<float> B_;        // packed weight scratch  (max K*N floats)
+    // Packed-int8 weight cache, one weight BO per (layer, wkey), packed on first
+    // use.  Without it every token re-dequantizes the Q4NX (load_layer) AND
+    // re-quantizes it (I8Ctx::packB) for every GEMM -- ~2 min/token on a 7B and
+    // ~8 min/token on the 8B, which makes a served answer impractical.  Only
+    // token-INVARIANT weights may be cached (dense projections, GDN, MLA, the
+    // shared expert); routed-expert weights vary per token and pass wkey < 0.
+    struct PackedWeight {
+        std::unique_ptr<xrt::bo> bo;
+        std::vector<float>       scales;
+    };
+    std::vector<std::unordered_map<int, PackedWeight>> wcache_;   // [layer][wkey]
+    bool weight_cache_ = true;    // NPU_NO_WCACHE=1 disables (A/B)
     int cur_layer_ = 0;           // layer currently being processed
 
     // ── Gemma4 (per-layer embeddings + KV sharing + extra norms) ──

npu/forward/include/q4nx_pack.h

@@ -67,6 +67,17 @@ typedef enum {
     Q4NX_DESIGN_O2,      /* 2nd attention type's O   (Gemma4 full layers)    */
     Q4NX_DESIGN_QKVLIN,  /* GDN linear attention's fused q|k|v: K=H, N=2*KD+VD */
     Q4NX_DESIGN_Z,       /* GDN in_proj_z (the gated-norm gate): K=H, N=VD   */
+    /* MLA (deepseek2, e.g. GLM-4.7-Flash).  These replace QKV; O/G/U/D are still
+     * used, with HD set to the per-head VALUE width so the O design's K = NH*HD =
+     * NH*v_head, and IM set to the per-expert intermediate so G/U/D are the
+     * expert MLP.  The dense leading layer gets its own DG/DU/DD.  Set
+     * dims->mla to derive this set. */
+    Q4NX_DESIGN_QA,      /* MLA q_a_proj:            K=H,       N=q_lora           */
+    Q4NX_DESIGN_QB,      /* MLA q_b_proj:            K=q_lora,  N=NH*(qk_nope+rope) */
+    Q4NX_DESIGN_KVA,     /* MLA kv_a_proj_with_mqa:  K=H,       N=kv_lora+rope     */
+    Q4NX_DESIGN_DG,      /* dense leading layer gate: K=H,       N=dense_im        */
+    Q4NX_DESIGN_DU,      /* dense leading layer up:   K=H,       N=dense_im        */
+    Q4NX_DESIGN_DD,      /* dense leading layer down: K=dense_im, N=H             */
     Q4NX_DESIGN_COUNT
 } Q4nxDesign;
 
@@ -76,6 +87,11 @@ typedef struct {
     char elf_name[Q4NX_ELF_NAME_MAX];  /* geometry-derived artifact name */
     int q_rows, k_rows, v_rows;        /* row counts along N for QKV */
     int gate_rows, up_rows;            /* row counts along N for GU/G/U */
+    /* Real N before the kernel's n-tile pad; 0 means "same as N".  Only KVA
+     * (kv_lora+rope_dim = 576, padded to 640) needs it: assemble_B scatters
+     * real_n columns and zero-fills the rest, so the weight is never read out
+     * of bounds. */
+    int real_n;
 } Q4nxDesignGeom;
 
 /* Architecture dims a design set is derived from. */
@@ -111,6 +127,17 @@ typedef struct {
      * [q(all heads) | gate(all heads)] (converter reorder).  Source:
      * attn_output_gate in the model's config.json. */
     int attn_output_gate;
+    /* MLA + MoE (deepseek2 / GLM-4.7-Flash).  mla == 0 for every other
+     * architecture, leaving the standard derivation untouched.  When set, HD must
+     * be the per-head VALUE width (so the O design's K = NH*HD = NH*v_head) and IM
+     * the per-expert intermediate (so G/U/D are the expert MLP). */
+    int mla;
+    int q_lora;      /* q_a_proj output width                       */
+    int kv_lora;     /* kv_a_proj latent width (excluding rope)      */
+    int rope_dim;    /* rope width inside kv_a and each q head       */
+    int qk_nope;     /* per-head qk width excluding rope             */
+    int dense_im;    /* leading dense layer's intermediate width     */
+    int mla_kva_pad; /* KVA design N pad: ceil((kv_lora+rope_dim)/128)*128 */
 } Q4nxModelDims;
 
 /* No model-specific geometry constant lives here by design: every caller