From 12196dac6634ee0c80de65e3289645f3b0930129 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 12:00:20 -0400 Subject: [PATCH] npu: a 17..128-row prefill chunk runs as 16-row launches when that is measured cheaper (fixes #1959) On a pack with both kernel families a chunk wider than 16 rows ran the 128-row design, which walks all of its baked rows whatever the chunk holds: on Qwen2.5-7B-Instruct_rai_1.9.0_npu a 16-row pass is 1073 ms and a 17-row pass was 6992 ms. npu_gemm_rows() now decides per projection, from the two launch costs this process has measured, whether such a chunk is one wide launch or is walked at the twin's width (batch_split.h, pure): split when ceil(M/16) x narrow cost is at least a tenth below the wide cost; keep the wide launch otherwise or while either side is unmeasured. Uncached weights (routed experts) are never split. NPU_PREFILL_SPLIT=0 keeps the old route, =1 forces the split. Measured, same binary, NPU_PREFILL_SPLIT=0 against the default, fresh serve per arm, one ticket (pass wall, ms): 17 rows 6948/6956 -> 2226/1938 24 rows 6973/8124 -> 1968/1955 48 rows 7041/7084 -> 3032/2892 96 rows 7295/7309 -> 5909/5874 112 rows 7323/7437 -> 6885/6872 128 rows 7421/7492 -> 7430/7343 16 rows unchanged (1077/1075 vs 1137/1060) Greedy text identical for all 14 requests between the arms; NPU_BATCH_CHECK=1 on a 17-id npu-run reports rows M=17 identical=17 differing=0 with the split forced and with it off. tests/npu_pack_test.cpp: 11 arms on the decision (crossover at the measured costs, unmeasured side, one family, a family where the wide launch wins). Co-Authored-By: Claude Fable 5.2 --- docs/npu.md | 40 +++++++++++++++++++ npu/forward/include/batch_split.h | 50 +++++++++++++++++++++++ npu/forward/include/q4nx_forward.h | 8 ++++ npu/forward/src/q4nx_forward.cpp | 64 +++++++++++++++++++++++++++++- tests/npu_pack_test.cpp | 29 ++++++++++++++ 5 files changed, 190 insertions(+), 1 deletion(-) create mode 100644 npu/forward/include/batch_split.h diff --git a/docs/npu.md b/docs/npu.md index 92d45cc7e..0f80a2c24 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -512,6 +512,46 @@ An **M≥256 t2s family** is what removes the 8x launch overhead; it is private- family passes, and the acceptance row is: both Llama3-8B prefill shapes (`K=4096 N=14336`, `K=14336 N=4096`) at M = 256 / 1024 / 16384, ≥ 4 TFLOP/s, byte-exact logits against the current path on a fixed prompt set, `amdxdna` counter 0. +#### A pack with both families: a 17-row chunk no longer pays for 128 rows (#1959) + +An imported int8 pack (the `b8` layout) carries a 16-row twin **and** a 128-row prefill design for each projection, and +`max_batch_rows()` is 128 there. A chunk of up to 16 rows runs the twin; anything wider ran the wide design, which walks +all 128 of its baked rows whatever the chunk holds. So one more prompt token cost 5.9 s: on +`Qwen2.5-7B-Instruct_rai_1.9.0_npu`, a 16-row pass is 1073 ms and a 17-row pass was 6992 ms. + +`npu_gemm_rows()` now decides per projection, from the launch costs this process has measured, whether a chunk of +17..128 rows runs as one wide launch or is walked at the twin's width (`npu/forward/include/batch_split.h`): it splits +when `ceil(M/16) x narrow cost` is at least a tenth below the wide cost, and keeps the wide launch otherwise or while +either side is unmeasured (the first wide launch of a projection is what measures it). The rule is measured, not a +constant, because the ratio belongs to the kernel family: here a wide launch costs about six narrow ones, and on a +family where it costs two the wide launch keeps every chunk above 32 rows. + +Same binary, `NPU_PREFILL_SPLIT=0` (the wide launch kept) against the default, fresh serve per arm, one ticket, +two prompts per length; the pass wall is the serve's own `[bs] M= start=0` line: + +| batched rows | wide launch kept | default (measured split) | launches per pass | +|---:|---:|---:|---| +| 16 | 1077 / 1075 ms | 1137 / 1060 ms | 112 / 112 | +| 17 | 6948 / 6956 ms | **2226 / 1938 ms** | 112 / 220 | +| 24 | 6973 / 8124 ms | **1968 / 1955 ms** | 112 / 224 | +| 48 | 7041 / 7084 ms | **3032 / 2892 ms** | 112 / 336 | +| 96 | 7295 / 7309 ms | **5909 / 5874 ms** | 112 / 392 | +| 112 | 7323 / 7437 ms | 6885 / 6872 ms | 112 / 328 | +| 128 | 7421 / 7492 ms | 7430 / 7343 ms | 112 / 308 | + +A request of 18 prompt ids and 3 output tokens went from 10.0 s to 5.0-5.3 s. The first 17-row pass of a serve is the +slower of its pair (2226 ms) because it contains the one wide launch per projection that takes the measurement. The +launch counts are not multiples of one factor because the decision is per projection: at 112 and 128 rows only the +projections whose wide launch costs more than eight or nine narrow ones are split. + +Nothing in the output moves. Every row of a batched GEMM is computed independently of the others, which is the +property the block walk for chunks wider than the design already relies on: the greedy text of all 14 requests is +identical between the two arms, and the per-row bitwise oracle (`NPU_BATCH_CHECK=1` on a 17-id `npu-run`) reports +`rows M=17 identical=17 differing=0` with the split forced (`NPU_PREFILL_SPLIT=1`) and with it off. A pack with one +family only (every t2s/pb16i container: `max_batch_rows()` is 16 there) never reaches the rule, and weights that are +re-packed per launch (routed MoE experts) are never split. The serve says once what it did: +`[batch] a 17-row chunk runs as 2 launch(es) of 16 rows instead of one 128-row launch: measured ...`. + #### Per-tile arithmetic ceiling, and where our GEMMs sit against it (AM027) AM027 (*Versal Adaptive SoC AIE-ML v2 Architecture Manual*, v1.1) **Appendix A, Table 14** gives 512 INT8×INT8 multipliers per compute tile and 256 diff --git a/npu/forward/include/batch_split.h b/npu/forward/include/batch_split.h new file mode 100644 index 000000000..80461cf86 --- /dev/null +++ b/npu/forward/include/batch_split.h @@ -0,0 +1,50 @@ +// 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. +// +// batch_split.h -- whether a prefill chunk runs as one wide launch or as several narrow ones. +// +// A pack that carries both kernel families has a narrow design (16 rows) and a wide one (128 rows) for the same +// projection. The wide design walks all of its baked rows whatever the chunk holds, so a 17-row chunk pays for 128: +// measured on an imported 7B pack, one 16-row pass is 1073 ms and one 17-row pass is 6992 ms. Two narrow launches +// would have cost a third of that. Which side wins depends on the family (on another family the wide launch is only +// twice a narrow one), so the choice is made from the two MEASURED launch costs of this process and not from a +// constant. +// +// Pure and header-only, like batch_width.h: the decision is the part that can be wrong, so it is pinned by a host +// test with no device. Row results do not depend on it -- every row of a batched GEMM is computed independently of +// the others, which is the property the blocked walk for chunks wider than the design already relies on. +#pragma once + +namespace onebit { +namespace npu { + +// The block width to walk an M-row chunk at, or 0 to keep it as one launch of the wide design. +// +// narrow, wide rows the two designs take (16 and 128); no split unless 2 <= narrow < wide +// narrow_ms, wide_ms measured wall of one launch of each for this projection; 0 or less means "not measured +// yet", and an unmeasured side keeps the one wide launch (which is how it gets measured) +// +// The split must save at least a tenth of the wide launch, so two costs that are level do not make the route depend +// on which one was measured in a noisier moment. +inline int batch_split_rows(int M, int narrow, int wide, double narrow_ms, double wide_ms) { + if (narrow < 2 || wide <= narrow) return 0; // one family only + if (M <= narrow || M > wide) return 0; // already a narrow launch, or the caller blocks it first + if (!(narrow_ms > 0.0) || !(wide_ms > 0.0)) return 0; + const int launches = (M + narrow - 1) / narrow; + return double(launches) * narrow_ms < 0.9 * wide_ms ? narrow : 0; +} + +} // namespace npu +} // namespace onebit diff --git a/npu/forward/include/q4nx_forward.h b/npu/forward/include/q4nx_forward.h index c64be7992..56e4b7117 100644 --- a/npu/forward/include/q4nx_forward.h +++ b/npu/forward/include/q4nx_forward.h @@ -385,6 +385,14 @@ class Q4nxNpuForward { // prompt-encoding forward (step_batch at start_pos 0) attends bidirectionally there and the GDN // layers keep their recurrence. Off unless the plan says so, so a causal model is untouched. bool noncausal_attn_ = false; + // engine#1959: measured wall of one batched launch per design, narrow (the 16-row twin) and wide (the prefill + // design), for cached weights only. npu_gemm_rows() reads them to decide whether a 17..wide-row chunk runs as one + // wide launch or as several narrow ones (batch_split.h). 0 = not measured yet. `split_mode_` is NPU_PREFILL_SPLIT: + // -1 unset (decide from the measurements), 0 never split (the behaviour before), 1 split whenever a twin exists. + double split_narrow_ms_[Q4NX_MAX_DESIGNS] = {0}; + double split_wide_ms_[Q4NX_MAX_DESIGNS] = {0}; + int split_mode_ = -2; // -2 = environment not read yet + bool split_said_ = false; // The refusal for a kernel a design needs but the directory does not have. With a plan it names // the plan file and the role, so "which description asked for this" needs no log reading. std::string missing_kernel(const std::string& dir, const char* elf_name, Q4nxDesign role) const; diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index a942bdf70..4eb90d288 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -24,6 +24,7 @@ // Numerics follow tools/q4nx_forward_ref.py so the two can be compared // element-wise. #include "q4nx_forward.h" +#include "batch_split.h" // one wide launch or several narrow ones, decided from measured costs #include "q4nx_state_planes.h" // the plane-copy half of a snapshot, device-free and unit-tested #include "model/op_dispatch.h" // the optional-op fallback seam (engine#803 A3) #include @@ -3228,6 +3229,52 @@ void Q4nxNpuForward::npu_gemm_rows(int role, const float* A, int M, int K, int N } return; } + // engine#1959: a chunk wider than the 16-row twin runs the wide (prefill) design, and that design walks all of + // its baked rows whatever the chunk holds -- a 17-row chunk pays for 128. Measured on an imported 7B pack: one + // 16-row pass 1073 ms, one 17-row pass 6992 ms. So when this projection has both designs and several narrow + // launches are measured cheaper than the one wide launch, the chunk is walked at the twin's width instead. Rows are + // independent, so the results are the same rows; only the launches differ. Which side wins is a property of the + // kernel family, so it is decided from this process's own two launch costs (batch_split.h), and an unmeasured side + // keeps the wide launch -- which is also how the wide cost gets measured. Uncached weights (wkey < 0, routed + // experts) are never split: each launch re-packs them, and that cost would be multiplied. + const bool in_range = role >= 0 && role < Q4NX_MAX_DESIGNS; + const int narrow = (xd_ && in_range && !decode_twin) ? xd_->rows_for(role, 1) : 0; + const int wide = (xd_ && in_range && narrow > 0) ? xd_->rows_for(role, narrow + 1) : 0; + if (split_mode_ == -2) { + const char* e = getenv("NPU_PREFILL_SPLIT"); + split_mode_ = !e ? -1 : (atoi(e) != 0 ? 1 : 0); + } + if (in_range && wkey >= 0 && split_mode_ != 0 && narrow >= 2 && wide > narrow && M > narrow && M <= wide) { + // Forced mode stands in two costs that always split, so the route can be A/B'd and checked bit for bit. + const double nms = split_mode_ == 1 ? 1.0 : split_narrow_ms_[role]; + const double wms = split_mode_ == 1 ? 1e9 : split_wide_ms_[role]; + const int sb = onebit::npu::batch_split_rows(M, narrow, wide, nms, wms); + if (sb > 0) { + if (!split_said_) { + split_said_ = true; + if (split_mode_ == 1) + std::fprintf(stderr, " [batch] a %d-row chunk runs as %d launch(es) of %d rows instead of one %d-row " + "launch (NPU_PREFILL_SPLIT=1)\n", M, (M + sb - 1) / sb, sb, wide); + else + std::fprintf(stderr, " [batch] a %d-row chunk runs as %d launch(es) of %d rows instead of one %d-row " + "launch: measured %.2f ms per %d-row launch against %.2f ms per %d-row launch " + "(%s; NPU_PREFILL_SPLIT=0 keeps the wide launch)\n", + M, (M + sb - 1) / sb, sb, wide, nms, sb, wms, wide, sc ? sc : "?"); + } + out.assign((size_t)M * N, 0.0f); + for (int b0 = 0; b0 < M && err_.empty(); b0 += sb) { + const int rows = std::min(sb, M - b0); + std::vector part; + npu_gemm_rows(role, A + (size_t)b0 * a_k, rows, K, N, Bmat, part, wkey, sc, a_k, decode_twin); + if (part.size() < (size_t)rows * N) { + if (err_.empty()) err_ = std::string("npu_gemm_rows: a split ") + (sc ? sc : "?") + " launch returned short"; + return; + } + std::copy(part.begin(), part.begin() + (size_t)rows * N, out.begin() + (size_t)b0 * N); + } + return; + } + } npu::target::GemvArgs a; a.role = role; a.layer = cur_layer_; a.M = M; a.K = K; a.N = N; a.x = A; a.Bmat = Bmat; a.wkey = wkey; a.tag = sc; @@ -3235,7 +3282,22 @@ void Q4nxNpuForward::npu_gemm_rows(int role, const float* A, int M, int K, int N a.rows = true; a.decode_twin = decode_twin; a.tiles = &Q4nxNpuForward::gemv_tile_cb; a.tiles_self = this; std::string err; - if (!xd_ || !xd_->gemv(a, out, err)) err_ = err.empty() ? "npu_gemm_rows: no target" : err; + // Time the launch only when its weights were already resident: a first touch packs or reads the weight image and + // would be recorded as the design's cost. + const bool timed = in_range && wkey >= 0 && narrow >= 2 && wide > narrow && weight_cached(wkey); + const auto t0 = timed ? std::chrono::steady_clock::now() : std::chrono::steady_clock::time_point{}; + if (!xd_ || !xd_->gemv(a, out, err)) { err_ = err.empty() ? "npu_gemm_rows: no target" : err; return; } + if (timed) { + const double ms = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + if (M <= narrow) { + // a running mean that follows the box's load: narrow launches happen on every decode step + split_narrow_ms_[role] = split_narrow_ms_[role] > 0.0 ? 0.75 * split_narrow_ms_[role] + 0.25 * ms : ms; + } else if (!(split_wide_ms_[role] > 0.0) || ms < split_wide_ms_[role]) { + // the BEST wide launch seen: once a chunk is split the wide design is not launched again, so a wide cost + // taken in a noisy moment must not be what keeps it out + split_wide_ms_[role] = ms; + } + } } // The folded container's sign table: magic "HAD1", u32 count, count x (u32 width, u32 offset), diff --git a/tests/npu_pack_test.cpp b/tests/npu_pack_test.cpp index e8f354a95..d50ce5760 100644 --- a/tests/npu_pack_test.cpp +++ b/tests/npu_pack_test.cpp @@ -23,6 +23,7 @@ #include "full_elf.h" #include "pack.h" #include "q4nx_pack.h" +#include "batch_split.h" #include "batch_width.h" #include "golden/host_gemm_i8_fixture.h" @@ -112,8 +113,36 @@ static void batch_width_arms() { std::printf(" batch width: 7 arms (min not first, empty slots, fallback)\n"); } +/* engine#1959: a 17-row chunk ran the 128-row design and paid for 128 rows (6992 ms against 1073 ms for 16 rows on an + * imported 7B pack). batch_split_rows() decides from the two MEASURED launch costs whether a chunk is walked at the + * twin's width. These arms pin the decision, including the direction that must not change anything: one family only, + * an unmeasured side, and a family where the wide launch is the cheaper one. */ +static void batch_split_arms() { + using onebit::npu::batch_split_rows; + const double n7b = 7.6, w7b = 56.3; // per-launch device wait measured on the 7B b8 family (ms) + check(batch_split_rows(17, 16, 128, n7b, w7b) == 16, "batch split: 17 rows run as 2 narrow launches (15.2 ms against 56.3)"); + check(batch_split_rows(96, 16, 128, n7b, w7b) == 16, "batch split: 96 rows run as 6 narrow launches (45.6 ms against 56.3)"); + check(batch_split_rows(97, 16, 128, n7b, w7b) == 0, + "batch split: 97 rows stay one wide launch (7 x 7.6 = 53.2 is not a tenth below 56.3)"); + check(batch_split_rows(128, 16, 128, n7b, w7b) == 0, "batch split: a full 128-row chunk stays one wide launch (60.8 > 56.3)"); + check(batch_split_rows(16, 16, 128, n7b, w7b) == 0 && batch_split_rows(1, 16, 128, n7b, w7b) == 0, + "batch split: a chunk the twin already takes is not split"); + check(batch_split_rows(129, 16, 128, n7b, w7b) == 0, "batch split: a chunk wider than the wide design is the block walk's, not this rule's"); + check(batch_split_rows(17, 16, 128, 0.0, w7b) == 0 && batch_split_rows(17, 16, 128, n7b, 0.0) == 0, + "batch split: an unmeasured side keeps the wide launch (that launch is how it gets measured)"); + check(batch_split_rows(17, 16, 16, n7b, w7b) == 0 && batch_split_rows(17, 128, 128, n7b, w7b) == 0 && + batch_split_rows(17, 0, 128, n7b, w7b) == 0 && batch_split_rows(17, 1, 128, n7b, w7b) == 0, + "batch split: one family only (no twin, or twin == wide) never splits"); + // a family whose wide launch costs twice a narrow one: the wide launch wins from 33 rows, and 17..32 is level + check(batch_split_rows(17, 16, 128, 0.59, 1.20) == 0, "batch split: 2 x 0.59 = 1.18 is not a tenth below 1.20, so 17 rows stay wide there"); + check(batch_split_rows(40, 16, 128, 0.59, 1.20) == 0, "batch split: 40 rows stay wide on a family where wide = 2 x narrow"); + check(batch_split_rows(17, 16, 128, 0.50, 1.20) == 16, "batch split: 2 x 0.50 = 1.00 is a tenth below 1.20, so it splits"); + std::printf(" batch split: 11 arms (crossover, unmeasured, one family, a family where wide wins)\n"); +} + int main(int argc, char** argv) { batch_width_arms(); + batch_split_arms(); quant_int4_arms(); for (int G : {8, 9, 16}) { for (int n : {G, 3 * G}) { // whole groups: a ragged tail is not a permutation