A one-token window commits itself

A one-token window always keeps its token, so its graph now advances the sequence state itself: the GDN step kernel
writes the state (n_keep = 1, the arithmetic the commit's replay used), the conv kernel writes the history it read
(gdn_conv_l2_multi's commit flag), and its indexer append and PLE history were final already.  Verifier::commit
launches no commit graph after it: the draft round graph then runs alone instead of beside the commit's 98 kernels,
which read and write each GDN layer's 3 MB state.  Window hashes on fixed text equal the previous build's (IQ3_S,
UD-Q4_K_XL; thinking and code text); cache_test.py's regenerated turn and stash round trip reproduce their output.

CUPTI, IQ3_S thinking text (--adapt-every 0), the 35% of rounds after one-token windows: the window 11,270 -> 11,380
us, the time from its end to the next window's start 3,706 -> 3,325 us (the round graph 655 -> 450 us, the draft
starting 150 us sooner): -270 us a round there, ~-95 us over all rounds.

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

include/strata/core/verify.hpp

@@ -315,6 +315,7 @@ private:
     void* hit_scratch_ = nullptr;
     float *head_mixed_ = nullptr, *head_inj_ = nullptr, *head_logits_ = nullptr;
     uint8_t* arg_scratch_ = nullptr;   ///< argmax_rows' partials and counters
+    int32_t* one_ = nullptr;           ///< device {1}: the n_keep of a one-token window, which commits itself
     float *sh_gate_ = nullptr, *sh_up_ = nullptr, *sh_g_ = nullptr;
     float* hist_snap_ = nullptr;                              // T * NG_HIST * NG_HC_DIM
     // the PLE block's projections (keys T * NG_HC_DIM, values T * n_embd) and the rest of the block's rows

include/strata/kernels/verify_kernels.hpp

@@ -21,10 +21,10 @@ namespace strata::kernels {
 inline constexpr int kVerifyMaxT = 8;
 
 /// For token t of T: conv over [history(3) | qkv_0 .. qkv_t] -> SiLU -> L2 norm of the q/k heads -> h[t].
-/// `history` is NOT written.  Bitwise `fused_gdn_conv_l2` per token.  Up to 65535 tokens (the prompt path runs a
-/// whole chunk through it).
+/// `history` is NOT written, but with `commit` (one token): it then keeps the token, as gdn_conv_commit.  Bitwise
+/// `fused_gdn_conv_l2` per token.  Up to 65535 tokens (the prompt path runs a whole chunk through it).
 void gdn_conv_l2_multi(const float* history, const float* qkv, const float* conv_w, float* h, int channels,
-                       int qk_heads, float eps, int n_tok, void* stream, int t_begin = 0);
+                       int qk_heads, float eps, int n_tok, void* stream, int t_begin = 0, bool commit = false);
 /// history <- the last 3 entries of [history | qkv_0 .. qkv_{n-1}], n = *n_keep (0 leaves it as it was).
 void gdn_conv_commit(float* history, const float* qkv, int channels, const int32_t* n_keep, void* stream);
 /// alpha/beta for T columns of x (T, n_embd): gate (T, h_v), beta (T, h_v).  Bitwise `fused_gdn_ab` per column.

src/core/verify.cpp

@@ -274,6 +274,7 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
         sh_up_ = b.take<float>(T * (uint64_t) g.n_ff); sh_g_ = b.take<float>(T + 4);
         head_logits_ = b.take<float>(T * (uint64_t) n_vocab_);
         arg_scratch_ = b.take<uint8_t>(strata::kernels::argmax_rows_scratch_bytes((int) T));
+        one_ = b.take<int32_t>(4);
         hist_snap_ = b.take<float>(T * HS);
         const uint64_t PD = (uint64_t) strata::kernels::NG_HC_DIM;
         ple_key_ = b.take<float>(T * PD); ple_val_ = b.take<float>(T * N); ple_nkey_ = b.take<float>(T * PD);
@@ -289,6 +290,13 @@ bool Verifier::init(const WeightTable& wt, const ModelGeometry& g, SessionState&
     Bump real;
     real.base = (uint8_t*) arena_;
     carve(real);
+    {
+        const int32_t one = 1;
+        if (cudaMemcpy(one_, &one, sizeof one, cudaMemcpyHostToDevice) != cudaSuccess) {
+            err = "verify: the arena could not be set";
+            return false;
+        }
+    }
     sink_.staging = (unsigned long long) staging_;
     sink_.staging_cap = kStagingBlobs;
     (void) TS;
@@ -615,11 +623,13 @@ bool Verifier::record_window(int T, cudaStream_t cs, std::string& err) {
                 proj(wqkv, qkv + (size_t) tb * C, N, C, cs);
                 proj(wg, z_ + (size_t) tb * ZV, N, ZV, shs_);
                 if (!mark(shs_, pjoin_)) return false;
-                gdn_conv_l2_multi(conv, qkv, (const float*) wc->data, hb, (int) C, (int) (2 * HK), EPS, n, cs, tb);
+                // a one-token window keeps its token: it advances the conv history and the state itself (commit)
+                gdn_conv_l2_multi(conv, qkv, (const float*) wc->data, hb, (int) C, (int) (2 * HK), EPS, n, cs, tb,
+                                  T == 1);
                 if (!wait(bjoin_) || !wait(pjoin_)) return false;
                 // the recurrence from the untouched state over tokens [0, te); outputs only for this group's
                 gdn_step_norm_multi(state, hb, (int) C, gate, beta, z_, (const float*) wnm->data, EPS, y_, (int) HK,
-                                    (int) HV, te, nullptr, cs, tb);
+                                    (int) HV, te, T == 1 ? one_ : nullptr, cs, tb);
                 quant(y_ + (size_t) tb * ZV, ZV);
                 proj(wout, bo_ + tb * N, ZV, N, cs);
             } else {
@@ -1403,11 +1413,13 @@ bool Verifier::commit(int n_keep, std::string& err, bool wait) {
     h_commit_[0] = n_keep;
     h_commit_[1] = n_keep - 1;
     for (int t = 0; t < max_t_; ++t) h_commit_[2 + t] = t < n_keep ? (int32_t) (last_pos0_ + t) : -1;
-    std::atomic_thread_fence(std::memory_order_seq_cst);
-    const cudaError_t le = cudaGraphLaunch(commit_exec_, cs_);
-    if (le != cudaSuccess) { err = std::string("verify: commit launch: ") + cudaGetErrorString(le); return false; }
-    (void) cudaStreamQuery(cs_);   // submit now, not at the next driver call
-    commit_pending_ = true;
+    if (last_t_ > 1) {   // a one-token window has advanced the state itself (record_window)
+        std::atomic_thread_fence(std::memory_order_seq_cst);
+        const cudaError_t le = cudaGraphLaunch(commit_exec_, cs_);
+        if (le != cudaSuccess) { err = std::string("verify: commit launch: ") + cudaGetErrorString(le); return false; }
+        (void) cudaStreamQuery(cs_);   // submit now, not at the next driver call
+        commit_pending_ = true;
+    }
     for (int t = 0; t < n_keep; ++t) {
         ss_->ple_prev[0] = ss_->ple_prev[1];
         ss_->ple_prev[1] = last_tokens_[t];

src/kernels/cuda/verify_kernels.cu

@@ -23,10 +23,10 @@ void check(const char* what) {
     if (e != cudaSuccess) { std::fprintf(stderr, "%s: %s\n", what, cudaGetErrorString(e)); std::exit(1); }
 }
 
-__global__ void __launch_bounds__(S) gdn_conv_l2_multi_kernel(const float* __restrict__ hist,
-                                                              const float* __restrict__ qkv,
+// `commit` (one token): the history then keeps it - [hist1, hist2, x_0], each thread its channel's, after reading them
+__global__ void __launch_bounds__(S) gdn_conv_l2_multi_kernel(float* hist, const float* __restrict__ qkv,
                                                               const float* __restrict__ w, float* __restrict__ h,
-                                                              int C, int qk_heads, float eps, int t_begin) {
+                                                              int C, int qk_heads, float eps, int t_begin, bool commit) {
     __shared__ float part[S / 32];
     const int t = t_begin + blockIdx.y;
     const int c = blockIdx.x * S + threadIdx.x;
@@ -38,6 +38,11 @@ __global__ void __launch_bounds__(S) gdn_conv_l2_multi_kernel(const float* __res
         win[j] = src < 3 ? hist[c * 3 + src] : qkv[(size_t) (src - 3) * C + c];
     }
     const float v0 = win[0], v1 = win[1], v2 = win[2], x = qkv[(size_t) t * C + c];
+    if (commit) {
+        hist[c * 3] = v1;
+        hist[c * 3 + 1] = v2;
+        hist[c * 3 + 2] = x;
+    }
     float sum = v0 * w[c * 4] + v1 * w[c * 4 + 1] + v2 * w[c * 4 + 2] + x * w[c * 4 + 3];
     float y = sum / (1.0f + __expf(-sum));
     if ((int) blockIdx.x < qk_heads) {
@@ -582,13 +587,14 @@ void dense_steps(const int32_t* cells, int n, int32_t* steps, void* stream) {
 }
 
 void gdn_conv_l2_multi(const float* history, const float* qkv, const float* conv_w, float* h, int channels,
-                       int qk_heads, float eps, int n_tok, void* stream, int t_begin) {
-    if (!history || !qkv || !conv_w || !h || channels % S != 0 || n_tok < 1 || n_tok > 65535) {
+                       int qk_heads, float eps, int n_tok, void* stream, int t_begin, bool commit) {
+    if (!history || !qkv || !conv_w || !h || channels % S != 0 || n_tok < 1 || n_tok > 65535 ||
+        (commit && (n_tok != 1 || t_begin != 0))) {
         std::fprintf(stderr, "gdn_conv_l2_multi: invalid arguments\n");
         std::exit(1);
     }
     gdn_conv_l2_multi_kernel<<<dim3((unsigned) (channels / S), (unsigned) n_tok), S, 0, (cudaStream_t) stream>>>(
-        history, qkv, conv_w, h, channels, qk_heads, eps, t_begin);
+        const_cast<float*>(history), qkv, conv_w, h, channels, qk_heads, eps, t_begin, commit);
     check("gdn_conv_l2_multi");
 }