[bong-water-water-bong]

Refs #1785 (decode speed; host side of a decode token).

The waste

On a hybrid (GDN + full-attention) model, step() calls load_layer() for every layer on every decode token, because its skip (dense_weights_cached) only applies to dense models. A GDN layer skips its own big dequants when its packed images are resident (proj_cached), and the MLP tail does too (ffn_tail_cached). The full-attention layers did not: q, k, v and o were dequantised into wq_/wk_/wv_/wo_ on every token.

Those four buffers have one reader, q4nx_assemble_B, and both paths already skip it once the layer's packed image is on the device (step(): !weight_cached(0) / !weight_cached(1); step_batch(): !cached). So the reload wrote buffers nothing read.

The change

load_layer() skips the four dequants when the layer's packed QKV (key 0) and O (key 1) images are both resident and the run is not NPU_INFER_HOST_GEMM. Both keys, because step_batch() guards its O assemble with the QKV key: a half-cached layer still loads everything. NPU_ATTN_RELOAD=1 keeps the old per-token reload (the A arm below).

Measured (NPU, one npu-hold ticket, serves ended with it)

FastFlowLM's Qwen3.5-2B-NPU2 (24 layers, 6 full-attention), this branch at 46bd1c2, NPU_PHASE_PROF=1, --dump-logits, 2 requests x 24 greedy tokens per arm, order A B B A (A = NPU_ATTN_RELOAD=1), box load 20 to 27. Medians over 46 decode tokens per arm:

arm	load_layer ms/token	layer loop host share	device wait	token wall
A1 reload (old)	12.96	23.9	241.7	347.4
B1 skip	0.04	10.6	240.8	331.9
B2 skip	0.03	11.1	247.3	341.4
A2 reload (old)	11.63	22.2	240.7	344.6

ra1 vs rb1: 48 of 48 byte-identical
ra1 vs rb2: 48 of 48 byte-identical
ra1 vs ra2: 48 of 48 byte-identical
rb1 vs dev1 (an earlier main binary, device head): 48 of 48

What this shows: the host share of the layer loop halves (about 12 ms per token) and no logit changes. What it does not show cleanly: the token wall. 347/345 against 332/341 ms points the right way, but B2 is within 4 ms of the A arms, so on this box I claim the host share, not a tok/s figure. The device wait (about 241 ms, 85 % of the token) is untouched.

Tests

No unit arm: the condition is "the device holds both packed images", which only exists with the device. The gate is the byte comparison above (old path against new, same binary, same ticket). ctest -R "npu_blocked_layout|npu_rmsnorm" still passes in the NPU build.

Models this reaches: any hybrid container with full-attention layers on the NPU route (Qwen3.5 family, MiniCPM-V 4.7's text side). Dense models never took this path on a cached layer. Only Qwen3.5-2B was measured.

Docs: docs/npu.md, next to the dense-tail cache paragraph.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at 46bd1c2

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1785 - Partially compliant Compliant requirements: The PR addresses the performance optimization for attention layer loading, which is a key part of the F5 flash serving mode's efficiency The PR introduces a mechanism to skip redundant dequantization of attention weights, which aligns with the goal of clean state per request and efficient token processing The PR includes documentation updates to reflect the new behavior and performance improvements Non-compliant requirements: The PR does not introduce the --flash flag or the specific single-turn flash serving mode as described in the ticket The PR does not include the full implementation of the fixed 1k context and clean state per request features The PR does not include the measured performance data against the specified bar (0.49 s TTFT / 33 tok/s) Requires further human verification: Verification that the --flash flag implementation is correctly integrated and functional Verification that the fixed 1k context and clean state per request behavior is correctly implemented Verification that the performance meets the specified bar (0.49 s TTFT / 33 tok/s)

⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Attention Reload Logic The change introduces a new condition to skip attention layer dequantization when both QKV and O images are cached. However, the logic assumes that step_batch() guards its O assemble with the QKV key, which may not always be true in all code paths. This could lead to inconsistent behavior in hybrid models where layers might be partially cached. const bool attn_cached = !attn_reload && !host_gemm_ && xd_ && xd_->weight_cached(l, 0) && xd_->weight_cached(l, 1); Host GEMM Compatibility The code explicitly checks for !host_gemm_ when deciding whether to skip attention reloads. However, there's no clear indication of how this interacts with the NPU_INFER_HOST_GEMM flag or whether this check is sufficient to prevent issues in host GEMM scenarios. This could be a potential source of inconsistency or performance regression. const bool attn_cached = !attn_reload && !host_gemm_ && xd_ && xd_->weight_cached(l, 0) && xd_->weight_cached(l, 1);

[bong-water-water-bong]

CI note for reviewers: build is red at 46bd1c2 for a reason outside the diff. Both attempts (13:58Z, 14:06Z) stop in the apt step with Could not resolve 'archive.ubuntu.com' on all four tries; ci runs on three other branches created 13:56Z to 14:05Z failed in the same window. Nothing was compiled. I will re-run it once the runners resolve again; the other 7 checks pass.

[bong-water-water-bong]

VERDICT: READY 46bd1c2

The change is small and the argument is the right kind — "these buffers have one reader, and that reader is already skipped" — so I checked the reader claim instead of the diff, because there are two assembly sites per path and a missed one would read a buffer nobody loaded.

Reader sites, by enclosing function (npu/forward/src/q4nx_forward.cpp):

line	function	what it reads	its guard
3764	step_batch	wq_/wk_/wv_	cached = weight_cached(0) (:3734)
3964	step_batch	wo_	the same cached flag
5021	step	wq_/wk_/wv_ (and wq_ twice for kv_shared)	!weight_cached(0) (:5026)
5100	step	wo_	!weight_cached(1) (:5100)

The guard the PR needs is exactly the one it wrote, and it closes in both directions: attn_cached requires weight_cached(l,0) && weight_cached(l,1). Every read site is guarded by one of those two keys, so any read implies at least one key is absent ⇒ attn_cached is false ⇒ the dequant ran. Conversely, with both keys resident, no site can read the four buffers. That is stronger than the body claims (it argues from step_batch's O guard alone), and it is why BOTH keys must be cached is not merely conservative but necessary — a QKV-only condition would have left step's O read unguarded while skipping its load.

Two things the change gets right that are easy to get wrong: it excludes host_gemm_ (the host-GEMM route does read these buffers by definition), and it keeps NPU_ATTN_RELOAD=1 as an A/B switch, so a surprising regression can be bisected in one env var without a rebuild.

The measurement is the honest kind: load_layer 13.0/11.6 ms per token with the reload against 0.04/0.03 without on Qwen3.5-2B-NPU2 (6 full-attention layers of 24), the layer loop's host share 23.9/22.2 → 10.6/11.1 ms, device wait unchanged (240.7→247.3 ms across all four arms) — and it states plainly that the whole-token move (347/345 → 332/341 ms) is inside the spread between two arms of one mode, so the host share is the number to read. Reporting a null on the headline number while the mechanism is clean is exactly right; 144 logits dumps byte-identical across the other three arms is what makes it safe.

No objection. The one thing worth a line in the docs if you touch it again: the docstring says the four buffers have one reader, q4nx_assemble_B, but the call sites are four (two per path, one per buffer group), which is what makes the two-key condition necessary — a future reader who takes "one reader" literally could simplify the condition and reintroduce the bug.

[bong-water-water-bong]

VERDICT at 46bd1c2: READY

It merges cleanly onto main.

Code: the skip is safe because of three facts, each checked at this head.

• The four buffers have exactly four readers. wq_/wk_/wv_/wo_ are read only at q4nx_forward.cpp:3895 and :4095 (step_batch, both behind !cached, which is QKV key 0) and at :5154 and :5233 (step, behind !weight_cached(0) and !weight_cached(1)). All four also sit behind !xd_->pb16i(). Requiring both keys covers step_batch's O guard, which uses key 0.
• A resident image cannot disappear. XdnaTarget::weight_cached() is wcache_[layer].find(key) != end(), and nothing in npu/ erases from wcache_. So load_layer()'s check cannot go stale before the layer's own O assemble.
• Host-GEMM runs are excluded (!host_gemm_), so the device-free paths that do read the f32 copies still load them.

Measurements, recomputed from the run directory (~/.cache/scratch-cl-npurt/hh):

• Dumps, by cmp: ra1 vs rb1 48 of 48, ra1 vs rb2 48 of 48, ra1 vs ra2 48 of 48, and rb1 vs dev1 48 of 48.
• load_layer from the [lh] lines, median over the 92 nonzero decode samples per arm: ra1 12.96, rb1 0.04, rb2 0.03, ra2 11.63 ms. This matches the table. The body also claims only the host share, not a tok/s figure, which is the right reading of B2 sitting within 4 ms of the A arms.

Build and tests at the head (ONEBIT_NPU=ON ONEBIT_NPU_HOST=ON): onebit_npu_forward builds, and the object carries NPU_ATTN_RELOAD. ctest -R "npu_blocked_layout|npu_rmsnorm" passes 2 of 2. One unrelated target, npu_rt_shim_test, does not link in my review build (an undefined npu_rt_api::ext::bo constructor against the runtime directory I configured with). This PR does not touch it.

CI: build at 46bd1c2 failed in Tools for checkout and the build at 14:51Z. That is the runner DNS outage, not this tree.

Non-blocking: there is no unit arm, as the body says, because the condition only exists with the device. NPU_ATTN_RELOAD is a measuring switch. Once the default has had a few windows, it could go, or move to the documented env list in docs/npu.md, which this PR already edits.

[bong-water-water-bong]

Closing under the lane PR cap (at most two open per lane). Nothing is dropped: this commit is cherry-picked unchanged onto #2209 as 27b3d92 (same diff as 46bd1c2 here, which held a READY), and #2209's body carries this PR's measurement and limits in full.