mtp: kill momus bugs #1-3 (recommit off-by-one, capture ownership, async stream)

#1: recommit verify wrote accept_n+2 KV slots vs fast-path's accept_n+1;
    runner's base_pos += accept_n+2 then re-verified bonus at its own slot,
    skipping every Dth token at γ≥3 partial-accept (reproduced AR/spec
    divergence in new test_recommit_byte_identical_to_ar).
#2: capture_topology_for_chain() virtual; runner owns the call.
    verify_batch no longer mutates last_tree_*.
#3: dedicated rollback CUDA stream + cudaMemcpy2DAsync batching in
    restore_kv_at_dfs (4× fewer launches per layer).

Bug #5 (step_sg_cache O(n_ctx)) deferred — needs ggml_set_rows refactor.

Author: dusterbloom

dflash/src/common/dflash_target.h

@@ -131,6 +131,12 @@ struct DFlashTarget {
     // without DeltaNet capture support.
     virtual void enable_chain_capture(bool /*on*/) {}
 
+    // Record linear-chain "tree" topology so a subsequent restore_kv_at_chain()
+    // knows the base_pos / n_tokens of the chain to roll back.  MUST be called
+    // by the chain runner BEFORE verify_batch on every iter that may need a
+    // rollback.  No-op for targets without chain capture support.
+    virtual void capture_topology_for_chain(int /*n_tokens*/, int /*base_pos*/) {}
+
     // ── Token utilities ─────────────────────────────────────────────
 
     // Check if a token is end-of-sequence for this model.

dflash/src/common/mtp_chain_runner.cpp

@@ -165,6 +165,9 @@ GenerateResult MtpChainRunner::run(const GenerateRequest & req,
             return result;
         }
 
+        // Caller-owned topology for restore_kv_at_chain (bug #2).
+        target_.capture_topology_for_chain((int)candidate.size(), base_pos);
+
         int last_argmax = -1;
         std::vector<int32_t> all_argmax;
         if (!target_.verify_batch(candidate, base_pos, last_argmax, &all_argmax)) {
@@ -192,50 +195,43 @@ GenerateResult MtpChainRunner::run(const GenerateRequest & req,
         const int total_this_iter = accept_n + 1;
 
         // ── KV reconciliation ──────────────────────────────────────────
-        // verify_batch advanced KV by 1+g_actual positions. We only want
-        // total_this_iter additional positions past base_pos.  Three paths:
-        //   1. accept_n == g_actual: KV is correct, nothing to do.
-        //   2. accept_n <  g_actual, fast rollback: target.restore_kv_at_chain
-        //      (accept_n) reuses the per-position SSM/conv intermediates
-        //      captured by the first verify_batch to undo the rejected tail
-        //      in-place.  No second backbone forward.  The bonus token's KV
-        //      is NOT written; next iter's verify_batch will process it as
-        //      the new cur_tok at slot base_pos+accept_n+1.  base_pos
-        //      advance and cur_tok bookkeeping match the accept-all path.
-        //   3. accept_n <  g_actual, fast path declined: fall back to the
-        //      legacy snapshot+recommit (verify_batch on [cur, accepted,
-        //      bonus]).  Used by targets that lack DeltaNet capture (CPU
-        //      stubs, layer-split shards) and any iter where chain capture
-        //      was skipped defensively.  base_pos advances by accept_n+2 to
-        //      skip past the recommit-written bonus slot; cur_tok_next is
-        //      taken from the recommit's argmax at the bonus position.
-        bool recommitted = false;
-        bool fast_rolled_back = false;
-        std::vector<int32_t> rc_argmax;
+        // Three paths converge on the SAME post-iter invariant:
+        //   base_pos advances by total_this_iter = accept_n + 1
+        //   cur_tok                 = bonus (= all_argmax[accept_n])
+        // and the bonus's KV is written by the NEXT iter's verify_batch.
+        //
+        //   1. accept-all (accept_n == g_actual): verify_batch wrote g+1 slots;
+        //      we treat the last (bonus) slot as uncommitted and let next iter
+        //      overwrite it.
+        //   2. fast rollback (restore_kv_at_chain succeeds): rolls cache.cur_pos
+        //      to base_pos + accept_n + 1, leaving the bonus slot unwritten.
+        //   3. recommit (fast path declined): snapshot+restore, then recommit
+        //      only [cur, accepted...] (accept_n+1 slots, NO bonus).  Bonus is
+        //      threaded via cur_tok like the other paths.  Advancing by
+        //      accept_n+2 here would skip a position every recommit iter and
+        //      diverge from AR — see test_recommit_byte_identical_to_ar.
         if (accept_n < g_actual) {
-            if (target_.restore_kv_at_chain(accept_n)) {
-                fast_rolled_back = true;
-            } else {
-                recommitted = true;
+            if (!target_.restore_kv_at_chain(accept_n)) {
+                // Slow path: snapshot rollback + commit ONLY [cur, accepted...]
+                // (accept_n+1 slots).  Bonus stays uncommitted, threaded via
+                // cur_tok like the fast path.
                 if (!target_.restore_kv()) {
                     result.ok = false;
                     result.error = "restore_kv";
                     return result;
                 }
                 std::vector<int32_t> commit_seq;
-                commit_seq.reserve((size_t)accept_n + 2);
+                commit_seq.reserve((size_t)accept_n + 1);
                 commit_seq.push_back(cur_tok);
                 for (int i = 0; i < accept_n; i++) commit_seq.push_back(drafts[i]);
-                commit_seq.push_back(all_argmax[accept_n]);
                 int discard = -1;
-                if (!target_.verify_batch(commit_seq, base_pos, discard, &rc_argmax)) {
+                if (!target_.verify_batch(commit_seq, base_pos, discard, nullptr)) {
                     result.ok = false;
                     result.error = "recommit";
                     return result;
                 }
             }
         }
-        (void)fast_rolled_back;  // bookkeeping is identical to accept-all path.
 
         // ── Emit accepted prefix + bonus, capped at n_gen ──────────────
         // emit_cap is the absolute ceiling on tokens that may be written
@@ -262,27 +258,8 @@ GenerateResult MtpChainRunner::run(const GenerateRequest & req,
             cur_tok = result.tokens.empty() ? cur_tok : result.tokens.back();
         }
 
-        // base_pos advance:
-        //   accept-all / fast_rolled_back:
-        //       verify_batch wrote g+1 KV slots; the fast rollback rolled
-        //       cache.cur_pos back to base_pos+accept_n+1 = base_pos +
-        //       total_this_iter, leaving the bonus token's KV unwritten.
-        //       Advance by total_this_iter and let next iter's verify_batch
-        //       process the bonus as the new cur_tok.
-        //   recommit:
-        //       recommit wrote accept_n+2 KV slots (cur_tok + accepted +
-        //       bonus). Advance by accept_n+2 so next iter's base_pos points
-        //       to a NEW position, not the bonus slot.  cur_tok_next is the
-        //       argmax at the bonus position from the recommit (a fresh
-        //       prediction, not yet in cache).
-        if (recommitted) {
-            base_pos += accept_n + 2;
-            if (!hit_eos && !rc_argmax.empty()) {
-                cur_tok = rc_argmax.back();
-            }
-        } else {
-            base_pos += total_this_iter;
-        }
+        // All paths share: base_pos += total_this_iter and cur_tok = bonus.
+        base_pos += total_this_iter;
 
         stats_.total_iters    += 1;
         stats_.total_accepted += accept_n;

dflash/src/device_runtime.h

@@ -36,6 +36,16 @@ using __nv_bfloat16 = __hip_bfloat16;
 #define cudaEventElapsedTime hipEventElapsedTime
 #define cudaEventDestroy hipEventDestroy
 #define cudaStreamSynchronize hipStreamSynchronize
+#define cudaStreamCreate hipStreamCreate
+#define cudaStreamDestroy hipStreamDestroy
+#define cudaMemcpy3DAsync hipMemcpy3DAsync
+#define cudaPitchedPtr hipPitchedPtr
+#define cudaPos hipPos
+#define cudaExtent hipExtent
+#define cudaMemcpy3DParms hipMemcpy3DParms
+#define make_cudaPitchedPtr make_hipPitchedPtr
+#define make_cudaPos make_hipPos
+#define make_cudaExtent make_hipExtent
 #define cudaGetLastError hipGetLastError
 #define cudaGetErrorString hipGetErrorString
 #define cudaDeviceSynchronize hipDeviceSynchronize

dflash/src/qwen35/qwen35_dflash_target.cpp

@@ -62,6 +62,7 @@ Qwen35DFlashTarget::~Qwen35DFlashTarget() {
             vprof_sum_get_hpre_, vprof_sum_get_argmax_, vprof_sum_total_);
     }
     step_graph_destroy(proj_sg_);
+    if (rollback_stream_) { cudaStreamDestroy(rollback_stream_); rollback_stream_ = nullptr; }
 }
 
 Qwen35DFlashTarget::Qwen35DFlashTarget(
@@ -264,27 +265,9 @@ bool Qwen35DFlashTarget::verify_batch(
 
     cache_.cur_pos = base_pos + n_tokens;
 
-    // Record linear-chain "tree" topology so restore_kv_at_chain (and
-    // restore_kv_at_dfs, for completeness) can roll back to any accepted
-    // prefix of this batch.  Only meaningful when capture actually fired —
-    // when it didn't, leave last_tree_base_pos_ at its prior value (could be
-    // -1 or a stale verify_tree value); the chain runner only calls
-    // restore_kv_at_chain on iters where it itself enabled capture, and a
-    // stale last_tree_* wouldn't be consulted by restore_kv_at_dfs in that
-    // window because verify_batch resets it here on every captured call.
-    if (capture_intermediate) {
-        last_tree_base_pos_ = base_pos;
-        last_tree_n_nodes_  = n_tokens - 1;
-        last_tree_parents_.resize(n_tokens);
-        last_tree_depths_.resize(n_tokens > 0 ? (size_t)(n_tokens - 1) : 0);
-        if (n_tokens > 0) last_tree_parents_[0] = -1;  // root marker
-        for (int i = 1; i < n_tokens; i++) {
-            last_tree_parents_[i] = i - 1;
-            last_tree_depths_[i - 1] = i;
-        }
-    } else {
-        // Invalidate so a stray restore_kv_at_chain() returns false instead of
-        // rolling back against stale capture buffers from an earlier iter.
+    // Topology is owned by the caller via capture_topology_for_chain(); if
+    // capture fired without prior topology, invalidate to be defensive.
+    if (!capture_intermediate) {
         last_tree_base_pos_ = -1;
     }
 
@@ -459,7 +442,14 @@ bool Qwen35DFlashTarget::restore_kv_at_dfs(const std::vector<int> & accepted_dfs
     }
 
     const int n_delta = (int)cache_.ssm_intermediate.size();
-    cudaStream_t stream = nullptr;  // default stream — serializes with next graph_compute
+    // Bug #3: dedicated stream so the rollback copies don't serialize with
+    // the default stream (e.g. ggml backend compute, host syncs).
+    if (!rollback_stream_) {
+        if (cudaStreamCreate(&rollback_stream_) != cudaSuccess) {
+            rollback_stream_ = nullptr;  // fall back to default
+        }
+    }
+    cudaStream_t stream = rollback_stream_;
     for (int il = 0; il < n_delta; il++) {
         ggml_tensor * ssm_inter = cache_.ssm_intermediate[il];
         ggml_tensor * conv_in   = cache_.conv_input_cache[il];
@@ -533,9 +523,9 @@ bool Qwen35DFlashTarget::restore_kv_at_dfs(const std::vector<int> & accepted_dfs
     }
 
     // Full-attention KV compaction: verify_tree wrote K/V at slots
-    // [base..base+N-1] in DFS order.  For the next iter's verify to see the
-    // correct committed prefix, slots [base..base+commit_n-1] must hold the
-    // K/V of the accepted path.  d==0 (root) is trivially aligned.
+    // [base..base+N-1] in DFS order.  Bug #3: collapse the per-head inner
+    // loop into one cudaMemcpy2DAsync (pitch=nb[2], height=n_kv) — saves
+    // 2*n_kv-2 launches per (layer, d) pair on a dedicated stream.
     if (walked_sibling) {
         const int base = last_tree_base_pos_;
         const int n_full_attn = (int)cache_.attn_k.size();
@@ -548,30 +538,46 @@ bool Qwen35DFlashTarget::restore_kv_at_dfs(const std::vector<int> & accepted_dfs
                 ggml_tensor * cv = cache_.attn_v[l];
                 if (!ck || !cv) continue;
                 const size_t slot_bytes = ck->nb[1];
-                const size_t src_off = (size_t)(base + src_dfs)  * slot_bytes;
-                const size_t dst_off = (size_t)(base + dst_slot) * slot_bytes;
-                const int n_kv = (int)ck->ne[2];
-                for (int h = 0; h < n_kv; h++) {
-                    const size_t head_src = src_off + (size_t)h * ck->nb[2];
-                    const size_t head_dst = dst_off + (size_t)h * ck->nb[2];
-                    cudaMemcpyAsync((char *)ck->data + head_dst,
-                                    (const char *)ck->data + head_src,
-                                    slot_bytes, cudaMemcpyDeviceToDevice, stream);
-                    cudaMemcpyAsync((char *)cv->data + head_dst,
-                                    (const char *)cv->data + head_src,
-                                    slot_bytes, cudaMemcpyDeviceToDevice, stream);
-                }
+                const int    n_kv       = (int)ck->ne[2];
+                const size_t pitch      = ck->nb[2];
+                const size_t src_off    = (size_t)(base + src_dfs)  * slot_bytes;
+                const size_t dst_off    = (size_t)(base + dst_slot) * slot_bytes;
+                cudaMemcpy2DAsync((char *)ck->data + dst_off, pitch,
+                                  (const char *)ck->data + src_off, pitch,
+                                  slot_bytes, n_kv,
+                                  cudaMemcpyDeviceToDevice, stream);
+                cudaMemcpy2DAsync((char *)cv->data + dst_off, cv->nb[2],
+                                  (const char *)cv->data + src_off, cv->nb[2],
+                                  slot_bytes, (int)cv->ne[2],
+                                  cudaMemcpyDeviceToDevice, stream);
             }
         }
     }
 
+    // Sync rollback stream so the next graph_compute (on default stream)
+    // sees a consistent KV/SSM state.
+    if (rollback_stream_) cudaStreamSynchronize(rollback_stream_);
+
     // Advance cur_pos to "just past the last committed slot" so the next
     // verify_batch's kv_start lines up.  root = dfs 0 lives at base, so
     // commit_n committed tokens occupy slots [base..base+commit_n-1].
     cache_.cur_pos = last_tree_base_pos_ + commit_n;
     return true;
 }
 
+void Qwen35DFlashTarget::capture_topology_for_chain(int n_tokens, int base_pos) {
+    if (n_tokens <= 0) { last_tree_base_pos_ = -1; return; }
+    last_tree_base_pos_ = base_pos;
+    last_tree_n_nodes_  = n_tokens - 1;
+    last_tree_parents_.resize(n_tokens);
+    last_tree_depths_.resize((size_t)(n_tokens - 1));
+    last_tree_parents_[0] = -1;
+    for (int i = 1; i < n_tokens; i++) {
+        last_tree_parents_[i] = i - 1;
+        last_tree_depths_[i - 1] = i;
+    }
+}
+
 bool Qwen35DFlashTarget::restore_kv_at_chain(int accept_n) {
     // A chain of N tokens recorded by verify_batch is the DFS spine
     // [0, 1, ..., N-1].  Roll back to slot accept_n: the first (accept_n + 1)

dflash/src/qwen35/qwen35_dflash_target.h

@@ -13,6 +13,7 @@
 
 #include "ggml.h"
 #include "ggml-backend.h"
+#include "device_runtime.h"  // cudaStream_t
 
 #include <vector>
 
@@ -125,6 +126,11 @@ class Qwen35DFlashTarget : public DFlashTarget {
     // recommit fallback.
     void enable_chain_capture(bool on) override { chain_capture_enabled_ = on; }
 
+    // Record a linear-chain "tree" (spine [0..n_tokens-1] at base_pos) so
+    // a subsequent restore_kv_at_chain() can locate the rollback slot.
+    // Owned by the caller (chain runner), separate from verify_batch state.
+    void capture_topology_for_chain(int n_tokens, int base_pos) override;
+
     bool is_eos(int token) const override;
 
     bool embed_tokens(const int32_t * tokens, int n,
@@ -255,6 +261,11 @@ class Qwen35DFlashTarget : public DFlashTarget {
     // n_tokens > max_verify_tokens, e.g. 512-token prefill chunks, where the
     // in-graph ggml_view_3d into the conv_input cache asserts).
     bool                       chain_capture_enabled_ = false;
+
+    // Dedicated CUDA stream for restore_kv_at_dfs copies (bug #3): avoids
+    // serializing ~384 per-head launches on the default stream.  Created on
+    // first use, destroyed in the dtor.
+    mutable cudaStream_t       rollback_stream_ = nullptr;
 };
 
 }  // namespace dflash27b