Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions docs/npu.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 11 additions & 1 deletion npu/forward/src/q4nx_forward.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Loading