The draft round graph runs the draft layer's front once

MtpDrafter::capture_round ran the draft layer's front (the embedding, the fc projections, the attention
hyper-connection read, the K/V projections and appends) over the window's T cells, then the whole layer again for row
a, repeating row a's front.  record_forward is now record_front (T rows) and record_rest (one row: the attention from
the q projection, the MLP, the final mixer, the head); the round graph runs the front once, copies row a's residual,
injection and mixed input to row 0 (copy_row_to_first, the row read on the device) and runs the rest there.  The
multi-row front is per row the one-row front's arithmetic: the window hashes on fixed text equal the previous build's
(IQ3_S, UD-Q4_K_XL; thinking and code text).

CUPTI, IQ3_S thinking text (--adapt-every 0): the round graph 721 / 747 / 789 / 784 / 898 -> 655 / 662 / 689 / 692 /
812 us after windows of 1 / 2 / 3 / 4 / 8 tokens.

Claude-Session: https://claude.ai/code/session_019XxRxQQNyk7Ycdr4ComSs1

include/strata/core/mtp.hpp

@@ -92,6 +92,13 @@ public:
 
 private:
     bool record_forward(int T, int step_row0, cudaStream_t cs, std::string& err);
+    /// The layer's front for T rows at step rows [row0, +T): the embedding, the fc projections, the attention
+    /// hyper-connection read (R_, inj_, mixed_) and the K/V appended.
+    bool record_front(int T, int row0, cudaStream_t cs, std::string& err);
+    /// The rest for row 0 at step row `step_row`: the attention from the q projection, the MLP, the final mixer,
+    /// the head, the draft and its probability.  Needs row 0's R_, inj_ and mixed_ (its q8_1 rows in xq_).
+    bool record_rest(int step_row, cudaStream_t cs, std::string& err);
+    void norm_rope(float* data, const float* gamma, int rows, int cols, const int32_t* p, cudaStream_t cs);
     bool capture_round(int T, std::string& err);
     bool capture_step(int j, std::string& err);
     cudaGraphExec_t step_exec_[9] = {};

include/strata/kernels/verify_kernels.hpp

@@ -80,6 +80,10 @@ void ident_hits(const int32_t* ids, int n, int32_t* slot, int32_t* dst, int32_t*
 void mtp_select(const float* R_src, int64_t R_stride, const int32_t* ids, const int32_t* row_dev, float* R_dst,
                 int32_t* tok_dst, int32_t* out, int j, void* stream, const float* probs = nullptr,
                 float* out_p = nullptr);
+/// Row *row_dev of a, b and c (a_n, b_n, c_n floats a row) copied to their row 0 (the draft layer's rest runs on
+/// row 0).  Graph-capturable: the row is read on the device.
+void copy_row_to_first(const int32_t* row_dev, float* a, int64_t a_n, float* b, int64_t b_n, float* c, int64_t c_n,
+                       void* stream);
 /// dst row i = src row ids[i] (row_bytes each, a multiple of 4), for n rows.
 void gather_rows(const uint8_t* src, int64_t row_bytes, const int32_t* ids, int64_t n, uint8_t* dst, void* stream);
 /// ids[t] = table[ids[t]] for n entries (a subset index back to a token id).

src/core/mtp.cpp

@@ -284,16 +284,18 @@ bool MtpDrafter::bind(const WeightTable& wt, const NativeHead* head, const float
 
 // The layer for T rows.  full = false stops after the K/V append (the prompt only needs the cache).
 bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::string& err) {
-    using namespace strata::kernels;
-    const ModelGeometry& g = *g_;
-    SessionState& ss = *ss_;
     // step_row0 >= 0: the full layer on step rows [step_row0, +T); step_row0 < 0: K/V only on rows [-1 - step_row0, +T)
     const bool full = step_row0 >= 0;
-    const int row0 = full ? step_row0 : -1 - step_row0;
     if (full && T != 1) { err = "mtp: the full layer runs one row at a time (its attention scratch is sized for one)"; return false; }
-    const int64_t N = g.n_embd, HC = g.hc, K = ss.k, NH = g.n_head, HD = g.head_dim, NKV = g.n_head_kv;
+    if (!record_front(T, full ? step_row0 : -1 - step_row0, cs, err)) return false;
+    return !full || record_rest(step_row0, cs, err);
+}
+
+bool MtpDrafter::record_front(int T, int row0, cudaStream_t cs, std::string& err) {
+    using namespace strata::kernels;
+    const ModelGeometry& g = *g_;
+    const int64_t N = g.n_embd, HC = g.hc, NH = g.n_head, HD = g.head_dim, NKV = g.n_head_kv;
     const QsaShapes s = shapes_of(g);
-    const GrShapes gs{g.n_embd, g.hc, g.hc_lr};
     const int32_t* step = step_ + row0 * 4;
     const int32_t* pos = pos_ + row0 * NH;
     try {
@@ -323,17 +325,12 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri
             }
             fused_gr_read_multi(fa, T, grs_, cs);
         }
-        // ---- attention: K/V into the layer's own cache, then (full) dense attention over every cell
-        auto norm_rope = [&](float* data, const float* gamma, int rows, int cols, const int32_t* p) {
-            native_qsa_rms_norm_weighted(data, gamma, data, cols, rows, EPS, cs);
-            if (native_rope_enabled()) native_rope_apply(data, data, rows, cols, (int) s.n_rot, (float) qsa_freq_base(), p, cs);
-            else rope_neox_apply(data, data, rows, cols, (int) s.n_rot, st_.cos_tab, st_.sin_tab, p, cs);
-        };
+        // ---- attention: K/V into the layer's own cache
         native_quantize_q8_1(mixed_, xq_, (int) N, T, cs);
         native_mmvq(GGML_Q8_0, q8("self_attn.k_proj.weight"), xq_, kcur_, (int) N, (int) (NKV * HD), T, cs);
         native_mmvq(GGML_Q8_0, q8("self_attn.v_proj.weight"), xq_, vcur_, (int) N, (int) (NKV * HD), T, cs);
         for (int t = 0; t < T; ++t) {
-            norm_rope(kcur_ + t * NKV * HD, f32("self_attn.k_norm.weight"), (int) NKV, (int) HD, pos + t * NH);
+            norm_rope(kcur_ + t * NKV * HD, f32("self_attn.k_norm.weight"), (int) NKV, (int) HD, pos + t * NH, cs);
             if (st_.kv_int8)
                 kv_append_q8_step(st_.k_q, st_.v_q, st_.k_scale, st_.v_scale, st_.page_table, step + t * 4,
                                   kcur_ + t * NKV * HD, vcur_ + t * NKV * HD, s, cs);
@@ -341,7 +338,33 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri
                 kv_append_step(st_.k_pool, st_.v_pool, st_.page_table, step + t * 4, kcur_ + t * NKV * HD,
                                vcur_ + t * NKV * HD, s, cs);
         }
-        if (!full) return true;
+    } catch (const std::exception& e) {
+        err = std::string("mtp: ") + e.what();
+        return false;
+    }
+    return true;
+}
+
+void MtpDrafter::norm_rope(float* data, const float* gamma, int rows, int cols, const int32_t* p, cudaStream_t cs) {
+    using namespace strata::kernels;
+    const QsaShapes s = shapes_of(*g_);
+    native_qsa_rms_norm_weighted(data, gamma, data, cols, rows, EPS, cs);
+    if (native_rope_enabled()) native_rope_apply(data, data, rows, cols, (int) s.n_rot, (float) qsa_freq_base(), p, cs);
+    else rope_neox_apply(data, data, rows, cols, (int) s.n_rot, st_.cos_tab, st_.sin_tab, p, cs);
+}
+
+bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) {
+    using namespace strata::kernels;
+    const ModelGeometry& g = *g_;
+    SessionState& ss = *ss_;
+    const int T = 1;
+    const int64_t N = g.n_embd, HC = g.hc, K = ss.k, NH = g.n_head, HD = g.head_dim;
+    const QsaShapes s = shapes_of(g);
+    const GrShapes gs{g.n_embd, g.hc, g.hc_lr};
+    const int32_t* step = step_ + step_row * 4;
+    const int32_t* pos = pos_ + step_row * NH;
+    try {
+        // ---- dense attention over every cell
         native_mmvq(GGML_Q8_0, q8("self_attn.q_proj.weight"), xq_, qfull_, (int) N, (int) (NH * 2 * HD), T, cs);
         for (int t = 0; t < T; ++t) {
             float* qc = qcur_ + t * NH * HD;
@@ -350,7 +373,7 @@ bool MtpDrafter::record_forward(int T, int step_row0, cudaStream_t cs, std::stri
                 err = "mtp: q split failed";
                 return false;
             }
-            norm_rope(qc, f32("self_attn.q_norm.weight"), (int) NH, (int) HD, pos + t * NH);
+            norm_rope(qc, f32("self_attn.q_norm.weight"), (int) NH, (int) HD, pos + t * NH, cs);
         }
         QsaAttnPools pools;
         pools.page_table = st_.page_table;
@@ -461,16 +484,20 @@ bool MtpDrafter::capture_round(int T, std::string& err) {
     copy_i32_from_mapped(pos_, m_pos_, (int64_t) 2 * T * g_->n_head, cs_);
     copy_i32_from_mapped(row_, m_row_, 2, cs_);
     copy_from_mapped(Rin_, window_R_, (int64_t) T * HCN, cs_);
-    // the catch-up: K/V for the window's T cells, then the full layer for row a only (its cell's K/V is written
-    // again, identically), staged by the host in step row 2*max_t - 1; the draft chain is one graph per step
-    // (`capture_step`) so the host can stop it when a draft is unlikely
+    // the catch-up: the layer's front for the window's T cells (their K/V), then its rest for row a only, on row
+    // a's intermediates copied to row 0, at the cell the host staged in step row 2*max_t - 1; the draft chain is
+    // one graph per step (`capture_step`) so the host can stop it when a draft is unlikely
     const int ra = 2 * max_t_ - 1;
-    ok = record_forward(T, -1, cs_, err);
-    if (ok) mtp_select(Rin_, HCN, tok_, row_, Rin_, tok_, nullptr, 0, cs_);
+    const int64_t N = g_->n_embd;
+    ok = record_front(T, 0, cs_, err);
     if (ok) {
+        if (T > 1) {
+            copy_row_to_first(row_, R_, HCN, inj_, g_->hc, mixed_, N, cs_);
+            native_quantize_q8_1(mixed_, xq_, (int) N, 1, cs_);
+        }
         copy_i32_from_mapped(step_ + ra * 4, m_step_ + ra * 4, 4, cs_);
         copy_i32_from_mapped(pos_ + ra * g_->n_head, m_pos_ + ra * g_->n_head, g_->n_head, cs_);
-        ok = record_forward(1, ra, cs_, err);
+        ok = record_rest(ra, cs_, err);
     }
     if (ok) mtp_select(R_, HCN, out_ids_, row_ + 1, Rin_, tok_, m_out_, 0, cs_, probs_, m_prob_);
     return finish_capture(cs_, ok, round_exec_[T], "round", err);

src/kernels/cuda/verify_kernels.cu

@@ -375,6 +375,24 @@ void ident_hits(const int32_t* ids, int n, int32_t* slot, int32_t* dst, int32_t*
     check("ident_hits");
 }
 
+namespace {
+__global__ void copy_row_to_first_kernel(const int32_t* __restrict__ row_dev, float* a, int64_t a_n, float* b,
+                                         int64_t b_n, float* c, int64_t c_n) {
+    const int64_t row = *row_dev;
+    if (row == 0) return;
+    const int64_t i0 = (int64_t) blockIdx.x * blockDim.x + threadIdx.x, st = (int64_t) gridDim.x * blockDim.x;
+    for (int64_t i = i0; i < a_n; i += st) a[i] = a[row * a_n + i];
+    for (int64_t i = i0; i < b_n; i += st) b[i] = b[row * b_n + i];
+    for (int64_t i = i0; i < c_n; i += st) c[i] = c[row * c_n + i];
+}
+}  // namespace
+
+void copy_row_to_first(const int32_t* row_dev, float* a, int64_t a_n, float* b, int64_t b_n, float* c, int64_t c_n,
+                       void* stream) {
+    copy_row_to_first_kernel<<<16, 256, 0, (cudaStream_t) stream>>>(row_dev, a, a_n, b, b_n, c, c_n);
+    check("copy_row_to_first");
+}
+
 void mtp_select(const float* R_src, int64_t R_stride, const int32_t* ids, const int32_t* row_dev, float* R_dst,
                 int32_t* tok_dst, int32_t* out, int j, void* stream, const float* probs, float* out_p) {
     mtp_select_kernel<<<16, 256, 0, (cudaStream_t) stream>>>(R_src, R_stride, ids, row_dev, R_dst, tok_dst, out, j,