From a34feeb0f9429db86bf941aeee9e7ef1fc9f1258 Mon Sep 17 00:00:00 2001
From: Marvin <nsanchez1980@gmail.com>
Date: Sat, 5 Sep 2026 21:13:04 -0300
Subject: [PATCH] cuda: chunked WY delta-net kernel v1 (opt-in DELTA_WY_CUDA=1)
 + harness --cuda mode

---
 ggml/src/ggml-cuda/delta-net.cu | 313 ++++++++++++++++++++++++++++++++
 tests/test-delta-chunk.cpp      |  15 +-
 2 files changed, 321 insertions(+), 7 deletions(-)

diff --git a/ggml/src/ggml-cuda/delta-net.cu b/ggml/src/ggml-cuda/delta-net.cu
index a7d509c1..c701a96d 100644
--- a/ggml/src/ggml-cuda/delta-net.cu
+++ b/ggml/src/ggml-cuda/delta-net.cu
@@ -1,5 +1,6 @@
 #include "common.cuh"
 #include "delta-net.cuh"
+#include <cstdio>
 #include <cstdlib>
 #include <cstring>
 
@@ -182,6 +183,266 @@ __global__ void delta_net_recurrent_f32(
     }
 }
 
+
+// Chunked WY parallel-scan delta net for long prefills (qwen4exp PP).
+// Validated math: tests/test-delta-chunk.cpp (candidate C) must stay green —
+// any change here needs a matching change there first. One CUDA block per
+// (seq, v-head); 64-token chunks (32-token fallback tier); per chunk:
+//   scalars b/d/logP per token (shared, 3*C floats; logP sequential by thread 0)
+//   L/KQ dot rows recomputed per use (no L/KQ storage; 64-way parallel dots,
+//     results broadcast via shared row, kh/qt stream from global/L1)
+//   triangular solve for E (E rows persist in shared; flat-partitioned updates)
+//   outputs + single-GEMM assembly of S_C (valid: dispatch guarantees
+//     saved_states == NULL, so no per-step states are needed)
+//   guard (max decay-ratio / max |e|, nan-aware): wild dynamics fall back to
+//     the exact sequential inner loop for that chunk (same block, S still
+//     intact in shared). Operating regime never trips it.
+// v1 is OPT-IN (DELTA_WY_CUDA=1) and correctness-first: SIMT code, no tensor
+// cores/cp.async (v2 after measurements). Shared budget d=128/C=64 is
+// (16384+8192+192+64)*4 = 99328 bytes (opt-in attr); denial or smaller
+// budgets drop to C=32 tier, then to the sequential kernel. All silent and
+// exact in every tier.
+#define DELTA_WY_CHUNK 64
+#define DELTA_WY_CHUNK_SMALL 32
+#define WY_GUARD_D 1e4f
+#define WY_GUARD_E 1e4f
+
+template <int HEAD_DIM>
+__global__ void delta_net_chunked_f32(
+    const float * __restrict__ q,
+    const float * __restrict__ k,
+    const float * __restrict__ v,
+    const float * __restrict__ g,
+    const float * __restrict__ beta_in,
+    const float * __restrict__ state_in,
+    float * __restrict__ dst,
+    float * __restrict__ state_out,
+    const int64_t n_heads,
+    const int64_t gqa_ratio,
+    const int repeat_type,
+    const int64_t n_tokens,
+    const int64_t n_seqs,
+    size_t vnb1, size_t vnb2, size_t vnb3,
+    const int chunk) {
+    const int batch_idx = (int)(blockIdx.x / n_heads);
+    const int head_idx  = (int)(blockIdx.x % n_heads);
+    const int head_idx_kq = repeat_type == 0 ? (int)(head_idx / gqa_ratio) : (int)(head_idx % (n_heads/gqa_ratio));
+    const int tid = (int)threadIdx.x;
+    const int nthreads = (int)blockDim.x;
+
+    // identical strides to delta_net_recurrent_f32
+    const int64_t qkv_stride_token = HEAD_DIM;
+    const int64_t qkv_stride_head = HEAD_DIM * n_tokens;
+    const int64_t qkv_stride_batch = HEAD_DIM * n_tokens * n_heads;
+    const int64_t qkv_stride_batch_kq = qkv_stride_batch / gqa_ratio;
+    const int64_t g_stride_batch = n_tokens * n_heads;
+    const int64_t state_head_offset = head_idx * HEAD_DIM * HEAD_DIM;
+    const int64_t state_batch_stride = HEAD_DIM * HEAD_DIM * n_heads;
+    const int64_t out_token_stride = HEAD_DIM * n_heads;
+
+    const float * q_base = q + batch_idx * qkv_stride_batch_kq + head_idx_kq * qkv_stride_head;
+    const float * k_base = k + batch_idx * qkv_stride_batch_kq + head_idx_kq * qkv_stride_head;
+    const float * v_base = v + batch_idx * vnb3 + head_idx * vnb2;
+    const float * g_base = g + batch_idx * g_stride_batch + head_idx;
+    const float * beta_base = beta_in + batch_idx * g_stride_batch + head_idx;
+    const float * state_src = state_in + batch_idx * state_batch_stride + state_head_offset;
+    float * out_base = dst + batch_idx * (HEAD_DIM * n_heads * n_tokens) + head_idx * HEAD_DIM;
+    float * state_dst = state_out + batch_idx * state_batch_stride + state_head_offset;
+
+    const float scale = rsqrtf((float)HEAD_DIM);
+
+    extern __shared__ float smem[];
+    float * sS = smem;                             // HEAD_DIM*HEAD_DIM (carry)
+    float * sE = sS + HEAD_DIM * HEAD_DIM;         // chunk*HEAD_DIM
+    float * sScl = sE + (int64_t)chunk * HEAD_DIM; // 3*chunk (b,d,logP)
+    float * sRow = sScl + 3 * chunk;               // chunk (dot rows; [0..7] reused by guard reduce)
+    float * sBv = sScl, * sDv = sScl + chunk, * sLogP = sScl + 2 * chunk;
+
+    for (int64_t x = tid; x < (int64_t)HEAD_DIM * HEAD_DIM; x += nthreads) sS[x] = state_src[x];
+    __syncthreads();
+
+    for (int64_t t0 = 0; t0 < n_tokens; t0 += chunk) {
+        const int64_t rem = n_tokens - t0;
+        const int C = (int)(rem < chunk ? rem : chunk);
+
+        // ---- scalars (logP sequential by thread 0: prefix dependency)
+        for (int i = tid; i < C; i += nthreads) {
+            const float g_val = g_base[(t0 + i) * n_heads];
+            const float b_raw = beta_base[(t0 + i) * n_heads];
+            sBv[i] = 1.0f / (1.0f + expf(-b_raw));
+            sDv[i] = expf(fminf(g_val, 50.0f));
+        }
+        if (tid == 0) {
+            float acc = 0.0f;
+            for (int i = 0; i < C; ++i) {
+                acc += fminf(g_base[(t0 + i) * n_heads], 50.0f);
+                sLogP[i] = acc;
+            }
+        }
+        __syncthreads();
+
+        // ---- solve for E (C sequential steps; flat-partitioned element updates)
+        float thr_maxD = 0.0f, thr_maxE = 0.0f;
+        for (int i = 0; i < C; ++i) {
+            // L[i][j] = b_i * exp(logP[i]-logP[j]) * (kh_i.kh_j), j<i; thread j
+            if (tid < i) {
+                const int j = tid;
+                float dot = 0.0f;
+                for (int d = 0; d < HEAD_DIM; ++d) {
+                    dot += k_base[(t0 + i) * qkv_stride_token + d]
+                         * k_base[(t0 + j) * qkv_stride_token + d];
+                }
+                const float Dij = expf(sLogP[i] - sLogP[j]);
+                if (!(Dij <= thr_maxD)) thr_maxD = Dij; // nan-aware max
+                sRow[j] = sBv[i] * Dij * dot;
+            }
+            __syncthreads();
+            // e_i[d] for owned d; ck recomputed (S rows in shared, kh streamed)
+            for (int dd = tid; dd < HEAD_DIM; dd += nthreads) {
+                float ck = 0.0f;
+                for (int e = 0; e < HEAD_DIM; ++e) ck += sS[dd + e * HEAD_DIM] * k_base[(t0 + i) * qkv_stride_token + e];
+                const float Pi = expf(sLogP[i]); // underflow to 0 is correct
+                const float * vv = v_base + (t0 + i) * vnb1;
+                float ev = sBv[i] * vv[dd] - sBv[i] * Pi * ck;
+                for (int j = 0; j < i; ++j) ev -= sRow[j] * sE[j * HEAD_DIM + dd];
+                sE[i * HEAD_DIM + dd] = ev;
+                const float ae = fabsf(ev);
+                if (!(ae <= thr_maxE)) thr_maxE = ae; // nan-aware max
+            }
+            __syncthreads();
+        }
+
+        // ---- guard reduce (warp shuffle -> sRow[0..7] -> scalar mailbox sRow[0])
+        {
+            float md = thr_maxD, me = thr_maxE;
+            for (int off = 16; off > 0; off >>= 1) {
+                const float od = __shfl_down_sync(0xffffffff, md, off);
+                const float oe = __shfl_down_sync(0xffffffff, me, off);
+                if (!(od <= md)) md = od;
+                if (!(oe <= me)) me = oe;
+            }
+            const int wid = tid / 32, lane = tid % 32;
+            if (lane == 0) { sRow[wid] = md; sRow[8 + wid] = me; }
+            __syncthreads();
+            if (tid == 0) {
+                float gmd = 0.0f, gme = 0.0f;
+                const int nwarps = nthreads / 32;
+                for (int w = 0; w < nwarps; ++w) {
+                    if (!(sRow[w] <= gmd)) gmd = sRow[w];
+                    if (!(sRow[8 + w] <= gme)) gme = sRow[8 + w];
+                }
+                sRow[0] = (gmd <= WY_GUARD_D && gme <= WY_GUARD_E) ? 1.0f : 0.0f;
+            }
+            __syncthreads();
+        }
+
+        if (sRow[0] < 0.5f) {
+            // ---- sequential fallback for this chunk (exact old math).
+            // Fully thread-local per owned row (no shared scratch): dots are
+            // per-row serial loops, only the scalar score needs broadcast.
+            for (int i = 0; i < C; ++i) {
+                float score = 0.0f;
+                if (tid == 0) {
+                    float s = 0.0f;
+                    for (int d = 0; d < HEAD_DIM; ++d) {
+                        s += k_base[(t0 + i) * qkv_stride_token + d] * q_base[(t0 + i) * qkv_stride_token + d];
+                    }
+                    score = s * scale;
+                    sRow[0] = score;
+                    sRow[1] = 1.0f / (1.0f + expf(-beta_base[(t0 + i) * n_heads]));
+                    sRow[2] = expf(fminf(g_base[(t0 + i) * n_heads], 50.0f));
+                }
+                __syncthreads();
+                {
+                    const float sc = sRow[0], bv = sRow[1], dv = sRow[2];
+                    float * out_t = out_base + (t0 + i) * out_token_stride;
+                    const float * vv = v_base + (t0 + i) * vnb1;
+                    const float * kk = k_base + (t0 + i) * qkv_stride_token;
+                    const float * qq = q_base + (t0 + i) * qkv_stride_token;
+                    for (int dd = tid; dd < HEAD_DIM; dd += nthreads) {
+                        float sk = 0.0f, sq = 0.0f;
+                        for (int e = 0; e < HEAD_DIM; ++e) {
+                            sk += sS[dd + e * HEAD_DIM] * kk[e];
+                            sq += sS[dd + e * HEAD_DIM] * qq[e];
+                        }
+                        const float vnew = bv * (vv[dd] - sk * dv);
+                        out_t[dd] = sq * dv * scale + vnew * sc;
+                        for (int cc = 0; cc < HEAD_DIM; ++cc) {
+                            float s = dv * sS[dd + cc * HEAD_DIM] + vnew * kk[cc];
+                            sS[dd + cc * HEAD_DIM] = fminf(fmaxf(s, -1e6f), 1e6f);
+                        }
+                    }
+                }
+                __syncthreads(); // S read/write hazard across threads per token
+            }
+        } else {
+            // ---- fast path outputs: KQ rows recomputed per step (no storage),
+            // aq recomputed per owned element, ci recomputed per thread (128
+            // MACs; cheaper than a broadcast round-trip).
+            for (int i = 0; i < C; ++i) {
+                // KQ[i][j] = kh_j . qt_i, j<i; thread j (idle beyond i)
+                if (tid < i) {
+                    const int j = tid;
+                    float kq = 0.0f;
+                    for (int e = 0; e < HEAD_DIM; ++e) {
+                        kq += k_base[(t0 + j) * qkv_stride_token + e]
+                            * q_base[(t0 + i) * qkv_stride_token + e] * scale;
+                    }
+                    // NOTE: qt carries *scale already (matches candidate C)
+                    sRow[j] = kq;
+                }
+                __syncthreads();
+                {
+                    // ci = kh_i . qt_i (recomputed per thread, no broadcast).
+                    // NOTE: pim/logPim/D[*] all come from the staged sLogP prefix
+                    // sums — bitwise identical to candidate C in the harness.
+                    float ci = 0.0f;
+                    for (int e = 0; e < HEAD_DIM; ++e) {
+                        ci += k_base[(t0 + i) * qkv_stride_token + e]
+                            * q_base[(t0 + i) * qkv_stride_token + e] * scale;
+                    }
+                    const float pim = i ? expf(sLogP[i - 1]) : 1.0f;
+                    const float dv = sDv[i];
+                    float * out_t = out_base + (t0 + i) * out_token_stride;
+                    for (int dd = tid; dd < HEAD_DIM; dd += nthreads) {
+                        float aq = 0.0f;
+                        for (int e = 0; e < HEAD_DIM; ++e) {
+                            aq += sS[dd + e * HEAD_DIM] * q_base[(t0 + i) * qkv_stride_token + e] * scale;
+                        }
+                        float a = pim * aq;
+                        for (int j = 0; j < i; ++j) {
+                            a += expf(sLogP[i - 1] - sLogP[j]) * sRow[j] * sE[j * HEAD_DIM + dd];
+                        }
+                        out_t[dd] = dv * a + ci * sE[i * HEAD_DIM + dd];
+                    }
+                }
+                __syncthreads();
+            }
+            // ---- GEMM assembly of S_C (single-GEMM form; valid: no per-step
+            // states needed since dispatch guarantees saved_states == NULL).
+            // S_C = P[C-1]*S_in + sum_i D[C-1][i] e_i kh_i^T, clamped.
+            {
+                const float pC = expf(sLogP[C - 1]);
+                for (int64_t x = tid; x < (int64_t)HEAD_DIM * HEAD_DIM; x += nthreads) {
+                    const int row = (int)(x % HEAD_DIM);
+                    const int col = (int)(x / HEAD_DIM);
+                    float s = pC * sS[x];
+                    for (int i = 0; i < C; ++i) {
+                        s += expf(sLogP[C - 1] - sLogP[i]) * sE[i * HEAD_DIM + row]
+                           * k_base[(t0 + i) * qkv_stride_token + col];
+                    }
+                    sS[x] = fminf(fmaxf(s, -1e6f), 1e6f);
+                }
+            }
+            __syncthreads();
+        }
+        __syncthreads();
+    }
+
+    // write back final state
+    for (int64_t x = tid; x < (int64_t)HEAD_DIM * HEAD_DIM; x += nthreads) state_dst[x] = sS[x];
+}
+
 static void delta_net_f32_cuda(
     const float * q,
     const float * k,
@@ -209,6 +470,58 @@ static void delta_net_f32_cuda(
         GGML_ABORT("Unsupported delta net head size");
     }
 
+    // Chunked WY path (opt-in via DELTA_WY_CUDA=1; validated math in
+    // tests/test-delta-chunk.cpp). Requirements: long enough prefill for at
+    // least 2 chunks, no per-step checkpoints (GEMM assembly form), and enough
+    // dynamic shared memory (else transparent fallback to sequential below).
+    // Falls back silently-exact in every other case: default behavior unchanged.
+    {
+        const char * wy_env = getenv("DELTA_WY_CUDA");
+        const bool wy_wanted = wy_env && wy_env[0] == '1';
+        if (wy_wanted && n_tokens > DELTA_WY_CHUNK && saved_states == nullptr) {
+            int chunk = DELTA_WY_CHUNK;
+            size_t need = ((size_t)head_dim * head_dim + (size_t)chunk * head_dim + 4 * chunk) * sizeof(float);
+            // shared tier selection: full chunks, then half chunks, then sequential.
+            // (cudaFuncSetAttribute form mirrors fattn-new-mma.cu usage.)
+            auto wy_attr_ok = [&](size_t bytes) -> bool {
+                cudaError_t err = (head_dim == 64)
+                    ? cudaFuncSetAttribute(delta_net_chunked_f32<64>,
+                            cudaFuncAttributeMaxDynamicSharedMemorySize, (int)bytes)
+                    : cudaFuncSetAttribute(delta_net_chunked_f32<128>,
+                            cudaFuncAttributeMaxDynamicSharedMemorySize, (int)bytes);
+                return err == cudaSuccess;
+            };
+            if (!wy_attr_ok(need)) {
+                chunk = DELTA_WY_CHUNK_SMALL;
+                need = ((size_t)head_dim * head_dim + (size_t)chunk * head_dim + 4 * chunk) * sizeof(float);
+                if (!wy_attr_ok(need)) {
+                    chunk = 0;
+                }
+            }
+            if (chunk > 0) {
+                static bool wy_note_done = false;
+                if (!wy_note_done) {
+                    wy_note_done = true;
+                    fprintf(stderr, "delta_net: chunked WY path, C=%d, shared %zu bytes\n", chunk, need);
+                }
+                constexpr int wy_threads = 256;
+                const int wy_blocks = (int)(n_seqs * n_heads);
+                if (head_dim == 64) {
+                    delta_net_chunked_f32<64><<<wy_blocks, wy_threads, need, stream>>>(
+                        q, k, v, g, beta, state_in, dst, state_out,
+                        n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, chunk);
+                } else {
+                    delta_net_chunked_f32<128><<<wy_blocks, wy_threads, need, stream>>>(
+                        q, k, v, g, beta, state_in, dst, state_out,
+                        n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, chunk);
+                }
+                CUDA_CHECK(cudaGetLastError());
+                return;
+            }
+            // else: fall through to the sequential kernel below
+        }
+    }
+
     GGML_ASSERT(head_dim % WARP_SIZE == 0);
     const int num_blocks = n_seqs * n_heads * (head_dim/WARP_SIZE);
     const size_t smem_size = 2 * head_dim * sizeof(float);
diff --git a/tests/test-delta-chunk.cpp b/tests/test-delta-chunk.cpp
index 35fd6fd9..9c18bc49 100644
--- a/tests/test-delta-chunk.cpp
+++ b/tests/test-delta-chunk.cpp
@@ -420,12 +420,14 @@ static void solve_chunk_wy(
     const int C = c1 - c0;
     const float scale = 1.0f / sqrtf((float)hd);
     // per-token ingredients (same factoring as run_token)
-    std::vector<float> kh(C * hd), qt(C * hd), bv(C), dv(C), pv(C), ck(C * hd), aq(C * hd);
+    // Absolute decay products are ALWAYS exp(logP[.]) (never progressive
+    // products): bitwise identical to what the CUDA kernel stages in shared.
+    std::vector<float> kh(C * hd), qt(C * hd), bv(C), dv(C), ck(C * hd), aq(C * hd);
     // log-decay prefix sums: every decay RATIO is evaluated as exp(logP[i]-logP[j]).
     // Raw products P[i] underflow to 0 over 64 steps with decay<1, turning later
     // P[i]/P[j] ratios into 0/0 = nan. exp-of-difference is exact-or-zero, which
-    // is the mathematically right answer (true ratio ~0). Absolute P[i] (=pv[i])
-    // may still underflow to 0 in f_i/S_C/output terms — also correct there.
+    // is the mathematically right answer (true ratio ~0). Absolute P[i]
+    // (=exp(logP[i])) may still underflow to 0 in f_i/S_C/output terms — also correct there.
     std::vector<float> logP(C);
     for (int i = 0; i < C; ++i) {
         const int tt = c0 + i;
@@ -441,7 +443,6 @@ static void solve_chunk_wy(
         kn = 1.0f / sqrtf(kn + 1e-12f);
         bv[i] = 1.0f / (1.0f + expf(-boff));
         dv[i] = expf(fminf(goff, 50.0f));
-        pv[i] = dv[i] * (i ? pv[i - 1] : 1.0f);
         logP[i] = fminf(goff, 50.0f) + (i ? logP[i - 1] : 0.0f);
         for (int d = 0; d < hd; ++d) {
             kh[i * hd + d] = k[d] * kn;
@@ -479,7 +480,7 @@ static void solve_chunk_wy(
         const int tt = c0 + i;
         const float * v = t.v.data() + (((size_t)b * c.hv + h) * c.nt + tt) * hd;
         for (int d = 0; d < hd; ++d) {
-            float e = bv[i] * v[d] - bv[i] * pv[i] * ck[i * hd + d];
+            float e = bv[i] * v[d] - bv[i] * expf(logP[i]) * ck[i * hd + d];
             for (int j = 0; j < i; ++j) e -= L[i * C + j] * E[j * hd + d];
             E[i * hd + d] = e;
             const float ae = fabsf(e);
@@ -511,7 +512,7 @@ static void solve_chunk_wy(
     memcpy(S.data(), s_in, stsz * sizeof(float));
     for (int i = 0; i < C; ++i) {
         const int tt = c0 + i;
-        const float pim = i ? pv[i - 1] : 1.0f;
+        const float pim = i ? expf(logP[i - 1]) : 1.0f;
         const float logPim = i ? logP[i - 1] : 0.0f;
         float ci = 0.0f;
         for (int d = 0; d < hd; ++d) ci += kh[i * hd + d] * qt[i * hd + d];
@@ -540,7 +541,7 @@ static void solve_chunk_wy(
         float md = 0.0f;
         for (int row = 0; row < hd; ++row) {
             for (int col = 0; col < hd; ++col) {
-                float s = pv[C - 1] * s_in[row + col * hd];
+                float s = expf(logP[C - 1]) * s_in[row + col * hd];
                 for (int i = 0; i < C; ++i) {
                     s += expf(logP[C - 1] - logP[i]) * E[i * hd + row] * kh[i * hd + col];
                 }