npu: the speculative loop end to end on the served 27B, and a fully accepted verify keeps its own state (refs #1781)

CMakeLists.txt

@@ -1981,6 +1981,15 @@ endif()
     # A DFlash draft scored against its target, both on the host (engine#1780): the target's tapped hidden states
     # into the draft's reference forward, proposals checked against the target's greedy tokens. A tool, not a ctest:
     # it needs two packages and minutes of host time.
+    # A DFlash 2 draft against the served 27B (engine#1781): acceptance in stages, and the speculative loop end to end
+    # (draft -> one batched verify pass -> commit). A tool, not a ctest: target pack, draft package and the device.
+    add_executable(dflash2_probe tools/dflash2_probe.cpp npu/dflash_forward.cpp npu/dflash_draft.cpp)
+    target_include_directories(dflash2_probe PRIVATE npu npu/forward/include)
+    target_link_libraries(dflash2_probe PRIVATE onebit_npu_forward onebit_npu)
+    find_package(OpenMP COMPONENTS CXX QUIET)
+    if(OpenMP_CXX_FOUND)
+        target_link_libraries(dflash2_probe PRIVATE OpenMP::OpenMP_CXX)
+    endif()
     add_executable(dflash_acceptance tools/dflash_acceptance.cpp npu/dflash_forward.cpp npu/dflash_draft.cpp)
     target_include_directories(dflash_acceptance PRIVATE npu npu/forward/include)
     target_link_libraries(dflash_acceptance PRIVATE onebit_npu_forward onebit_npu)

docs/amd-drafters.md

@@ -189,6 +189,23 @@ Measured with these pieces against the served 27B, in a local loop: on a code pr
 tok/s against 4.2 tok/s plain decoding; output identical to plain greedy decoding. The loop is not part of this change;
 the numbers are on engine#1781.
 
+### The speculative loop on the served 27B (`tools/dflash2_probe loop`)
+
+`dflash2_probe loop` runs the pair end to end. The target is `bonsai2-27b-t2s2` on the NPU. The draft's projections can
+run on the NPU through `set_gemm_hook()`; its attention, convolutions and selector stay on the host. Each round, the
+draft proposes 7 tokens, the target verifies the 8 rows in one `step_batch_all()` pass, and `commit_verify_rows()` keeps
+the accepted prefix. When every row is accepted, the commit keeps the batch's own recurrent state instead of restoring
+the snapshot and replaying the rows: `NPU_REPLAY_CHECK=1` holds the two equal, and `NPU_COMMIT_FULL_REPLAY=1` keeps the
+old path. Measured on engine#1781, with output identical to plain greedy decoding:
+
+| prompt | plain decode | the loop | tokens per round |
+|---|---|---|---|
+| a Python function | 4.24-4.27 tok/s | 7.5-7.8 tok/s | 6.86 |
+| prose (with the fallback to plain decoding) | 4.17-4.21 | 4.08-4.12 | 2.82-3.0 |
+
+A loop has to be compared with the serve's own batched prefill. The per-token prefill is not bit-identical on this route,
+and against it the loop's output parts after 23 tokens.
+
 ## PARD-2's repos are weights-only — borrow the target's tokenizer
 
 `amd/PARD2-*` ship **`README.md`, `config.json`, `model.safetensors`, `warp_model.bin`** and nothing else: **no tokenizer

npu/forward/src/q4nx_forward.cpp

@@ -2505,8 +2505,15 @@ bool Q4nxNpuForward::commit_verify_rows(int accepted) {
         if (!kcache_[l].empty() && kcache_[l].size() > keep) kcache_[l].resize(keep);
         if (!vcache_[l].empty() && vcache_[l].size() > keep) vcache_[l].resize(keep);
     }
-    for (size_t l = 0; l < conv_state_.size(); l++) conv_state_[l] = snap_conv_[l];
-    for (size_t l = 0; l < rec_state_.size(); l++) rec_state_[l] = snap_rec_[l];
+    // A full accept keeps the state the batch's own recurrence left: that is the state the restore + replay would
+    // rebuild (NPU_REPLAY_CHECK holds the two equal), so the restore and the replay are skipped. On the served 27B
+    // they are 11 ms + 7 ms per row, 67 ms of an 8-row round. NPU_COMMIT_FULL_REPLAY=1 keeps them, for an A/B.
+    static const bool always_replay = std::getenv("NPU_COMMIT_FULL_REPLAY") != nullptr;
+    const bool keep_batch_state = accepted == verify_rows_ && accepted > 0 && !always_replay && !replay_check;
+    if (!keep_batch_state) {
+        for (size_t l = 0; l < conv_state_.size(); l++) conv_state_[l] = snap_conv_[l];
+        for (size_t l = 0; l < rec_state_.size(); l++) rec_state_[l] = snap_rec_[l];
+    }
     // The accepted prefix is the only thing a drafter may advance over: drop the rejected rows' hidden here, so
     // that last_hidden_rows() is exactly `accepted` rows starting at verify_pos_.  Leaving all verify_rows_ made
     // the draft loop's own size check fail, and it skipped the drafter advance without a word.
@@ -2516,8 +2523,8 @@ bool Q4nxNpuForward::commit_verify_rows(int accepted) {
     }
     const auto cv_t1 = std::chrono::steady_clock::now();
     verify_open_ = false;                        // the replay below must not re-record
-    if (accepted > 0 && shortconv_) conv_apply_rows(accepted);
-    if (accepted > 0 && !gdn_apply_rows(0, accepted)) return false;
+    if (!keep_batch_state && accepted > 0 && shortconv_) conv_apply_rows(accepted);
+    if (!keep_batch_state && accepted > 0 && !gdn_apply_rows(0, accepted)) return false;
     if (cv_prof) {
         const auto cv_t2 = std::chrono::steady_clock::now();
         const double restore = std::chrono::duration<double, std::milli>(cv_t1 - cv_t0).count();