From 473fdf4e81a13c3c98724f3a561c8522817b2665 Mon Sep 17 00:00:00 2001 From: t Date: Tue, 10 Nov 2026 07:42:25 -0400 Subject: [PATCH] npu: the batch-check oracle compares from the same state it claims to, and names the positions (fixes #1905) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `NPU_BATCH_CHECK` ran a sequential replay over `start_pos .. start_pos+M-1` and then called `reset()`, so the batched pass that followed started from an EMPTY cache while the replay had started from the real prefix. For any chunk with `start_pos > 0` the oracle therefore compared "one token at position 1 from scratch" against "position 1 continuing the previous chunk" and reported MISMATCH — measured on `bonsai2-27b-NPU-b8`, width 1, four ids: [batchcheck] M=1 start=0 seq_argmax=198 batch_argmax=198 MATCH [batchcheck] M=1 start=1 seq_argmax=220 batch_argmax=653 MISMATCH The first row matching and every later row not is the signature of a state mismatch in the debug path, not of a row-dependence bug in the kernels -- and it sends the reader to the kernels. Latent until now because the packs on this box make a single pass (#1904); it goes live as soon as an M=128 prefill chunks a prompt, where every chunk after the first would "diverge". **The fix uses what the forward already has:** `snapshot_state(start_pos, snap)` before the replay and `restore_state(snap, start_pos)` after it -- the pair the prefix cache uses, whose plane-copy half (`q4nx_state_planes.h`) is device-free and unit-tested. `reset()` stays for the `start_pos == 0` case, where there is nothing to restore and it is exactly right. If the snapshot cannot be taken the oracle REFUSES with a message rather than comparing two different states again. **And the line names the positions compared:** `positions=%d..%d` instead of only `start=%d`, so the range is in the output rather than inferred from `start` and `M`. --- npu/forward/src/q4nx_forward.cpp | 30 +++++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index a8e503162..f487b66cb 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -3368,6 +3368,16 @@ bool Q4nxNpuForward::step_batch(const std::vector& tokens, int start_pos, // Not with injected rows: the oracle's own sequential pass runs step(tokens[m]) on the placeholder ids this // entry point supplies, so it would compare against a zero-token pass and report a divergence that is not // there. A debug tool that lies is worse than one that says "not applicable". + // engine#1905: the oracle's sequential replay must start from the SAME state the batched pass will, or the + // comparison is between different prefixes. It used to walk `start_pos .. start_pos+M-1` and then `reset()`, so for + // any chunk with `start_pos > 0` the batched pass ran from an empty cache and the replay from the real prefix: the + // oracle reported MISMATCH for every chunk but the first (measured on `bonsai2-27b-NPU-b8`, width 1, four ids: + // start=0 MATCH, start=1 seq=220 bat=653). That reads as a row-dependence bug in the kernels and sends the reader + // to fix the wrong thing. Snapshot the state the prefix established and restore it after the replay -- the same + // pair the prefix cache uses (`snapshot_state`/`restore_state`, and `q4nx_state_planes.h` is device-free). + StateSnapshot oracle_snap; + const bool oracle_snapped = batch_check && !inject_embeds_ && start_pos > 0 && + snapshot_state(start_pos, oracle_snap); if (batch_check && !inject_embeds_) { std::vector seq; if (c_.NV > 0) seq_rows.reserve((size_t)M * (size_t)c_.NV); @@ -3378,7 +3388,18 @@ bool Q4nxNpuForward::step_batch(const std::vector& tokens, int start_pos, for (size_t i = 0; i < seq.size(); i++) if (seq_argmax < 0 || seq[i] > seq[(size_t)seq_argmax]) seq_argmax = (int)i; } - reset(); + // Restore instead of reset() when the replay ran over a real prefix: `reset()` would leave the batched pass + // starting from nothing while the replay started from `start_pos` rows, which is the false mismatch above. With + // `start_pos == 0` there is nothing to restore and `reset()` is exactly right. + if (oracle_snapped) { + if (!restore_state(oracle_snap, start_pos)) { + err_ = "batch check: the sequential replay's prefix could not be restored, so the batched pass would " + "not be comparable with it"; + return false; + } + } else { + reset(); + } } // Batch phase attribution: the same partition step() prints, for an M-row verify pass, so the @@ -3817,8 +3838,11 @@ bool Q4nxNpuForward::step_batch(const std::vector& tokens, int start_pos, int ba = 0; for (size_t i = 1; i < logits.size(); i++) if (logits[i] > logits[(size_t)ba]) ba = (int)i; - std::fprintf(stderr, "[batchcheck] M=%d start=%d seq_argmax=%d batch_argmax=%d %s\n", - M, start_pos, seq_argmax, ba, seq_argmax == ba ? "MATCH" : "MISMATCH"); + // engine#1905: name the positions compared, not only the start. Before the snapshot fix this line said + // `start=1 MATCH`-looking numbers for a comparison that was really "position 1 from scratch" against "position + // 1 continuing a prefix", and the reader had to infer the range from `start` and `M`. + std::fprintf(stderr, "[batchcheck] M=%d positions=%d..%d seq_argmax=%d batch_argmax=%d %s\n", + M, start_pos, start_pos + M - 1, seq_argmax, ba, seq_argmax == ba ? "MATCH" : "MISMATCH"); // engine#1781: and the comparison the gate actually asks for -- every row, bitwise. if (!seq_rows.empty() && bat_rows.size() == seq_rows.size() && c_.NV > 0) { onebit::npu::RowMismatch mm;