Skip to content

npu plan P7: int8 KV with per-head Hadamard rotation, and NPU prefill in 1024/2048-token chunks with block-fp expansion #806

Description

@bong-water-water-bong

plan kv.dtype int8 + per-head rotation (KLD vs f16 KV <= 1e-3); prefill chunks on the NPU with block-fp expansion of t2 weights (public class: 110-150 tok/s on the Max+ 395). Lane: pi-attn for the design note, then the xdna/hexagon lanes. Depends on P4.

Part of #799 (the NPU plan layer). Gate and lane per the plan table in #799. Not a c20 blocker (bucket C: sequenced after the release freeze unless marked additive).

Activity

  1. bong-water-water-bong commented on Nov 2, 2026

    @bong-water-water-bong
    CollaboratorAuthor

    P7 design note delivered (design only, per #799's sequencing — nothing started, no device time): ~/briefs/DESIGN-P7-int8kv-hadamard-npu-prefill.md (285 lines, both halves, written against P4's NpuTarget/NpuModel split).

    Two things in it the plan should act on, before P4 lands.

    1. A conflict with a standing rule, and a proposed split

    The plan parks M>16 (decode geometry M=16 rows). Part B of P7 cannot exist under that rule — prefill in 1024/2048-token chunks is M≫16. Proposal:

    keep M>16 parked for decode (gemv, attn_step, gdn_step), and gate a distinct prefill GEMM path separately — its own kernel, its own byte-exact oracle, its own permission. Nothing in P7b asks for M>16 in the decode path.

    That decision is the prerequisite for every number in Part B, so it is the one thing I need an answer on.

    2. The model numbers in my earlier note were wrong, so here are the parsed ones

    I had claimed Bonsai-2-27B is "all layers full attention" and d=128. Both were wrong. Parsed from the GGUF headers themselves (Ternary-Bonsai-2-27B-PTQ1_0.gguf, Flash-Next shard 1):

    Bonsai-2-27B (qwen35) Flash-Next (qwen4exp)
    blocks / full_attention_interval 64 / 4 48 / 4
    layers holding a KV cache 16 of 64 12 of 48
    head_count_kv 4 2
    key_length = value_length (head_dim) 256 256
    KV per token 64 KiB f16 → 32 KiB int8 24 KiB f16 → 12 KiB int8

    So on Bonsai the int8 cache is 2 GiB → 1 GiB at 32k and 8 → 4 GiB at 128k — on an NPU whose t2 model is 7.2 GB that is the 32k-vs-64k arena difference, and it is where this half of P7 actually pays. (The Flash-Next 24 KiB/token reproduces the 462 MB/step figure already recorded in my state file at 18,782 tokens, so the arithmetic agrees with a measured row.)

    With d=256 the byte math improves slightly: int8 + one f16 scale per (head, token) = 1.008 B/element = 1.98x, and the Hadamard is 8 stages / 2048 adds per head per token.

    What the note specifies

    • Part A (kv.dtype): rotation is free on the K side because (QH)(KH)ᵀ = QKᵀ — rotate Q, store K rotated, no un-rotation. V needs one Hᵀ on the output or stays f16. Four variants costed; variant 4 (int8 K, V f16, 1.33x) goes first because it proves rotation + sidecar + snapshot accounting with the smallest diff, and it does not depend on t2, so it can run beside P3. plan kv {dtype, hadamard} + caps().supports.int8_kv, with a named refusal when the plan asks and the target can't — never a silent f16 fallback. The KV is part of the prefix-cache snapshot payload (app/npu_prefix_cache.h), so the snapshot byte accounting changes with it.
    • Part B (prefill): the win is that expansion amortises — decode expands a weight tile for 16 rows, a 1024-row chunk uses it 64x more, which is why the t2-on-core expansion (P3) makes prefill cheap. Chunk = 1024, tile = 128 rows (1024×5120 bf16 is 10.5 MB per activation tensor, so a chunk must be row-tiled), one layer at a time, KV written by the prefill path in the same int8+rotated contract the decode path reads.
    • Where their 110-150 tok/s goes (arithmetic, labelled as such in the note): one 27B layer pass ≈ 7.2 GB / 47 GB/s ≈ 153 ms streaming bound per 1024-token chunk; the dot work at int8 peak would be ~1.1 s; onwx's measured 6.8-9.3 s is ~45-60x off streaming and ~6-8x off compute. So the gap is engine occupancy per layer, not FLOPs or bandwidth — which is why the note's first measurable is "% of the streaming bound per layer", with tok/s beside it.
    • Gates: npu plan P7: int8 KV with per-head Hadamard rotation, and NPU prefill in 1024/2048-token chunks with block-fp expansion #806's KLD ≤ 1e-3 vs f16 KV (onwx's 0.0003 is the stretch target, labelled theirs), needle ≥ 8k, window11 identical-text, byte-exact oracle for the new kernel before any model run, and prefill tok/s + % of streaming bound per target.
    • One caveat made precise instead of hand-waved: the sparse-attention risk is narrower than it sounds. Flash-Next's indexer carries its own key stream (attention.indexer.key_length = 128, separate from key_length 256), so quantizing the attention K leaves top-k selection untouched; quantizing the indexer K would change it, and only that needs the 2051-cliff/mask-rebuild re-measurement.
    • Risks: outlier heads survive a per-head rotation (fallback: per-32-group scales at a small byte cost), RoPE/Hadamard ordering is silent-wrong-answer, Q pays an FHT per step, the scale sidecar is a separate DMA descriptor per token not just bytes, the 32-token int8 path for <256-token prompts is a third code path, and the 128-vs-64 tile is a one-week measurement on the sealed family (if the largest projection's expanded K-slice plus activations does not fit AIE local memory, amortisation halves).

    Sequencing: P7a (Part A variant 4) after P4; P7b (prefill) after P4 + P3 + the M>16 decision; P7c (int8 V) after P7a is measured.

  2. bong-water-water-bong commented on Nov 2, 2026

    @bong-water-water-bong
    CollaboratorAuthor

    exp-path audit done (on the two AIE2P traps you passed on), written into Part C of the P7 note: ~/briefs/DESIGN-P7-int8kv-hadamard-npu-prefill.md.

    Result: we are not exposed to either trap, and the reason is worth keeping.

    Every exp in our AIE attention/softmax path is our own range-reduced polynomial, and we never call the library's silu/tanh. Evidence (private npu-kernels, npu/dx/kernels/, read at 59ee2fd):

    site what it does
    attn.cc:74 exp_f32 x·log2e → int n + f∈[0,1) → degree-8 Taylor for 2^f (~1e-7) → scaled by 2^n from the exponent bits, explicit −87 cutoff. Not native exp2.
    attn.cc:104 exp32 same range reduction and polynomial, evaluated at bf16 width — the per-block softmax weight path (16/32 lanes).
    attn.cc:169,227 online-softmax rescales and per-score weights use those two.
    silu.cc:29,45,87 (+ gemvs.cc:24, layer.cc:26 which #include "silu.cc") silu(g)·u with its own exp_f32 — no aie::silu, no vtanh, no aie::tanh.
    npu/forward/src/q4nx_forward.cpp:3299-3320 host attention is std::exp(score − max) f32 — the NPU route's attention is still host-side, so no AIE exp there at all today.

    So: no native exp2 anywhere (trap i inapplicable) and no aie::silu/vtanh anywhere (trap ii inapplicable — our silu is hand-written exactly so it cannot be). The mlir-aie aie2p/silu.cc that carries the vtanh path is not in our build.

    What the audit does surface, and it belongs in P7's numbers:

    1. We already pay the "+26 % attention time" class of cost, deliberately. The corrected exp is the expensive one and our kernel is on that side — so when P7 moves attention onto the AIE (attn_step), that cost has to be in the estimate from the start, not discovered later.
    2. exp32 is bf16-limited, not 1e-7. Its polynomial is evaluated in bf16 (bf16 Horner, cast back), so ~0.4 % typical, up to ~0.8 % on the 2^f factor — the same order as the +6 % the trap warns about, but unbiased rounding rather than a one-sided bias. It is a documented, deliberate speed/precision trade; it is the first place to look if a sharp/low-temperature row ever misbehaves, with scalar exp_f32 as the bounded-error fallback.
    3. So the trap to guard is a future optimisation, not a bug we have. The failure mode is someone "speeding up" exp32 by swapping in the native exp2 (+6 % high on every probability, silently) or replacing our silu with aie::silu. Both look like wins; both move logits.

    Added to P7's gates: the plan may not carry an exp: native option (a target that wants it takes a named accuracy waiver, same shape as int8_kv's named refusal); whoever implements attn_step checks softmax weights against a f32 std::exp reference reporting max relative error (not mean), at temperature 1 and at a low temperature; and silu(g)·u is checked across |g| ≤ 1 as well as the tails, because the vtanh trap is only visible in the |x| ≤ 0.5 band where a spot check on large values would miss it.

    I did not edit addons/npu-dx/forge/gen.py's TRAPS list — it is another lane's repo and that list is its own record. Both items are suggested there verbatim in the note for whoever owns it: "native exp2 is piecewise linear, up to +6 % high; the corrected path costs ~26 % of attention time — budget it, do not optimise it away" and "mlir-aie silu_vec's vtanh returns x for |x| <= 0.5; do not use the library silu".

  3. bong-water-water-bong commented on Nov 6, 2026

    @bong-water-water-bong
    CollaboratorAuthor

    Parked: unstarted, and dependent on P4 plus kernel work.

    Audit of origin/main (19486dcc): no kv.dtype/int8 KV/block-fp in npu/plan/plan.h or anywhere in the tree; the design audit for P7 lives in a private brief (~/briefs/DESIGN-P7-int8kv-hadamard-npu-prefill.md), not in this repository.

    Unpark: P4 (open PR #808, waiting on the release/c20 freeze) merges so the NpuModel/target layer exists, then the int8-KV + Hadamard and 1024/2048-token NPU-prefill kernel work can be specified against the plan.

  4. bong-water-water-bong commented on Nov 6, 2026

    @bong-water-water-bong
    CollaboratorAuthor

    Scoped first step landed as PR #1273 — the kv contract, plus the arithmetic it rests on measured in CI

    This issue is a roadmap item, so per #799's rule it gets a scoped first PR rather than a close. #1273 does not close #806; this comment is the list of what remains.

    The gap this issue's own last comment named was that the design audit "lives in a private brief, not in this repository", and that docs/npu-plan.md says — its own words under Not yet — "the plan has no kv.dtype field yet". Both are fixed for the contract half, and the half that can be wrong silently is now gated on every build with no device at all.

    What #1273 lands

    kv {key, value, hadamard {block, signs}} in npu/plan/plan.h, defaulting to f16/no rotation so every plan written before the field keeps loading; int8 = int8 with one fp32 scale per (head, token) taken after a per-head Hadamard rotation. Refusals name the field (unknown dtype, rotation block not 0 or a power of two in [2, 65536], signs that is a path instead of a bare name, signs with no rotation block, non-string dtype). kv_int8() is the question P4's target-capability check will ask. tools/npu_plan.py states the contract it writes; docs/npu-plan.md carries the schema, the refusals, the measurements and a rewritten P7 entry.

    Measured (the two claims part A rests on, no model, no device)

    Attention-weight KLD against the f16 cache — the statistic this issue's gate names — head width 256, one scale per (head, token), 512 keys, 8 draws per cell:

    KV vectors f16 + rotation int8 K, unrotated int8 K + rotation rotation's worth
    uniform (control) 3.9e-15 2.5e-05 2.4e-05 1.0x
    8x outlier channels 8.6e-15 4.5e-04 4.7e-05 9.6x
    30x outlier channels 4.5e-14 8.2e-03 1.5e-04 54.5x
    • The rotation is free on the K side, exactly as designed: (QH)(KH)ᵀ = QKᵀ to 1.6e-09 relative, which is 0.01x one f32 dot product's own rounding error — a rotated f16 cache is indistinguishable from today's, not merely close.
    • The rotation is what keeps int8 inside the 1e-3 gate. At 30× outliers the unrotated cache sits an order of magnitude outside it (8.2e-03) and the rotated one sits inside (1.5e-04).
    • And the control says where it does not matter: at 8× the unrotated cache already passes, so on a model with mild channels the rotation buys accuracy this gate does not ask for. On uniform data it measures 1.0×, which is why that row exists — so nobody can cite this PR as evidence for a win that was not measured.

    What remains (the honest list)

    1. The cache writer/reader in the forward — an int8 K store plus its fp32 scale sidecar, and the prefix-cache snapshot's byte accounting moving with it (app/npu_prefix_cache.h). Needs P4's NpuModel (PR npu: extract the XDNA launch layer behind NpuTarget (P4 of #799, A1) #808) to have a place to live.
    2. The gate itself: KLD ≤ 1e-3 on real KV from a real model on the device. The CI test bounds the mechanism on synthetic vectors generated to the outlier structure; it does not and cannot decide the model. This is the number that closes this issue, and it needs an NPU ticket and the kernel-side scale path (private npu-kernels).
    3. The target-side refusal — caps().supports.int8_kv and a named refusal when a plan asks for int8 on a target that cannot do it, never a silent f16 fallback. Same dependency: P4.
    4. Part B (prefill in 1024/2048-token chunks with block-fp expansion) — unstarted, and correctly so: it needs P3's on-core t2 expansion to amortise against, P4's interface, and the M>16 decision, since prefill is M≫16. Per the design note's proposal that parked rule applies to the decode path only (gemv, attn_step, gdn_step) and a distinct prefill GEMM path gets its own kernel, byte-exact oracle and permission — that decision is the prerequisite for every number in part B.
    5. P7c (int8 V) after P7a is measured, since an int8 V needs an inverse transform on the output that K does not.

    One correction to the design note, from doing the arithmetic in code

    The note proposed variant 4 as the smallest diff and implied the rotation is what makes int8 affordable at all. Measured, that is true at high outlier severity and not true at moderate severity (8× already passes unrotated). So the accurate statement is: the rotation is what keeps int8 inside the gate when a model's channels are severe, and on a mild model it is accuracy margin the gate does not require. That distinction only showed up because the control row was run; it is now in docs/npu-plan.md and in the test.

    Device cost of this whole issue tonight: 0 s GPU, 0 s NPU, no tickets taken — the c31 gate chain had the box, and none of this needed a device.

  5. bong-water-water-bong commented on Nov 7, 2026

    @bong-water-water-bong
    CollaboratorAuthor

    #1320 — item 3 of the remaining list is now enforced, not just documented

    P4 landed (#808, A1) with Caps::int8_kv, and nothing read it: an int8 plan loaded and opened clean, so it would have been answered out of an f16 cache. #1320 closes that before the writer exists.

    The refusal, as a user now reads it:

    plan asks for an int8 KV cache (kv.key="int8", kv.value="f16") but this target has no int8 KV path
    (caps().int8_kv is false): refusing rather than storing f16 in its place
    

    One shared helper (check_kv_caps, npu/target/target.h) called from xdna_target.cpp::open() before the device checks, so P5's Hexagon target inherits the rule instead of re-remembering it. Gated device-free end-to-end from a plan file: ctest -R npu_target_caps, which also asserts the other direction — a plan with no kv block (every container shipped so far) is accepted, since refusing a working model is the worse regression.

    Remaining for this issue is unchanged and still device-bound:

    1. the cache writer/reader in the forward (int8 K store + scale sidecar, with the prefix-cache snapshot's byte accounting);
    2. the KLD ≤ 1e-3 gate on real KV from a real model on the device;
    3. part B (1024/2048-token prefill chunks with block-fp expansion), which still waits on P3's on-core expansion, P4's A2, and the M>16 decision — the decode path keeps M≤16, a prefill GEMM path gets its own permission.

    Device cost of this pass: 0 s GPU, 0 s NPU, no tickets.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions