metal: gdn_step partial commit (commit = 0x100 | k): the state and conv history after the first k window tokens

- gdn.metal: with commit = 0x100 | k (1 <= k <= GDN_NTOK) the state is written in place and the conv history to
  hist_out right after token k - 1, while every token's y is still computed; commit = 1 / 0 unchanged (the full
  commit now writes at the last token inside the loop - same values, same speed)
- gdn_parity: the partial commits are bitwise a k-token commit's state and history, and y bitwise the verify window's
  (k = 1 .. ntok - 1, ntok 2..4); gr_parity GR_DOWN=SG_KS picks gr_down's variant (sweep: 10.1-11.8 us at T 2, none
  better than 8_8 by more than 7%)

metal/kernels/gdn.metal

@@ -19,7 +19,9 @@ struct GdnArgs {
     uint h_k;        // 16
     uint h_v;        // 48
     uint channels;   // 10240 = (2 h_k + h_v) * 128
-    uint commit;
+    uint commit;     // 0: verify (no write); 1: commit after the last token; 0x100 | k (1 <= k <= GDN_NTOK): commit
+                     // the state and the conv history after the first k tokens, in place, while every token's y is
+                     // still computed (a window whose first k rows are kept) - bitwise the commit of a k-token step
     float eps;       // 1e-6
     uint plane;      // 0: y only; else y is the raw plane of prepared activations (xprep layout, n = h_v * 128,
                      // `plane` floats between planes) and the other planes and the per-16 sums are written too
@@ -60,6 +62,7 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b
     float st[RPT];
     for (int r = 0; r < RPT; ++r) st[r] = state[((i0 + r) * hv + h) * S + j];
 
+    const int commit_at = a.commit == 1u ? GDN_NTOK : (a.commit >> 8) ? int(a.commit & 0xFFu) : 0;
     for (int t = 0; t < GDN_NTOK; ++t) {
         // conv + SiLU
         float cv = 0.0f;
@@ -139,6 +142,16 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b
                 ys[j] = yv;
             }
         }
+        if (t + 1 == commit_at) {
+            // the state and the conv history after token t (registers: st after this token's update, hst after its
+            // shift); the state is read only before the loop, so in place is safe.  q/k history: the threadgroups
+            // h < h_k own their q/k head (h == hk); v history: every head its own - to hist_out, NOT in place: a q/k
+            // channel is read by the three threadgroups h % h_k at their start, and a threadgroup may start after its
+            // owner has finished (seen on the 10-core M4); the caller swaps the two
+            for (int r = 0; r < RPT; ++r) state[((i0 + r) * hv + h) * S + j] = st[r];
+            if (ch >= 0 && (tid >= 256 || h < int(a.h_k)))
+                for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k];
+        }
         threadgroup_barrier(mem_flags::mem_threadgroup);
         if (a.plane && ib == 0 && (j & 15) == 0) {
             float s0 = 0.0f, s1 = 0.0f;
@@ -146,14 +159,6 @@ kernel void gdn_step_t(device float* state [[buffer(0)]], device float* hist [[b
             y[4 * a.plane + t * (hv * S / 16) + (h * S + j) / 16] = s0 + s1;
         }
     }
-    if (a.commit) {
-        for (int r = 0; r < RPT; ++r) state[((i0 + r) * hv + h) * S + j] = st[r];
-        // q/k history: the threadgroups h < h_k own their q/k head (h == hk); v history: every head its own
-        // to hist_out, NOT in place: a q/k channel is read by the three threadgroups h % h_k at their start, and
-        // a threadgroup may start after its owner has finished (seen on the 10-core M4); the caller swaps the two
-        if (ch >= 0 && (tid >= 256 || h < int(a.h_k)))
-            for (int k = 0; k < 3; ++k) hist_out[ch * 3 + k] = hst[k];
-    }
 }
 
 template [[host_name("gdn_step")]] kernel decltype(gdn_step_t<4>) gdn_step_t<4>;

metal/tests/gdn_parity.cpp

@@ -170,11 +170,31 @@ int main(int argc, char** argv) {
             cx.run({disp(p2, 0, 0, y2)});
             prefix = std::memcmp(y.ptr, y2.ptr, 4 * HV * S * nt) == 0;
         }
-        const bool ok = untouched && same_y && prefix && ey.l2() < 1e-5 && es.l2() < 1e-5 && eh.l2() == 0;
+        // partial commit (commit = 0x100 | k): the state and history after the first k tokens, bitwise those of a
+        // k-token commit; y bitwise the verify window's
+        bool partial = true;
+        for (int k = 1; k < nt; ++k) {
+            mtl::Pipeline pk = cx.pipeline(lib, std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step", mtl::Constants().i(2, k));
+            mtl::Pipeline pn = cx.pipeline(lib, std::getenv("GDN_IB8") ? "gdn_step8" : "gdn_step", mtl::Constants().i(2, nt));
+            std::memcpy(state.ptr, st0.ptr, 4 * st_elems);
+            std::memcpy(hist.ptr, h0.ptr, 4 * C * 3);
+            cx.run({disp(pk, 0, 1, y2)});
+            std::vector<float> sk(state.as<float>(), state.as<float>() + st_elems), hk(hist2.as<float>(), hist2.as<float>() + C * 3);
+            std::memcpy(state.ptr, st0.ptr, 4 * st_elems);
+            cx.run({disp(pn, 0, 0, y)});
+            std::vector<float> yv(y.as<float>(), y.as<float>() + HV * S * nt);
+            std::memcpy(state.ptr, st0.ptr, 4 * st_elems);
+            std::memset(hist2.ptr, 0xFF, 4 * C * 3);
+            cx.run({disp(pn, 0, 0x100 | k, y2)});
+            partial = partial && std::memcmp(state.ptr, sk.data(), 4 * st_elems) == 0 &&
+                      std::memcmp(hist2.ptr, hk.data(), 4 * C * 3) == 0 && std::memcmp(y2.ptr, yv.data(), 4 * HV * S * nt) == 0;
+        }
+        const bool ok = untouched && same_y && prefix && partial && ey.l2() < 1e-5 && es.l2() < 1e-5 && eh.l2() == 0;
         fails += !ok;
         std::printf("gdn parity ntok %d: y rel L2 %.2e, state %.2e, history %.1e; verify leaves state: %s, verify y == "
-                    "commit y: %s, first %d of a longer window bitwise: %s  %s\n", nt, ey.l2(), es.l2(), eh.l2(),
-                    untouched ? "yes" : "NO", same_y ? "yes" : "NO", nt, prefix ? "yes" : "NO", ok ? "PASS" : "FAIL");
+                    "commit y: %s, first %d of a longer window bitwise: %s, partial commits bitwise: %s  %s\n", nt, ey.l2(), es.l2(),
+                    eh.l2(), untouched ? "yes" : "NO", same_y ? "yes" : "NO", nt, prefix ? "yes" : "NO", partial ? "yes" : "NO",
+                    ok ? "PASS" : "FAIL");
     }
     std::printf("%s\n", fails ? "SOME FAILED" : "ALL PASS");
     return fails ? 1 : 0;

metal/tests/gr_parity.cpp

@@ -62,6 +62,11 @@ int main(int argc, char** argv) {
         }
     }
     const int MAXT = 8;
+    // GR_DOWN=SG_KS selects gr_down's variant (default 8_8: 8 simdgroups split one row)
+    const std::string grd = std::getenv("GR_DOWN") ? std::getenv("GR_DOWN") : "8_8";
+    int grd_sg = 8, grd_ks = 8;
+    std::sscanf(grd.c_str(), "%d_%d", &grd_sg, &grd_ks);
+    const int grd_rs = grd_sg / grd_ks;
     mtl::Buffer R = cx.alloc(4 * MAXT * NIN), R0 = cx.alloc(4 * MAXT * NIN), bo = cx.alloc(4 * MAXT * N),
                 inj = cx.alloc(4 * MAXT * HC), xn = cx.alloc(4 * MAXT * NIN), lo = cx.alloc(4 * MAXT * LR),
                 injo = cx.alloc(4 * MAXT * HC), xp = cx.alloc(4 * (4 * MAXT * N + MAXT * N / 16)),
@@ -73,7 +78,7 @@ int main(int argc, char** argv) {
     int fails = 0;
     for (int nt : ntoks) {
         mtl::Pipeline pn = cx.pipeline(lib, "gr_norm"),
-                      pd = cx.pipeline(lib, q8 ? "gr_down_q8_8_8" : "gr_down_8_8", mtl::Constants().i(1, nt)),
+                      pd = cx.pipeline(lib, std::string(q8 ? "gr_down_q8_" : "gr_down_") + grd, mtl::Constants().i(1, nt)),
                       pu = cx.pipeline(lib, q8 ? "gr_up_q8" : "gr_up", mtl::Constants().i(1, nt));
         for (int apply = 0; apply < 2; ++apply) {
             GrArgs a{N, HC, LR, (uint32_t)nt, (uint32_t)apply, 1e-6f};
@@ -95,8 +100,8 @@ int main(int argc, char** argv) {
                     e.p = &pd;
                     e.args = {mtl::Arg::B(W, o), mtl::Arg::B(W, o + wd_bytes), mtl::Arg::B(xn), mtl::Arg::B(lo),
                               mtl::Arg::B(injo), mtl::Arg::V(a)};
-                    e.grid[0] = LR + HC;   // RS = 1 row per threadgroup (8 simdgroups split it)
-                    e.tg[0] = 256;
+                    e.grid[0] = (LR + HC + grd_rs - 1) / grd_rs;   // RS = SG / KS rows per threadgroup
+                    e.tg[0] = 32 * grd_sg;
                     ds.push_back(e);
                 }
                 if (which & 4) {