diff --git a/docs/npu.md b/docs/npu.md index baedf4d8e..8256de551 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -2058,6 +2058,19 @@ On FastFlowLM's Qwen3.5-2B-NPU2 (IM 6144, H 2048, a 151 MB tail that the earlier run in one ticket measured decode at 362–367 ms per token with the cache on (both arms) against 485–520 ms with it off in the quieter of the two off arms and 478–707 ms in the other, same text in all twelve runs (engine#1785). +The full-attention layers of such a model had the same reload, for their own four buffers: `load_layer()` +dequantised `q`, `k`, `v` and `o` into `wq_`/`wk_`/`wv_`/`wo_` on every decode token, and the only reader of +those buffers (`q4nx_assemble_B`) is skipped once the layer's packed QKV and O images are on the device. The +reload is now skipped when **both** images are resident (the batched path guards its O assemble with the QKV +key, so a half-cached layer still loads everything) and the run is not `NPU_INFER_HOST_GEMM`. +`NPU_ATTN_RELOAD=1` keeps the old per-token reload. On Qwen3.5-2B-NPU2 (6 full-attention layers of 24), an +A B B A run in one ticket with `NPU_PHASE_PROF=1`, 46 decode tokens per arm, medians: `load_layer` 13.0 and +11.6 ms per token with the reload against 0.04 and 0.03 ms without; the layer loop's host share 23.9 and 22.2 +ms against 10.6 and 11.1 ms; device wait unchanged (240.7 to 247.3 ms in all four). All 144 logits dumps of +the other three arms are byte-identical to the first arm's, so this removes work and changes no value. The +whole token moved from 347 and 345 ms to 332 and 341 ms, which on that box (load 20 to 27) is inside the +spread between two arms of one mode; the host share is the number to read (engine#1785). + ### The prefix cache's snapshot: two kinds of plane `snapshot_state`/`restore_state` (`npu/forward/include/q4nx_forward.h`) hand a prefill's state to the next diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index 9dca11690..e74b47369 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -2041,12 +2041,22 @@ bool Q4nxNpuForward::load_layer(int l) { // copies of the big projections are never needed: skipping them here is the // difference between a first token that dequantises the whole 17.9 GB model and // one that opens 13.4 GB of images. The norms and every small member still load. + // The four attention f32 buffers (wq_/wk_/wv_/wo_) are read only by q4nx_assemble_B(), which step() and + // step_batch() skip once this layer's packed QKV (key 0) and O (key 1) images are on the device -- the + // same argument as ffn_tail_cached() for the MLP and proj_cached for a GDN layer. A dense model never + // gets here on a cached layer (step()'s `dense_weights_cached`), but a hybrid one does on every token: + // its full-attention layers re-dequantised q, k, v and o each decode step, 17.5 ms of a 285 ms token on + // Qwen3.5-2B (6 of 24 layers), for buffers nothing read (engine#1785). BOTH keys must be cached: + // step_batch() guards its O assemble with the QKV key, so a half-cached layer still loads everything. + // NPU_ATTN_RELOAD=1 keeps the old per-token reload, for the A/B. + static const bool attn_reload = getenv("NPU_ATTN_RELOAD") != nullptr; + const bool attn_cached = !attn_reload && !host_gemm_ && xd_ && xd_->weight_cached(l, 0) && xd_->weight_cached(l, 1); if (is_conv_layer(l)) { // LFM2 conv layer: two projections where an attention layer has four. Same f32-until-packed rule. LhScope _t_p(LH_LOAD_PROJ); if (!deq(&lw->shortconv_in_proj_weight, wqkvlin_.data(), 3 * H, H)) return false; if (!deq(&lw->shortconv_out_proj_weight, woutlin_.data(), H, q4nx_pad_k(H))) return false; - } else if (!xd_->pb16i()) { + } else if (!xd_->pb16i() && !attn_cached) { if (!deq(&lw->q_proj_weight, wq_.data(), NH * hd * (attn_gate_ ? 2 : 1), kpad_)) return false; // q|gate rows if (!kv_shared) {