diff --git a/CMakeLists.txt b/CMakeLists.txt index 807965496..0cba4b4aa 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -856,6 +856,12 @@ add_dependencies(npu_sampling_test npu_rope_llama3_test) add_executable(hadamard_rotate_test tests/hadamard_rotate_test.cpp) target_include_directories(hadamard_rotate_test PRIVATE npu/forward/include) add_test(NAME hadamard_rotate COMMAND hadamard_rotate_test) +# The arithmetic P7 (#806) part A rests on: (QH)(KH)^T = QK^T to f32 rounding, and what the rotation is +# worth against an f16 KV cache (attention-weight KLD) with and without outlier channels. Header only, no +# device, no model: runs in CI. +add_executable(npu_kv_int8_test tests/npu_kv_int8_test.cpp) +target_include_directories(npu_kv_int8_test PRIVATE npu/forward/include) +add_test(NAME npu_kv_int8 COMMAND npu_kv_int8_test) # The A-row quantiser every NPU GEMM runs on the host before its launch (npu/forward/include/quant_i8.h): its # AVX2 path must be byte-identical to the shipped scalar loop, including the .5 boundaries, the NaN/inf inputs # and the values past 2^31 that the (int) cast saturates. Header only, no device, no model: runs in CI. @@ -873,7 +879,7 @@ add_test(NAME attn_math COMMAND attn_math_test) # ci.yml's build step names its targets explicitly and needs the `workflow` scope to extend, so bind the new # binary to a target the list already builds (same reason as npu_moe_test above); drop this line when the # workflow is next touched with that scope and add hadamard_rotate_test to the list itself. -add_dependencies(npu_forward_loader_test hadamard_rotate_test quant_i8_row_test attn_math_test) +add_dependencies(npu_forward_loader_test hadamard_rotate_test quant_i8_row_test attn_math_test npu_kv_int8_test) # The window11 gate's comparison policy (engine#834): no tolerance without a measurement behind it, the # statistic named, and a control pair setting the tolerance from its own spread. Shell + the fixtures it # writes itself, so there is no target for the build list to miss. diff --git a/docs/npu-plan.md b/docs/npu-plan.md index 172dafb62..8388ce183 100644 --- a/docs/npu-plan.md +++ b/docs/npu-plan.md @@ -61,6 +61,8 @@ tools/npu_plan.py [--kernels DIR] [--out PATH] [--check] [--print] | `decode.sealed_elf` | pb16i only: the sealed ELF the decode roles run | | `kernels.tag` | artifact-name suffix: `""`, `_b8`, `_pb16i`, `_t2s` or `_t2s2`; must agree with the container layout and `container.pb16i` | | `attention.mode` | `causal` (default) or `noncausal_full_attention`. The latter is a prompt-encoding mode: every full-attention layer sees the whole prompt, and the forward refuses to generate with it | +| `kv.key`, `kv.value` | what the KV cache is stored as, per side: `f16` (the meaning of an absent `kv` block, and of every plan written before P7) or `int8`. `int8` is int8 with **one fp32 scale per (head, token)**, taken *after* a per-head Hadamard rotation — the rotation is what makes one scale affordable, because it collapses the per-channel magnitude spread the scale would otherwise have to cover. `key` is the side that costs no un-rotation (`(QH)(KH)ᵀ = QKᵀ`), so storing K rotated and rotating Q needs no inverse transform; `value` is a separate choice because an int8 V needs one on the output, which is why K-only is the cheap variant to prove the path with. This is a description: a target without the capability must refuse `int8` **by name**, never serve f16 in its place | +| `kv.hadamard.block`, `kv.hadamard.signs` | the rotation itself: `block` is `0` (no rotation) or a power of two in `[2, 65536]` — the head width the rotation runs over — and `signs` is an optional **bare file name** (no directory part, under the model dir) for per-channel signs. A `signs` file with `block: 0` is refused rather than ignored | | `moe.batched_expert_gu` | `{enabled, K, N}`; when enabled, K and N must be named | | `layers[]` | one entry per layer, in order: `index`, `role` (`attn`/`gdn`/`moe`/`mla`), and projections `qkv {K, N, fused}`, `attn_gate {K, N}`, `gate_up {K, N}`, `down {K, N}`. The count must equal `model.layers`. The projection shapes are recorded for P4’s `NpuModel`; the current forward reads only the count (for its [plan] log line) and derives designs from the container | @@ -81,6 +83,9 @@ The loader never substitutes a default for something it cannot honour; every ref `_t2s`/`_t2s2` require `t2s`); - `decode.dirs` with an empty array, a `root` other than `model`/`kernels`, or a path that is empty, absolute, contains `..` or a backslash; +- `kv.key`/`kv.value` outside `f16`/`int8`; a `kv.hadamard.block` that is not `0` or a power of two in + `[2, 65536]`; a `kv.hadamard.signs` that is a path rather than a bare name, or that names a file while + `kv.hadamard.block` is `0`; - `moe.batched_expert_gu.enabled` without positive K and N; - `layers` missing/out of order, a role outside the four, a projection naming only one of K and N, or a layer count that disagrees with `model.layers`; @@ -90,6 +95,30 @@ The loader never substitutes a default for something it cannot honour; every ref The forward adds one more check on load: `kernels.tag` must equal the tag the container itself implies (`q4nx_forward.cpp:510-539`), so a plan and a container cannot disagree about the artifact family. +## What the int8 KV contract is worth, measured + +The two claims `kv` rests on are arithmetic, not device behaviour, so they are gated in CI with no model and +no device (`ctest -R npu_kv_int8`, `tests/npu_kv_int8_test.cpp`). Attention-weight **KLD against the f16 +cache**, 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** — nothing to concentrate, nothing to claim | +| 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** | + +Two readings matter. The rotation is **free**: rotating both sides and taking the dot product preserves the +score to 1.6e-09 relative, which is *smaller* than one f32 dot product's own rounding error (0.01x of it), so +a rotated f16 cache is indistinguishable from today's. And the rotation is what **keeps int8 inside the +plan's gate**: at 30x outliers an unrotated int8 K cache lands at 8.2e-03, an order of magnitude outside +1e-03, while the rotated one stays at 1.5e-04. At 8x outliers the unrotated cache is already inside the gate +— so on a model whose channels are mild, the rotation buys accuracy this gate does not ask for, and on one +whose channels are not, it is the difference between passing and failing. + +What this does **not** establish is the gate itself: the vectors are synthetic and generated to the outlier +structure this page describes, and #806's KLD is taken on real KV from a real model on the device. The test +bounds the mechanism; the model run decides. + ## How the forward consumes it At init the forward looks for `/plan.json`. If present it loads it and uses the plan for the @@ -141,4 +170,10 @@ layers[0] {"index": 0, "role": "gdn", "qkv": {"K": 5120, "N": 10240, "fused" until that interface lands. - **P5** (#804): Hexagon target (t2 gemv on HVX/HMX, QEMU v73 oracle) has no engine rows yet. - **P6** (#805): the shared drafter protocol is on main; tau=0.6 and Hexagon acceptance rows are still missing. -- **P7** (#806): int8 KV with per-head Hadamard and NPU prefill chunks; the plan has no `kv.dtype` field yet. +- **P7** (#806): the `kv` contract above is in the plan and its arithmetic is gated in CI (`ctest -R npu_kv_int8`, + which needs no device: see below). Still open, in the order it can be taken: the cache **writer/reader** in the + forward (an int8 K store plus its scale sidecar, the prefix-cache snapshot's byte accounting with it), the **KLD + gate on real KV from a model on the device** — the CI test bounds the mechanism on synthetic vectors, not the + model — and **part B**, prefill in 1024/2048-token chunks with block-fp expansion, which waits on P3's on-core + expansion, P4's `NpuModel`, and the `M>16` decision (prefill *is* M≫16; nothing in P7b asks for M>16 in the + decode path). diff --git a/npu/plan/plan.h b/npu/plan/plan.h index 8b8d14a66..2bb9df51c 100644 --- a/npu/plan/plan.h +++ b/npu/plan/plan.h @@ -138,6 +138,21 @@ struct Plan { std::string tile_order = "engine-1"; std::string attention_mode = "causal"; bool noncausal() const { return attention_mode == "noncausal_full_attention"; } + // P7 (#806) part A: what the KV cache is stored as. `f16` is every plan written before this field + // existed, and stays the default for that reason. `int8` stores that side as int8 with one fp32 + // scale per (head, token), taken AFTER a per-head Hadamard rotation -- the rotation is what makes one + // scale affordable, because it collapses the per-channel magnitude spread the quantizer would + // otherwise have to cover (docs/npu-plan.md). `key` is the side that pays no un-rotation, since + // (QH)(KH)^T = QK^T; `value` is a separate choice because an int8 V needs its own inverse transform + // on the output, so K-only is the cheapest variant to prove the path with. + // + // This is a description, not an implementation: a target that cannot honour `int8` must refuse it by + // name. Falling back to f16 would answer with the right text and a different cache contract, which + // is the failure mode the plan layer exists to remove. + std::string kv_key = "f16", kv_value = "f16"; + int kv_hadamard_block = 0; // 0 = no rotation + std::string kv_hadamard_signs; // optional per-channel sign file, bare name under the model dir + bool kv_int8() const { return kv_key == "int8" || kv_value == "int8"; } // Where the decode family is looked for, in order: which given path it hangs off, and the // relative directory under it ("." = that directory itself). struct DecodeDir { bool under_model = false; std::string path; }; @@ -211,6 +226,40 @@ struct Plan { throw PlanError("plan: " + path + " attention.mode \"" + p.attention_mode + "\" is not causal or noncausal_full_attention"); } + // Optional and defaulted on purpose, for the same reason as attention.mode: a plan written + // before P7 existed describes f16 KV, which is what every plan meant then. + if (j.contains("kv")) { + const auto& kv = obj(j, "kv", path); + // Type before value: a number where a dtype belongs is a refusal that names the field, not a + // JSON exception escaping the loader. + for (const char* side : {"key", "value"}) + if (kv.contains(side) && !kv.at(side).is_string()) + throw PlanError("plan: " + path + " kv." + side + " is not a string"); + p.kv_key = kv.value("key", "f16"); + p.kv_value = kv.value("value", "f16"); + if (p.kv_key != "f16" && p.kv_key != "int8") + throw PlanError("plan: " + path + " kv.key \"" + p.kv_key + "\" is not f16 or int8"); + if (p.kv_value != "f16" && p.kv_value != "int8") + throw PlanError("plan: " + path + " kv.value \"" + p.kv_value + "\" is not f16 or int8"); + if (kv.contains("hadamard")) { + const auto& h = obj(kv, "hadamard", path); + p.kv_hadamard_block = positive_or_zero(h, "block", path); + const int b = p.kv_hadamard_block; + if (b != 0 && (b < 2 || b > 65536 || (b & (b - 1)) != 0)) + throw PlanError("plan: " + path + " kv.hadamard.block " + std::to_string(b) + + " is not 0 or a power of two in [2, 65536]"); + p.kv_hadamard_signs = h.value("signs", ""); + // A bare file name, never a path: a plan must not be able to point the loader at a file + // the command line did not name (the same rule decode.dirs follows). + if (p.kv_hadamard_signs.find('/') != std::string::npos || + p.kv_hadamard_signs.find("..") != std::string::npos) + throw PlanError("plan: " + path + " kv.hadamard.signs \"" + p.kv_hadamard_signs + + "\" must be a bare file name (no directory part)"); + if (b == 0 && !p.kv_hadamard_signs.empty()) + throw PlanError("plan: " + path + + " kv.hadamard.signs names a file but kv.hadamard.block is 0 (no rotation)"); + } + } if (j.contains("kernels")) { const auto& k = obj(j, "kernels", path); p.tag = k.value("tag", ""); diff --git a/tests/npu_kv_int8_test.cpp b/tests/npu_kv_int8_test.cpp new file mode 100644 index 000000000..8c4bf70bd --- /dev/null +++ b/tests/npu_kv_int8_test.cpp @@ -0,0 +1,278 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// npu_kv_int8_test: the arithmetic P7 (#806) part A rests on, measured on this host with no model and no +// device. The plan's `kv {key, value, hadamard}` field (npu/plan/plan.h) is a description; this is the +// part that decides whether the description is worth implementing, and it is the part that can be wrong +// silently, so it runs on every build. +// +// Claim 1 -- the rotation is free on the K side. For an orthonormal Hadamard H, (QH)(KH)^T = QK^T, so a +// cache can store K rotated and the scores are the same scores. If that were only approximately true the +// variant would be a different model that answers the same text most of the time, which is the worst +// failure shape there is. Reported as a multiple of one f32 dot product's own rounding error, so the +// number means something instead of being compared to a tolerance somebody picked. +// +// Claim 2 -- what the rotation buys. One int8 scale per (head, token) has to cover that vector's dynamic +// range, and real KV vectors carry a few large channels; rotating spreads them, so the same 8 bits cover +// the same vector with less error. Measured as the attention-weight KL divergence against the f16 cache +// the plan calls the reference (P7's gate is KLD <= 1e-3), on an outlier-channel distribution and -- as +// the control that keeps this honest -- on a uniform one, where a rotation has nothing to concentrate and +// must not be credited with a win it did not earn. +// +// What this does NOT measure: model quality. The vectors are synthetic, generated to the shapes and the +// outlier structure docs/npu-plan.md describes; the KLD that gates #806 is taken on real KV from a real +// model on the device. This bounds the mechanism, not the model. +#include "hadamard.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +int failures = 0; +void check(bool ok, const char* what) { + if (ok) return; + ++failures; + std::printf("FAIL %s\n", what); +} + +// The attention head width both models this phase targets carry (docs/npu-plan.md: key_length = 256). +constexpr int kHead = 256; +constexpr int kKeys = 512; + +// float -> binary16 -> float, round-to-nearest-even. Written out rather than reached for with a compiler +// extension (_Float16 / std::float16_t), because this test is gated in CI on more than one compiler and a +// baseline is only a baseline if its rounding is the defined one. Only the normal range this test uses is +// implemented; the early returns are the cases these magnitudes cannot reach. +float half_roundtrip(float f) { + uint32_t x; + std::memcpy(&x, &f, sizeof x); + const uint32_t sign = x & 0x80000000u; + const int32_t exp = (int32_t)((x >> 23) & 0xFFu); + const uint32_t mant = x & 0x7FFFFFu; + if (exp == 0xFF || exp == 0) return f; + int32_t e = exp - 127; + if (e < -14) return sign ? -0.0f : 0.0f; + if (e > 15) return f; + const uint32_t discarded = mant & 0x1FFFu; // the 13 bits binary16 does not keep + uint32_t keep = mant >> 13; + constexpr uint32_t kHalfway = 0x1000u; + if (discarded > kHalfway || (discarded == kHalfway && (keep & 1u))) keep += 1; + if (keep == 0x400u) { // rounding carried into the exponent + keep = 0; + if (++e > 15) return sign ? -INFINITY : INFINITY; + } + const uint32_t out = sign | ((uint32_t)(e + 127) << 23) | (keep << 13); + float r; + std::memcpy(&r, &out, sizeof r); + return r; +} + +// One quantization cell per row: scale = max|x| / 127, round to nearest, clamp, dequantize. The contract +// variant 4 specifies (K int8, one scale per head per token). +void quant_int8(const float* x, int n, float* out) { + float m = 0; + for (int i = 0; i < n; i++) m = std::max(m, std::fabs(x[i])); + const float s = m > 0 ? m / 127.0f : 1.0f; + for (int i = 0; i < n; i++) { + float q = std::round(x[i] / s); + q = std::max(-127.0f, std::min(127.0f, q)); + out[i] = q * s; + } +} + +// Softmax over keys in double, so the reference is not itself a source of error. +std::vector softmax(const std::vector& s) { + double m = -INFINITY; + for (float v : s) m = std::max(m, v); + std::vector e(s.size()); + double sum = 0; + for (size_t i = 0; i < s.size(); i++) { e[i] = std::exp((double)s[i] - m); sum += e[i]; } + for (double& v : e) v /= sum; + return e; +} + +// KL(p || q): the arm's attention weights against the reference's, over the same keys. +double kl(const std::vector& p, const std::vector& q) { + double d = 0; + for (size_t i = 0; i < p.size(); i++) { + if (p[i] <= 0) continue; + d += p[i] * std::log(p[i] / std::max(q[i], 1e-300)); + } + return d; +} + +// A KV matrix with the structure real KV vectors have: every channel Gaussian, but a handful scaled up -- +// the outlier-channel effect the rotation exists to remove. `gain` of 1.0 is the control: nothing to +// concentrate, so nothing for a rotation to win. +std::vector make_kv(std::mt19937& rng, double outlier_gain) { + std::normal_distribution g(0.0f, 1.0f); + std::vector gain(kHead, 1.0); + if (outlier_gain > 1.0) { + std::uniform_int_distribution chan(0, kHead - 1), n_out(4, 8); + for (int i = 0, n = n_out(rng); i < n; i++) gain[chan(rng)] = outlier_gain; + } + std::vector k((size_t)kKeys * kHead); + for (int r = 0; r < kKeys; r++) + for (int c = 0; c < kHead; c++) k[(size_t)r * kHead + c] = g(rng) * (float)gain[c]; + return k; +} + +struct KlRow { + double mean = 0, max = 0; +}; + +// One cache arm, scored against the f16 unrotated cache that ships today. `rotate` applies the per-head +// Hadamard to Q and to every K row; `int8_k` stores K as int8 with one scale per row. The reference is +// always the f16 cache, unrotated, so every row is "how far is this cache from today's". +KlRow measure(std::mt19937& rng, const std::vector& K, const std::vector& k_f16, bool rotate, + bool int8_k) { + std::vector signs((size_t)kHead, 1.0f); + std::vector k = k_f16; + if (rotate) + for (int r = 0; r < kKeys; r++) + q4nx_hadamard::rotate_scalar(signs.data(), &k[(size_t)r * kHead], kHead, kHead); + if (int8_k) + for (int r = 0; r < kKeys; r++) quant_int8(&k[(size_t)r * kHead], kHead, &k[(size_t)r * kHead]); + (void)K; + + std::vector klds; + std::normal_distribution g(0.0f, 1.0f); + const double inv = 1.0 / std::sqrt((double)kHead); // as attention scales + for (int rep = 0; rep < 8; rep++) { + std::vector q((size_t)kHead); + for (float& v : q) v = g(rng); + std::vector q_arm = q; + if (rotate) q4nx_hadamard::rotate_scalar(signs.data(), q_arm.data(), kHead, kHead); + + std::vector s_ref(kKeys), s_arm(kKeys); + for (int r = 0; r < kKeys; r++) { + double a = 0, b = 0; + for (int c = 0; c < kHead; c++) { + a += (double)q[c] * (double)k_f16[(size_t)r * kHead + c]; + b += (double)q_arm[c] * (double)k[(size_t)r * kHead + c]; + } + s_ref[r] = (float)(a * inv); + s_arm[r] = (float)(b * inv); + } + klds.push_back(kl(softmax(s_ref), softmax(s_arm))); + } + KlRow row; + for (double v : klds) { row.mean += v; row.max = std::max(row.max, v); } + row.mean /= (double)klds.size(); + return row; +} + +// Claim 1: (QH)(KH)^T = QK^T up to f32 rounding, read as a multiple of one f32 dot product's own error. +void claim_rotation_is_free(std::mt19937& rng) { + std::vector signs((size_t)kHead, 1.0f); + std::normal_distribution g(0.0f, 1.0f); + std::vector q((size_t)kHead), k((size_t)kHead); + for (float& v : q) v = g(rng); + for (float& v : k) v = g(rng); + + std::vector qr = q, kr = k; + q4nx_hadamard::rotate_scalar(signs.data(), qr.data(), kHead, kHead); + q4nx_hadamard::rotate_scalar(signs.data(), kr.data(), kHead, kHead); + + double plain = 0, rot = 0; + for (int c = 0; c < kHead; c++) { plain += (double)q[c] * (double)k[c]; rot += (double)qr[c] * (double)kr[c]; } + float plain_f32 = 0; + for (int c = 0; c < kHead; c++) plain_f32 += q[c] * k[c]; + + const double floor = std::max(std::fabs(plain - (double)plain_f32), 1e-12); + const double rel = std::fabs(plain - rot) / std::max(std::fabs(plain), 1e-12); + std::printf("claim 1 scores after the rotation: |QK - (QH)(KH)| = %.3e absolute, %.3e relative -- %.2fx one " + "f32 dot product's own rounding error, i.e. smaller than not rotating at all\n", + std::fabs(plain - rot), rel, std::fabs(plain - rot) / floor); + check(rel < 1e-5, "the rotation preserves the score to f32 rounding"); +} + +} // namespace + +int main() { + std::mt19937 rng(20261106); + claim_rotation_is_free(rng); + + double factor_at[3] = {0, 0, 0}; + int gi = 0; + for (double gain : {1.0, 8.0, 30.0}) { + std::vector K = make_kv(rng, gain); + std::vector k_f16 = K; + for (float& v : k_f16) v = half_roundtrip(v); + + const KlRow f16_rot = measure(rng, K, k_f16, /*rotate=*/true, /*int8_k=*/false); + const KlRow i8_plain = measure(rng, K, k_f16, false, true); + const KlRow i8_rot = measure(rng, K, k_f16, true, true); + + std::printf("\nKV distribution: %s (outlier channel gain %.0fx)\n", + gain > 1.0 ? "outlier channels, the shape real KV has" : "uniform, the control", gain); + std::printf(" f16 cache, rotated KLD mean %.3e max %.3e\n", f16_rot.mean, f16_rot.max); + std::printf(" int8 K, not rotated KLD mean %.3e max %.3e\n", i8_plain.mean, i8_plain.max); + std::printf(" int8 K, rotated (variant 4) KLD mean %.3e max %.3e %.1fx better than unrotated\n", + i8_rot.mean, i8_rot.max, i8_plain.mean / std::max(i8_rot.mean, 1e-300)); + + factor_at[gi++] = i8_plain.mean / std::max(i8_rot.mean, 1e-300); + if (gain > 1.0) { + check(i8_rot.mean < 1e-3, "variant 4 lands inside the plan's 1e-3 gate on this distribution"); + } else { + // The control. Uniform magnitudes have nothing to concentrate, so a rotation claiming a big win + // here would mean this test is measuring something other than what it says it measures. + check(f16_rot.max < 1e-5, "a rotated f16 cache still matches the f16 cache (claim 1, at the weights)"); + } + } + + // What the sweep above is for. Both thresholds are set from the measured columns in this test, not + // picked: the rotation's value has to *scale with outlier severity*, and it has to be worth nothing + // where there are no outliers (otherwise this test is measuring something else). + std::printf("\nrotation's KLD improvement: uniform %.1fx, 8x outliers %.1fx, 30x outliers %.1fx\n", + factor_at[0], factor_at[1], factor_at[2]); + check(factor_at[2] > factor_at[1] && factor_at[1] > factor_at[0], + "the rotation's benefit grows with outlier severity (measured 54.5x / 9.6x / 1.0x)"); + check(factor_at[0] < 1.5, "on uniform data the rotation claims no win (measured 1.0x)"); + + // The decisive one for part A: at 30x outliers an unrotated int8 K cache leaves the plan's gate and a + // rotated one stays inside it. That is the whole reason the rotation is in the contract. + { + std::vector K = make_kv(rng, 30.0); + std::vector k_f16 = K; + for (float& v : k_f16) v = half_roundtrip(v); + const KlRow plain = measure(rng, K, k_f16, false, true); + const KlRow rot = measure(rng, K, k_f16, true, true); + std::printf("30x outliers, same question as the gate: unrotated int8 KLD %.3e %s 1e-3, rotated %.3e %s 1e-3\n", + plain.mean, plain.mean > 1e-3 ? "OUTSIDE" : "inside", rot.mean, rot.mean > 1e-3 ? "OUTSIDE" : "inside"); + check(plain.mean > 1e-3, "at 30x outliers an unrotated int8 K cache is outside the 1e-3 gate"); + check(rot.mean < 1e-3, "at 30x outliers the rotated int8 K cache stays inside it"); + } + + // The byte accounting the plan's variants are quoted from (docs/npu-plan.md). + const double f16_bytes = (double)kHead * 2.0; + const double int8_bytes = (double)kHead * 1.0 + 4.0; // one f32 scale per (head, token) + std::printf("\nbytes per (head, token): f16 K %.0f, int8 K + f32 scale %.0f (%.3fx smaller); " + "with V int8 as well, %.3fx smaller on the whole KV\n", + f16_bytes, int8_bytes, f16_bytes / int8_bytes, (2.0 * f16_bytes) / (2.0 * int8_bytes)); + + if (failures) { + std::printf("\n%d check(s) failed\n", failures); + return 1; + } + std::printf("npu_kv_int8: ok\n"); + return 0; +} diff --git a/tests/npu_plan_test.cpp b/tests/npu_plan_test.cpp index aad885943..d05b0a150 100644 --- a/tests/npu_plan_test.cpp +++ b/tests/npu_plan_test.cpp @@ -292,8 +292,77 @@ void test_refusals() { } // namespace + +// P7 (#806) part A: the KV contract. Absent means f16 with no rotation -- which is what every plan written +// before this field meant -- and anything this build cannot honour is a refusal that names the field, so a +// plan asking for an int8 cache on a target without one cannot be served as f16 by accident. +void test_kv() { + { + const Plan p = Plan::load(write_plan("kv_absent", base_plan())); + check(p.kv_key == "f16" && p.kv_value == "f16" && p.kv_hadamard_block == 0, + "a plan with no kv block means f16 KV with no rotation"); + check(!p.kv_int8(), "and reports itself as not int8"); + } + { + auto j = base_plan(); + j["kv"] = {{"key", "int8"}, {"value", "f16"}, {"hadamard", {{"block", 256}, {"signs", "had.signs"}}}}; + const Plan p = Plan::load(write_plan("kv_variant4", j)); + check(p.kv_key == "int8" && p.kv_value == "f16", "variant 4 (int8 K, f16 V) loads as written"); + check(p.kv_int8(), "and reports itself as int8"); + check(p.kv_hadamard_block == 256 && p.kv_hadamard_signs == "had.signs", + "the per-head rotation block and its sign file are read"); + } + { + auto j = base_plan(); + j["kv"] = {{"key", "f16"}, {"value", "f16"}, {"hadamard", {{"block", 0}}}}; + const Plan p = Plan::load(write_plan("kv_plain", j)); + check(!p.kv_int8() && p.kv_hadamard_block == 0, "an explicit f16/no-rotation block loads"); + } + { + auto j = base_plan(); + j["kv"] = {{"key", "int4"}}; + const std::string why = refusal(write_plan("kv_bad_dtype", j)); + check(has(why, "kv.key") && has(why, "int4"), "an unknown kv.key is refused, naming the field and value"); + } + { + auto j = base_plan(); + j["kv"] = {{"key", 8}}; + const std::string why = refusal(write_plan("kv_key_not_string", j)); + check(has(why, "kv.key") && has(why, "not a string"), + "a non-string kv.key is a named refusal, not a JSON exception escaping the loader"); + } + { + auto j = base_plan(); + j["kv"] = {{"key", "int8"}, {"value", "q8_0"}}; + const std::string why = refusal(write_plan("kv_bad_value", j)); + check(has(why, "kv.value"), "an unknown kv.value is refused, naming the field"); + } + { + auto j = base_plan(); + j["kv"] = {{"key", "int8"}, {"hadamard", {{"block", 250}}}}; + const std::string why = refusal(write_plan("kv_bad_block", j)); + check(has(why, "kv.hadamard.block"), "a rotation block that is not a power of two is refused"); + } + { + auto j = base_plan(); + j["kv"] = {{"key", "int8"}, {"hadamard", {{"block", 256}, {"signs", "../elsewhere.signs"}}}}; + const std::string why = refusal(write_plan("kv_bad_signs", j)); + check(has(why, "kv.hadamard.signs"), + "a sign file that is a path rather than a bare name is refused (a plan cannot point the " + "loader outside the dirs the command line named)"); + } + { + auto j = base_plan(); + j["kv"] = {{"key", "int8"}, {"hadamard", {{"block", 0}, {"signs", "had.signs"}}}}; + const std::string why = refusal(write_plan("kv_signs_no_block", j)); + check(has(why, "kv.hadamard.signs"), + "a sign file with no rotation block is refused as contradictory rather than ignored"); + } +} + int main() { test_valid(); + test_kv(); test_tile_order(); test_refusals(); std::printf("%s\n", failures ? "FAILED" : "npu_plan: ok"); diff --git a/tools/npu_plan.py b/tools/npu_plan.py index 7400a5d69..f5891b41b 100644 --- a/tools/npu_plan.py +++ b/tools/npu_plan.py @@ -294,6 +294,12 @@ def generate(model_dir: str, kernels: str) -> dict: }, "kernels": {"tag": tag}, "attention": {"mode": attention_mode(model_dir)}, + # P7 (#806) part A: the KV contract. f16 with no rotation is what every container shipped so + # far means; the generator states it rather than leaving it to a key's absence, so a plan reads + # as what its cache is. Turning int8 on is a per-target decision (its rotation and scale + # sidecar are kernel-adjacent), so nothing here enables it -- a plan that asks for int8 is + # either written by hand or by a target-aware generator, and the plan loader validates it. + "kv": {"key": "f16", "value": "f16", "hadamard": {"block": 0, "signs": ""}}, "moe": {"batched_expert_gu": moe_batched(cfg, hidden, model_dir, kernels)}, "layers": layers, }