手撕KDA

前置知识：KDA算法公式、triton、cuda、c++

从 FLA（flash-linear-attention）仓库的代码入手

仓库链接：fla-org/flash-linear-attention: 🚀 Efficient implementations for emerging model architectures

0 公式回顾

1 Naive

代码链接：flash-linear-attention/fla/ops/kda/naive.py at main · fla-org/flash-linear-attention

# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
#   https://github.com/fla-org/flash-linear-attention/graphs/contributors

import torch
from einops import rearrange


def naive_chunk_kda(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    scale: float | None = None,
    initial_state: torch.Tensor | None = None,
    output_final_state: bool = False,
    chunk_size: int = 64,
):
    r"""
    Args:
        q (torch.Tensor):
            Queries of shape ``[B, T, H, K]``.
        k (torch.Tensor):
            Keys of shape ``[B, T, H, K]``.
        v (torch.Tensor):
            Values of shape ``[B, T, HV, V]``. ``HV`` must be divisible by ``H``.
        g (torch.Tensor):
            Per-dimension decay gates (log-space) of shape ``[B, T, HV, K]``.
        beta (torch.Tensor):
            Beta scalars of shape ``[B, T, HV]``.
        scale (Optional[float]):
            Scale factor. Defaults to ``1 / sqrt(K)``.
        initial_state (Optional[torch.Tensor]):
            Initial state of shape ``[B, HV, K, V]``.
        output_final_state (bool):
            Whether to return the final state.
        chunk_size (int):
            Chunk size for the chunked computation. Default: 64.

    Returns:
        A tuple ``(o, S)`` where ``o`` has shape ``[B, T, HV, V]`` and
        ``S`` has shape ``[B, HV, K, V]`` if ``output_final_state`` else ``None``.
    """
    dtype = v.dtype
    B, T, H, K, HV, V = *q.shape, v.shape[2], v.shape[-1]
    G = HV // H
    BT = chunk_size
    NT = T // BT
    if scale is None:
        scale = K ** -0.5
    assert T % BT == 0

    # Rearrange into chunks: [B, head, NT, BT, ...]
    q, k = [rearrange(x, 'b (n c) h ... -> b h n c ...', c=BT).to(torch.float) for x in [q, k]]
    v, g, beta = [rearrange(x, 'b (n c) h ... -> b h n c ...', c=BT).to(torch.float) for x in [v, g, beta]]
    # Expand q/k to value head dim for GVA: [B, H, ...] -> [B, HV, ...]
    q = q.repeat_interleave(G, dim=1) * scale  # [B, HV, NT, BT, K]
    k = k.repeat_interleave(G, dim=1)          # [B, HV, NT, BT, K]
    g = g.cumsum(-2)

    # note that diagonal is masked.
    mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0)

    # Akk uses k (expanded to HV) and g (per value head)
    A = torch.zeros(*g.shape[:-1], BT, dtype=torch.float, device=q.device)
    for i in range(BT):
        k_i = k[..., i, :]
        g_i = g[..., i:i+1, :]
        A[..., i] = torch.einsum('... c d, ... d -> ... c', k * (g - g_i).exp(), k_i)
    A = A * beta[..., None]

    A = -A.masked_fill(mask, 0)
    for i in range(1, BT):
        A[..., i, :i] = A[..., i, :i].clone() + (A[..., i, :, None].clone() * A[..., :, :i].clone()).sum(-2)
    A = (A + torch.eye(BT, dtype=torch.float, device=q.device)) * beta[..., None, :]

    w = A @ (g.exp() * k)
    u = A @ v

    S = k.new_zeros(B, HV, K, V).to(q)
    if initial_state is not None:
        S += initial_state
    o = torch.zeros_like(v)
    mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1)
    for i in range(0, NT):
        # [B, HV, BT, ...]
        q_i = q[:, :, i]      # [B, HV, BT, K]
        k_i = k[:, :, i]      # [B, HV, BT, K]
        u_i = u[:, :, i]        # [B, HV, BT, V]
        g_i = g[:, :, i]        # [B, HV, BT, K]
        w_i = w[:, :, i]        # [B, HV, BT, K]
        # Aqk: per value head (q from qk head, g from value head, k from qk head)
        Aqk = torch.zeros(B, HV, BT, BT, dtype=torch.float, device=q.device)
        for j in range(BT):
            k_j = k[:, :, i, j]
            g_j = g[:, :, i, j:j+1, :]
            Aqk[..., j] = torch.einsum('... c d, ... d -> ... c', q_i * (g_i - g_j).exp(), k_j)
        Aqk = Aqk.masked_fill(mask, 0)
        v_i = u_i - w_i @ S
        o[:, :, i] = (q_i * g_i.exp()) @ S + Aqk @ v_i
        S = S * rearrange(g_i[:, :, -1].exp(), 'b h k -> b h k 1')
        S += rearrange((g_i[:, :, -1:] - g_i).exp() * k_i, 'b h c k -> b h k c') @ v_i
    if not output_final_state:
        S = None
    return rearrange(o, 'b h n c d -> b (n c) h d').to(dtype), S

1.1 pytorch 充满着向量化的计算，将其翻译成C++，代码如下

/**
 * Pure C++ implementation of naive_chunk_kda
 *
 * Key-Delta-Attention (KDA) chunked forward pass.
 * This is a reference implementation matching the Python naive_chunk_kda function.
 *
 * Uses h (per-chunk hidden state) instead of S (single final state) for clarity.
 * h[n] stores the hidden state BEFORE processing chunk n, matching FLA's convention.
 *
 * Data layout (all row-major):
 *   q:     [B][T][H][K]
 *   k:     [B][T][H][K]
 *   v:     [B][T][HV][V]
 *   g:     [B][T][HV][K]
 *   beta:  [B][T][HV]
 *   o:     [B][T][HV][V]
 *   h:     [B][HV][NT][K][V]   (hidden state per chunk)
 *   h0:    [B][HV][K][V]       (initial state)
 *   ht:    [B][HV][K][V]       (final state)
 */

#include <vector>
#include <cmath>
#include <cassert>
#include <algorithm>
#include <iostream>

// ============================================================================
// Tensor helper: a simple row-major multidimensional array
// ============================================================================
template <typename T>
struct Tensor {
    std::vector<int> shape;
    std::vector<T> data;

    Tensor() = default;

    // Create an uninitialized tensor with given shape
    Tensor(std::vector<int> shape_) : shape(std::move(shape_)) {
        int total = 1;
        for (int s : shape) total *= s;
        data.resize(total, T(0));
    }

    // Create a tensor filled with a value
    Tensor(std::vector<int> shape_, T val) : shape(std::move(shape_)) {
        int total = 1;
        for (int s : shape) total *= s;
        data.resize(total, val);
    }

    int size() const {
        int total = 1;
        for (int s : shape) total *= s;
        return total;
    }

    // Compute flat index from multi-dim indices
    template <typename... Idx>
    int index(Idx... idxs) const {
        std::vector<int> v = {idxs...};
        int flat = 0;
        int stride = 1;
        for (int d = (int)shape.size() - 1; d >= 0; --d) {
            flat += v[d] * stride;
            stride *= shape[d];
        }
        return flat;
    }

    template <typename... Idx>
    T& operator()(Idx... idxs) { return data[index(idxs...)]; }

    template <typename... Idx>
    const T& operator()(Idx... idxs) const { return data[index(idxs...)]; }
};

// ============================================================================
// Helper: compute exp2(x) = 2^x
// ============================================================================
static inline float exp2f_fast(float x) {
    return exp2f(x);  // or powf(2.0f, x)
}

// ============================================================================
// naive_chunk_kda in pure C++
// ============================================================================
struct ChunkKdaResult {
    Tensor<float> o;      // [B][T][HV][V]
    Tensor<float> h;      // [B][HV][NT][K][V]  (hidden state per chunk)
    Tensor<float> ht;     // [B][HV][K][V]      (final state) or empty
};

ChunkKdaResult naive_chunk_kda(
    const Tensor<float>& q,      // [B][T][H][K]
    const Tensor<float>& k,      // [B][T][H][K]
    const Tensor<float>& v,      // [B][T][HV][V]
    const Tensor<float>& g,      // [B][T][HV][K]
    const Tensor<float>& beta,   // [B][T][HV]
    float scale,
    const Tensor<float>* h0,     // nullable, [B][HV][K][V]  (initial state)
    bool output_final_state,
    int chunk_size
) {
    int B  = q.shape[0];
    int T  = q.shape[1];
    int H  = q.shape[2];
    int K  = q.shape[3];
    int HV = v.shape[2];
    int V  = v.shape[3];
    int G  = HV / H;
    int BT = chunk_size;
    int NT = T / BT;

    assert(T % BT == 0);
    assert(HV % H == 0);

    // -----------------------------------------------------------------------
    // Step 0: Expand q, k from H heads to HV heads (GQA repeat)
    //         Rearrange from [B][T][H][K] -> [B][HV][NT][BT][K]
    //         q is also scaled by `scale`
    // -----------------------------------------------------------------------

    // qe: [B][HV][NT][BT][K]  (expanded q with scale)
    // ke: [B][HV][NT][BT][K]  (expanded k)
    // ve: [B][HV][NT][BT][V]  (rearranged v)
    // ge: [B][HV][NT][BT][K]  (rearranged g, will be cumsum'd)
    // be: [B][HV][NT][BT]     (rearranged beta)
    Tensor<float> qe({B, HV, NT, BT, K}, 0.0f);
    Tensor<float> ke({B, HV, NT, BT, K}, 0.0f);
    Tensor<float> ve({B, HV, NT, BT, V}, 0.0f);
    Tensor<float> ge({B, HV, NT, BT, K}, 0.0f);
    Tensor<float> be({B, HV, NT, BT},    0.0f);

    for (int b = 0; b < B; b++) {
        for (int t = 0; t < T; t++) {
            int n = t / BT;       // chunk index
            int c = t % BT;       // position within chunk
            for (int hv = 0; hv < HV; hv++) {
                int h = hv / G;   // corresponding q/k head
                for (int kk = 0; kk < K; kk++) {
                    qe(b, hv, n, c, kk) = q(b, t, h, kk) * scale;
                    ke(b, hv, n, c, kk) = k(b, t, h, kk);
                }
                for (int vv = 0; vv < V; vv++) {
                    ve(b, hv, n, c, vv) = v(b, t, hv, vv);
                }
                for (int kk = 0; kk < K; kk++) {
                    ge(b, hv, n, c, kk) = g(b, t, hv, kk);
                }
                be(b, hv, n, c) = beta(b, t, hv);
            }
        }
    }

    // -----------------------------------------------------------------------
    // Step 1: Compute g cumsum along the BT dimension (within each chunk)
    //         ge(b, hv, n, c, k) = sum_{j=0}^{c} g(b, hv, n, j, k)
    // -----------------------------------------------------------------------
    for (int b = 0; b < B; b++) {
        for (int hv = 0; hv < HV; hv++) {
            for (int n = 0; n < NT; n++) {
                for (int kk = 0; kk < K; kk++) {
                    float cumsum = 0.0f;
                    for (int c = 0; c < BT; c++) {
                        cumsum += ge(b, hv, n, c, kk);
                        ge(b, hv, n, c, kk) = cumsum;
                    }
                }
            }
        }
    }

    // -----------------------------------------------------------------------
    // Step 2: Compute Akk matrix (strictly lower triangular, then solve)
    //
    // Akk[i][j] = beta[i] * sum_k ( k[i][k] * 2^{g[i][k]-g[j][k]} * k[j][k] )
    // for i > j
    //
    // Then: A = -(I - M) where M = Akk (strict lower tri)
    //       Forward substitution: A becomes (I-M)^{-1}
    //       Then multiply by beta column-wise
    // -----------------------------------------------------------------------

    // Akk: [B][HV][NT][BT][BT]
    Tensor<float> Akk({B, HV, NT, BT, BT}, 0.0f);

    for (int b = 0; b < B; b++) {
        for (int hv = 0; hv < HV; hv++) {
            for (int n = 0; n < NT; n++) {
                // Compute raw Akk
                for (int i = 0; i < BT; i++) {
                    for (int j = 0; j < BT; j++) {
                        float val = 0.0f;
                        for (int kk = 0; kk < K; kk++) {
                            // k[i] * 2^{g[i]-g[j]} * k[j]
                            float decay = exp2f_fast(ge(b, hv, n, i, kk) - ge(b, hv, n, j, kk));
                            val += ke(b, hv, n, i, kk) * decay * ke(b, hv, n, j, kk);
                        }
                        // Upper triangular (i < j) is zero, diagonal (i==j) is also zero
                        if (i > j) {
                            Akk(b, hv, n, i, j) = -val * be(b, hv, n, i);
                        }
                        // i <= j: stays 0
                    }
                }

                // Forward substitution: A becomes (I - M)^{-1}
                // for i in range(1, BT):
                //     A[i, :i] += (A[i, :, None] * A[:, :i]).sum(dim=-2)
                for (int i = 1; i < BT; i++) {
                    for (int j = 0; j < i; j++) {
                        float sum = 0.0f;
                        for (int m = 0; m < i; m++) {
                            sum += Akk(b, hv, n, i, m) * Akk(b, hv, n, m, j);
                        }
                        Akk(b, hv, n, i, j) += sum;
                    }
                }

                // Add identity and multiply by beta (column-wise)
                for (int i = 0; i < BT; i++) {
                    for (int j = 0; j < BT; j++) {
                        if (i == j) {
                            Akk(b, hv, n, i, j) += 1.0f;
                        }
                        Akk(b, hv, n, i, j) *= be(b, hv, n, j);
                    }
                }
            }
        }
    }

    // -----------------------------------------------------------------------
    // Step 3: Compute w and u from Akk
    //   w = Akk @ (2^g * k)    -- delta rule weight
    //   u = Akk @ v             -- WY corrected value
    // -----------------------------------------------------------------------

    // w: [B][HV][NT][BT][K]
    // u: [B][HV][NT][BT][V]
    Tensor<float> w({B, HV, NT, BT, K}, 0.0f);
    Tensor<float> u({B, HV, NT, BT, V}, 0.0f);

    for (int b = 0; b < B; b++) {
        for (int hv = 0; hv < HV; hv++) {
            for (int n = 0; n < NT; n++) {
                // w = Akk @ (exp2(g) * k)
                for (int i = 0; i < BT; i++) {
                    for (int kk = 0; kk < K; kk++) {
                        float val = 0.0f;
                        for (int j = 0; j < BT; j++) {
                            float k_gated = ke(b, hv, n, j, kk) * exp2f_fast(ge(b, hv, n, j, kk));
                            val += Akk(b, hv, n, i, j) * k_gated;
                        }
                        w(b, hv, n, i, kk) = val;
                    }
                }

                // u = Akk @ v
                for (int i = 0; i < BT; i++) {
                    for (int vv = 0; vv < V; vv++) {
                        float val = 0.0f;
                        for (int j = 0; j < BT; j++) {
                            val += Akk(b, hv, n, i, j) * ve(b, hv, n, j, vv);
                        }
                        u(b, hv, n, i, vv) = val;
                    }
                }
            }
        }
    }

    // -----------------------------------------------------------------------
    // Step 4: Inter-chunk recurrence + output computation
    //
    // h[n] is the hidden state BEFORE processing chunk n.
    //   h[0] = h0 (initial state) or zeros
    //   h[n+1] = h[n] * 2^{g_last[n]} + kg[n]^T @ v_new[n]
    //
    // For each chunk n:
    //   1. Compute Aqk (Q-Key attention, lower triangular incl. diagonal)
    //   2. v_new = u - w @ h[n]    (Delta Rule: subtract key-state interaction)
    //   3. o_inter = (q * 2^g) @ h[n]
    //   4. o_intra = Aqk @ v_new
    //   5. o = o_inter + o_intra
    //   6. Compute h[n+1] from h[n] and chunk n's data
    // -----------------------------------------------------------------------

    // h: [B][HV][NT][K][V]  — per-chunk hidden state
    // h[n] = state at the BEGINNING of chunk n
    Tensor<float> h({B, HV, NT, K, V}, 0.0f);

    // Initialize h[0] from h0 (initial state)
    if (h0 != nullptr) {
        for (int b = 0; b < B; b++) {
            for (int hv = 0; hv < HV; hv++) {
                for (int kk = 0; kk < K; kk++) {
                    for (int vv = 0; vv < V; vv++) {
                        h(b, hv, 0, kk, vv) = (*h0)(b, hv, kk, vv);
                    }
                }
            }
        }
    }

    // o: [B][HV][NT][BT][V]
    Tensor<float> o({B, HV, NT, BT, V}, 0.0f);

    for (int n = 0; n < NT; n++) {
        for (int b = 0; b < B; b++) {
            for (int hv = 0; hv < HV; hv++) {

                // ---- Step 4a: Compute Aqk for this chunk ----
                // Aqk[i][j] = sum_k ( q[i][k] * 2^{g[i][k]-g[j][k]} * k[j][k] )
                // for i >= j (lower triangular including diagonal)
                std::vector<float> Aqk_vec(BT * BT, 0.0f);

                for (int i = 0; i < BT; i++) {
                    for (int j = 0; j <= i; j++) {  // only compute lower tri
                        float val = 0.0f;
                        for (int kk = 0; kk < K; kk++) {
                            float decay = exp2f_fast(ge(b, hv, n, i, kk) - ge(b, hv, n, j, kk));
                            val += qe(b, hv, n, i, kk) * decay * ke(b, hv, n, j, kk);
                        }
                        Aqk_vec[i * BT + j] = val;
                    }
                }

                // ---- Step 4b: Compute v_new = u - w @ h[n] ----
                // v_new[i][v] = u[i][v] - sum_k w[i][k] * h[n][k][v]
                std::vector<float> v_new(BT * V, 0.0f);

                for (int i = 0; i < BT; i++) {
                    for (int vv = 0; vv < V; vv++) {
                        float wh = 0.0f;
                        for (int kk = 0; kk < K; kk++) {
                            wh += w(b, hv, n, i, kk) * h(b, hv, n, kk, vv);
                        }
                        v_new[i * V + vv] = u(b, hv, n, i, vv) - wh;
                    }
                }

                // ---- Step 4c: Compute output ----
                // o = (q * 2^g) @ h[n] + Aqk @ v_new

                for (int i = 0; i < BT; i++) {
                    for (int vv = 0; vv < V; vv++) {
                        // inter-chunk: o_inter = (q * 2^g) @ h[n]
                        float o_inter = 0.0f;
                        for (int kk = 0; kk < K; kk++) {
                            float q_gated = qe(b, hv, n, i, kk) * exp2f_fast(ge(b, hv, n, i, kk));
                            o_inter += q_gated * h(b, hv, n, kk, vv);
                        }

                        // intra-chunk: o_intra = Aqk @ v_new
                        float o_intra = 0.0f;
                        for (int j = 0; j <= i; j++) {
                            o_intra += Aqk_vec[i * BT + j] * v_new[j * V + vv];
                        }

                        o(b, hv, n, i, vv) = o_inter + o_intra;
                    }
                }

                // ---- Step 4d: Compute h[n+1] from h[n] ----
                // h[n+1] = h[n] * 2^{g_last} + kg^T @ v_new
                // where kg[c][k] = k[c][k] * 2^{g_last[k] - g[c][k]}
                // and g_last = ge[..., BT-1, :]  (last token's cumulated gate in chunk)

                if (n + 1 < NT) {
                    for (int kk = 0; kk < K; kk++) {
                        float g_last_k = ge(b, hv, n, BT - 1, kk);
                        float decay = exp2f_fast(g_last_k);
                        for (int vv = 0; vv < V; vv++) {
                            // h[n+1][k][v] = h[n][k][v] * 2^{g_last[k]} + sum_c kg[c][k] * v_new[c][v]
                            float kg_dot_vnew = 0.0f;
                            for (int c = 0; c < BT; c++) {
                                float kg = ke(b, hv, n, c, kk) * exp2f_fast(g_last_k - ge(b, hv, n, c, kk));
                                kg_dot_vnew += kg * v_new[c * V + vv];
                            }
                            h(b, hv, n + 1, kk, vv) = h(b, hv, n, kk, vv) * decay + kg_dot_vnew;
                        }
                    }
                }

            } // hv
        } // b
    } // n (chunk)

    // -----------------------------------------------------------------------
    // Step 5: Rearrange output from [B][HV][NT][BT][V] -> [B][T][HV][V]
    // -----------------------------------------------------------------------
    Tensor<float> o_out({B, T, HV, V}, 0.0f);
    for (int b = 0; b < B; b++) {
        for (int n = 0; n < NT; n++) {
            for (int c = 0; c < BT; c++) {
                int t = n * BT + c;
                for (int hv = 0; hv < HV; hv++) {
                    for (int vv = 0; vv < V; vv++) {
                        o_out(b, t, hv, vv) = o(b, hv, n, c, vv);
                    }
                }
            }
        }
    }

    ChunkKdaResult result;
    result.o = std::move(o_out);
    result.h = std::move(h);

    // Compute final state ht = h[NT-1] after applying last chunk's gate + kg update
    // In the loop above, h[NT] is not stored (n+1 < NT check), so we compute ht separately
    if (output_final_state) {
        result.ht = Tensor<float>({B, HV, K, V}, 0.0f);
        int n_last = NT - 1;
        for (int b = 0; b < B; b++) {
            for (int hv = 0; hv < HV; hv++) {
                // We need to recompute the v_new for the last chunk
                // and then compute ht = h[n_last] * 2^{g_last} + kg^T @ v_new
                std::vector<float> v_new_last(BT * V, 0.0f);
                for (int i = 0; i < BT; i++) {
                    for (int vv = 0; vv < V; vv++) {
                        float wh = 0.0f;
                        for (int kk = 0; kk < K; kk++) {
                            wh += w(b, hv, n_last, i, kk) * h(b, hv, n_last, kk, vv);
                        }
                        v_new_last[i * V + vv] = u(b, hv, n_last, i, vv) - wh;
                    }
                }
                for (int kk = 0; kk < K; kk++) {
                    float g_last_k = ge(b, hv, n_last, BT - 1, kk);
                    float decay = exp2f_fast(g_last_k);
                    for (int vv = 0; vv < V; vv++) {
                        float kg_dot_vnew = 0.0f;
                        for (int c = 0; c < BT; c++) {
                            float kg = ke(b, hv, n_last, c, kk) * exp2f_fast(g_last_k - ge(b, hv, n_last, c, kk));
                            kg_dot_vnew += kg * v_new_last[c * V + vv];
                        }
                        result.ht(b, hv, kk, vv) = h(b, hv, n_last, kk, vv) * decay + kg_dot_vnew;
                    }
                }
            }
        }
    }

    return result;
}


// ============================================================================
// Test / example usage
// ============================================================================
#ifdef NAIVE_CHUNK_KDA_MAIN

#include <random>
#include <cstdio>

int main() {
    // Small test case
    int B = 2, T = 64, H = 4, HV = 8, K = 16, V = 32;
    int BT = 16;  // chunk_size
    float scale = 1.0f / sqrtf((float)K);

    std::mt19937 rng(42);
    std::normal_distribution<float> dist(0.0f, 1.0f);

    // Create random inputs
    Tensor<float> q({B, T, H, K});
    Tensor<float> k({B, T, H, K});
    Tensor<float> v({B, T, HV, V});
    Tensor<float> g({B, T, HV, K});
    Tensor<float> beta({B, T, HV});

    for (auto& x : q.data) x = dist(rng) * 0.1f;
    for (auto& x : k.data) x = dist(rng) * 0.1f;
    for (auto& x : v.data) x = dist(rng) * 0.1f;
    // g should be negative (decay), typically from sigmoid/log
    for (auto& x : g.data) x = -fabs(dist(rng)) * 0.5f;
    // beta should be in (0, 1)
    for (auto& x : beta.data) x = 0.5f + 0.5f * tanhf(dist(rng));

    auto result = naive_chunk_kda(q, k, v, g, beta, scale,
                                  nullptr /* h0 */,
                                  true /* output_final_state */,
                                  BT);

    printf("Output shape: [%d, %d, %d, %d]\n",
           result.o.shape[0], result.o.shape[1],
           result.o.shape[2], result.o.shape[3]);
    printf("Hidden state shape: [%d, %d, %d, %d, %d]\n",
           result.h.shape[0], result.h.shape[1],
           result.h.shape[2], result.h.shape[3],
           result.h.shape[4]);
    printf("Final state shape:  [%d, %d, %d, %d]\n",
           result.ht.shape[0], result.ht.shape[1],
           result.ht.shape[2], result.ht.shape[3]);

    // Print a few output values for sanity check
    printf("o[0][0][0][0..3] = %.6f %.6f %.6f %.6f\n",
           result.o(0, 0, 0, 0), result.o(0, 0, 0, 1),
           result.o(0, 0, 0, 2), result.o(0, 0, 0, 3));
    printf("o[0][63][7][0..3] = %.6f %.6f %.6f %.6f\n",
           result.o(0, 63, 7, 0), result.o(0, 63, 7, 1),
           result.o(0, 63, 7, 2), result.o(0, 63, 7, 3));

    printf("Done.\n");
    return 0;
}

#endif // NAIVE_CHUNK_KDA_MAIN

下面开始手撕 KDA

2 baseline

先给出一套 c++ 版本 和 cuda 版本的 baseline

2.1 c++ 版本

// KDA chunk reference implementation (CPU, fp32).
//
// Forward + backward port of the KDA chunk algorithm, structurally faithful to
// the Triton `chunk_kda` pipeline but written as plain nested loops for clarity.
// Algorithmic template: fla/ops/kda/naive.py::naive_chunk_kda (+ gate / l2norm /
// beta-sigmoid pre-stages). Assumes HV == H and fixed sequence length.
//
// All math is done in fp32. Natural exp is used instead of the exp2/RCP_LN2
// folding of the Triton path (mathematically equivalent).

#include <torch/extension.h>

#include <cmath>
#include <vector>

namespace {

inline float softplusf(float x) {
  // numerically stable log(1+exp(x))
  if (x > 20.0f) return x;
  if (x < -20.0f) return std::exp(x);
  return std::log1p(std::exp(x));
}

inline float sigmoidf(float x) { return 1.0f / (1.0f + std::exp(-x)); }

constexpr float L2_EPS = 1e-6f;

// Per (batch, head) workspace holding every intermediate the backward pass
// needs. Recomputed by running the forward; backward then sweeps in reverse.
struct Workspace {
  int T, K, V, BT, NT;
  // pre-stage (per token)
  std::vector<float> qn, kn;     // [T,K] l2-normed q,k
  std::vector<float> rq, rk;     // [T]   l2norm rstd
  std::vector<float> beta_s;     // [T]   sigmoid(beta_raw)
  std::vector<float> ga;         // [T,K] gate activation (pre-cumsum)
  std::vector<float> gc;         // [T,K] cumsum(ga) within chunk
  std::vector<float> qs;         // [T,K] qn * scale
  std::vector<float> vv;         // [T,V] copy of v
  // chunk-level
  std::vector<float> X;          // [NT,BT,BT] (I-Araw)^{-1}
  std::vector<float> Afull;      // [NT,BT,BT]
  std::vector<float> Aqk;        // [NT,BT,BT]
  std::vector<float> w;          // [T,K]
  std::vector<float> u;          // [T,V]
  std::vector<float> vcorr;      // [T,V]
  std::vector<float> Sentry;     // [NT+1, K, V]  state at entry of each chunk
};

// Forward for a single (b,h). Pointers already offset to this head's data.
// q_raw,k_raw,g_raw,beta_raw are the *raw* inputs for this head.
void forward_bh(
    const float* q_raw, const float* k_raw, const float* v_in,
    const float* g_raw, const float* beta_raw,
    float A_log_h, const float* dt_bias_h,
    const float* h0,  // [K,V]
    float scale, int T, int K, int V, int BT,
    float* o_out,       // [T,V]
    float* fs_out,      // [K,V]
    Workspace* ws) {
  const int NT = T / BT;
  const float aexp = std::exp(A_log_h);

  ws->T = T; ws->K = K; ws->V = V; ws->BT = BT; ws->NT = NT;
  ws->qn.assign(T * K, 0.f); ws->kn.assign(T * K, 0.f);
  ws->rq.assign(T, 0.f); ws->rk.assign(T, 0.f);
  ws->beta_s.assign(T, 0.f);
  ws->ga.assign(T * K, 0.f); ws->gc.assign(T * K, 0.f);
  ws->qs.assign(T * K, 0.f); ws->vv.assign(T * V, 0.f);
  ws->X.assign(NT * BT * BT, 0.f); ws->Afull.assign(NT * BT * BT, 0.f);
  ws->Aqk.assign(NT * BT * BT, 0.f);
  ws->w.assign(T * K, 0.f); ws->u.assign(T * V, 0.f);
  ws->vcorr.assign(T * V, 0.f);
  ws->Sentry.assign((NT + 1) * K * V, 0.f);

  // ---- pre-stage: l2norm, beta sigmoid, gate activation ----
  for (int t = 0; t < T; ++t) {
    float sq = 0.f, sk = 0.f;
    for (int d = 0; d < K; ++d) { float a = q_raw[t * K + d]; sq += a * a; }
    for (int d = 0; d < K; ++d) { float a = k_raw[t * K + d]; sk += a * a; }
    float rq = 1.0f / std::sqrt(sq + L2_EPS);
    float rk = 1.0f / std::sqrt(sk + L2_EPS);
    ws->rq[t] = rq; ws->rk[t] = rk;
    for (int d = 0; d < K; ++d) {
      float qn = q_raw[t * K + d] * rq;
      float kn = k_raw[t * K + d] * rk;
      ws->qn[t * K + d] = qn;
      ws->kn[t * K + d] = kn;
      ws->qs[t * K + d] = qn * scale;
      float x = g_raw[t * K + d] + dt_bias_h[d];
      ws->ga[t * K + d] = -aexp * softplusf(x);
    }
    ws->beta_s[t] = sigmoidf(beta_raw[t]);
    for (int vd = 0; vd < V; ++vd) ws->vv[t * V + vd] = v_in[t * V + vd];
  }
  // cumsum of ga within each chunk -> gc
  for (int n = 0; n < NT; ++n) {
    for (int d = 0; d < K; ++d) {
      float acc = 0.f;
      for (int c = 0; c < BT; ++c) {
        int t = n * BT + c;
        acc += ws->ga[t * K + d];
        ws->gc[t * K + d] = acc;
      }
    }
  }

  // ---- per-chunk WY / UT transform ----
  for (int n = 0; n < NT; ++n) {
    float* X = &ws->X[n * BT * BT];
    float* Af = &ws->Afull[n * BT * BT];
    const int base = n * BT;
    // Araw (strict lower) stored directly into X as starting matrix; X holds
    // the running inverse during forward substitution.
    for (int c = 0; c < BT; ++c) {
      for (int i = 0; i < BT; ++i) X[c * BT + i] = 0.f;
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int i = 0; i < c; ++i) {
        int ti = base + i;
        float inner = 0.f;
        for (int d = 0; d < K; ++d) {
          float f = std::exp(ws->gc[tc * K + d] - ws->gc[ti * K + d]);
          inner += ws->kn[tc * K + d] * f * ws->kn[ti * K + d];
        }
        X[c * BT + i] = -ws->beta_s[tc] * inner;  // Araw
      }
    }
    // forward substitution: X becomes (I - Araw)^{-1} strict-lower part.
    for (int c = 1; c < BT; ++c) {
      for (int j = 0; j < c; ++j) {
        float acc = X[c * BT + j];  // Araw[c,j]
        for (int m = j + 1; m < c; ++m) acc += X[c * BT + m] * X[m * BT + j];
        X[c * BT + j] = acc;
      }
    }
    // Afull[c,i] = (X[c,i] + delta) * beta_s[i]
    for (int c = 0; c < BT; ++c) {
      for (int i = 0; i < BT; ++i) {
        float xci = X[c * BT + i] + (c == i ? 1.0f : 0.0f);
        Af[c * BT + i] = xci * ws->beta_s[base + i];
      }
    }
    // w = Afull @ (exp(gc) * kn) ; u = Afull @ v
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float acc = 0.f;
        for (int i = 0; i < BT; ++i) {
          int ti = base + i;
          float kg = std::exp(ws->gc[ti * K + d]) * ws->kn[ti * K + d];
          acc += Af[c * BT + i] * kg;
        }
        ws->w[tc * K + d] = acc;
      }
      for (int vd = 0; vd < V; ++vd) {
        float acc = 0.f;
        for (int i = 0; i < BT; ++i)
          acc += Af[c * BT + i] * ws->vv[(base + i) * V + vd];
        ws->u[tc * V + vd] = acc;
      }
    }
  }

  // ---- inter-chunk state recurrence + output ----
  std::vector<float> S(K * V, 0.f);
  for (int i = 0; i < K * V; ++i) S[i] = h0 ? h0[i] : 0.f;
  for (int kv = 0; kv < K * V; ++kv) ws->Sentry[0 * K * V + kv] = S[kv];

  for (int n = 0; n < NT; ++n) {
    const int base = n * BT;
    float* Aqk = &ws->Aqk[n * BT * BT];
    // Aqk[c,j] for c>=j
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int j = 0; j <= c; ++j) {
        int tj = base + j;
        float acc = 0.f;
        for (int d = 0; d < K; ++d) {
          float f = std::exp(ws->gc[tc * K + d] - ws->gc[tj * K + d]);
          acc += ws->qs[tc * K + d] * f * ws->kn[tj * K + d];
        }
        Aqk[c * BT + j] = acc;
      }
    }
    // vcorr = u - w @ S
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float acc = ws->u[tc * V + vd];
        for (int d = 0; d < K; ++d) acc -= ws->w[tc * K + d] * S[d * V + vd];
        ws->vcorr[tc * V + vd] = acc;
      }
    }
    // o = (qs*exp(gc)) @ S + Aqk @ vcorr
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float acc = 0.f;
        for (int d = 0; d < K; ++d)
          acc += ws->qs[tc * K + d] * std::exp(ws->gc[tc * K + d]) * S[d * V + vd];
        for (int j = 0; j <= c; ++j)
          acc += Aqk[c * BT + j] * ws->vcorr[(base + j) * V + vd];
        o_out[tc * V + vd] = acc;
      }
    }
    // S update
    const int tlast = base + BT - 1;
    std::vector<float> Snew(K * V, 0.f);
    for (int d = 0; d < K; ++d) {
      float glast = ws->gc[tlast * K + d];
      float EL = std::exp(glast);
      for (int vd = 0; vd < V; ++vd) {
        float acc = S[d * V + vd] * EL;
        for (int c = 0; c < BT; ++c) {
          int tc = base + c;
          float P = std::exp(glast - ws->gc[tc * K + d]) * ws->kn[tc * K + d];
          acc += P * ws->vcorr[tc * V + vd];
        }
        Snew[d * V + vd] = acc;
      }
    }
    S.swap(Snew);
    for (int kv = 0; kv < K * V; ++kv) ws->Sentry[(n + 1) * K * V + kv] = S[kv];
  }
  for (int kv = 0; kv < K * V; ++kv) fs_out[kv] = S[kv];
}

// Backward for a single (b,h). Accumulates grads into the d* output buffers
// (which are head-local). dA_log_h / ddt_bias_h are accumulated.
void backward_bh(
    const float* q_raw, const float* k_raw,
    float A_log_h, const float* dt_bias_h,
    const float* h0, float scale, int T, int K, int V, int BT,
    const float* do_in,   // [T,V]
    const float* dht,     // [K,V] grad wrt final state
    const Workspace* ws,
    float* dq_raw, float* dk_raw, float* dv_out, float* dg_raw,
    float* dbeta_raw, float* dA_log_h, float* ddt_bias_h, float* dh0_out) {
  const int NT = T / BT;
  const float aexp = std::exp(A_log_h);

  // head-local accumulators
  std::vector<float> dgc(T * K, 0.f), dqs(T * K, 0.f), dkn(T * K, 0.f);
  std::vector<float> dqn(T * K, 0.f), dbeta_s(T, 0.f), dga(T * K, 0.f);
  std::vector<float> dv(T * V, 0.f);

  // dS flowing backward, init from dht (grad of final state)
  std::vector<float> dS(K * V, 0.f);
  for (int i = 0; i < K * V; ++i) dS[i] = dht ? dht[i] : 0.f;

  for (int n = NT - 1; n >= 0; --n) {
    const int base = n * BT;
    const float* Aqk = &ws->Aqk[n * BT * BT];
    const float* Af = &ws->Afull[n * BT * BT];
    const float* X = &ws->X[n * BT * BT];
    const float* S = &ws->Sentry[n * K * V];  // S_entry of chunk n
    const int tlast = base + BT - 1;

    std::vector<float> dvcorr(BT * V, 0.f);
    std::vector<float> dAqk(BT * BT, 0.f);
    std::vector<float> dQE(BT * K, 0.f);      // grad wrt (qs*exp(gc))
    std::vector<float> dP(BT * K, 0.f);       // grad wrt P[c,d]
    std::vector<float> dEL(K, 0.f);           // grad wrt EL[d]
    std::vector<float> dS_entry(K * V, 0.f);

    // ---- adjoint of S update (R3): S_new = S*EL + sum_c P[c]⊗vcorr[c] ----
    // incoming dS is grad wrt S_new
    for (int d = 0; d < K; ++d) {
      for (int vd = 0; vd < V; ++vd) {
        float g = dS[d * V + vd];
        float EL = std::exp(ws->gc[tlast * K + d]);
        dS_entry[d * V + vd] += g * EL;
        dEL[d] += g * S[d * V + vd];
        for (int c = 0; c < BT; ++c) {
          int tc = base + c;
          float P = std::exp(ws->gc[tlast * K + d] - ws->gc[tc * K + d]) * ws->kn[tc * K + d];
          dvcorr[c * V + vd] += P * g;
          dP[c * K + d] += g * ws->vcorr[tc * V + vd];
        }
      }
    }
    // ---- adjoint of output (R2): o = (qs*E)@S + Aqk@vcorr ----
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float god = do_in[tc * V + vd];
        for (int d = 0; d < K; ++d) {
          dQE[c * K + d] += god * S[d * V + vd];
          dS_entry[d * V + vd] += ws->qs[tc * K + d] * std::exp(ws->gc[tc * K + d]) * god;
        }
        for (int j = 0; j <= c; ++j) {
          dAqk[c * BT + j] += god * ws->vcorr[(base + j) * V + vd];
          dvcorr[j * V + vd] += Aqk[c * BT + j] * god;
        }
      }
    }
    // ---- adjoint of vcorr (R1): vcorr = u - w@S ----
    std::vector<float> du(BT * V, 0.f), dw(BT * K, 0.f);
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float g = dvcorr[c * V + vd];
        du[c * V + vd] += g;
        for (int d = 0; d < K; ++d) {
          dw[c * K + d] += -g * S[d * V + vd];
          dS_entry[d * V + vd] += -ws->w[tc * K + d] * g;
        }
      }
    }
    // ---- adjoint of P[c,d] = exp(glast-gc[c,d])*kn[c,d] ----
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float factor = std::exp(ws->gc[tlast * K + d] - ws->gc[tc * K + d]);
        float P = factor * ws->kn[tc * K + d];
        float g = dP[c * K + d];
        dkn[tc * K + d] += g * factor;
        dgc[tlast * K + d] += g * P;
        dgc[tc * K + d] += -g * P;
      }
    }
    // ---- adjoint of EL[d] = exp(gc[tlast,d]) ----
    for (int d = 0; d < K; ++d)
      dgc[tlast * K + d] += dEL[d] * std::exp(ws->gc[tlast * K + d]);

    // ---- adjoint of Aqk[c,j] = sum_d qs[c,d]*exp(gc[c]-gc[j])*kn[j,d] ----
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int j = 0; j <= c; ++j) {
        int tj = base + j;
        float g = dAqk[c * BT + j];
        if (g == 0.f) continue;
        for (int d = 0; d < K; ++d) {
          float f = std::exp(ws->gc[tc * K + d] - ws->gc[tj * K + d]);
          float term = ws->qs[tc * K + d] * f * ws->kn[tj * K + d];
          dqs[tc * K + d] += g * f * ws->kn[tj * K + d];
          dkn[tj * K + d] += g * ws->qs[tc * K + d] * f;
          dgc[tc * K + d] += g * term;
          dgc[tj * K + d] += -g * term;
        }
      }
    }
    // ---- adjoint of dQE: (qs*E) ----
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float E = std::exp(ws->gc[tc * K + d]);
        float g = dQE[c * K + d];
        dqs[tc * K + d] += g * E;
        dgc[tc * K + d] += g * ws->qs[tc * K + d] * E;
      }
    }

    // ---- adjoint of w = Afull @ kg_decay, u = Afull @ v ----
    std::vector<float> dAfull(BT * BT, 0.f);
    std::vector<float> dkg(BT * K, 0.f);  // grad wrt kg_decay[i,d]
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float gwd = dw[c * K + d];
        if (gwd == 0.f) continue;
        for (int i = 0; i < BT; ++i) {
          int ti = base + i;
          float kg = std::exp(ws->gc[ti * K + d]) * ws->kn[ti * K + d];
          dAfull[c * BT + i] += gwd * kg;
          dkg[i * K + d] += gwd * Af[c * BT + i];
        }
      }
      for (int vd = 0; vd < V; ++vd) {
        float gud = du[c * V + vd];
        for (int i = 0; i < BT; ++i) {
          dAfull[c * BT + i] += gud * ws->vv[(base + i) * V + vd];
          dv[(base + i) * V + vd] += gud * Af[c * BT + i];
        }
      }
    }
    // ---- adjoint of kg_decay[i,d] = exp(gc[i,d])*kn[i,d] ----
    for (int i = 0; i < BT; ++i) {
      int ti = base + i;
      for (int d = 0; d < K; ++d) {
        float E = std::exp(ws->gc[ti * K + d]);
        float g = dkg[i * K + d];
        dkn[ti * K + d] += g * E;
        dgc[ti * K + d] += g * ws->kn[ti * K + d] * E;  // = g*kg_decay
      }
    }
    // ---- adjoint of Afull[c,i] = X[c,i]*beta_s[i] (+ diag) ----
    std::vector<float> dX(BT * BT, 0.f);
    for (int c = 0; c < BT; ++c) {
      for (int i = 0; i < BT; ++i) {
        float xci = X[c * BT + i] + (c == i ? 1.0f : 0.0f);
        float g = dAfull[c * BT + i];
        dbeta_s[base + i] += g * xci;
        dX[c * BT + i] += g * ws->beta_s[base + i];
      }
    }
    // ---- matrix-inverse VJP: X=(I-Araw)^{-1}; dAraw = Xt @ dX @ Xt ----
    // Build full X (with unit diagonal) Xf.
    std::vector<float> Xf(BT * BT, 0.f);
    for (int c = 0; c < BT; ++c)
      for (int i = 0; i < BT; ++i)
        Xf[c * BT + i] = X[c * BT + i] + (c == i ? 1.0f : 0.0f);
    // tmp = Xf^T @ dX  -> [BT,BT]
    std::vector<float> tmp(BT * BT, 0.f);
    for (int a = 0; a < BT; ++a)
      for (int b = 0; b < BT; ++b) {
        float acc = 0.f;
        for (int m = 0; m < BT; ++m) acc += Xf[m * BT + a] * dX[m * BT + b];
        tmp[a * BT + b] = acc;
      }
    // dAraw = tmp @ Xf^T
    std::vector<float> dAraw(BT * BT, 0.f);
    for (int a = 0; a < BT; ++a)
      for (int b = 0; b < BT; ++b) {
        float acc = 0.f;
        for (int m = 0; m < BT; ++m) acc += tmp[a * BT + m] * Xf[b * BT + m];
        dAraw[a * BT + b] = acc;
      }
    // ---- adjoint of Araw[c,i] = -beta_s[c]*inner[c,i], c>i ----
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int i = 0; i < c; ++i) {
        int ti = base + i;
        float g = dAraw[c * BT + i];
        // recompute inner
        float inner = 0.f;
        for (int d = 0; d < K; ++d) {
          float f = std::exp(ws->gc[tc * K + d] - ws->gc[ti * K + d]);
          inner += ws->kn[tc * K + d] * f * ws->kn[ti * K + d];
        }
        dbeta_s[tc] += g * (-inner);
        float dinner = g * (-ws->beta_s[tc]);
        for (int d = 0; d < K; ++d) {
          float f = std::exp(ws->gc[tc * K + d] - ws->gc[ti * K + d]);
          float term = ws->kn[tc * K + d] * f * ws->kn[ti * K + d];
          dkn[tc * K + d] += dinner * f * ws->kn[ti * K + d];
          dkn[ti * K + d] += dinner * ws->kn[tc * K + d] * f;
          dgc[tc * K + d] += dinner * term;
          dgc[ti * K + d] += -dinner * term;
        }
      }
    }
    // pass dS_entry to previous chunk
    dS.swap(dS_entry);
  }
  // dS now holds grad wrt S_entry of chunk 0 == grad wrt h0
  for (int i = 0; i < K * V; ++i) dh0_out[i] = dS[i];

  // ---- adjoint of qs = qn*scale ----
  for (int t = 0; t < T; ++t)
    for (int d = 0; d < K; ++d) dqn[t * K + d] += dqs[t * K + d] * scale;

  // ---- adjoint of gc = cumsum(ga): reverse cumsum ----
  for (int n = 0; n < NT; ++n) {
    for (int d = 0; d < K; ++d) {
      float acc = 0.f;
      for (int c = BT - 1; c >= 0; --c) {
        int t = n * BT + c;
        acc += dgc[t * K + d];
        dga[t * K + d] = acc;
      }
    }
  }
  // ---- adjoint of ga = -aexp*softplus(graw+dt_bias) ----
  // ga = -aexp*softplus(x); softplus(x) = -ga/aexp; sigmoid(x) = 1-exp(-sp).
  float dA_acc = 0.f;
  for (int t = 0; t < T; ++t) {
    for (int d = 0; d < K; ++d) {
      float ga = ws->ga[t * K + d];
      // sp = -ga/aexp ; x = inverse-softplus not needed; sigmoid(x) derived:
      // ga = -aexp*softplus(x) -> softplus(x) = -ga/aexp
      // d softplus = sigmoid(x); need sigmoid(x). sigmoid(x)=1-exp(-softplus(x)).
      float sp = -ga / aexp;
      float sig = 1.0f - std::exp(-sp);  // sigmoid(x)
      float dxin = dga[t * K + d] * (-aexp) * sig;
      dg_raw[t * K + d] = dxin;
      ddt_bias_h[d] += dxin;
      dA_acc += dga[t * K + d] * ga;  // d/dA_log: ga
    }
  }
  *dA_log_h += dA_acc;

  // ---- adjoint of beta_s = sigmoid(beta_raw) ----
  for (int t = 0; t < T; ++t) {
    float b = ws->beta_s[t];
    dbeta_raw[t] = dbeta_s[t] * b * (1.0f - b);
  }

  // ---- adjoint of l2norm for q and k ----
  for (int t = 0; t < T; ++t) {
    float rq = ws->rq[t], rk = ws->rk[t];
    float sdq = 0.f, sdk = 0.f;
    for (int d = 0; d < K; ++d) {
      sdq += dqn[t * K + d] * ws->qn[t * K + d];
      sdk += dkn[t * K + d] * ws->kn[t * K + d];
    }
    for (int d = 0; d < K; ++d) {
      dq_raw[t * K + d] = dqn[t * K + d] * rq - sdq * ws->qn[t * K + d] * rq;
      dk_raw[t * K + d] = dkn[t * K + d] * rk - sdk * ws->kn[t * K + d] * rk;
    }
  }
  // copy dv
  for (int i = 0; i < T * V; ++i) dv_out[i] = dv[i];
}

}  // namespace

// ---------------- Python-facing entry points ----------------

std::vector<torch::Tensor> kda_chunk_fwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v,
    torch::Tensor g, torch::Tensor beta,
    torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size) {
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);

  const int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3);
  const int Vd = v.size(3);
  const int BT = chunk_size;

  auto o = torch::zeros({B, T, H, Vd}, torch::kFloat32);
  auto fs = torch::zeros({B, H, K, Vd}, torch::kFloat32);

  const float* qp = q.data_ptr<float>();
  const float* kp = k.data_ptr<float>();
  const float* vp = v.data_ptr<float>();
  const float* gp = g.data_ptr<float>();
  const float* bp = beta.data_ptr<float>();
  const float* alp = A_log.data_ptr<float>();
  const float* dtp = dt_bias.data_ptr<float>();
  const float* h0p = initial_state.data_ptr<float>();
  float* op = o.data_ptr<float>();
  float* fsp = fs.data_ptr<float>();

  for (int b = 0; b < B; ++b) {
    for (int h = 0; h < H; ++h) {
      // gather this head's data into contiguous [T,*] buffers
      std::vector<float> qh(T * K), kh(T * K), vh(T * Vd), gh(T * K), bh(T);
      for (int t = 0; t < T; ++t) {
        for (int d = 0; d < K; ++d) {
          int idx = ((b * T + t) * H + h) * K + d;
          qh[t * K + d] = qp[idx];
          kh[t * K + d] = kp[idx];
          gh[t * K + d] = gp[idx];
        }
        for (int vd = 0; vd < Vd; ++vd)
          vh[t * Vd + vd] = vp[((b * T + t) * H + h) * Vd + vd];
        bh[t] = bp[(b * T + t) * H + h];
      }
      std::vector<float> oh(T * Vd), fsh(K * Vd);
      Workspace ws;
      forward_bh(qh.data(), kh.data(), vh.data(), gh.data(), bh.data(),
                 alp[h], dtp + h * K, h0p + (b * H + h) * K * Vd,
                 (float)scale, T, K, Vd, BT, oh.data(), fsh.data(), &ws);
      for (int t = 0; t < T; ++t)
        for (int vd = 0; vd < Vd; ++vd)
          op[((b * T + t) * H + h) * Vd + vd] = oh[t * Vd + vd];
      for (int i = 0; i < K * Vd; ++i)
        fsp[(b * H + h) * K * Vd + i] = fsh[i];
    }
  }
  return {o, fs};
}

std::vector<torch::Tensor> kda_chunk_bwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v,
    torch::Tensor g, torch::Tensor beta,
    torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size,
    torch::Tensor do_, torch::Tensor dht) {
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);
  do_ = do_.contiguous().to(torch::kFloat32);
  dht = dht.contiguous().to(torch::kFloat32);

  const int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3);
  const int Vd = v.size(3);
  const int BT = chunk_size;

  auto dq = torch::zeros_like(q);
  auto dk = torch::zeros_like(k);
  auto dv = torch::zeros_like(v);
  auto dg = torch::zeros_like(g);
  auto dbeta = torch::zeros_like(beta);
  auto dA_log = torch::zeros_like(A_log);
  auto ddt_bias = torch::zeros_like(dt_bias);
  auto dh0 = torch::zeros_like(initial_state);

  const float* qp = q.data_ptr<float>();
  const float* kp = k.data_ptr<float>();
  const float* vp = v.data_ptr<float>();
  const float* gp = g.data_ptr<float>();
  const float* bp = beta.data_ptr<float>();
  const float* alp = A_log.data_ptr<float>();
  const float* dtp = dt_bias.data_ptr<float>();
  const float* h0p = initial_state.data_ptr<float>();
  const float* dop = do_.data_ptr<float>();
  const float* dhtp = dht.data_ptr<float>();

  float* dqp = dq.data_ptr<float>();
  float* dkp = dk.data_ptr<float>();
  float* dvp = dv.data_ptr<float>();
  float* dgp = dg.data_ptr<float>();
  float* dbp = dbeta.data_ptr<float>();
  float* dalp = dA_log.data_ptr<float>();
  float* ddtp = ddt_bias.data_ptr<float>();
  float* dh0p = dh0.data_ptr<float>();

  for (int b = 0; b < B; ++b) {
    for (int h = 0; h < H; ++h) {
      std::vector<float> qh(T * K), kh(T * K), vh(T * Vd), gh(T * K), bh(T),
          doh(T * Vd);
      for (int t = 0; t < T; ++t) {
        for (int d = 0; d < K; ++d) {
          int idx = ((b * T + t) * H + h) * K + d;
          qh[t * K + d] = qp[idx];
          kh[t * K + d] = kp[idx];
          gh[t * K + d] = gp[idx];
        }
        for (int vd = 0; vd < Vd; ++vd) {
          vh[t * Vd + vd] = vp[((b * T + t) * H + h) * Vd + vd];
          doh[t * Vd + vd] = dop[((b * T + t) * H + h) * Vd + vd];
        }
        bh[t] = bp[(b * T + t) * H + h];
      }
      std::vector<float> oh(T * Vd), fsh(K * Vd);
      Workspace ws;
      forward_bh(qh.data(), kh.data(), vh.data(), gh.data(), bh.data(),
                 alp[h], dtp + h * K, h0p + (b * H + h) * K * Vd,
                 (float)scale, T, K, Vd, BT, oh.data(), fsh.data(), &ws);

      std::vector<float> dqh(T * K, 0.f), dkh(T * K, 0.f), dvh(T * Vd, 0.f),
          dgh(T * K, 0.f), dbh(T, 0.f), dh0h(K * Vd, 0.f),
          ddth(K, 0.f);
      float dAlogh = 0.f;
      backward_bh(qh.data(), kh.data(), alp[h], dtp + h * K,
                  h0p + (b * H + h) * K * Vd, (float)scale, T, K, Vd, BT,
                  doh.data(), dhtp + (b * H + h) * K * Vd, &ws,
                  dqh.data(), dkh.data(), dvh.data(), dgh.data(), dbh.data(),
                  &dAlogh, ddth.data(), dh0h.data());

      for (int t = 0; t < T; ++t) {
        for (int d = 0; d < K; ++d) {
          int idx = ((b * T + t) * H + h) * K + d;
          dqp[idx] = dqh[t * K + d];
          dkp[idx] = dkh[t * K + d];
          dgp[idx] = dgh[t * K + d];
        }
        for (int vd = 0; vd < Vd; ++vd)
          dvp[((b * T + t) * H + h) * Vd + vd] = dvh[t * Vd + vd];
        dbp[(b * T + t) * H + h] = dbh[t];
      }
      for (int i = 0; i < K * Vd; ++i)
        dh0p[(b * H + h) * K * Vd + i] = dh0h[i];
      dalp[h] += dAlogh;
      for (int d = 0; d < K; ++d) ddtp[h * K + d] += ddth[d];
    }
  }
  return {dq, dk, dv, dg, dbeta, dA_log, ddt_bias, dh0};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fwd", &kda_chunk_fwd, "KDA chunk forward (CPU)");
  m.def("bwd", &kda_chunk_bwd, "KDA chunk backward (CPU)");
}

2.1.2 cuda 版本

KDA chunk 的 CUDA 参考实现（fp32），逐 stage 对应 CPU 参考实现，每个 stage 单独 launch 一个 kernel。

线程排布：一个per (batch, head)只分配一个线程去计算所有chunk

// KDA chunk reference implementation (CUDA, fp32).
//
// Stage-by-stage port mirroring the CPU reference (kda_chunk_cpu.cpp). Each
// stage is a separate kernel launched in sequence; every kernel uses one thread
// per (batch, head) so the body is a direct transcription of the CPU loops.
// Per-head intermediates live in global scratch buffers (allocated by the host
// dispatcher) to avoid large per-thread local memory.
//
// Assumes HV == H, fixed sequence length, fp32 throughout.

#include <cuda_runtime.h>
#include <torch/extension.h>

#include <vector>

namespace {

constexpr float L2_EPS = 1e-6f;

__device__ inline float softplusf(float x) {
  if (x > 20.0f) return x;
  if (x < -20.0f) return expf(x);
  return log1pf(expf(x));
}
__device__ inline float sigmoidf(float x) { return 1.0f / (1.0f + expf(-x)); }

// inputs: q/k/g [B,T,H,K], v/o/do/dv [B,T,H,V], beta [B,T,H],
//         h0/fs/dht/dh0 [B,H,K,V] == [BH,K,V], A_log [H], dt_bias [H*K]
__device__ inline int iTK(int bh, int t, int d, int T, int K) { return bh * T * K + t * K + d; }
__device__ inline int iTV(int bh, int t, int vd, int T, int V) { return bh * T * V + t * V + vd; }
__device__ inline int inTK(int b, int t, int h, int d, int T, int H, int K) {
  return ((b * T + t) * H + h) * K + d;
}
__device__ inline int inTV(int b, int t, int h, int vd, int T, int H, int V) {
  return ((b * T + t) * H + h) * V + vd;
}

// ================= FORWARD =================

__global__ void k_pre(
    const float* q_raw, const float* k_raw, const float* g_raw,
    const float* beta_raw, const float* A_log, const float* dt_bias,
    float* qn, float* kn, float* qs, float* rq, float* rk, float* beta_s,
    float* ga, float* gc, float scale, int B, int T, int H, int K, int BT) {
  int bh = blockIdx.x * blockDim.x + threadIdx.x;
  if (bh >= B * H) return;
  int b = bh / H, h = bh % H, NT = T / BT;
  float aexp = expf(A_log[h]);
  for (int t = 0; t < T; ++t) {
    float sq = 0.f, sk = 0.f;
    for (int d = 0; d < K; ++d) {
      float a = q_raw[inTK(b, t, h, d, T, H, K)]; sq += a * a;
      float c = k_raw[inTK(b, t, h, d, T, H, K)]; sk += c * c;
    }
    float r_q = rsqrtf(sq + L2_EPS), r_k = rsqrtf(sk + L2_EPS);
    rq[bh * T + t] = r_q; rk[bh * T + t] = r_k;
    for (int d = 0; d < K; ++d) {
      float qv = q_raw[inTK(b, t, h, d, T, H, K)] * r_q;
      float kv = k_raw[inTK(b, t, h, d, T, H, K)] * r_k;
      qn[iTK(bh, t, d, T, K)] = qv;
      kn[iTK(bh, t, d, T, K)] = kv;
      qs[iTK(bh, t, d, T, K)] = qv * scale;
      float x = g_raw[inTK(b, t, h, d, T, H, K)] + dt_bias[h * K + d];
      ga[iTK(bh, t, d, T, K)] = -aexp * softplusf(x);
    }
    beta_s[bh * T + t] = sigmoidf(beta_raw[(b * T + t) * H + h]);
  }
  for (int n = 0; n < NT; ++n)
    for (int d = 0; d < K; ++d) {
      float acc = 0.f;
      for (int c = 0; c < BT; ++c) {
        int t = n * BT + c;
        acc += ga[iTK(bh, t, d, T, K)];
        gc[iTK(bh, t, d, T, K)] = acc;
      }
    }
}

__global__ void k_wy(
    const float* v_in, const float* kn, const float* gc, const float* beta_s,
    float* X, float* Afull, float* w, float* u,
    int B, int T, int H, int K, int V, int BT) {
  int bh = blockIdx.x * blockDim.x + threadIdx.x;
  if (bh >= B * H) return;
  int b = bh / H, h = bh % H, NT = T / BT;
  for (int n = 0; n < NT; ++n) {
    float* Xn = X + (bh * NT + n) * BT * BT;
    float* Af = Afull + (bh * NT + n) * BT * BT;
    int base = n * BT;
    for (int c = 0; c < BT * BT; ++c) Xn[c] = 0.f;
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int i = 0; i < c; ++i) {
        int ti = base + i;
        float inner = 0.f;
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, ti, d, T, K)]);
          inner += kn[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, ti, d, T, K)];
        }
        Xn[c * BT + i] = -beta_s[bh * T + tc] * inner;
      }
    }
    for (int c = 1; c < BT; ++c)
      for (int j = 0; j < c; ++j) {
        float acc = Xn[c * BT + j];
        for (int m = j + 1; m < c; ++m) acc += Xn[c * BT + m] * Xn[m * BT + j];
        Xn[c * BT + j] = acc;
      }
    for (int c = 0; c < BT; ++c)
      for (int i = 0; i < BT; ++i) {
        float xci = Xn[c * BT + i] + (c == i ? 1.0f : 0.0f);
        Af[c * BT + i] = xci * beta_s[bh * T + base + i];
      }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float acc = 0.f;
        for (int i = 0; i < BT; ++i) {
          int ti = base + i;
          float kg = expf(gc[iTK(bh, ti, d, T, K)]) * kn[iTK(bh, ti, d, T, K)];
          acc += Af[c * BT + i] * kg;
        }
        w[iTK(bh, tc, d, T, K)] = acc;
      }
      for (int vd = 0; vd < V; ++vd) {
        float acc = 0.f;
        for (int i = 0; i < BT; ++i)
          acc += Af[c * BT + i] * v_in[inTV(b, base + i, h, vd, T, H, V)];
        u[iTV(bh, tc, vd, T, V)] = acc;
      }
    }
  }
}

__global__ void k_recur(
    const float* qs, const float* kn, const float* gc, const float* w,
    const float* u, const float* h0,
    float* Aqk, float* vcorr, float* Sentry, float* o_out, float* fs_out,
    int B, int T, int H, int K, int V, int BT) {
  int bh = blockIdx.x * blockDim.x + threadIdx.x;
  if (bh >= B * H) return;
  int b = bh / H, h = bh % H, NT = T / BT;
  float* Se = Sentry + bh * (NT + 1) * K * V;
  for (int i = 0; i < K * V; ++i) Se[i] = h0[bh * K * V + i];
  for (int n = 0; n < NT; ++n) {
    int base = n * BT;
    float* S = Se + n * K * V;
    float* Snew = Se + (n + 1) * K * V;
    float* Aq = Aqk + (bh * NT + n) * BT * BT;
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int j = 0; j <= c; ++j) {
        int tj = base + j;
        float acc = 0.f;
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, tj, d, T, K)]);
          acc += qs[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, tj, d, T, K)];
        }
        Aq[c * BT + j] = acc;
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float acc = u[iTV(bh, tc, vd, T, V)];
        for (int d = 0; d < K; ++d) acc -= w[iTK(bh, tc, d, T, K)] * S[d * V + vd];
        vcorr[iTV(bh, tc, vd, T, V)] = acc;
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float acc = 0.f;
        for (int d = 0; d < K; ++d)
          acc += qs[iTK(bh, tc, d, T, K)] * expf(gc[iTK(bh, tc, d, T, K)]) * S[d * V + vd];
        for (int j = 0; j <= c; ++j)
          acc += Aq[c * BT + j] * vcorr[iTV(bh, base + j, vd, T, V)];
        o_out[inTV(b, tc, h, vd, T, H, V)] = acc;
      }
    }
    int tlast = base + BT - 1;
    for (int d = 0; d < K; ++d) {
      float glast = gc[iTK(bh, tlast, d, T, K)], EL = expf(glast);
      for (int vd = 0; vd < V; ++vd) {
        float acc = S[d * V + vd] * EL;
        for (int c = 0; c < BT; ++c) {
          int tc = base + c;
          float P = expf(glast - gc[iTK(bh, tc, d, T, K)]) * kn[iTK(bh, tc, d, T, K)];
          acc += P * vcorr[iTV(bh, tc, vd, T, V)];
        }
        Snew[d * V + vd] = acc;
      }
    }
  }
  for (int i = 0; i < K * V; ++i) fs_out[bh * K * V + i] = Se[NT * K * V + i];
}

// ================= BACKWARD =================

struct BwdScratch {
  float *dgc, *dqs, *dkn, *dbeta_s;
  float *dS, *dS_entry, *dEL;
  float *dvcorr, *du;
  float *dAqk, *dAfull, *dX, *Xf, *tmp, *dAraw;
  float *dQE, *dP, *dw, *dkg;
};

__global__ void k_bwd_recur(
    const float* qs, const float* kn, const float* gc, const float* v_in,
    const float* w, const float* Afull, const float* X, const float* Aqk,
    const float* vcorr, const float* Sentry, const float* beta_s,
    const float* do_in, const float* dht,
    BwdScratch sc, float* dv_out, float* dh0_out,
    int B, int T, int H, int K, int V, int BT) {
  int bh = blockIdx.x * blockDim.x + threadIdx.x;
  if (bh >= B * H) return;
  int b = bh / H, h = bh % H, NT = T / BT;

  // dgc/dqs/dkn use the global base + iTK(bh,...) convention (matching
  // k_bwd_pre); do NOT pre-offset by bh*T*K or the iTK index double-counts it.
  float* dgc = sc.dgc;
  float* dqs = sc.dqs;
  float* dkn = sc.dkn;
  float* dbeta_s = sc.dbeta_s + bh * T;
  float* dS = sc.dS + bh * K * V;
  float* dS_entry = sc.dS_entry + bh * K * V;
  float* dEL = sc.dEL + bh * K;
  float* dvcorr = sc.dvcorr + bh * BT * V;
  float* du = sc.du + bh * BT * V;
  float* dAqk = sc.dAqk + bh * BT * BT;
  float* dAfull = sc.dAfull + bh * BT * BT;
  float* dX = sc.dX + bh * BT * BT;
  float* Xf = sc.Xf + bh * BT * BT;
  float* tmpm = sc.tmp + bh * BT * BT;
  float* dAraw = sc.dAraw + bh * BT * BT;
  float* dQE = sc.dQE + bh * BT * K;
  float* dP = sc.dP + bh * BT * K;
  float* dw = sc.dw + bh * BT * K;
  float* dkg = sc.dkg + bh * BT * K;

  for (int i = 0; i < T * K; ++i) {
    int idx = bh * T * K + i;
    dgc[idx] = 0.f; dqs[idx] = 0.f; dkn[idx] = 0.f;
  }
  for (int i = 0; i < T; ++i) dbeta_s[i] = 0.f;
  for (int t = 0; t < T; ++t)
    for (int vd = 0; vd < V; ++vd) dv_out[inTV(b, t, h, vd, T, H, V)] = 0.f;
  for (int i = 0; i < K * V; ++i) dS[i] = dht[bh * K * V + i];

  const float* Sentry_bh = Sentry + bh * (NT + 1) * K * V;

  for (int n = NT - 1; n >= 0; --n) {
    int base = n * BT, tlast = base + BT - 1;
    const float* Aq = Aqk + (bh * NT + n) * BT * BT;
    const float* Af = Afull + (bh * NT + n) * BT * BT;
    const float* Xn = X + (bh * NT + n) * BT * BT;
    const float* S = Sentry_bh + n * K * V;

    for (int i = 0; i < BT * V; ++i) { dvcorr[i] = 0.f; du[i] = 0.f; }
    for (int i = 0; i < BT * BT; ++i) { dAqk[i] = 0.f; dAfull[i] = 0.f; dX[i] = 0.f; }
    for (int i = 0; i < BT * K; ++i) { dQE[i] = 0.f; dP[i] = 0.f; dw[i] = 0.f; dkg[i] = 0.f; }
    for (int i = 0; i < K; ++i) dEL[i] = 0.f;
    for (int i = 0; i < K * V; ++i) dS_entry[i] = 0.f;

    // R3
    for (int d = 0; d < K; ++d) {
      float EL = expf(gc[iTK(bh, tlast, d, T, K)]);
      for (int vd = 0; vd < V; ++vd) {
        float g = dS[d * V + vd];
        dS_entry[d * V + vd] += g * EL;
        dEL[d] += g * S[d * V + vd];
        for (int c = 0; c < BT; ++c) {
          int tc = base + c;
          float P = expf(gc[iTK(bh, tlast, d, T, K)] - gc[iTK(bh, tc, d, T, K)]) * kn[iTK(bh, tc, d, T, K)];
          dvcorr[c * V + vd] += P * g;
          dP[c * K + d] += g * vcorr[iTV(bh, tc, vd, T, V)];
        }
      }
    }
    // R2
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float god = do_in[inTV(b, tc, h, vd, T, H, V)];
        for (int d = 0; d < K; ++d) {
          dQE[c * K + d] += god * S[d * V + vd];
          dS_entry[d * V + vd] += qs[iTK(bh, tc, d, T, K)] * expf(gc[iTK(bh, tc, d, T, K)]) * god;
        }
        for (int j = 0; j <= c; ++j) {
          dAqk[c * BT + j] += god * vcorr[iTV(bh, base + j, vd, T, V)];
          dvcorr[j * V + vd] += Aq[c * BT + j] * god;
        }
      }
    }
    // R1
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float g = dvcorr[c * V + vd];
        du[c * V + vd] += g;
        for (int d = 0; d < K; ++d) {
          dw[c * K + d] += -g * S[d * V + vd];
          dS_entry[d * V + vd] += -w[iTK(bh, tc, d, T, K)] * g;
        }
      }
    }
    // adjoint P
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float factor = expf(gc[iTK(bh, tlast, d, T, K)] - gc[iTK(bh, tc, d, T, K)]);
        float P = factor * kn[iTK(bh, tc, d, T, K)];
        float g = dP[c * K + d];
        dkn[iTK(bh, tc, d, T, K)] += g * factor;
        dgc[iTK(bh, tlast, d, T, K)] += g * P;
        dgc[iTK(bh, tc, d, T, K)] += -g * P;
      }
    }
    // adjoint EL
    for (int d = 0; d < K; ++d)
      dgc[iTK(bh, tlast, d, T, K)] += dEL[d] * expf(gc[iTK(bh, tlast, d, T, K)]);
    // adjoint Aqk
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int j = 0; j <= c; ++j) {
        int tj = base + j;
        float g = dAqk[c * BT + j];
        if (g == 0.f) continue;
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, tj, d, T, K)]);
          float term = qs[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, tj, d, T, K)];
          dqs[iTK(bh, tc, d, T, K)] += g * f * kn[iTK(bh, tj, d, T, K)];
          dkn[iTK(bh, tj, d, T, K)] += g * qs[iTK(bh, tc, d, T, K)] * f;
          dgc[iTK(bh, tc, d, T, K)] += g * term;
          dgc[iTK(bh, tj, d, T, K)] += -g * term;
        }
      }
    }
    // adjoint dQE
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float E = expf(gc[iTK(bh, tc, d, T, K)]);
        float g = dQE[c * K + d];
        dqs[iTK(bh, tc, d, T, K)] += g * E;
        dgc[iTK(bh, tc, d, T, K)] += g * qs[iTK(bh, tc, d, T, K)] * E;
      }
    }
    // adjoint w=Af@kg, u=Af@v
    for (int c = 0; c < BT; ++c) {
      for (int d = 0; d < K; ++d) {
        float gwd = dw[c * K + d];
        if (gwd != 0.f)
          for (int i = 0; i < BT; ++i) {
            int ti = base + i;
            float kg = expf(gc[iTK(bh, ti, d, T, K)]) * kn[iTK(bh, ti, d, T, K)];
            dAfull[c * BT + i] += gwd * kg;
            dkg[i * K + d] += gwd * Af[c * BT + i];
          }
      }
      for (int vd = 0; vd < V; ++vd) {
        float gud = du[c * V + vd];
        for (int i = 0; i < BT; ++i) {
          dAfull[c * BT + i] += gud * v_in[inTV(b, base + i, h, vd, T, H, V)];
          dv_out[inTV(b, base + i, h, vd, T, H, V)] += gud * Af[c * BT + i];
        }
      }
    }
    // adjoint kg_decay
    for (int i = 0; i < BT; ++i) {
      int ti = base + i;
      for (int d = 0; d < K; ++d) {
        float E = expf(gc[iTK(bh, ti, d, T, K)]);
        float g = dkg[i * K + d];
        dkn[iTK(bh, ti, d, T, K)] += g * E;
        dgc[iTK(bh, ti, d, T, K)] += g * kn[iTK(bh, ti, d, T, K)] * E;
      }
    }
    // adjoint Afull = X*beta (+diag)
    for (int c = 0; c < BT; ++c)
      for (int i = 0; i < BT; ++i) {
        float xci = Xn[c * BT + i] + (c == i ? 1.0f : 0.0f);
        float g = dAfull[c * BT + i];
        dbeta_s[base + i] += g * xci;
        dX[c * BT + i] = g * beta_s[bh * T + base + i];
      }
    // matrix-inverse VJP: dAraw = Xf^T @ dX @ Xf^T
    for (int c = 0; c < BT; ++c)
      for (int i = 0; i < BT; ++i)
        Xf[c * BT + i] = Xn[c * BT + i] + (c == i ? 1.0f : 0.0f);
    for (int a = 0; a < BT; ++a)
      for (int bb = 0; bb < BT; ++bb) {
        float acc = 0.f;
        for (int m = 0; m < BT; ++m) acc += Xf[m * BT + a] * dX[m * BT + bb];
        tmpm[a * BT + bb] = acc;
      }
    for (int a = 0; a < BT; ++a)
      for (int bb = 0; bb < BT; ++bb) {
        float acc = 0.f;
        for (int m = 0; m < BT; ++m) acc += tmpm[a * BT + m] * Xf[bb * BT + m];
        dAraw[a * BT + bb] = acc;
      }
    // adjoint Araw[c,i] = -beta_s[c]*inner (c>i)
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int i = 0; i < c; ++i) {
        int ti = base + i;
        float g = dAraw[c * BT + i];
        float inner = 0.f;
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, ti, d, T, K)]);
          inner += kn[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, ti, d, T, K)];
        }
        dbeta_s[tc] += g * (-inner);
        float dinner = g * (-beta_s[bh * T + tc]);
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, ti, d, T, K)]);
          float term = kn[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, ti, d, T, K)];
          dkn[iTK(bh, tc, d, T, K)] += dinner * f * kn[iTK(bh, ti, d, T, K)];
          dkn[iTK(bh, ti, d, T, K)] += dinner * kn[iTK(bh, tc, d, T, K)] * f;
          dgc[iTK(bh, tc, d, T, K)] += dinner * term;
          dgc[iTK(bh, ti, d, T, K)] += -dinner * term;
        }
      }
    }
    for (int i = 0; i < K * V; ++i) dS[i] = dS_entry[i];
  }
  for (int i = 0; i < K * V; ++i) dh0_out[bh * K * V + i] = dS[i];
}

__global__ void k_bwd_pre(
    const float* q_raw, const float* k_raw, const float* qn, const float* kn,
    const float* rq, const float* rk, const float* beta_s, const float* ga,
    const float* A_log,
    const float* dgc, const float* dqs, const float* dkn, const float* dbeta_s,
    float* dq_raw, float* dk_raw, float* dg_raw, float* dbeta_raw,
    float* dA_log, float* ddt_bias, float scale,
    int B, int T, int H, int K, int BT) {
  int bh = blockIdx.x * blockDim.x + threadIdx.x;
  if (bh >= B * H) return;
  int b = bh / H, h = bh % H, NT = T / BT;
  float aexp = expf(A_log[h]);

  // reverse cumsum of dgc -> dga, then gate bwd
  float dA_acc = 0.f;
  for (int n = 0; n < NT; ++n)
    for (int d = 0; d < K; ++d) {
      float acc = 0.f;
      for (int c = BT - 1; c >= 0; --c) {
        int t = n * BT + c;
        acc += dgc[iTK(bh, t, d, T, K)];
        // dga = acc ; gate bwd
        float gav = ga[iTK(bh, t, d, T, K)];
        float sp = -gav / aexp;
        float sig = 1.0f - expf(-sp);
        float dxin = acc * (-aexp) * sig;
        dg_raw[inTK(b, t, h, d, T, H, K)] = dxin;
        atomicAdd(&ddt_bias[h * K + d], dxin);
        dA_acc += acc * gav;
      }
    }
  atomicAdd(&dA_log[h], dA_acc);

  // beta bwd
  for (int t = 0; t < T; ++t) {
    float be = beta_s[bh * T + t];
    dbeta_raw[(b * T + t) * H + h] = dbeta_s[bh * T + t] * be * (1.0f - be);
  }

  // l2norm bwd for q (dqn = dqs*scale) and k (dkn from recur)
  for (int t = 0; t < T; ++t) {
    float r_q = rq[bh * T + t], r_k = rk[bh * T + t];
    float sdq = 0.f, sdk = 0.f;
    for (int d = 0; d < K; ++d) {
      float dqn = dqs[iTK(bh, t, d, T, K)] * scale;
      sdq += dqn * qn[iTK(bh, t, d, T, K)];
      sdk += dkn[iTK(bh, t, d, T, K)] * kn[iTK(bh, t, d, T, K)];
    }
    for (int d = 0; d < K; ++d) {
      float dqn = dqs[iTK(bh, t, d, T, K)] * scale;
      dq_raw[inTK(b, t, h, d, T, H, K)] = dqn * r_q - sdq * qn[iTK(bh, t, d, T, K)] * r_q;
      dk_raw[inTK(b, t, h, d, T, H, K)] =
          dkn[iTK(bh, t, d, T, K)] * r_k - sdk * kn[iTK(bh, t, d, T, K)] * r_k;
    }
  }
}

// ---------------- host launchers ----------------

static inline float* fp(torch::Tensor& t) { return t.data_ptr<float>(); }

}  // namespace

std::vector<torch::Tensor> kda_cuda_fwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  int BT = chunk_size, NT = T / BT, BH = B * H;

  auto o = torch::zeros({B, T, H, V}, opt);
  auto fs = torch::zeros({B, H, K, V}, opt);
  auto qn = torch::zeros({BH, T, K}, opt), kn = torch::zeros({BH, T, K}, opt);
  auto qs = torch::zeros({BH, T, K}, opt), ga = torch::zeros({BH, T, K}, opt);
  auto gc = torch::zeros({BH, T, K}, opt);
  auto rq = torch::zeros({BH, T}, opt), rk = torch::zeros({BH, T}, opt);
  auto beta_s = torch::zeros({BH, T}, opt);
  auto X = torch::zeros({BH, NT, BT, BT}, opt);
  auto Afull = torch::zeros({BH, NT, BT, BT}, opt);
  auto Aqk = torch::zeros({BH, NT, BT, BT}, opt);
  auto w = torch::zeros({BH, T, K}, opt), u = torch::zeros({BH, T, V}, opt);
  auto vcorr = torch::zeros({BH, T, V}, opt);
  auto Sentry = torch::zeros({BH, NT + 1, K, V}, opt);

  int threads = 64, blocks = (BH + threads - 1) / threads;
  k_pre<<<blocks, threads>>>(fp(q), fp(k), fp(g), fp(beta), fp(A_log),
      fp(dt_bias), fp(qn), fp(kn), fp(qs), fp(rq), fp(rk), fp(beta_s),
      fp(ga), fp(gc), (float)scale, B, T, H, K, BT);
  k_wy<<<blocks, threads>>>(fp(v), fp(kn), fp(gc), fp(beta_s), fp(X),
      fp(Afull), fp(w), fp(u), B, T, H, K, V, BT);
  k_recur<<<blocks, threads>>>(fp(qs), fp(kn), fp(gc), fp(w), fp(u),
      fp(initial_state), fp(Aqk), fp(vcorr), fp(Sentry), fp(o), fp(fs),
      B, T, H, K, V, BT);
  cudaDeviceSynchronize();
  return {o, fs};
}

std::vector<torch::Tensor> kda_cuda_bwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size,
    torch::Tensor do_, torch::Tensor dht) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);
  do_ = do_.contiguous().to(torch::kFloat32);
  dht = dht.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  int BT = chunk_size, NT = T / BT, BH = B * H;

  // forward scratch (recompute)
  auto qn = torch::zeros({BH, T, K}, opt), kn = torch::zeros({BH, T, K}, opt);
  auto qs = torch::zeros({BH, T, K}, opt), ga = torch::zeros({BH, T, K}, opt);
  auto gc = torch::zeros({BH, T, K}, opt);
  auto rq = torch::zeros({BH, T}, opt), rk = torch::zeros({BH, T}, opt);
  auto beta_s = torch::zeros({BH, T}, opt);
  auto X = torch::zeros({BH, NT, BT, BT}, opt);
  auto Afull = torch::zeros({BH, NT, BT, BT}, opt);
  auto Aqk = torch::zeros({BH, NT, BT, BT}, opt);
  auto w = torch::zeros({BH, T, K}, opt), u = torch::zeros({BH, T, V}, opt);
  auto vcorr = torch::zeros({BH, T, V}, opt);
  auto Sentry = torch::zeros({BH, NT + 1, K, V}, opt);
  auto o = torch::zeros({B, T, H, V}, opt), fs = torch::zeros({B, H, K, V}, opt);

  int threads = 64, blocks = (BH + threads - 1) / threads;
  k_pre<<<blocks, threads>>>(fp(q), fp(k), fp(g), fp(beta), fp(A_log),
      fp(dt_bias), fp(qn), fp(kn), fp(qs), fp(rq), fp(rk), fp(beta_s),
      fp(ga), fp(gc), (float)scale, B, T, H, K, BT);
  k_wy<<<blocks, threads>>>(fp(v), fp(kn), fp(gc), fp(beta_s), fp(X),
      fp(Afull), fp(w), fp(u), B, T, H, K, V, BT);
  k_recur<<<blocks, threads>>>(fp(qs), fp(kn), fp(gc), fp(w), fp(u),
      fp(initial_state), fp(Aqk), fp(vcorr), fp(Sentry), fp(o), fp(fs),
      B, T, H, K, V, BT);

  // outputs
  auto dq = torch::zeros_like(q), dk = torch::zeros_like(k);
  auto dv = torch::zeros_like(v), dgr = torch::zeros_like(g);
  auto dbeta = torch::zeros_like(beta);
  auto dA_log = torch::zeros_like(A_log), ddt_bias = torch::zeros_like(dt_bias);
  auto dh0 = torch::zeros_like(initial_state);

  // backward scratch
  BwdScratch sc;
  auto s_dgc = torch::zeros({BH, T, K}, opt); sc.dgc = fp(s_dgc);
  auto s_dqs = torch::zeros({BH, T, K}, opt); sc.dqs = fp(s_dqs);
  auto s_dkn = torch::zeros({BH, T, K}, opt); sc.dkn = fp(s_dkn);
  auto s_dbeta_s = torch::zeros({BH, T}, opt); sc.dbeta_s = fp(s_dbeta_s);
  auto s_dS = torch::zeros({BH, K, V}, opt); sc.dS = fp(s_dS);
  auto s_dS_entry = torch::zeros({BH, K, V}, opt); sc.dS_entry = fp(s_dS_entry);
  auto s_dEL = torch::zeros({BH, K}, opt); sc.dEL = fp(s_dEL);
  auto s_dvcorr = torch::zeros({BH, BT, V}, opt); sc.dvcorr = fp(s_dvcorr);
  auto s_du = torch::zeros({BH, BT, V}, opt); sc.du = fp(s_du);
  auto s_dAqk = torch::zeros({BH, BT, BT}, opt); sc.dAqk = fp(s_dAqk);
  auto s_dAfull = torch::zeros({BH, BT, BT}, opt); sc.dAfull = fp(s_dAfull);
  auto s_dX = torch::zeros({BH, BT, BT}, opt); sc.dX = fp(s_dX);
  auto s_Xf = torch::zeros({BH, BT, BT}, opt); sc.Xf = fp(s_Xf);
  auto s_tmp = torch::zeros({BH, BT, BT}, opt); sc.tmp = fp(s_tmp);
  auto s_dAraw = torch::zeros({BH, BT, BT}, opt); sc.dAraw = fp(s_dAraw);
  auto s_dQE = torch::zeros({BH, BT, K}, opt); sc.dQE = fp(s_dQE);
  auto s_dP = torch::zeros({BH, BT, K}, opt); sc.dP = fp(s_dP);
  auto s_dw = torch::zeros({BH, BT, K}, opt); sc.dw = fp(s_dw);
  auto s_dkg = torch::zeros({BH, BT, K}, opt); sc.dkg = fp(s_dkg);

  k_bwd_recur<<<blocks, threads>>>(fp(qs), fp(kn), fp(gc), fp(v), fp(w),
      fp(Afull), fp(X), fp(Aqk), fp(vcorr), fp(Sentry), fp(beta_s),
      fp(do_), fp(dht), sc, fp(dv), fp(dh0), B, T, H, K, V, BT);
  k_bwd_pre<<<blocks, threads>>>(fp(q), fp(k), fp(qn), fp(kn), fp(rq), fp(rk),
      fp(beta_s), fp(ga), fp(A_log), sc.dgc, sc.dqs, sc.dkn, sc.dbeta_s,
      fp(dq), fp(dk), fp(dgr), fp(dbeta), fp(dA_log), fp(ddt_bias),
      (float)scale, B, T, H, K, BT);
  cudaDeviceSynchronize();
  return {dq, dk, dv, dgr, dbeta, dA_log, ddt_bias, dh0};
}

3 优化 1

一种直接的优化是：k_pre、k_wy都是chunk间独立的，那么可以一个block处理一个per (batch, head)，block内的每个线程处理一个chunk。k_recur里的Aq也是chunk独立的，剩余部分chunk有依赖，只能按照当前单线程处理。

// KDA chunk reference (CUDA, fp32) — chunk-parallel optimized variant.
//
// Optimization over kda_chunk_cuda.cu (which used one thread per (batch,head)):
//   grid  = B*H blocks, one block per (batch, head)
//   block = NT threads, one thread per chunk
//
// Stage dependency analysis:
//   * k_pre : per-time work (l2norm/beta/gate) is time-independent and the
//             cumsum is within-chunk only  -> chunks independent -> thread/chunk.
//   * k_wy  : each chunk's WY/UT transform depends only on that chunk's data
//             -> chunks independent -> thread/chunk.
//   * k_recur: Aqk[n] depends only on chunk n (parallel), but vcorr/o and the
//             S recurrence carry state across chunks -> all threads compute
//             their Aqk in parallel, __syncthreads(), then thread 0 runs the
//             serial chunk sweep.
//   * k_bwd_pre : reverse-cumsum/gate/beta/l2norm are per-chunk independent
//             -> thread/chunk (atomicAdd into dA_log/ddt_bias).
//   * k_bwd_recur : the reverse sweep carries dS backward across chunks
//             -> inherently serial -> thread 0 only.
//
// Assumes HV == H, fixed sequence length, fp32 throughout.

#include <cuda_runtime.h>
#include <torch/extension.h>

#include <vector>

namespace {

constexpr float L2_EPS = 1e-6f;

__device__ inline float softplusf(float x) {
  if (x > 20.0f) return x;
  if (x < -20.0f) return expf(x);
  return log1pf(expf(x));
}
__device__ inline float sigmoidf(float x) { return 1.0f / (1.0f + expf(-x)); }

__device__ inline int iTK(int bh, int t, int d, int T, int K) { return bh * T * K + t * K + d; }
__device__ inline int iTV(int bh, int t, int vd, int T, int V) { return bh * T * V + t * V + vd; }
__device__ inline int inTK(int b, int t, int h, int d, int T, int H, int K) {
  return ((b * T + t) * H + h) * K + d;
}
__device__ inline int inTV(int b, int t, int h, int vd, int T, int H, int V) {
  return ((b * T + t) * H + h) * V + vd;
}

// ================= FORWARD =================

// grid = B*H, block = NT threads (thread n owns chunk n).
__global__ void k_pre_opt(
    const float* q_raw, const float* k_raw, const float* g_raw,
    const float* beta_raw, const float* A_log, const float* dt_bias,
    float* qn, float* kn, float* qs, float* rq, float* rk, float* beta_s,
    float* ga, float* gc, float scale, int B, int T, int H, int K, int BT, int NT) {
  int bh = blockIdx.x, n = threadIdx.x;
  if (bh >= B * H || n >= NT) return;
  int b = bh / H, h = bh % H;
  float aexp = expf(A_log[h]);
  int base = n * BT;
  for (int c = 0; c < BT; ++c) {
    int t = base + c;
    float sq = 0.f, sk = 0.f;
    for (int d = 0; d < K; ++d) {
      float a = q_raw[inTK(b, t, h, d, T, H, K)]; sq += a * a;
      float cc = k_raw[inTK(b, t, h, d, T, H, K)]; sk += cc * cc;
    }
    float r_q = rsqrtf(sq + L2_EPS), r_k = rsqrtf(sk + L2_EPS);
    rq[bh * T + t] = r_q; rk[bh * T + t] = r_k;
    for (int d = 0; d < K; ++d) {
      float qv = q_raw[inTK(b, t, h, d, T, H, K)] * r_q;
      float kv = k_raw[inTK(b, t, h, d, T, H, K)] * r_k;
      qn[iTK(bh, t, d, T, K)] = qv;
      kn[iTK(bh, t, d, T, K)] = kv;
      qs[iTK(bh, t, d, T, K)] = qv * scale;
      float x = g_raw[inTK(b, t, h, d, T, H, K)] + dt_bias[h * K + d];
      ga[iTK(bh, t, d, T, K)] = -aexp * softplusf(x);
    }
    beta_s[bh * T + t] = sigmoidf(beta_raw[(b * T + t) * H + h]);
  }
  for (int d = 0; d < K; ++d) {
    float acc = 0.f;
    for (int c = 0; c < BT; ++c) {
      int t = base + c;
      acc += ga[iTK(bh, t, d, T, K)];
      gc[iTK(bh, t, d, T, K)] = acc;
    }
  }
}

__global__ void k_wy_opt(
    const float* v_in, const float* kn, const float* gc, const float* beta_s,
    float* X, float* Afull, float* w, float* u,
    int B, int T, int H, int K, int V, int BT, int NT) {
  int bh = blockIdx.x, n = threadIdx.x;
  if (bh >= B * H || n >= NT) return;
  int b = bh / H, h = bh % H;
  float* Xn = X + (bh * NT + n) * BT * BT;
  float* Af = Afull + (bh * NT + n) * BT * BT;
  int base = n * BT;
  for (int c = 0; c < BT * BT; ++c) Xn[c] = 0.f;
  for (int c = 0; c < BT; ++c) {
    int tc = base + c;
    for (int i = 0; i < c; ++i) {
      int ti = base + i;
      float inner = 0.f;
      for (int d = 0; d < K; ++d) {
        float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, ti, d, T, K)]);
        inner += kn[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, ti, d, T, K)];
      }
      Xn[c * BT + i] = -beta_s[bh * T + tc] * inner;
    }
  }
  for (int c = 1; c < BT; ++c)
    for (int j = 0; j < c; ++j) {
      float acc = Xn[c * BT + j];
      for (int m = j + 1; m < c; ++m) acc += Xn[c * BT + m] * Xn[m * BT + j];
      Xn[c * BT + j] = acc;
    }
  for (int c = 0; c < BT; ++c)
    for (int i = 0; i < BT; ++i) {
      float xci = Xn[c * BT + i] + (c == i ? 1.0f : 0.0f);
      Af[c * BT + i] = xci * beta_s[bh * T + base + i];
    }
  for (int c = 0; c < BT; ++c) {
    int tc = base + c;
    for (int d = 0; d < K; ++d) {
      float acc = 0.f;
      for (int i = 0; i < BT; ++i) {
        int ti = base + i;
        float kg = expf(gc[iTK(bh, ti, d, T, K)]) * kn[iTK(bh, ti, d, T, K)];
        acc += Af[c * BT + i] * kg;
      }
      w[iTK(bh, tc, d, T, K)] = acc;
    }
    for (int vd = 0; vd < V; ++vd) {
      float acc = 0.f;
      for (int i = 0; i < BT; ++i)
        acc += Af[c * BT + i] * v_in[inTV(b, base + i, h, vd, T, H, V)];
      u[iTV(bh, tc, vd, T, V)] = acc;
    }
  }
}

// block = NT threads: thread n computes Aqk[n] (parallel), then thread 0 runs
// the serial recurrence over chunks.
__global__ void k_recur_opt(
    const float* qs, const float* kn, const float* gc, const float* w,
    const float* u, const float* h0,
    float* Aqk, float* vcorr, float* Sentry, float* o_out, float* fs_out,
    int B, int T, int H, int K, int V, int BT, int NT) {
  int bh = blockIdx.x, n = threadIdx.x;
  if (bh >= B * H) return;
  int b = bh / H, h = bh % H;
  // parallel: each thread builds Aqk for its chunk
  if (n < NT) {
    int base = n * BT;
    float* Aq = Aqk + (bh * NT + n) * BT * BT;
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int j = 0; j <= c; ++j) {
        int tj = base + j;
        float acc = 0.f;
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, tj, d, T, K)]);
          acc += qs[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, tj, d, T, K)];
        }
        Aq[c * BT + j] = acc;
      }
    }
  }
  __syncthreads();
  if (n != 0) return;
  // serial recurrence (thread 0)
  float* Se = Sentry + bh * (NT + 1) * K * V;
  for (int i = 0; i < K * V; ++i) Se[i] = h0[bh * K * V + i];
  for (int nn = 0; nn < NT; ++nn) {
    int base = nn * BT;
    float* S = Se + nn * K * V;
    float* Snew = Se + (nn + 1) * K * V;
    const float* Aq = Aqk + (bh * NT + nn) * BT * BT;
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float acc = u[iTV(bh, tc, vd, T, V)];
        for (int d = 0; d < K; ++d) acc -= w[iTK(bh, tc, d, T, K)] * S[d * V + vd];
        vcorr[iTV(bh, tc, vd, T, V)] = acc;
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float acc = 0.f;
        for (int d = 0; d < K; ++d)
          acc += qs[iTK(bh, tc, d, T, K)] * expf(gc[iTK(bh, tc, d, T, K)]) * S[d * V + vd];
        for (int j = 0; j <= c; ++j)
          acc += Aq[c * BT + j] * vcorr[iTV(bh, base + j, vd, T, V)];
        o_out[inTV(b, tc, h, vd, T, H, V)] = acc;
      }
    }
    int tlast = base + BT - 1;
    for (int d = 0; d < K; ++d) {
      float glast = gc[iTK(bh, tlast, d, T, K)], EL = expf(glast);
      for (int vd = 0; vd < V; ++vd) {
        float acc = S[d * V + vd] * EL;
        for (int c = 0; c < BT; ++c) {
          int tc = base + c;
          float P = expf(glast - gc[iTK(bh, tc, d, T, K)]) * kn[iTK(bh, tc, d, T, K)];
          acc += P * vcorr[iTV(bh, tc, vd, T, V)];
        }
        Snew[d * V + vd] = acc;
      }
    }
  }
  for (int i = 0; i < K * V; ++i) fs_out[bh * K * V + i] = Se[NT * K * V + i];
}

// ================= BACKWARD =================

struct BwdScratch {
  float *dgc, *dqs, *dkn, *dbeta_s;
  float *dS, *dS_entry, *dEL;
  float *dvcorr, *du;
  float *dAqk, *dAfull, *dX, *Xf, *tmp, *dAraw;
  float *dQE, *dP, *dw, *dkg;
};

// serial reverse sweep (thread 0). Carries dS across chunks.
__global__ void k_bwd_recur_opt(
    const float* qs, const float* kn, const float* gc, const float* v_in,
    const float* w, const float* Afull, const float* X, const float* Aqk,
    const float* vcorr, const float* Sentry, const float* beta_s,
    const float* do_in, const float* dht,
    BwdScratch sc, float* dv_out, float* dh0_out,
    int B, int T, int H, int K, int V, int BT) {
  int bh = blockIdx.x;
  if (bh >= B * H || threadIdx.x != 0) return;
  int b = bh / H, h = bh % H, NT = T / BT;

  float* dgc = sc.dgc;
  float* dqs = sc.dqs;
  float* dkn = sc.dkn;
  float* dbeta_s = sc.dbeta_s + bh * T;
  float* dS = sc.dS + bh * K * V;
  float* dS_entry = sc.dS_entry + bh * K * V;
  float* dEL = sc.dEL + bh * K;
  float* dvcorr = sc.dvcorr + bh * BT * V;
  float* du = sc.du + bh * BT * V;
  float* dAqk = sc.dAqk + bh * BT * BT;
  float* dAfull = sc.dAfull + bh * BT * BT;
  float* dX = sc.dX + bh * BT * BT;
  float* Xf = sc.Xf + bh * BT * BT;
  float* tmpm = sc.tmp + bh * BT * BT;
  float* dAraw = sc.dAraw + bh * BT * BT;
  float* dQE = sc.dQE + bh * BT * K;
  float* dP = sc.dP + bh * BT * K;
  float* dw = sc.dw + bh * BT * K;
  float* dkg = sc.dkg + bh * BT * K;

  for (int i = 0; i < T * K; ++i) {
    int idx = bh * T * K + i;
    dgc[idx] = 0.f; dqs[idx] = 0.f; dkn[idx] = 0.f;
  }
  for (int i = 0; i < T; ++i) dbeta_s[i] = 0.f;
  for (int t = 0; t < T; ++t)
    for (int vd = 0; vd < V; ++vd) dv_out[inTV(b, t, h, vd, T, H, V)] = 0.f;
  for (int i = 0; i < K * V; ++i) dS[i] = dht[bh * K * V + i];

  const float* Sentry_bh = Sentry + bh * (NT + 1) * K * V;

  for (int n = NT - 1; n >= 0; --n) {
    int base = n * BT, tlast = base + BT - 1;
    const float* Aq = Aqk + (bh * NT + n) * BT * BT;
    const float* Af = Afull + (bh * NT + n) * BT * BT;
    const float* Xn = X + (bh * NT + n) * BT * BT;
    const float* S = Sentry_bh + n * K * V;

    for (int i = 0; i < BT * V; ++i) { dvcorr[i] = 0.f; du[i] = 0.f; }
    for (int i = 0; i < BT * BT; ++i) { dAqk[i] = 0.f; dAfull[i] = 0.f; dX[i] = 0.f; }
    for (int i = 0; i < BT * K; ++i) { dQE[i] = 0.f; dP[i] = 0.f; dw[i] = 0.f; dkg[i] = 0.f; }
    for (int i = 0; i < K; ++i) dEL[i] = 0.f;
    for (int i = 0; i < K * V; ++i) dS_entry[i] = 0.f;

    for (int d = 0; d < K; ++d) {
      float EL = expf(gc[iTK(bh, tlast, d, T, K)]);
      for (int vd = 0; vd < V; ++vd) {
        float g = dS[d * V + vd];
        dS_entry[d * V + vd] += g * EL;
        dEL[d] += g * S[d * V + vd];
        for (int c = 0; c < BT; ++c) {
          int tc = base + c;
          float P = expf(gc[iTK(bh, tlast, d, T, K)] - gc[iTK(bh, tc, d, T, K)]) * kn[iTK(bh, tc, d, T, K)];
          dvcorr[c * V + vd] += P * g;
          dP[c * K + d] += g * vcorr[iTV(bh, tc, vd, T, V)];
        }
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float god = do_in[inTV(b, tc, h, vd, T, H, V)];
        for (int d = 0; d < K; ++d) {
          dQE[c * K + d] += god * S[d * V + vd];
          dS_entry[d * V + vd] += qs[iTK(bh, tc, d, T, K)] * expf(gc[iTK(bh, tc, d, T, K)]) * god;
        }
        for (int j = 0; j <= c; ++j) {
          dAqk[c * BT + j] += god * vcorr[iTV(bh, base + j, vd, T, V)];
          dvcorr[j * V + vd] += Aq[c * BT + j] * god;
        }
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float g = dvcorr[c * V + vd];
        du[c * V + vd] += g;
        for (int d = 0; d < K; ++d) {
          dw[c * K + d] += -g * S[d * V + vd];
          dS_entry[d * V + vd] += -w[iTK(bh, tc, d, T, K)] * g;
        }
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float factor = expf(gc[iTK(bh, tlast, d, T, K)] - gc[iTK(bh, tc, d, T, K)]);
        float P = factor * kn[iTK(bh, tc, d, T, K)];
        float g = dP[c * K + d];
        dkn[iTK(bh, tc, d, T, K)] += g * factor;
        dgc[iTK(bh, tlast, d, T, K)] += g * P;
        dgc[iTK(bh, tc, d, T, K)] += -g * P;
      }
    }
    for (int d = 0; d < K; ++d)
      dgc[iTK(bh, tlast, d, T, K)] += dEL[d] * expf(gc[iTK(bh, tlast, d, T, K)]);
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int j = 0; j <= c; ++j) {
        int tj = base + j;
        float g = dAqk[c * BT + j];
        if (g == 0.f) continue;
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, tj, d, T, K)]);
          float term = qs[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, tj, d, T, K)];
          dqs[iTK(bh, tc, d, T, K)] += g * f * kn[iTK(bh, tj, d, T, K)];
          dkn[iTK(bh, tj, d, T, K)] += g * qs[iTK(bh, tc, d, T, K)] * f;
          dgc[iTK(bh, tc, d, T, K)] += g * term;
          dgc[iTK(bh, tj, d, T, K)] += -g * term;
        }
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float E = expf(gc[iTK(bh, tc, d, T, K)]);
        float g = dQE[c * K + d];
        dqs[iTK(bh, tc, d, T, K)] += g * E;
        dgc[iTK(bh, tc, d, T, K)] += g * qs[iTK(bh, tc, d, T, K)] * E;
      }
    }
    for (int c = 0; c < BT; ++c) {
      for (int d = 0; d < K; ++d) {
        float gwd = dw[c * K + d];
        if (gwd != 0.f)
          for (int i = 0; i < BT; ++i) {
            int ti = base + i;
            float kg = expf(gc[iTK(bh, ti, d, T, K)]) * kn[iTK(bh, ti, d, T, K)];
            dAfull[c * BT + i] += gwd * kg;
            dkg[i * K + d] += gwd * Af[c * BT + i];
          }
      }
      for (int vd = 0; vd < V; ++vd) {
        float gud = du[c * V + vd];
        for (int i = 0; i < BT; ++i) {
          dAfull[c * BT + i] += gud * v_in[inTV(b, base + i, h, vd, T, H, V)];
          dv_out[inTV(b, base + i, h, vd, T, H, V)] += gud * Af[c * BT + i];
        }
      }
    }
    for (int i = 0; i < BT; ++i) {
      int ti = base + i;
      for (int d = 0; d < K; ++d) {
        float E = expf(gc[iTK(bh, ti, d, T, K)]);
        float g = dkg[i * K + d];
        dkn[iTK(bh, ti, d, T, K)] += g * E;
        dgc[iTK(bh, ti, d, T, K)] += g * kn[iTK(bh, ti, d, T, K)] * E;
      }
    }
    for (int c = 0; c < BT; ++c)
      for (int i = 0; i < BT; ++i) {
        float xci = Xn[c * BT + i] + (c == i ? 1.0f : 0.0f);
        float g = dAfull[c * BT + i];
        dbeta_s[base + i] += g * xci;
        dX[c * BT + i] = g * beta_s[bh * T + base + i];
      }
    for (int c = 0; c < BT; ++c)
      for (int i = 0; i < BT; ++i)
        Xf[c * BT + i] = Xn[c * BT + i] + (c == i ? 1.0f : 0.0f);
    for (int a = 0; a < BT; ++a)
      for (int bb = 0; bb < BT; ++bb) {
        float acc = 0.f;
        for (int m = 0; m < BT; ++m) acc += Xf[m * BT + a] * dX[m * BT + bb];
        tmpm[a * BT + bb] = acc;
      }
    for (int a = 0; a < BT; ++a)
      for (int bb = 0; bb < BT; ++bb) {
        float acc = 0.f;
        for (int m = 0; m < BT; ++m) acc += tmpm[a * BT + m] * Xf[bb * BT + m];
        dAraw[a * BT + bb] = acc;
      }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int i = 0; i < c; ++i) {
        int ti = base + i;
        float g = dAraw[c * BT + i];
        float inner = 0.f;
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, ti, d, T, K)]);
          inner += kn[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, ti, d, T, K)];
        }
        dbeta_s[tc] += g * (-inner);
        float dinner = g * (-beta_s[bh * T + tc]);
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, ti, d, T, K)]);
          float term = kn[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, ti, d, T, K)];
          dkn[iTK(bh, tc, d, T, K)] += dinner * f * kn[iTK(bh, ti, d, T, K)];
          dkn[iTK(bh, ti, d, T, K)] += dinner * kn[iTK(bh, tc, d, T, K)] * f;
          dgc[iTK(bh, tc, d, T, K)] += dinner * term;
          dgc[iTK(bh, ti, d, T, K)] += -dinner * term;
        }
      }
    }
    for (int i = 0; i < K * V; ++i) dS[i] = dS_entry[i];
  }
  for (int i = 0; i < K * V; ++i) dh0_out[bh * K * V + i] = dS[i];
}

// per-chunk independent: thread n owns chunk n.
__global__ void k_bwd_pre_opt(
    const float* q_raw, const float* k_raw, const float* qn, const float* kn,
    const float* rq, const float* rk, const float* beta_s, const float* ga,
    const float* A_log,
    const float* dgc, const float* dqs, const float* dkn, const float* dbeta_s,
    float* dq_raw, float* dk_raw, float* dg_raw, float* dbeta_raw,
    float* dA_log, float* ddt_bias, float scale,
    int B, int T, int H, int K, int BT, int NT) {
  int bh = blockIdx.x, n = threadIdx.x;
  if (bh >= B * H || n >= NT) return;
  int b = bh / H, h = bh % H;
  float aexp = expf(A_log[h]);
  int base = n * BT;

  float dA_acc = 0.f;
  for (int d = 0; d < K; ++d) {
    float acc = 0.f;
    for (int c = BT - 1; c >= 0; --c) {
      int t = base + c;
      acc += dgc[iTK(bh, t, d, T, K)];
      float gav = ga[iTK(bh, t, d, T, K)];
      float sp = -gav / aexp;
      float sig = 1.0f - expf(-sp);
      float dxin = acc * (-aexp) * sig;
      dg_raw[inTK(b, t, h, d, T, H, K)] = dxin;
      atomicAdd(&ddt_bias[h * K + d], dxin);
      dA_acc += acc * gav;
    }
  }
  atomicAdd(&dA_log[h], dA_acc);

  for (int c = 0; c < BT; ++c) {
    int t = base + c;
    float be = beta_s[bh * T + t];
    dbeta_raw[(b * T + t) * H + h] = dbeta_s[bh * T + t] * be * (1.0f - be);
  }

  for (int c = 0; c < BT; ++c) {
    int t = base + c;
    float r_q = rq[bh * T + t], r_k = rk[bh * T + t];
    float sdq = 0.f, sdk = 0.f;
    for (int d = 0; d < K; ++d) {
      float dqn = dqs[iTK(bh, t, d, T, K)] * scale;
      sdq += dqn * qn[iTK(bh, t, d, T, K)];
      sdk += dkn[iTK(bh, t, d, T, K)] * kn[iTK(bh, t, d, T, K)];
    }
    for (int d = 0; d < K; ++d) {
      float dqn = dqs[iTK(bh, t, d, T, K)] * scale;
      dq_raw[inTK(b, t, h, d, T, H, K)] = dqn * r_q - sdq * qn[iTK(bh, t, d, T, K)] * r_q;
      dk_raw[inTK(b, t, h, d, T, H, K)] =
          dkn[iTK(bh, t, d, T, K)] * r_k - sdk * kn[iTK(bh, t, d, T, K)] * r_k;
    }
  }
}

static inline float* fp(torch::Tensor& t) { return t.data_ptr<float>(); }

}  // namespace

std::vector<torch::Tensor> kda_cuda_fwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  int BT = chunk_size, NT = T / BT, BH = B * H;

  auto o = torch::zeros({B, T, H, V}, opt);
  auto fs = torch::zeros({B, H, K, V}, opt);
  auto qn = torch::zeros({BH, T, K}, opt), kn = torch::zeros({BH, T, K}, opt);
  auto qs = torch::zeros({BH, T, K}, opt), ga = torch::zeros({BH, T, K}, opt);
  auto gc = torch::zeros({BH, T, K}, opt);
  auto rq = torch::zeros({BH, T}, opt), rk = torch::zeros({BH, T}, opt);
  auto beta_s = torch::zeros({BH, T}, opt);
  auto X = torch::zeros({BH, NT, BT, BT}, opt);
  auto Afull = torch::zeros({BH, NT, BT, BT}, opt);
  auto Aqk = torch::zeros({BH, NT, BT, BT}, opt);
  auto w = torch::zeros({BH, T, K}, opt), u = torch::zeros({BH, T, V}, opt);
  auto vcorr = torch::zeros({BH, T, V}, opt);
  auto Sentry = torch::zeros({BH, NT + 1, K, V}, opt);

  dim3 grid(BH), block(NT);
  k_pre_opt<<<grid, block>>>(fp(q), fp(k), fp(g), fp(beta), fp(A_log),
      fp(dt_bias), fp(qn), fp(kn), fp(qs), fp(rq), fp(rk), fp(beta_s),
      fp(ga), fp(gc), (float)scale, B, T, H, K, BT, NT);
  k_wy_opt<<<grid, block>>>(fp(v), fp(kn), fp(gc), fp(beta_s), fp(X),
      fp(Afull), fp(w), fp(u), B, T, H, K, V, BT, NT);
  k_recur_opt<<<grid, block>>>(fp(qs), fp(kn), fp(gc), fp(w), fp(u),
      fp(initial_state), fp(Aqk), fp(vcorr), fp(Sentry), fp(o), fp(fs),
      B, T, H, K, V, BT, NT);
  cudaDeviceSynchronize();
  return {o, fs};
}

std::vector<torch::Tensor> kda_cuda_bwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size,
    torch::Tensor do_, torch::Tensor dht) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);
  do_ = do_.contiguous().to(torch::kFloat32);
  dht = dht.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  int BT = chunk_size, NT = T / BT, BH = B * H;

  auto qn = torch::zeros({BH, T, K}, opt), kn = torch::zeros({BH, T, K}, opt);
  auto qs = torch::zeros({BH, T, K}, opt), ga = torch::zeros({BH, T, K}, opt);
  auto gc = torch::zeros({BH, T, K}, opt);
  auto rq = torch::zeros({BH, T}, opt), rk = torch::zeros({BH, T}, opt);
  auto beta_s = torch::zeros({BH, T}, opt);
  auto X = torch::zeros({BH, NT, BT, BT}, opt);
  auto Afull = torch::zeros({BH, NT, BT, BT}, opt);
  auto Aqk = torch::zeros({BH, NT, BT, BT}, opt);
  auto w = torch::zeros({BH, T, K}, opt), u = torch::zeros({BH, T, V}, opt);
  auto vcorr = torch::zeros({BH, T, V}, opt);
  auto Sentry = torch::zeros({BH, NT + 1, K, V}, opt);
  auto o = torch::zeros({B, T, H, V}, opt), fs = torch::zeros({B, H, K, V}, opt);

  dim3 grid(BH), block(NT);
  k_pre_opt<<<grid, block>>>(fp(q), fp(k), fp(g), fp(beta), fp(A_log),
      fp(dt_bias), fp(qn), fp(kn), fp(qs), fp(rq), fp(rk), fp(beta_s),
      fp(ga), fp(gc), (float)scale, B, T, H, K, BT, NT);
  k_wy_opt<<<grid, block>>>(fp(v), fp(kn), fp(gc), fp(beta_s), fp(X),
      fp(Afull), fp(w), fp(u), B, T, H, K, V, BT, NT);
  k_recur_opt<<<grid, block>>>(fp(qs), fp(kn), fp(gc), fp(w), fp(u),
      fp(initial_state), fp(Aqk), fp(vcorr), fp(Sentry), fp(o), fp(fs),
      B, T, H, K, V, BT, NT);

  auto dq = torch::zeros_like(q), dk = torch::zeros_like(k);
  auto dv = torch::zeros_like(v), dgr = torch::zeros_like(g);
  auto dbeta = torch::zeros_like(beta);
  auto dA_log = torch::zeros_like(A_log), ddt_bias = torch::zeros_like(dt_bias);
  auto dh0 = torch::zeros_like(initial_state);

  BwdScratch sc;
  auto s_dgc = torch::zeros({BH, T, K}, opt); sc.dgc = fp(s_dgc);
  auto s_dqs = torch::zeros({BH, T, K}, opt); sc.dqs = fp(s_dqs);
  auto s_dkn = torch::zeros({BH, T, K}, opt); sc.dkn = fp(s_dkn);
  auto s_dbeta_s = torch::zeros({BH, T}, opt); sc.dbeta_s = fp(s_dbeta_s);
  auto s_dS = torch::zeros({BH, K, V}, opt); sc.dS = fp(s_dS);
  auto s_dS_entry = torch::zeros({BH, K, V}, opt); sc.dS_entry = fp(s_dS_entry);
  auto s_dEL = torch::zeros({BH, K}, opt); sc.dEL = fp(s_dEL);
  auto s_dvcorr = torch::zeros({BH, BT, V}, opt); sc.dvcorr = fp(s_dvcorr);
  auto s_du = torch::zeros({BH, BT, V}, opt); sc.du = fp(s_du);
  auto s_dAqk = torch::zeros({BH, BT, BT}, opt); sc.dAqk = fp(s_dAqk);
  auto s_dAfull = torch::zeros({BH, BT, BT}, opt); sc.dAfull = fp(s_dAfull);
  auto s_dX = torch::zeros({BH, BT, BT}, opt); sc.dX = fp(s_dX);
  auto s_Xf = torch::zeros({BH, BT, BT}, opt); sc.Xf = fp(s_Xf);
  auto s_tmp = torch::zeros({BH, BT, BT}, opt); sc.tmp = fp(s_tmp);
  auto s_dAraw = torch::zeros({BH, BT, BT}, opt); sc.dAraw = fp(s_dAraw);
  auto s_dQE = torch::zeros({BH, BT, K}, opt); sc.dQE = fp(s_dQE);
  auto s_dP = torch::zeros({BH, BT, K}, opt); sc.dP = fp(s_dP);
  auto s_dw = torch::zeros({BH, BT, K}, opt); sc.dw = fp(s_dw);
  auto s_dkg = torch::zeros({BH, BT, K}, opt); sc.dkg = fp(s_dkg);

  k_bwd_recur_opt<<<grid, block>>>(fp(qs), fp(kn), fp(gc), fp(v), fp(w),
      fp(Afull), fp(X), fp(Aqk), fp(vcorr), fp(Sentry), fp(beta_s),
      fp(do_), fp(dht), sc, fp(dv), fp(dh0), B, T, H, K, V, BT);
  k_bwd_pre_opt<<<grid, block>>>(fp(q), fp(k), fp(qn), fp(kn), fp(rq), fp(rk),
      fp(beta_s), fp(ga), fp(A_log), sc.dgc, sc.dqs, sc.dkn, sc.dbeta_s,
      fp(dq), fp(dk), fp(dgr), fp(dbeta), fp(dA_log), fp(ddt_bias),
      (float)scale, B, T, H, K, BT, NT);
  cudaDeviceSynchronize();
  return {dq, dk, dv, dgr, dbeta, dA_log, ddt_bias, dh0};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fwd", &kda_cuda_fwd, "KDA chunk forward (CUDA, chunk-parallel)");
  m.def("bwd", &kda_cuda_bwd, "KDA chunk backward (CUDA, chunk-parallel)");
}

4 优化 2

k_pre、k_wy都是chunk间独立的，那么可以一个block处理一个per (batch, head)，block内的每个warp（32线程）处理一个chunk，chunk内K或者V维度分块（K/warp），BT维度不分，warp依次处理每K/V维度的个分块，warp内线程顺序处理即可。k_recur，有chunk依赖，用一个block处理一个chunk，block内每个线程处理一个k/v维度，block直接串行。

// KDA chunk reference (CUDA, fp32) — warp/block-parallel optimized variant (v2).
//
// Over v1 (kda_chunk_cuda_opt.cu, 1 thread/chunk):
//   * k_pre / k_wy / k_recur_aqk : grid = B*H blocks; one WARP (32 lanes) per
//     chunk (warps stride over chunks). Lanes split the feature dim (K or V);
//     BT is looped. K-reductions use warp-shuffle. -> 32x more threads/chunk,
//     coalesced loads.
//   * k_recur_seq : one BLOCK per (batch,head); thread = one v-dim (vd). The
//     block loops chunks SERIALLY (the S carry stays in the block), threads
//     parallelize the K/V inner loops. NOTE: CUDA does not order blocks, so the
//     chunk dependency is serialized *inside* one block, not across blocks.
//   * backward : same structure as v1 (chunk-parallel k_bwd_pre; serial
//     k_bwd_recur on thread 0, since dS carries backward across chunks).
//
// Assumes HV == H, fixed length, fp32, chunk_size BT <= 64.

#include <cuda_runtime.h>
#include <torch/extension.h>

#include <vector>

namespace {

constexpr float L2_EPS = 1e-6f;

__device__ inline float softplusf(float x) {
  if (x > 20.0f) return x;
  if (x < -20.0f) return expf(x);
  return log1pf(expf(x));
}
__device__ inline float sigmoidf(float x) { return 1.0f / (1.0f + expf(-x)); }
__device__ inline float warp_sum(float v) {
  for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o);
  return v;
}

__device__ inline int iTK(int bh, int t, int d, int T, int K) { return bh * T * K + t * K + d; }
__device__ inline int iTV(int bh, int t, int vd, int T, int V) { return bh * T * V + t * V + vd; }
__device__ inline int inTK(int b, int t, int h, int d, int T, int H, int K) {
  return ((b * T + t) * H + h) * K + d;
}
__device__ inline int inTV(int b, int t, int h, int vd, int T, int H, int V) {
  return ((b * T + t) * H + h) * V + vd;
}

// ================= FORWARD =================

// grid = B*H ; block = nwarps*32 ; warp w owns chunks {w, w+nwarps, ...}.
__global__ void k_pre_v2(
    const float* q_raw, const float* k_raw, const float* g_raw,
    const float* beta_raw, const float* A_log, const float* dt_bias,
    float* qn, float* kn, float* qs, float* rq, float* rk, float* beta_s,
    float* ga, float* gc, float scale, int B, int T, int H, int K, int BT, int NT) {
  int bh = blockIdx.x;
  if (bh >= B * H) return;
  int b = bh / H, h = bh % H;
  int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, nwarps = blockDim.x >> 5;
  float aexp = expf(A_log[h]);
  for (int n = warp; n < NT; n += nwarps) {
    int base = n * BT;
    for (int c = 0; c < BT; ++c) {
      int t = base + c;
      float sq = 0.f, sk = 0.f;
      for (int d = lane; d < K; d += 32) {
        float a = q_raw[inTK(b, t, h, d, T, H, K)]; sq += a * a;
        float cc = k_raw[inTK(b, t, h, d, T, H, K)]; sk += cc * cc;
      }
      sq = warp_sum(sq); sk = warp_sum(sk);
      float r_q = rsqrtf(sq + L2_EPS), r_k = rsqrtf(sk + L2_EPS);
      if (lane == 0) { rq[bh * T + t] = r_q; rk[bh * T + t] = r_k; }
      for (int d = lane; d < K; d += 32) {
        float qv = q_raw[inTK(b, t, h, d, T, H, K)] * r_q;
        float kv = k_raw[inTK(b, t, h, d, T, H, K)] * r_k;
        qn[iTK(bh, t, d, T, K)] = qv;
        kn[iTK(bh, t, d, T, K)] = kv;
        qs[iTK(bh, t, d, T, K)] = qv * scale;
        float x = g_raw[inTK(b, t, h, d, T, H, K)] + dt_bias[h * K + d];
        ga[iTK(bh, t, d, T, K)] = -aexp * softplusf(x);
      }
      if (lane == 0) beta_s[bh * T + t] = sigmoidf(beta_raw[(b * T + t) * H + h]);
    }
    for (int d = lane; d < K; d += 32) {
      float acc = 0.f;
      for (int c = 0; c < BT; ++c) {
        int t = base + c;
        acc += ga[iTK(bh, t, d, T, K)];
        gc[iTK(bh, t, d, T, K)] = acc;
      }
    }
  }
}

// warp per chunk. K-reductions over lanes; triangular inverse parallel over j
// (rows sequential) using register snapshot to avoid intra-row races (BT<=64).
__global__ void k_wy_v2(
    const float* v_in, const float* kn, const float* gc, const float* beta_s,
    float* X, float* Afull, float* w, float* u,
    int B, int T, int H, int K, int V, int BT, int NT) {
  int bh = blockIdx.x;
  if (bh >= B * H) return;
  int b = bh / H, h = bh % H;
  int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, nwarps = blockDim.x >> 5;
  for (int n = warp; n < NT; n += nwarps) {
    float* Xn = X + (bh * NT + n) * BT * BT;
    float* Af = Afull + (bh * NT + n) * BT * BT;
    int base = n * BT;
    for (int idx = lane; idx < BT * BT; idx += 32) Xn[idx] = 0.f;
    __syncwarp();
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int i = 0; i < c; ++i) {
        int ti = base + i;
        float inner = 0.f;
        for (int d = lane; d < K; d += 32) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, ti, d, T, K)]);
          inner += kn[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, ti, d, T, K)];
        }
        inner = warp_sum(inner);
        if (lane == 0) Xn[c * BT + i] = -beta_s[bh * T + tc] * inner;
      }
    }
    __syncwarp();
    // forward substitution: rows c sequential; lanes own j (and j+32).
    for (int c = 1; c < BT; ++c) {
      int j0 = lane, j1 = lane + 32;
      float a0 = 0.f, a1 = 0.f;
      if (j0 < c) {
        a0 = Xn[c * BT + j0];
        for (int m = j0 + 1; m < c; ++m) a0 += Xn[c * BT + m] * Xn[m * BT + j0];
      }
      if (j1 < c) {
        a1 = Xn[c * BT + j1];
        for (int m = j1 + 1; m < c; ++m) a1 += Xn[c * BT + m] * Xn[m * BT + j1];
      }
      __syncwarp();
      if (j0 < c) Xn[c * BT + j0] = a0;
      if (j1 < c) Xn[c * BT + j1] = a1;
      __syncwarp();
    }
    for (int idx = lane; idx < BT * BT; idx += 32) {
      int c = idx / BT, i = idx % BT;
      float xci = Xn[c * BT + i] + (c == i ? 1.0f : 0.0f);
      Af[idx] = xci * beta_s[bh * T + base + i];
    }
    __syncwarp();
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = lane; d < K; d += 32) {
        float acc = 0.f;
        for (int i = 0; i < BT; ++i) {
          int ti = base + i;
          float kg = expf(gc[iTK(bh, ti, d, T, K)]) * kn[iTK(bh, ti, d, T, K)];
          acc += Af[c * BT + i] * kg;
        }
        w[iTK(bh, tc, d, T, K)] = acc;
      }
      for (int vd = lane; vd < V; vd += 32) {
        float acc = 0.f;
        for (int i = 0; i < BT; ++i)
          acc += Af[c * BT + i] * v_in[inTV(b, base + i, h, vd, T, H, V)];
        u[iTV(bh, tc, vd, T, V)] = acc;
      }
    }
  }
}

// warp per chunk: Aqk[c,j] = sum_d qs*exp(gc[c]-gc[j])*kn (j<=c).
__global__ void k_aqk_v2(
    const float* qs, const float* kn, const float* gc, float* Aqk,
    int B, int T, int H, int K, int BT, int NT) {
  int bh = blockIdx.x;
  if (bh >= B * H) return;
  int warp = threadIdx.x >> 5, lane = threadIdx.x & 31, nwarps = blockDim.x >> 5;
  for (int n = warp; n < NT; n += nwarps) {
    int base = n * BT;
    float* Aq = Aqk + (bh * NT + n) * BT * BT;
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int j = 0; j <= c; ++j) {
        int tj = base + j;
        float acc = 0.f;
        for (int d = lane; d < K; d += 32) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, tj, d, T, K)]);
          acc += qs[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, tj, d, T, K)];
        }
        acc = warp_sum(acc);
        if (lane == 0) Aq[c * BT + j] = acc;
      }
    }
  }
}

// one block per (b,h); thread = vd. Loops chunks serially (S carry in block).
__global__ void k_recur_seq_v2(
    const float* qs, const float* kn, const float* gc, const float* w,
    const float* u, const float* h0, const float* Aqk,
    float* vcorr, float* Sentry, float* o_out, float* fs_out,
    int B, int T, int H, int K, int V, int BT, int NT) {
  int bh = blockIdx.x;
  if (bh >= B * H) return;
  int b = bh / H, h = bh % H, vd = threadIdx.x;
  float* Se = Sentry + bh * (NT + 1) * K * V;
  if (vd < V)
    for (int d = 0; d < K; ++d) Se[d * V + vd] = h0[bh * K * V + d * V + vd];
  __syncthreads();
  for (int n = 0; n < NT; ++n) {
    int base = n * BT;
    float* S = Se + n * K * V;
    float* Snew = Se + (n + 1) * K * V;
    const float* Aq = Aqk + (bh * NT + n) * BT * BT;
    if (vd < V)
      for (int c = 0; c < BT; ++c) {
        int tc = base + c;
        float acc = u[iTV(bh, tc, vd, T, V)];
        for (int d = 0; d < K; ++d) acc -= w[iTK(bh, tc, d, T, K)] * S[d * V + vd];
        vcorr[iTV(bh, tc, vd, T, V)] = acc;
      }
    __syncthreads();
    if (vd < V)
      for (int c = 0; c < BT; ++c) {
        int tc = base + c;
        float acc = 0.f;
        for (int d = 0; d < K; ++d)
          acc += qs[iTK(bh, tc, d, T, K)] * expf(gc[iTK(bh, tc, d, T, K)]) * S[d * V + vd];
        for (int j = 0; j <= c; ++j)
          acc += Aq[c * BT + j] * vcorr[iTV(bh, base + j, vd, T, V)];
        o_out[inTV(b, tc, h, vd, T, H, V)] = acc;
      }
    int tlast = base + BT - 1;
    if (vd < V)
      for (int d = 0; d < K; ++d) {
        float glast = gc[iTK(bh, tlast, d, T, K)], EL = expf(glast);
        float acc = S[d * V + vd] * EL;
        for (int c = 0; c < BT; ++c) {
          int tc = base + c;
          float P = expf(glast - gc[iTK(bh, tc, d, T, K)]) * kn[iTK(bh, tc, d, T, K)];
          acc += P * vcorr[iTV(bh, tc, vd, T, V)];
        }
        Snew[d * V + vd] = acc;
      }
    __syncthreads();
  }
  if (vd < V)
    for (int d = 0; d < K; ++d) fs_out[bh * K * V + d * V + vd] = Se[NT * K * V + d * V + vd];
}

// ================= BACKWARD (same structure as v1) =================

struct BwdScratch {
  float *dgc, *dqs, *dkn, *dbeta_s;
  float *dS, *dS_entry, *dEL;
  float *dvcorr, *du;
  float *dAqk, *dAfull, *dX, *Xf, *tmp, *dAraw;
  float *dQE, *dP, *dw, *dkg;
};

// serial reverse sweep (thread 0). Carries dS across chunks.
__global__ void k_bwd_recur_v2(
    const float* qs, const float* kn, const float* gc, const float* v_in,
    const float* w, const float* Afull, const float* X, const float* Aqk,
    const float* vcorr, const float* Sentry, const float* beta_s,
    const float* do_in, const float* dht,
    BwdScratch sc, float* dv_out, float* dh0_out,
    int B, int T, int H, int K, int V, int BT) {
  int bh = blockIdx.x;
  if (bh >= B * H || threadIdx.x != 0) return;
  int b = bh / H, h = bh % H, NT = T / BT;

  float* dgc = sc.dgc;
  float* dqs = sc.dqs;
  float* dkn = sc.dkn;
  float* dbeta_s = sc.dbeta_s + bh * T;
  float* dS = sc.dS + bh * K * V;
  float* dS_entry = sc.dS_entry + bh * K * V;
  float* dEL = sc.dEL + bh * K;
  float* dvcorr = sc.dvcorr + bh * BT * V;
  float* du = sc.du + bh * BT * V;
  float* dAqk = sc.dAqk + bh * BT * BT;
  float* dAfull = sc.dAfull + bh * BT * BT;
  float* dX = sc.dX + bh * BT * BT;
  float* Xf = sc.Xf + bh * BT * BT;
  float* tmpm = sc.tmp + bh * BT * BT;
  float* dAraw = sc.dAraw + bh * BT * BT;
  float* dQE = sc.dQE + bh * BT * K;
  float* dP = sc.dP + bh * BT * K;
  float* dw = sc.dw + bh * BT * K;
  float* dkg = sc.dkg + bh * BT * K;

  for (int i = 0; i < T * K; ++i) {
    int idx = bh * T * K + i;
    dgc[idx] = 0.f; dqs[idx] = 0.f; dkn[idx] = 0.f;
  }
  for (int i = 0; i < T; ++i) dbeta_s[i] = 0.f;
  for (int t = 0; t < T; ++t)
    for (int vd = 0; vd < V; ++vd) dv_out[inTV(b, t, h, vd, T, H, V)] = 0.f;
  for (int i = 0; i < K * V; ++i) dS[i] = dht[bh * K * V + i];

  const float* Sentry_bh = Sentry + bh * (NT + 1) * K * V;

  for (int n = NT - 1; n >= 0; --n) {
    int base = n * BT, tlast = base + BT - 1;
    const float* Aq = Aqk + (bh * NT + n) * BT * BT;
    const float* Af = Afull + (bh * NT + n) * BT * BT;
    const float* Xn = X + (bh * NT + n) * BT * BT;
    const float* S = Sentry_bh + n * K * V;

    for (int i = 0; i < BT * V; ++i) { dvcorr[i] = 0.f; du[i] = 0.f; }
    for (int i = 0; i < BT * BT; ++i) { dAqk[i] = 0.f; dAfull[i] = 0.f; dX[i] = 0.f; }
    for (int i = 0; i < BT * K; ++i) { dQE[i] = 0.f; dP[i] = 0.f; dw[i] = 0.f; dkg[i] = 0.f; }
    for (int i = 0; i < K; ++i) dEL[i] = 0.f;
    for (int i = 0; i < K * V; ++i) dS_entry[i] = 0.f;

    for (int d = 0; d < K; ++d) {
      float EL = expf(gc[iTK(bh, tlast, d, T, K)]);
      for (int vd = 0; vd < V; ++vd) {
        float g = dS[d * V + vd];
        dS_entry[d * V + vd] += g * EL;
        dEL[d] += g * S[d * V + vd];
        for (int c = 0; c < BT; ++c) {
          int tc = base + c;
          float P = expf(gc[iTK(bh, tlast, d, T, K)] - gc[iTK(bh, tc, d, T, K)]) * kn[iTK(bh, tc, d, T, K)];
          dvcorr[c * V + vd] += P * g;
          dP[c * K + d] += g * vcorr[iTV(bh, tc, vd, T, V)];
        }
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float god = do_in[inTV(b, tc, h, vd, T, H, V)];
        for (int d = 0; d < K; ++d) {
          dQE[c * K + d] += god * S[d * V + vd];
          dS_entry[d * V + vd] += qs[iTK(bh, tc, d, T, K)] * expf(gc[iTK(bh, tc, d, T, K)]) * god;
        }
        for (int j = 0; j <= c; ++j) {
          dAqk[c * BT + j] += god * vcorr[iTV(bh, base + j, vd, T, V)];
          dvcorr[j * V + vd] += Aq[c * BT + j] * god;
        }
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int vd = 0; vd < V; ++vd) {
        float g = dvcorr[c * V + vd];
        du[c * V + vd] += g;
        for (int d = 0; d < K; ++d) {
          dw[c * K + d] += -g * S[d * V + vd];
          dS_entry[d * V + vd] += -w[iTK(bh, tc, d, T, K)] * g;
        }
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float factor = expf(gc[iTK(bh, tlast, d, T, K)] - gc[iTK(bh, tc, d, T, K)]);
        float P = factor * kn[iTK(bh, tc, d, T, K)];
        float g = dP[c * K + d];
        dkn[iTK(bh, tc, d, T, K)] += g * factor;
        dgc[iTK(bh, tlast, d, T, K)] += g * P;
        dgc[iTK(bh, tc, d, T, K)] += -g * P;
      }
    }
    for (int d = 0; d < K; ++d)
      dgc[iTK(bh, tlast, d, T, K)] += dEL[d] * expf(gc[iTK(bh, tlast, d, T, K)]);
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int j = 0; j <= c; ++j) {
        int tj = base + j;
        float g = dAqk[c * BT + j];
        if (g == 0.f) continue;
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, tj, d, T, K)]);
          float term = qs[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, tj, d, T, K)];
          dqs[iTK(bh, tc, d, T, K)] += g * f * kn[iTK(bh, tj, d, T, K)];
          dkn[iTK(bh, tj, d, T, K)] += g * qs[iTK(bh, tc, d, T, K)] * f;
          dgc[iTK(bh, tc, d, T, K)] += g * term;
          dgc[iTK(bh, tj, d, T, K)] += -g * term;
        }
      }
    }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int d = 0; d < K; ++d) {
        float E = expf(gc[iTK(bh, tc, d, T, K)]);
        float g = dQE[c * K + d];
        dqs[iTK(bh, tc, d, T, K)] += g * E;
        dgc[iTK(bh, tc, d, T, K)] += g * qs[iTK(bh, tc, d, T, K)] * E;
      }
    }
    for (int c = 0; c < BT; ++c) {
      for (int d = 0; d < K; ++d) {
        float gwd = dw[c * K + d];
        if (gwd != 0.f)
          for (int i = 0; i < BT; ++i) {
            int ti = base + i;
            float kg = expf(gc[iTK(bh, ti, d, T, K)]) * kn[iTK(bh, ti, d, T, K)];
            dAfull[c * BT + i] += gwd * kg;
            dkg[i * K + d] += gwd * Af[c * BT + i];
          }
      }
      for (int vd = 0; vd < V; ++vd) {
        float gud = du[c * V + vd];
        for (int i = 0; i < BT; ++i) {
          dAfull[c * BT + i] += gud * v_in[inTV(b, base + i, h, vd, T, H, V)];
          dv_out[inTV(b, base + i, h, vd, T, H, V)] += gud * Af[c * BT + i];
        }
      }
    }
    for (int i = 0; i < BT; ++i) {
      int ti = base + i;
      for (int d = 0; d < K; ++d) {
        float E = expf(gc[iTK(bh, ti, d, T, K)]);
        float g = dkg[i * K + d];
        dkn[iTK(bh, ti, d, T, K)] += g * E;
        dgc[iTK(bh, ti, d, T, K)] += g * kn[iTK(bh, ti, d, T, K)] * E;
      }
    }
    for (int c = 0; c < BT; ++c)
      for (int i = 0; i < BT; ++i) {
        float xci = Xn[c * BT + i] + (c == i ? 1.0f : 0.0f);
        float g = dAfull[c * BT + i];
        dbeta_s[base + i] += g * xci;
        dX[c * BT + i] = g * beta_s[bh * T + base + i];
      }
    for (int c = 0; c < BT; ++c)
      for (int i = 0; i < BT; ++i)
        Xf[c * BT + i] = Xn[c * BT + i] + (c == i ? 1.0f : 0.0f);
    for (int a = 0; a < BT; ++a)
      for (int bb = 0; bb < BT; ++bb) {
        float acc = 0.f;
        for (int m = 0; m < BT; ++m) acc += Xf[m * BT + a] * dX[m * BT + bb];
        tmpm[a * BT + bb] = acc;
      }
    for (int a = 0; a < BT; ++a)
      for (int bb = 0; bb < BT; ++bb) {
        float acc = 0.f;
        for (int m = 0; m < BT; ++m) acc += tmpm[a * BT + m] * Xf[bb * BT + m];
        dAraw[a * BT + bb] = acc;
      }
    for (int c = 0; c < BT; ++c) {
      int tc = base + c;
      for (int i = 0; i < c; ++i) {
        int ti = base + i;
        float g = dAraw[c * BT + i];
        float inner = 0.f;
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, ti, d, T, K)]);
          inner += kn[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, ti, d, T, K)];
        }
        dbeta_s[tc] += g * (-inner);
        float dinner = g * (-beta_s[bh * T + tc]);
        for (int d = 0; d < K; ++d) {
          float f = expf(gc[iTK(bh, tc, d, T, K)] - gc[iTK(bh, ti, d, T, K)]);
          float term = kn[iTK(bh, tc, d, T, K)] * f * kn[iTK(bh, ti, d, T, K)];
          dkn[iTK(bh, tc, d, T, K)] += dinner * f * kn[iTK(bh, ti, d, T, K)];
          dkn[iTK(bh, ti, d, T, K)] += dinner * kn[iTK(bh, tc, d, T, K)] * f;
          dgc[iTK(bh, tc, d, T, K)] += dinner * term;
          dgc[iTK(bh, ti, d, T, K)] += -dinner * term;
        }
      }
    }
    for (int i = 0; i < K * V; ++i) dS[i] = dS_entry[i];
  }
  for (int i = 0; i < K * V; ++i) dh0_out[bh * K * V + i] = dS[i];
}

// per-chunk independent: thread n owns chunk n.
__global__ void k_bwd_pre_v2(
    const float* q_raw, const float* k_raw, const float* qn, const float* kn,
    const float* rq, const float* rk, const float* beta_s, const float* ga,
    const float* A_log,
    const float* dgc, const float* dqs, const float* dkn, const float* dbeta_s,
    float* dq_raw, float* dk_raw, float* dg_raw, float* dbeta_raw,
    float* dA_log, float* ddt_bias, float scale,
    int B, int T, int H, int K, int BT, int NT) {
  int bh = blockIdx.x, n = threadIdx.x;
  if (bh >= B * H || n >= NT) return;
  int b = bh / H, h = bh % H;
  float aexp = expf(A_log[h]);
  int base = n * BT;

  float dA_acc = 0.f;
  for (int d = 0; d < K; ++d) {
    float acc = 0.f;
    for (int c = BT - 1; c >= 0; --c) {
      int t = base + c;
      acc += dgc[iTK(bh, t, d, T, K)];
      float gav = ga[iTK(bh, t, d, T, K)];
      float sp = -gav / aexp;
      float sig = 1.0f - expf(-sp);
      float dxin = acc * (-aexp) * sig;
      dg_raw[inTK(b, t, h, d, T, H, K)] = dxin;
      atomicAdd(&ddt_bias[h * K + d], dxin);
      dA_acc += acc * gav;
    }
  }
  atomicAdd(&dA_log[h], dA_acc);

  for (int c = 0; c < BT; ++c) {
    int t = base + c;
    float be = beta_s[bh * T + t];
    dbeta_raw[(b * T + t) * H + h] = dbeta_s[bh * T + t] * be * (1.0f - be);
  }

  for (int c = 0; c < BT; ++c) {
    int t = base + c;
    float r_q = rq[bh * T + t], r_k = rk[bh * T + t];
    float sdq = 0.f, sdk = 0.f;
    for (int d = 0; d < K; ++d) {
      float dqn = dqs[iTK(bh, t, d, T, K)] * scale;
      sdq += dqn * qn[iTK(bh, t, d, T, K)];
      sdk += dkn[iTK(bh, t, d, T, K)] * kn[iTK(bh, t, d, T, K)];
    }
    for (int d = 0; d < K; ++d) {
      float dqn = dqs[iTK(bh, t, d, T, K)] * scale;
      dq_raw[inTK(b, t, h, d, T, H, K)] = dqn * r_q - sdq * qn[iTK(bh, t, d, T, K)] * r_q;
      dk_raw[inTK(b, t, h, d, T, H, K)] =
          dkn[iTK(bh, t, d, T, K)] * r_k - sdk * kn[iTK(bh, t, d, T, K)] * r_k;
    }
  }
}

static inline float* fp(torch::Tensor& t) { return t.data_ptr<float>(); }

}  // namespace

// ---------------- host launchers ----------------

static int fwd_block_threads(int NT) {
  int warps = NT < 32 ? NT : 32;  // up to 32 warps; warps stride over chunks
  if (warps < 1) warps = 1;
  return warps * 32;
}

std::vector<torch::Tensor> kda_cuda_fwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  int BT = chunk_size, NT = T / BT, BH = B * H;

  auto o = torch::zeros({B, T, H, V}, opt);
  auto fs = torch::zeros({B, H, K, V}, opt);
  auto qn = torch::zeros({BH, T, K}, opt), kn = torch::zeros({BH, T, K}, opt);
  auto qs = torch::zeros({BH, T, K}, opt), ga = torch::zeros({BH, T, K}, opt);
  auto gc = torch::zeros({BH, T, K}, opt);
  auto rq = torch::zeros({BH, T}, opt), rk = torch::zeros({BH, T}, opt);
  auto beta_s = torch::zeros({BH, T}, opt);
  auto X = torch::zeros({BH, NT, BT, BT}, opt);
  auto Afull = torch::zeros({BH, NT, BT, BT}, opt);
  auto Aqk = torch::zeros({BH, NT, BT, BT}, opt);
  auto w = torch::zeros({BH, T, K}, opt), u = torch::zeros({BH, T, V}, opt);
  auto vcorr = torch::zeros({BH, T, V}, opt);
  auto Sentry = torch::zeros({BH, NT + 1, K, V}, opt);

  int wth = fwd_block_threads(NT);
  k_pre_v2<<<BH, wth>>>(fp(q), fp(k), fp(g), fp(beta), fp(A_log),
      fp(dt_bias), fp(qn), fp(kn), fp(qs), fp(rq), fp(rk), fp(beta_s),
      fp(ga), fp(gc), (float)scale, B, T, H, K, BT, NT);
  k_wy_v2<<<BH, wth>>>(fp(v), fp(kn), fp(gc), fp(beta_s), fp(X),
      fp(Afull), fp(w), fp(u), B, T, H, K, V, BT, NT);
  k_aqk_v2<<<BH, wth>>>(fp(qs), fp(kn), fp(gc), fp(Aqk), B, T, H, K, BT, NT);
  k_recur_seq_v2<<<BH, V>>>(fp(qs), fp(kn), fp(gc), fp(w), fp(u),
      fp(initial_state), fp(Aqk), fp(vcorr), fp(Sentry), fp(o), fp(fs),
      B, T, H, K, V, BT, NT);
  cudaDeviceSynchronize();
  return {o, fs};
}

std::vector<torch::Tensor> kda_cuda_bwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size,
    torch::Tensor do_, torch::Tensor dht) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);
  do_ = do_.contiguous().to(torch::kFloat32);
  dht = dht.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  int BT = chunk_size, NT = T / BT, BH = B * H;

  auto qn = torch::zeros({BH, T, K}, opt), kn = torch::zeros({BH, T, K}, opt);
  auto qs = torch::zeros({BH, T, K}, opt), ga = torch::zeros({BH, T, K}, opt);
  auto gc = torch::zeros({BH, T, K}, opt);
  auto rq = torch::zeros({BH, T}, opt), rk = torch::zeros({BH, T}, opt);
  auto beta_s = torch::zeros({BH, T}, opt);
  auto X = torch::zeros({BH, NT, BT, BT}, opt);
  auto Afull = torch::zeros({BH, NT, BT, BT}, opt);
  auto Aqk = torch::zeros({BH, NT, BT, BT}, opt);
  auto w = torch::zeros({BH, T, K}, opt), u = torch::zeros({BH, T, V}, opt);
  auto vcorr = torch::zeros({BH, T, V}, opt);
  auto Sentry = torch::zeros({BH, NT + 1, K, V}, opt);
  auto o = torch::zeros({B, T, H, V}, opt), fs = torch::zeros({B, H, K, V}, opt);

  int wth = fwd_block_threads(NT);
  k_pre_v2<<<BH, wth>>>(fp(q), fp(k), fp(g), fp(beta), fp(A_log),
      fp(dt_bias), fp(qn), fp(kn), fp(qs), fp(rq), fp(rk), fp(beta_s),
      fp(ga), fp(gc), (float)scale, B, T, H, K, BT, NT);
  k_wy_v2<<<BH, wth>>>(fp(v), fp(kn), fp(gc), fp(beta_s), fp(X),
      fp(Afull), fp(w), fp(u), B, T, H, K, V, BT, NT);
  k_aqk_v2<<<BH, wth>>>(fp(qs), fp(kn), fp(gc), fp(Aqk), B, T, H, K, BT, NT);
  k_recur_seq_v2<<<BH, V>>>(fp(qs), fp(kn), fp(gc), fp(w), fp(u),
      fp(initial_state), fp(Aqk), fp(vcorr), fp(Sentry), fp(o), fp(fs),
      B, T, H, K, V, BT, NT);

  auto dq = torch::zeros_like(q), dk = torch::zeros_like(k);
  auto dv = torch::zeros_like(v), dgr = torch::zeros_like(g);
  auto dbeta = torch::zeros_like(beta);
  auto dA_log = torch::zeros_like(A_log), ddt_bias = torch::zeros_like(dt_bias);
  auto dh0 = torch::zeros_like(initial_state);

  BwdScratch sc;
  auto s_dgc = torch::zeros({BH, T, K}, opt); sc.dgc = fp(s_dgc);
  auto s_dqs = torch::zeros({BH, T, K}, opt); sc.dqs = fp(s_dqs);
  auto s_dkn = torch::zeros({BH, T, K}, opt); sc.dkn = fp(s_dkn);
  auto s_dbeta_s = torch::zeros({BH, T}, opt); sc.dbeta_s = fp(s_dbeta_s);
  auto s_dS = torch::zeros({BH, K, V}, opt); sc.dS = fp(s_dS);
  auto s_dS_entry = torch::zeros({BH, K, V}, opt); sc.dS_entry = fp(s_dS_entry);
  auto s_dEL = torch::zeros({BH, K}, opt); sc.dEL = fp(s_dEL);
  auto s_dvcorr = torch::zeros({BH, BT, V}, opt); sc.dvcorr = fp(s_dvcorr);
  auto s_du = torch::zeros({BH, BT, V}, opt); sc.du = fp(s_du);
  auto s_dAqk = torch::zeros({BH, BT, BT}, opt); sc.dAqk = fp(s_dAqk);
  auto s_dAfull = torch::zeros({BH, BT, BT}, opt); sc.dAfull = fp(s_dAfull);
  auto s_dX = torch::zeros({BH, BT, BT}, opt); sc.dX = fp(s_dX);
  auto s_Xf = torch::zeros({BH, BT, BT}, opt); sc.Xf = fp(s_Xf);
  auto s_tmp = torch::zeros({BH, BT, BT}, opt); sc.tmp = fp(s_tmp);
  auto s_dAraw = torch::zeros({BH, BT, BT}, opt); sc.dAraw = fp(s_dAraw);
  auto s_dQE = torch::zeros({BH, BT, K}, opt); sc.dQE = fp(s_dQE);
  auto s_dP = torch::zeros({BH, BT, K}, opt); sc.dP = fp(s_dP);
  auto s_dw = torch::zeros({BH, BT, K}, opt); sc.dw = fp(s_dw);
  auto s_dkg = torch::zeros({BH, BT, K}, opt); sc.dkg = fp(s_dkg);

  k_bwd_recur_v2<<<BH, 32>>>(fp(qs), fp(kn), fp(gc), fp(v), fp(w),
      fp(Afull), fp(X), fp(Aqk), fp(vcorr), fp(Sentry), fp(beta_s),
      fp(do_), fp(dht), sc, fp(dv), fp(dh0), B, T, H, K, V, BT);
  k_bwd_pre_v2<<<BH, NT>>>(fp(q), fp(k), fp(qn), fp(kn), fp(rq), fp(rk),
      fp(beta_s), fp(ga), fp(A_log), sc.dgc, sc.dqs, sc.dkn, sc.dbeta_s,
      fp(dq), fp(dk), fp(dgr), fp(dbeta), fp(dA_log), fp(ddt_bias),
      (float)scale, B, T, H, K, BT, NT);
  cudaDeviceSynchronize();
  return {dq, dk, dv, dgr, dbeta, dA_log, ddt_bias, dh0};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fwd", &kda_cuda_fwd, "KDA chunk forward (CUDA v2, warp/block-parallel)");
  m.def("bwd", &kda_cuda_bwd, "KDA chunk backward (CUDA v2)");
}

实现要点

核心 Kernel 设计

k_pre / k_wy / k_aqk

• Grid: B×H 个 block
• Warp 调度: 一个 warp 处理一个 chunk，warps 跨 chunk 步进
• Lane 分工: lane 切分 K/V 维度
• BT 循环: 块内串行处理时间步

关键优化:

• K-reduction (L2 norm、Araw 内积、Aqk 内积) 使用 __shfl_xor warp-shuffle 归约
• Global load 沿 lane 连续 → 合并访存
• WY 三角求逆: 行 c 串行，但行内对 j 并行。已证明行内各 j 互不依赖，只需将整行快照进寄存器再写回即可避免竞争。BT ≤ 64 时每 lane 最多处理 2 个 j

k_recur_seq

• 调度策略: 一个 block 处理一个 (b, h) 对
• 线程映射: 每个 thread = 一个 vd (维度 V)，即 block 包含 V 个线程
• 串行递推: block 内串行遍历所有 chunk (S 进位保留在 block 内部)
• 并行内循环: 线程并行执行 K/V 内层循环，子步间使用 __syncthreads 同步

⚠️ 重要说明: CUDA 中不同 block 之间没有执行顺序保证，因此不能依赖"一个 block 处理一个 chunk、blocks 串行执行"的假设。正确做法是将 chunk 依赖的串行化放在单个 (b, h) 的 block 内部，语义等价且安全。

---

5 优化 3

v2 基线 — 已有特性

kda_chunk_cuda_v2.cu（load_cuda_v2）。在 v1 基础上的前向重设计：

• grid=BH；k_pre/k_wy/k_aqk 使用 每块 1 个 warp（warp 按 chunk 步进， lane 拆分 K/V，warp_sum 规约）；WY 三角逆采用并行-over-j + 寄存器快照。
• k_recur_seq = 每个 (b,h) 一个 block，在块内串行循环 chunk （S 状态驻留），子步骤间 __syncthreads。
• 所有计算为标量 fp32：softplus 门控导致 ki = kn·exp(−gc)、L 构建、 三角逆和 Aqk 无法使用张量核心（一旦 cumsum(g) < −88，exp(−gc) 在 fp32 中溢出）。
• 循环层仅有 BH 个 block → H20（约 78 SMs × 多 block 能力）利用率严重不足； 串行的 chunk 间 S 依赖 + 子步骤间 barrier 占据主导。

构建于 v2 的 warp/block 并行前向传播基础之上，但将所有不会溢出的矩阵乘法都路由到可复用的 CuTe 块级 GEMM 辅助函数中处理。

下面是cute的gemm封装

// CuTe block-level shared-memory tile-gemm helper for KDA v8 (CHUNK=16).
//
//   cute_gemm16<M,N,Kk>(sA, sB, sC)
//     C[M,N] = A[M,Kk] @ B[N,Kk]^T   (TN: A row-major [M,Kk], B row-major [N,Kk],
//     C row-major [M,N]; fp32 accum).  Operands already in shared (plain
//     row-major).  Uses an AtomLayout<1,2,1> over SM80_16x8x16 => 64 threads,
//     MMA tile 16x16x16, so it accepts M%16==0, N%16==0, Kk%16==0 — in
//     particular the BT=16 shapes that the 2x2/32-tile v7 helper cannot do.
//
//   cute_gemm16_cg<M,N,Kk>(sA, sB, gC, ldc)
//     same but C written straight to GLOBAL with row stride ldc.
//
// Only the first 64 threads participate; callers __syncthreads() around these.

#pragma once

#include <cuda_bf16.h>
#include <cute/tensor.hpp>
#include <cute/atom/mma_atom.hpp>
#include <cute/atom/copy_atom.hpp>

namespace v8 {

using bf16 = __nv_bfloat16;
using namespace cute;

static constexpr int GEMM16_THREADS = 64;

// AtomLayout<1,2,1>: 1 warp in M, 2 warps in N over the 16x8x16 atom =>
// natural TiledMMA tile (16, 16, 16), 64 threads.
template <int M, int N, int Kk>
using MMA16 = decltype(make_tiled_mma(
    MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>{},
    make_layout(make_shape(_1{}, _2{}, _1{}))));

// ---- C in shared ----
template <int M, int N, int Kk>
__device__ inline void cute_gemm16_idx(const bf16* sAp, const bf16* sBp,
                                       float* sCp, int idx) {
  Tensor sA = make_tensor(make_smem_ptr(sAp),
      make_layout(make_shape(Int<M>{}, Int<Kk>{}), make_stride(Int<Kk>{}, _1{})));
  Tensor sB = make_tensor(make_smem_ptr(sBp),
      make_layout(make_shape(Int<N>{}, Int<Kk>{}), make_stride(Int<Kk>{}, _1{})));
  Tensor sC = make_tensor(make_smem_ptr(sCp),
      make_layout(make_shape(Int<M>{}, Int<N>{}), make_stride(Int<N>{}, _1{})));

  MMA16<M, N, Kk> tiled_mma;
  auto thr_mma = tiled_mma.get_slice(idx);
  auto tCsA = thr_mma.partition_A(sA);
  auto tCsB = thr_mma.partition_B(sB);
  auto tCrA = thr_mma.make_fragment_A(tCsA);  // (MMA,MMA_M,MMA_K)
  auto tCrB = thr_mma.make_fragment_B(tCsB);  // (MMA,MMA_N,MMA_K)
  auto tCrC = thr_mma.partition_fragment_C(sC);
  clear(tCrC);

  int nk = size<2>(tCrA);
#pragma unroll
  for (int ik = 0; ik < nk; ++ik) {
    copy(tCsA(_, _, ik), tCrA(_, _, ik));
    copy(tCsB(_, _, ik), tCrB(_, _, ik));
    gemm(tiled_mma, tCrC, tCrA(_, _, ik), tCrB(_, _, ik), tCrC);
  }
  copy(tCrC, thr_mma.partition_C(sC));
}
template <int M, int N, int Kk>
__device__ inline void cute_gemm16(const bf16* sAp, const bf16* sBp, float* sCp) {
  cute_gemm16_idx<M, N, Kk>(sAp, sBp, sCp, threadIdx.x);
}

// ---- C straight to global (row stride ldc) ----
template <int M, int N, int Kk>
__device__ inline void cute_gemm16_cg(const bf16* sAp, const bf16* sBp,
                                      float* gCp, int ldc) {
  int idx = threadIdx.x;
  Tensor sA = make_tensor(make_smem_ptr(sAp),
      make_layout(make_shape(Int<M>{}, Int<Kk>{}), make_stride(Int<Kk>{}, _1{})));
  Tensor sB = make_tensor(make_smem_ptr(sBp),
      make_layout(make_shape(Int<N>{}, Int<Kk>{}), make_stride(Int<Kk>{}, _1{})));
  Tensor gC = make_tensor(make_gmem_ptr(gCp),
      make_layout(make_shape(Int<M>{}, Int<N>{}), make_stride(ldc, _1{})));

  MMA16<M, N, Kk> tiled_mma;
  auto thr_mma = tiled_mma.get_slice(idx);
  auto tCsA = thr_mma.partition_A(sA);
  auto tCsB = thr_mma.partition_B(sB);
  auto tCrA = thr_mma.make_fragment_A(tCsA);
  auto tCrB = thr_mma.make_fragment_B(tCsB);
  auto tCrC = thr_mma.partition_fragment_C(gC);
  clear(tCrC);

  int nk = size<2>(tCrA);
#pragma unroll
  for (int ik = 0; ik < nk; ++ik) {
    copy(tCsA(_, _, ik), tCrA(_, _, ik));
    copy(tCsB(_, _, ik), tCrB(_, _, ik));
    gemm(tiled_mma, tCrC, tCrA(_, _, ik), tCrB(_, _, ik), tCrC);
  }
  copy(tCrC, thr_mma.partition_C(gC));
}

// ---- C straight to global as bf16 (row stride ldc) ----
// Same TN gemm but the fp32 accumulator is converted to bf16 on store, so the
// per-chunk operands land in global as bf16 (lets K2 use 128-bit cp.async with
// no fp32->bf16 conversion in the load warp).
template <int M, int N, int Kk>
__device__ inline void cute_gemm16_cg_bf(const bf16* sAp, const bf16* sBp,
                                         bf16* gCp, int ldc) {
  int idx = threadIdx.x;
  Tensor sA = make_tensor(make_smem_ptr(sAp),
      make_layout(make_shape(Int<M>{}, Int<Kk>{}), make_stride(Int<Kk>{}, _1{})));
  Tensor sB = make_tensor(make_smem_ptr(sBp),
      make_layout(make_shape(Int<N>{}, Int<Kk>{}), make_stride(Int<Kk>{}, _1{})));
  Tensor gC = make_tensor(make_gmem_ptr(gCp),
      make_layout(make_shape(Int<M>{}, Int<N>{}), make_stride(ldc, _1{})));

  MMA16<M, N, Kk> tiled_mma;
  auto thr_mma = tiled_mma.get_slice(idx);
  auto tCsA = thr_mma.partition_A(sA);
  auto tCsB = thr_mma.partition_B(sB);
  auto tCrA = thr_mma.make_fragment_A(tCsA);
  auto tCrB = thr_mma.make_fragment_B(tCsB);
  auto tCrC = thr_mma.partition_fragment_C(gC);  // fp32 accum
  clear(tCrC);

  int nk = size<2>(tCrA);
#pragma unroll
  for (int ik = 0; ik < nk; ++ik) {
    copy(tCsA(_, _, ik), tCrA(_, _, ik));
    copy(tCsB(_, _, ik), tCrB(_, _, ik));
    gemm(tiled_mma, tCrC, tCrA(_, _, ik), tCrB(_, _, ik), tCrC);
  }
  // convert fp32 fragment -> bf16 fragment, then store to bf16 global.
  auto tCrC_bf = make_fragment_like<bf16>(tCrC);
#pragma unroll
  for (int i = 0; i < size(tCrC); ++i)
    tCrC_bf(i) = __float2bfloat16(tCrC(i));
  copy(tCrC_bf, thr_mma.partition_C(gC));
}

// ---- runtime-dimension dispatcher over the BT=16 shapes the v8 kernels use ----
// C[M,N] = A[M,Kk] @ B[N,Kk]^T.  (K==V in {64,128}; VC=32; BT=16.)
#define V8_GEMM_CASE(MM, NN, KK) \
  if (M == (MM) && N == (NN) && Kk == (KK)) { cute_gemm16<MM, NN, KK>(sAp, sBp, sCp); return; }
__device__ inline void cute_gemm16_d(const bf16* sAp, const bf16* sBp,
                                     float* sCp, int M, int N, int Kk) {
  // L = kd[16,K] @ ki[16,K]^T -> [16,16];  Aqk = qd@ki^T -> [16,16].
  V8_GEMM_CASE(16, 16, 64)
  V8_GEMM_CASE(16, 16, 128)
  // recur wS/qS = A[16,K] @ St[VC,K]^T -> [16,VC].
  V8_GEMM_CASE(16, 32, 64)
  V8_GEMM_CASE(16, 32, 128)
  // recur AV = Aqk[16,16] @ vcorrT[VC,16]^T -> [16,VC].
  V8_GEMM_CASE(16, 32, 16)
  // recur newS = vcorrT[VC,16] @ P[K,16]^T -> [VC,K].
  V8_GEMM_CASE(32, 64, 16)
  V8_GEMM_CASE(32, 128, 16)
}
#undef V8_GEMM_CASE

// ---- same dispatcher, explicit thread index (lets two independent gemms run on
// two disjoint 64-thread warp-groups in parallel: group0 idx=tid, group1
// idx=tid-64). ----
#define V8_GEMM_IDX_CASE(MM, NN, KK) \
  if (M == (MM) && N == (NN) && Kk == (KK)) { cute_gemm16_idx<MM, NN, KK>(sAp, sBp, sCp, idx); return; }
__device__ inline void cute_gemm16_d_idx(const bf16* sAp, const bf16* sBp,
                                         float* sCp, int M, int N, int Kk, int idx) {
  V8_GEMM_IDX_CASE(16, 16, 64)
  V8_GEMM_IDX_CASE(16, 16, 128)
  V8_GEMM_IDX_CASE(16, 32, 64)
  V8_GEMM_IDX_CASE(16, 32, 128)
  V8_GEMM_IDX_CASE(16, 32, 16)
  V8_GEMM_IDX_CASE(32, 64, 16)
  V8_GEMM_IDX_CASE(32, 128, 16)
}
#undef V8_GEMM_IDX_CASE

// ---- same dispatcher, C straight to global (row stride ldc). ----
#define V8_GEMM_CG_CASE(MM, NN, KK) \
  if (M == (MM) && N == (NN) && Kk == (KK)) { cute_gemm16_cg<MM, NN, KK>(sAp, sBp, gCp, ldc); return; }
__device__ inline void cute_gemm16_cg_d(const bf16* sAp, const bf16* sBp,
                                        float* gCp, int ldc, int M, int N, int Kk) {
  // w = Af[16,16] @ kd[K,16... -> actually Af@kg]: M=16,N=K,Kk=16; u likewise.
  V8_GEMM_CG_CASE(16, 64, 16)
  V8_GEMM_CG_CASE(16, 128, 16)
}
#undef V8_GEMM_CG_CASE

// ---- bf16-output cg dispatcher (w/u written bf16 to global). ----
#define V8_GEMM_CGBF_CASE(MM, NN, KK) \
  if (M == (MM) && N == (NN) && Kk == (KK)) { cute_gemm16_cg_bf<MM, NN, KK>(sAp, sBp, gCp, ldc); return; }
__device__ inline void cute_gemm16_cg_bf_d(const bf16* sAp, const bf16* sBp,
                                           bf16* gCp, int ldc, int M, int N, int Kk) {
  V8_GEMM_CGBF_CASE(16, 64, 16)
  V8_GEMM_CGBF_CASE(16, 128, 16)
}
#undef V8_GEMM_CGBF_CASE

}  // namespace v8

// KDA chunk reference (CUDA, bf16) — v8 CHUNK=16 + bounded-gate redesign.
//
// FORWARD-ONLY. This is a structural redesign over v7 that adopts FlashKDA's
// two key choices so that the three FlashKDA-inspired optimizations actually
// pay off (in the BT=64 + softplus build they were neutral/regressive):
//
//   * CHUNK = 16 (BT=16) and the BOUNDED gate
//        ga = lower_bound * sigmoid(exp(A_log) * (g_raw + dt_bias))   in [lb, 0)
//     With BT=16 and lb=-5 the worst within-chunk cumsum span is 15*(-5)=-75,
//     so exp(-gc) <= e^75 ~ 3.7e32 < bf16 max 3.4e38 => EVERY matmul (incl. the
//     L matrix and Aqk, which v7 had to keep scalar fp32 because ki=kn*exp(-gc)
//     overflowed) is now overflow-safe and goes to bf16 tensor cores.
//
// Two-kernel structure mirroring FlashKDA (forward only):
//   K1  k_prepare_v8 (opt2 = fusion): grid=BH*NT, 1 CTA per (chunk,head), 16
//       rows. l2norm + bounded gate + cumsum, then on TC (cute_gemm16):
//         L = kd@ki^T (strict-lower, *(-beta)), triangular inverse (I+L)^-1,
//         Aqk = qd@ki^T (lower-tri), w = Af@kd, u = Af@v.
//       kd=kn*e^{gc}, ki=kn*e^{-gc}, qd=qn*scale*e^{gc}. Stores w,u,qd,Aqk,kr,
//       glast to a per-chunk workspace (former k_pre+k_aqk+k_wy fused into one).
//   K2  k_recur_v8 (opt1 + opt3): grid=BH*VS, V-split (VCOLS=32), S^T[VC,K]
//       resident in shared, serial chunk loop. Per chunk on TC:
//         vcorr = u - w@S^T ; o = qd@S^T + Aqk@vcorr ; S^T = S^T*e^{glast} + vcorr^T@kr^T.
//       opt3: forward skips per-chunk Sentry/vcorr global writes (only o and
//       final_state written).  opt1: with BT=16 the recur kernel is ~50KB (not
//       ~100KB), so an occupancy-safe double-buffer of the St-INDEPENDENT
//       operands (w/qd/Aqk/kr) is feasible (PIPELINE template flag) — these are
//       prepared in K1 and don't depend on the running state, so chunk t+1 can
//       be staged while chunk t computes.
//
// Validation: chunk_kda rejects chunk_size=16, so v8 (BT=16) is validated vs
//   chunk_kda(safe_gate=True, lower_bound=-5, chunk_size=64) — the forward o /
//   final_state are mathematically chunk-size-independent — and cross-checked vs
//   fused_recurrent_kda(lower_bound=-5).
//
// Assumes HV==H, fixed length, K==V multiple of 16.
//
// NOTE: the BACKWARD path below is left as the v7 (softplus / BT in {32,64})
// kernels verbatim and is intentionally NOT updated or validated — v8 scope is
// forward-only. kda_cuda_bwd still launches the old *_v7 kernels so the module
// keeps compiling, but it does not match v8's CHUNK=16 + bounded-gate forward.

#include <cuda_runtime.h>
#include <torch/extension.h>

#include <vector>

#include "cute_gemm_v7.cuh"
#include "cute_gemm_v8.cuh"

namespace {

using v7::bf16;
constexpr float L2_EPS = 1e-6f;

__device__ inline float softplusf(float x) {
  if (x > 20.0f) return x;
  if (x < -20.0f) return expf(x);
  return log1pf(expf(x));
}
__device__ inline float sigmoidf(float x) { return 1.0f / (1.0f + expf(-x)); }
// store a float into an output slot of type OT (float passthrough / bf16 cast).
__device__ __forceinline__ void ot_store(float* p, float v) { *p = v; }
__device__ __forceinline__ void ot_store(bf16* p, float v) { *p = __float2bfloat16(v); }
__device__ inline float warp_sum(float v) {
  for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o);
  return v;
}
// Named barrier: synchronize exactly `cnt` threads on logical barrier `id`
// (id 1..15). Lets a warp-specialized kernel sync only its compute warps
// (cnt=COMPUTE_THREADS) while the dedicated load warp stays off that barrier.
__device__ __forceinline__ void bar_sync(int id, int cnt) {
  asm volatile("barrier.sync %0, %1;" :: "r"(id), "r"(cnt) : "memory");
}
// Arrive (no wait) on named barrier `id` expecting `cnt` participants.
__device__ __forceinline__ void bar_arrive(int id, int cnt) {
  asm volatile("barrier.arrive %0, %1;" :: "r"(id), "r"(cnt) : "memory");
}
// 16-byte cp.async global->shared (8 bf16). dst must be a shared address.
__device__ __forceinline__ void cp_async16(bf16* dst, const bf16* src) {
  unsigned s = (unsigned)__cvta_generic_to_shared(dst);
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: "r"(s), "l"(src));
}
__device__ __forceinline__ void cp_async_commit() {
  asm volatile("cp.async.commit_group;\n");
}
__device__ __forceinline__ void cp_async_wait_all() {
  asm volatile("cp.async.wait_group 0;\n");
}

// ---- optional in-kernel per-region timing (clock64). OFF by default: when
// V8_PROFILE is not defined every macro below expands to nothing, so the
// production kernels are byte-identical. ON: thread 0 accumulates cycle deltas
// per region into a __device__ global, read back via the v8 pybind. ----
#ifdef V8_PROFILE
__device__ unsigned long long g_prof_k1[16];
__device__ unsigned long long g_prof_k2[16];
#define PROF_DECL long long _pc = 0;
#define PROF_START() do { if (threadIdx.x == 0) _pc = clock64(); } while (0)
#define PROF_ADD(arr, i) do { if (threadIdx.x == 0) { \
    long long _n = clock64(); atomicAdd(&arr[i], (unsigned long long)(_n - _pc)); _pc = _n; } } while (0)
#else
#define PROF_DECL
#define PROF_START()
#define PROF_ADD(arr, i)
#endif

__device__ inline int iTK(int bh, int t, int d, int T, int K) { return bh * T * K + t * K + d; }
__device__ inline int iTV(int bh, int t, int vd, int T, int V) { return bh * T * V + t * V + vd; }
__device__ inline int inTK(int b, int t, int h, int d, int T, int H, int K) {
  return ((b * T + t) * H + h) * K + d;
}
__device__ inline int inTV(int b, int t, int h, int vd, int T, int H, int V) {
  return ((b * T + t) * H + h) * V + vd;
}

// ================= v8 FORWARD (CHUNK=16, bounded gate, all-TC) =================
//
// Per-chunk workspace produced by k_prepare_v8 (layout [BH, NT, ...]):
//   w8   [BH,NT,16,K]   = Af @ kd            (TC)
//   u8   [BH,NT,16,V]   = Af @ v             (TC)
//   qd8  [BH,NT,16,K]   = qn*scale*exp(gc)   (q_decayed, ready for q@S)
//   kr8  [BH,NT,16,K]   = kn*exp(glast-gc)   (k_restored, for S update)
//   Aqk8 [BH,NT,16,16]  = qd@ki^T lower-tri  (TC)
//   gl8  [BH,NT,K]      = gc_last per dim    (decay EL = exp(gl8))
constexpr int BT16 = 16;

extern __shared__ char prep_smem[];
// OT = workspace output dtype for the per-chunk operands (w8/u8/qd8/kr8/Aqk8).
// float = baseline/pipeline K2; bf16 = warp_spec K2 (lets it cp.async-load
// directly with no fp32->bf16 conversion). gl8 stays fp32 either way.
template <typename OT>
__global__ void k_prepare_v8(
    const float* q_raw, const float* k_raw, const float* g_raw,
    const float* beta_raw, const float* A_log, const float* dt_bias,
    const float* v_in,
    OT* w8, OT* u8, OT* qd8, OT* kr8, OT* Aqk8, float* gl8,
    float scale, float lower_bound, int B, int T, int H, int K, int V, int NT) {
  int blk = blockIdx.x;
  if (blk >= B * H * NT) return;
  int bh = blk / NT, n = blk % NT;
  int b = bh / H, h = bh % H;
  int base = n * BT16, tid = threadIdx.x;
  int warp = tid >> 5, lane = tid & 31, nwarps = blockDim.x >> 5;
  float aexp = expf(A_log[h]);
  PROF_DECL PROF_START();

  // shared layout (fp32 first for 16B align, then bf16):
  //   gc[16,K] | qn[16,K] | kn[16,K] | L/Af fp32[16,16]
  //   | kd_bf[16,K] | ki_bf[16,K] | v_bf[16,V] | qd_bf[16,K] | Af_bf[16,16]
  float* gc_s = (float*)prep_smem;                 // [16,K] (ga then cumsum)
  float* qn_s = gc_s + BT16 * K;                   // [16,K]
  float* kn_s = qn_s + BT16 * K;                   // [16,K]
  float* Lf = kn_s + BT16 * K;                     // [16,16] fp32 (L then Af)
  bf16* kd_bf = (bf16*)(Lf + BT16 * BT16);         // [16,K]
  bf16* ki_bf = kd_bf + BT16 * K;                  // [16,K]
  bf16* v_bf = ki_bf + BT16 * K;                   // [16,V]
  bf16* qd_bf = v_bf + BT16 * V;                   // [16,K]
  bf16* Af_bf = qd_bf + BT16 * K;                  // [16,16]
  __shared__ float beta_sh[BT16];
  __shared__ float gl_sh[256];  // glast per dim (K<=256)

  // ---- l2norm q,k + bounded gate ; cumsum done below ----
  // one warp per row of the 16-row chunk.
  for (int c = warp; c < BT16; c += nwarps) {
    int t = base + c;
    float sq = 0.f, sk = 0.f;
    for (int d = lane; d < K; d += 32) {
      float a = q_raw[inTK(b, t, h, d, T, H, K)]; sq += a * a;
      float cc = k_raw[inTK(b, t, h, d, T, H, K)]; sk += cc * cc;
    }
    sq = warp_sum(sq); sk = warp_sum(sk);
    float r_q = rsqrtf(sq + L2_EPS), r_k = rsqrtf(sk + L2_EPS);
    // TODO：优化q_raw、k_raw的 global 访存
    for (int d = lane; d < K; d += 32) {
      qn_s[c * K + d] = q_raw[inTK(b, t, h, d, T, H, K)] * r_q;
      kn_s[c * K + d] = k_raw[inTK(b, t, h, d, T, H, K)] * r_k;
      float x = g_raw[inTK(b, t, h, d, T, H, K)] + dt_bias[h * K + d];
      gc_s[c * K + d] = lower_bound * sigmoidf(aexp * x);  // bounded gate in [lb,0)
      v_bf[c * V + d] = __float2bfloat16(v_in[inTV(b, t, h, d, T, H, V)]);
    }
    if (lane == 0) beta_sh[c] = sigmoidf(beta_raw[(b * T + t) * H + h]);
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 0);  // l2norm + gate

  // cumsum gc over the 16 rows (per dim d); one thread per dim handles column.
  // TODO：128个thread，每个线程负责一列
  for (int d = tid; d < K; d += blockDim.x) {
    float acc = 0.f;
    for (int c = 0; c < BT16; ++c) { acc += gc_s[c * K + d]; gc_s[c * K + d] = acc; }
    gl_sh[d] = acc;                 // glast for this dim
    gl8[(bh * NT + n) * K + d] = acc;
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 1);  // cumsum

  // kd = kn*exp(gc), ki = kn*exp(-gc), qd = qn*scale*exp(gc), kr = kn*exp(glast-gc).
  // kr8 is stored TRANSPOSED as [K,16] (krT[d,c]) so K2 loads it with a straight
  // copy / cp.async (no transpose needed at consume time).
  for (int idx = tid; idx < BT16 * K; idx += blockDim.x) {
    int c = idx / K, d = idx % K;
    float gcv = gc_s[idx];
    float kn = kn_s[idx], qn = qn_s[idx];
    float egc = expf(gcv), eneg = expf(-gcv);
    float qd = qn * scale * egc;
    kd_bf[idx] = __float2bfloat16(kn * egc);
    ki_bf[idx] = __float2bfloat16(kn * eneg);
    qd_bf[idx] = __float2bfloat16(qd);
    ot_store(&qd8[(bh * NT + n) * BT16 * K + idx], qd);            // q_decayed -> global
    ot_store(&kr8[(bh * NT + n) * BT16 * K + d * BT16 + c], kn * expf(gl_sh[d] - gcv));// krT[d,c]
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 2);  // kd/ki/qd/kr emit

  // ---- L = kd @ ki^T  (TC) ; strict-lower, store Araw = -beta*L into Lf ----
  v8::cute_gemm16_d(kd_bf, ki_bf, Lf, BT16, BT16, K);
  __syncthreads();
  PROF_ADD(g_prof_k1, 3);  // L gemm
  for (int idx = tid; idx < BT16 * BT16; idx += blockDim.x) {
    int c = idx / BT16, i = idx % BT16;
    Lf[idx] = (c > i) ? (-beta_sh[c] * Lf[idx]) : 0.f;
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 4);  // L mask

  // ---- triangular inverse: X = (I - Araw)^{-1} - I via forward substitution ----
  // (identical recurrence to v7/v2: X[c,i] = Araw[c,i] + sum_{i<m<c} Araw[c,m]*X[m,i],
  //  with Araw = -beta*tril(L) already in Lf; result is (I+L)^{-1} off-diagonal.)
  if (tid < 32) {
    int ln = tid;
    for (int c = 1; c < BT16; ++c) {
      int i = ln;
      float a = 0.f;
      if (i < c) {
        a = Lf[c * BT16 + i];
        for (int mm = i + 1; mm < c; ++mm) a += Lf[c * BT16 + mm] * Lf[mm * BT16 + i];
      }
      __syncwarp();
      if (i < c) Lf[c * BT16 + i] = a;
      __syncwarp();
    }
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 5);  // triangular inverse
  // Af = (X + I) * beta_s[i]  (fold beta into columns), and bf16 copy.
  for (int idx = tid; idx < BT16 * BT16; idx += blockDim.x) {
    int c = idx / BT16, i = idx % BT16;
    float xci = Lf[idx] + (c == i ? 1.0f : 0.0f);
    Af_bf[idx] = __float2bfloat16(xci * beta_sh[i]);
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 6);  // Af build

  // ---- Aqk = qd @ ki^T  (TC) ; lower-triangular (j<=c) ----
  // reuse Lf as fp32 output scratch for Aqk.
  v8::cute_gemm16_d(qd_bf, ki_bf, Lf, BT16, BT16, K);
  __syncthreads();
  PROF_ADD(g_prof_k1, 7);  // Aqk gemm
  for (int idx = tid; idx < BT16 * BT16; idx += blockDim.x) {
    int c = idx / BT16, j = idx % BT16;
    ot_store(&Aqk8[(bh * NT + n) * BT16 * BT16 + idx], (j <= c) ? Lf[idx] : 0.f);
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 8);  // Aqk store

  // NOTE: the __syncthreads that used to sit here (between the Aqk8 store and
  // the w-staging below) is REMOVED: the Aqk store only reads Lf + writes
  // global, while w/u staging overwrites ki_bf/qd_bf, whose last reader was the
  // Aqk gemm — already separated by the sync right after that gemm. (answer to
  // "算aqk和w、v也没有依赖，不用__syncthreads")
  // ---- w = Af @ kd (->w8) ; u = Af @ v (->u8) ----
  // w[16,K]=Af[16,16]@kd[16,K]; TN form C=A@B^T needs B[K,16]=kd^T. Give w and
  // u SEPARATE transpose buffers (kd^T in ki_bf, v^T in qd_bf) so the two
  // cg-gemms need NO sync between them (different inputs + different outputs).
  // (answer to "算w和u可以并行，不需要中间__syncthreads")
  bf16* tBw = ki_bf;  // kd^T [K,16]
  bf16* tBu = qd_bf;  // v^T  [V,16]
  for (int idx = tid; idx < K * BT16; idx += blockDim.x) {
    int d = idx / BT16, c = idx % BT16;
    tBw[d * BT16 + c] = kd_bf[c * K + d];
  }
  for (int idx = tid; idx < V * BT16; idx += blockDim.x) {
    int vd = idx / BT16, c = idx % BT16;
    tBu[vd * BT16 + c] = v_bf[c * V + vd];
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 9);  // w/u transpose
  if constexpr (cute::is_same<OT, bf16>::value) {
    v8::cute_gemm16_cg_bf_d(Af_bf, tBw, w8 + (bh * NT + n) * BT16 * K, K, BT16, K, BT16);
    v8::cute_gemm16_cg_bf_d(Af_bf, tBu, u8 + (bh * NT + n) * BT16 * V, V, BT16, V, BT16);
  } else {
    v8::cute_gemm16_cg_d(Af_bf, tBw, w8 + (bh * NT + n) * BT16 * K, K, BT16, K, BT16);
    v8::cute_gemm16_cg_d(Af_bf, tBu, u8 + (bh * NT + n) * BT16 * V, V, BT16, V, BT16);
  }
  PROF_ADD(g_prof_k1, 10);  // w/u gemm
}

// V-split serial recurrence (CHUNK=16). grid=BH*VS, block=GEMM16_THREADS.
// PIPELINE: when true, prefetch chunk t+1's St-independent operands
// (w/qd/Aqk/kr) into the back buffer while computing chunk t.
constexpr int VCOLS8 = 32;
extern __shared__ char recur8_smem[];
template <bool PIPELINE>
__global__ void k_recur_v8(
    const float* w8, const float* u8, const float* qd8, const float* kr8,
    const float* Aqk8, const float* gl8, const float* h0,
    float* o_out, float* fs_out,
    int B, int T, int H, int K, int V, int NT) {
  int VS = V / VCOLS8;
  int blk = blockIdx.x;
  if (blk >= B * H * VS) return;
  int bh = blk / VS, vsblk = blk % VS, col0 = vsblk * VCOLS8;
  int b = bh / H, h = bh % H, tid = threadIdx.x;

  // shared layout: St_f fp32[VC,K] | Cbuf fp32[max(2*16*VC, VC*K)]
  //   | St_bf[VC,K] | double-buffered operands x2:
  //     w_bf[16,K] qd_bf[16,K] Aqk_bf[16,16] krT_bf[K,16]
  //   | vcorrT_bf[VC,16]
  int CBUF = 2 * BT16 * VCOLS8;
  if (VCOLS8 * K > CBUF) CBUF = VCOLS8 * K;
  float* St_f = (float*)recur8_smem;             // [VC,K]
  float* Cbuf = St_f + VCOLS8 * K;               // reused gemm scratch
  bf16* St_bf = (bf16*)(Cbuf + CBUF);            // [VC,K]
  int OPSZ = BT16 * K + BT16 * K + BT16 * BT16 + K * BT16;  // per-buffer operands
  int NBUF = PIPELINE ? 2 : 1;
  bf16* opbuf = St_bf + VCOLS8 * K;              // [NBUF * OPSZ]
  bf16* vcorrT_bf = opbuf + NBUF * OPSZ;         // [VC,16]
  float* C1 = Cbuf;                              // wS [16,VC]
  float* C2 = Cbuf;                              // qS [16,VC]
  float* C3 = Cbuf + BT16 * VCOLS8;              // AV [16,VC]
  float* C4 = Cbuf;                              // newS [VC,K]
  __shared__ float EL_sh[256];                   // exp(glast) per dim (K<=256)

  // operand sub-pointers within a buffer slot
  auto W_OF = [&](int slot) { return opbuf + slot * OPSZ; };
  auto QD_OF = [&](int slot) { return opbuf + slot * OPSZ + BT16 * K; };
  auto AQ_OF = [&](int slot) { return opbuf + slot * OPSZ + 2 * BT16 * K; };
  auto KR_OF = [&](int slot) { return opbuf + slot * OPSZ + 2 * BT16 * K + BT16 * BT16; };

  // stage chunk n's St-independent operands into buffer `slot`.
  auto stage = [&](int n, int slot) {
    const float* w = w8 + (bh * NT + n) * BT16 * K;
    const float* qd = qd8 + (bh * NT + n) * BT16 * K;
    const float* Aq = Aqk8 + (bh * NT + n) * BT16 * BT16;
    const float* kr = kr8 + (bh * NT + n) * BT16 * K;
    bf16* w_bf = W_OF(slot); bf16* qd_bf = QD_OF(slot);
    bf16* Aqk_bf = AQ_OF(slot); bf16* krT_bf = KR_OF(slot);
    for (int idx = tid; idx < BT16 * K; idx += blockDim.x) {
      w_bf[idx] = __float2bfloat16(w[idx]);
      qd_bf[idx] = __float2bfloat16(qd[idx]);
    }
    for (int idx = tid; idx < BT16 * BT16; idx += blockDim.x)
      Aqk_bf[idx] = __float2bfloat16(Aq[idx]);
    // kr8 is already stored transposed as krT[K,16] by K1 -> straight copy.
    for (int idx = tid; idx < K * BT16; idx += blockDim.x)
      krT_bf[idx] = __float2bfloat16(kr[idx]);
  };

  // seed S^T from initial_state: St_f[vd,d] = h0[d, col0+vd]
  for (int idx = tid; idx < VCOLS8 * K; idx += blockDim.x) {
    int vd = idx / K, d = idx % K;
    St_f[idx] = h0[bh * K * V + d * V + (col0 + vd)];
  }
  if (PIPELINE) { stage(0, 0); }
  __syncthreads();

  for (int n = 0; n < NT; ++n) {
    int slot = PIPELINE ? (n & 1) : 0;
    if (!PIPELINE) { stage(n, 0); __syncthreads(); }
    bf16* w_bf = W_OF(slot); bf16* qd_bf = QD_OF(slot);
    bf16* Aqk_bf = AQ_OF(slot); bf16* krT_bf = KR_OF(slot);
    const float* u = u8 + (bh * NT + n) * BT16 * V;
    const float* gl = gl8 + (bh * NT + n) * K;

    // St_bf from current St_f
    for (int idx = tid; idx < VCOLS8 * K; idx += blockDim.x)
      St_bf[idx] = __float2bfloat16(St_f[idx]);
    __syncthreads();

    // PIPELINE: prefetch next chunk's operands into the other buffer now;
    // it only reads global w8/qd8/Aqk8/kr8 (St-independent) so it overlaps.
    if (PIPELINE && n + 1 < NT) stage(n + 1, (n + 1) & 1);

    // gemm1: wS = w @ S^T ; vcorr = u - wS  -> vcorrT[VC,16]
    v8::cute_gemm16_d(w_bf, St_bf, C1, BT16, VCOLS8, K);
    __syncthreads();
    for (int idx = tid; idx < BT16 * VCOLS8; idx += blockDim.x) {
      int c = idx / VCOLS8, vd = idx % VCOLS8;
      float vc = u[c * V + (col0 + vd)] - C1[idx];
      vcorrT_bf[vd * BT16 + c] = __float2bfloat16(vc);
    }
    __syncthreads();

    // gemm2: qS = qd @ S^T ; gemm3: AV = Aqk @ vcorr ; o = qS + AV
    v8::cute_gemm16_d(qd_bf, St_bf, C2, BT16, VCOLS8, K);
    v8::cute_gemm16_d(Aqk_bf, vcorrT_bf, C3, BT16, VCOLS8, BT16);
    __syncthreads();
    for (int idx = tid; idx < BT16 * VCOLS8; idx += blockDim.x) {
      int c = idx / VCOLS8, vd = idx % VCOLS8;
      o_out[inTV(b, n * BT16 + c, h, col0 + vd, T, H, V)] = C2[idx] + C3[idx];
    }
    __syncthreads();

    // gemm4: newS[VC,K] = vcorrT[VC,16] @ krT[K,16]^T ; S^T = S^T*EL + newS
    v8::cute_gemm16_d(vcorrT_bf, krT_bf, C4, VCOLS8, K, BT16);
    __syncthreads();
    for (int d = tid; d < K; d += blockDim.x) EL_sh[d] = expf(gl[d]);
    __syncthreads();
    for (int idx = tid; idx < VCOLS8 * K; idx += blockDim.x) {
      int d = idx % K;
      St_f[idx] = St_f[idx] * EL_sh[d] + C4[idx];
    }
    __syncthreads();
  }

  for (int idx = tid; idx < VCOLS8 * K; idx += blockDim.x) {
    int vd = idx / K, d = idx % K;
    fs_out[bh * K * V + d * V + (col0 + vd)] = St_f[idx];
  }
}

// ---- warp-specialized recurrence (real load/compute overlap) ----
// blockDim = WS_COMPUTE (96 = 3 compute warps) + 32 (1 dedicated load warp).
// Operands now arrive as BF16 in global (K1 emits bf16 for this path) and kr8
// is pre-transposed [K,16], so the load warp uses 128-bit cp.async straight
// into shared with NO conversion/transpose. Double-buffered: load warp stages
// chunk n+1 while the compute warps run chunk n's 4 gemms (FlashKDA load-warp).
// The 4 cute gemms use exactly 64 threads, so they are guarded to tid<64; the
// extra compute warp (tid 64..95) only accelerates the elementwise loops.
//   vcorr stays in shared (CuTe helper hides fragment layout; vcorr is a
//   B-operand in AV but A-operand in newS -> incompatible frags).
constexpr int WS_COMPUTE = 96;                  // 3 compute warps
constexpr int WS_GEMM = v8::GEMM16_THREADS;     // 64 threads do the MMA
// barrier ids: 1 = compute-only; 2/3 = operand-ready slot0/1; 4/5 = free slot0/1
extern __shared__ char recur8ws_smem[];
__global__ void k_recur_v8_ws(
    const bf16* w8, const bf16* u8, const bf16* qd8, const bf16* kr8,
    const bf16* Aqk8, const float* gl8, const float* h0,
    float* o_out, float* fs_out,
    int B, int T, int H, int K, int V, int NT) {
  int VS = V / VCOLS8;
  int blk = blockIdx.x;
  if (blk >= B * H * VS) return;
  int bh = blk / VS, vsblk = blk % VS, col0 = vsblk * VCOLS8;
  int b = bh / H, h = bh % H, tid = threadIdx.x;
  bool is_load = tid >= WS_COMPUTE;          // load warp = last 32 threads
  int ll = tid - WS_COMPUTE;                 // load-warp lane 0..31
  int WSALL = WS_COMPUTE + 32;               // participants on ready/free barriers

  int CBUF = 2 * BT16 * VCOLS8;
  if (VCOLS8 * K > CBUF) CBUF = VCOLS8 * K;
  float* St_f = (float*)recur8ws_smem;        // [VC,K]
  float* Cbuf = St_f + VCOLS8 * K;            // gemm scratch
  bf16* St_bf = (bf16*)(Cbuf + CBUF);         // [VC,K]
  int OPSZ = BT16 * K + BT16 * K + BT16 * BT16 + K * BT16;
  bf16* opbuf = St_bf + VCOLS8 * K;           // [2 * OPSZ]
  bf16* vcorrT_bf = opbuf + 2 * OPSZ;         // [VC,16]
  float* C1 = Cbuf;
  float* C2 = Cbuf;
  float* C3 = Cbuf + BT16 * VCOLS8;
  float* C4 = Cbuf;
  __shared__ float EL_sh[256];  // exp(glast) per dim, hoisted per chunk (K<=256)

  auto W_OF = [&](int slot) { return opbuf + slot * OPSZ; };
  auto QD_OF = [&](int slot) { return opbuf + slot * OPSZ + BT16 * K; };
  auto AQ_OF = [&](int slot) { return opbuf + slot * OPSZ + 2 * BT16 * K; };
  auto KR_OF = [&](int slot) { return opbuf + slot * OPSZ + 2 * BT16 * K + BT16 * BT16; };

  if (is_load) {
    // -------- dedicated load warp: cp.async chunk operands ahead of compute --
    // producer: wait FREE(slot) -> cp.async -> commit+wait -> signal READY(slot).
    for (int n = 0; n < NT; ++n) {
      int slot = n & 1;
      bar_sync(4 + slot, WSALL);               // wait slot freed (primed by compute)
      const bf16* w = w8 + (bh * NT + n) * BT16 * K;
      const bf16* qd = qd8 + (bh * NT + n) * BT16 * K;
      const bf16* Aq = Aqk8 + (bh * NT + n) * BT16 * BT16;
      const bf16* kr = kr8 + (bh * NT + n) * BT16 * K;  // pre-transposed [K,16]
      bf16* w_bf = W_OF(slot); bf16* qd_bf = QD_OF(slot);
      bf16* Aqk_bf = AQ_OF(slot); bf16* krT_bf = KR_OF(slot);
      // each cp.async moves 8 bf16 (16B). all these counts are multiples of 8.
      for (int v = ll; v < (BT16 * K) / 8; v += 32) cp_async16(w_bf + v * 8, w + v * 8);
      for (int v = ll; v < (BT16 * K) / 8; v += 32) cp_async16(qd_bf + v * 8, qd + v * 8);
      for (int v = ll; v < (BT16 * BT16) / 8; v += 32) cp_async16(Aqk_bf + v * 8, Aq + v * 8);
      for (int v = ll; v < (K * BT16) / 8; v += 32) cp_async16(krT_bf + v * 8, kr + v * 8);
      cp_async_commit();
      cp_async_wait_all();                     // ensure landed before compute reads
      bar_arrive(2 + slot, WSALL);             // operands ready for compute
    }
    return;
  }

  // -------- compute warps (96): carry St_f, run the 4 gemms per chunk --------
  for (int idx = tid; idx < VCOLS8 * K; idx += WS_COMPUTE) {
    int vd = idx / K, d = idx % K;
    St_f[idx] = h0[bh * K * V + d * V + (col0 + vd)];
  }
  bar_sync(1, WS_COMPUTE);
  // prime FREE(slot): both buffers start empty so the load warp may fill them.
  bar_arrive(4 + 0, WSALL);
  bar_arrive(4 + 1, WSALL);
  PROF_DECL PROF_START();

  for (int n = 0; n < NT; ++n) {
    int slot = n & 1;
    bar_sync(2 + slot, WSALL);                 // wait operands staged by load warp
    PROF_ADD(g_prof_k2, 7);  // stall: wait operands
    bf16* w_bf = W_OF(slot); bf16* qd_bf = QD_OF(slot);
    bf16* Aqk_bf = AQ_OF(slot); bf16* krT_bf = KR_OF(slot);
    const bf16* u = u8 + (bh * NT + n) * BT16 * V;
    const float* gl = gl8 + (bh * NT + n) * K;

    for (int idx = tid; idx < VCOLS8 * K; idx += WS_COMPUTE)
      St_bf[idx] = __float2bfloat16(St_f[idx]);
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 0);  // St->bf16

    // gemm1: wS = w @ S^T ; vcorr = u - wS -> vcorrT[VC,16]  (gemm: 64 threads)
    if (tid < WS_GEMM) v8::cute_gemm16_d(w_bf, St_bf, C1, BT16, VCOLS8, K);
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 1);  // gemm1 wS
    for (int idx = tid; idx < BT16 * VCOLS8; idx += WS_COMPUTE) {
      int c = idx / VCOLS8, vd = idx % VCOLS8;
      float vc = __bfloat162float(u[c * V + (col0 + vd)]) - C1[idx];
      vcorrT_bf[vd * BT16 + c] = __float2bfloat16(vc);
    }
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 2);  // vcorr

    // gemm2/3: qS = qd @ S^T ; AV = Aqk @ vcorr ; o = qS + AV
    if (tid < WS_GEMM) {
      v8::cute_gemm16_d(qd_bf, St_bf, C2, BT16, VCOLS8, K);
      v8::cute_gemm16_d(Aqk_bf, vcorrT_bf, C3, BT16, VCOLS8, BT16);
    }
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 3);  // gemm2/3 qS,AV
    for (int idx = tid; idx < BT16 * VCOLS8; idx += WS_COMPUTE) {
      int c = idx / VCOLS8, vd = idx % VCOLS8;
      o_out[inTV(b, n * BT16 + c, h, col0 + vd, T, H, V)] = C2[idx] + C3[idx];
    }
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 4);  // o write

    // gemm4: newS[VC,K] = vcorrT @ krT^T ; S^T = S^T*EL + newS
    if (tid < WS_GEMM) v8::cute_gemm16_d(vcorrT_bf, krT_bf, C4, VCOLS8, K, BT16);
    bar_arrive(4 + slot, WSALL);               // slot fully consumed (incl. krT)
    // hoist exp(glast[d]) once per dim (K calls) instead of VC*K times.
    for (int d = tid; d < K; d += WS_COMPUTE) EL_sh[d] = expf(gl[d]);
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 5);  // gemm4 newS
    for (int idx = tid; idx < VCOLS8 * K; idx += WS_COMPUTE) {
      int d = idx % K;
      St_f[idx] = St_f[idx] * EL_sh[d] + C4[idx];
    }
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 6);  // S update
  }

  for (int idx = tid; idx < VCOLS8 * K; idx += WS_COMPUTE) {
    int vd = idx / K, d = idx % K;
    fs_out[bh * K * V + d * V + (col0 + vd)] = St_f[idx];
  }
}


static int fwd_block_threads(int NT) {
  int warps = NT < 32 ? NT : 32;
  if (warps < 1) warps = 1;
  return warps * 32;
}

std::vector<torch::Tensor> kda_cuda_fwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size,
    double lower_bound, bool pipeline, bool warp_spec) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  // v8 forces CHUNK=16 (bounded-gate redesign); chunk_size arg is ignored.
  const int BT = 16;
  int NT = T / BT, BH = B * H;

  auto o = torch::zeros({B, T, H, V}, opt);
  auto fs = torch::zeros({B, H, K, V}, opt);
  // per-chunk workspaces. warp_spec consumes them as bf16 (cp.async, no
  // conversion); baseline/pipeline consume them as fp32. gl8 is always fp32.
  auto wopt = warp_spec ? torch::TensorOptions().dtype(torch::kBFloat16).device(q.device()) : opt;
  auto w8 = torch::zeros({BH, NT, BT, K}, wopt);
  auto u8 = torch::zeros({BH, NT, BT, V}, wopt);
  auto qd8 = torch::zeros({BH, NT, BT, K}, wopt);
  auto kr8 = torch::zeros({BH, NT, BT, K}, wopt);
  auto Aqk8 = torch::zeros({BH, NT, BT, BT}, wopt);
  auto gl8 = torch::zeros({BH, NT, K}, opt);

  // ---- K1: fused prepare (l2norm + bounded gate + cumsum + L/inv/Aqk/w/u) ----
  int mxKV = K > V ? K : V;
  int prep_sh = (BT * K + BT * K + BT * K + BT * BT) * sizeof(float)
              + (BT * K + BT * K + BT * V + BT * K + BT * BT) * sizeof(bf16);
  int prep_threads = 128;
  if (warp_spec) {
    cudaFuncSetAttribute(k_prepare_v8<bf16>, cudaFuncAttributeMaxDynamicSharedMemorySize, prep_sh);
    k_prepare_v8<bf16><<<BH * NT, prep_threads, prep_sh>>>(
        fp(q), fp(k), fp(g), fp(beta), fp(A_log), fp(dt_bias), fp(v),
        bp(w8), bp(u8), bp(qd8), bp(kr8), bp(Aqk8), fp(gl8),
        (float)scale, (float)lower_bound, B, T, H, K, V, NT);
  } else {
    cudaFuncSetAttribute(k_prepare_v8<float>, cudaFuncAttributeMaxDynamicSharedMemorySize, prep_sh);
    k_prepare_v8<float><<<BH * NT, prep_threads, prep_sh>>>(
        fp(q), fp(k), fp(g), fp(beta), fp(A_log), fp(dt_bias), fp(v),
        fp(w8), fp(u8), fp(qd8), fp(kr8), fp(Aqk8), fp(gl8),
        (float)scale, (float)lower_bound, B, T, H, K, V, NT);
  }

  // ---- K2: V-split serial recurrence ----
  {
    int VC = VCOLS8, VS = V / VC;
    int cbuf = 2 * BT * VC; if (VC * K > cbuf) cbuf = VC * K;
    int opsz = BT * K + BT * K + BT * BT + K * BT;
    if (warp_spec) {
      // warp-specialized: dedicated load warp (cp.async) + 2 operand buffers.
      int rsh = (VC * K + cbuf) * sizeof(float)
              + (VC * K + 2 * opsz + VC * BT) * sizeof(bf16);
      int thr = WS_COMPUTE + 32;  // 96 compute + 1 load warp
      cudaFuncSetAttribute(k_recur_v8_ws, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
      k_recur_v8_ws<<<BH * VS, thr, rsh>>>(
          bp(w8), bp(u8), bp(qd8), bp(kr8), bp(Aqk8), fp(gl8),
          fp(initial_state), fp(o), fp(fs), B, T, H, K, V, NT);
    } else {
      int nbuf = pipeline ? 2 : 1;
      int rsh = (VC * K + cbuf) * sizeof(float)
              + (VC * K + nbuf * opsz + VC * BT) * sizeof(bf16);
      if (pipeline) {
        cudaFuncSetAttribute(k_recur_v8<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
        k_recur_v8<true><<<BH * VS, v8::GEMM16_THREADS, rsh>>>(
            fp(w8), fp(u8), fp(qd8), fp(kr8), fp(Aqk8), fp(gl8),
            fp(initial_state), fp(o), fp(fs), B, T, H, K, V, NT);
      } else {
        cudaFuncSetAttribute(k_recur_v8<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
        k_recur_v8<false><<<BH * VS, v8::GEMM16_THREADS, rsh>>>(
            fp(w8), fp(u8), fp(qd8), fp(kr8), fp(Aqk8), fp(gl8),
            fp(initial_state), fp(o), fp(fs), B, T, H, K, V, NT);
      }
    }
  }
  (void)mxKV;
  cudaDeviceSynchronize();
  return {o, fs};
}

std::vector<torch::Tensor> kda_cuda_bwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size,
    torch::Tensor do_, torch::Tensor dht) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);
  do_ = do_.contiguous().to(torch::kFloat32);
  dht = dht.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  int BT = chunk_size, NT = T / BT, BH = B * H;

  auto qn = torch::zeros({BH, T, K}, opt), kn = torch::zeros({BH, T, K}, opt);
  auto qs = torch::zeros({BH, T, K}, opt), ga = torch::zeros({BH, T, K}, opt);
  auto gc = torch::zeros({BH, T, K}, opt);
  auto rq = torch::zeros({BH, T}, opt), rk = torch::zeros({BH, T}, opt);
  auto beta_s = torch::zeros({BH, T}, opt);
  auto X = torch::zeros({BH, NT, BT, BT}, opt);
  auto Afull = torch::zeros({BH, NT, BT, BT}, opt);
  auto Aqk = torch::zeros({BH, NT, BT, BT}, opt);
  auto w = torch::zeros({BH, T, K}, opt), u = torch::zeros({BH, T, V}, opt);
  auto vcorr = torch::zeros({BH, T, V}, opt);
  auto Sentry = torch::zeros({BH, NT + 1, K, V}, opt);
  auto o = torch::zeros({B, T, H, V}, opt), fs = torch::zeros({B, H, K, V}, opt);

  int wth = fwd_block_threads(NT);
  int mxKV = K > V ? K : V;
  int wy_sh = BT * BT * sizeof(float) + BT * BT * sizeof(bf16)
            + mxKV * BT * sizeof(bf16);
  cudaFuncSetAttribute(k_wy_v7, cudaFuncAttributeMaxDynamicSharedMemorySize, wy_sh);

  k_pre_v7<<<BH, wth>>>(fp(q), fp(k), fp(g), fp(beta), fp(A_log),
      fp(dt_bias), fp(qn), fp(kn), fp(qs), fp(rq), fp(rk), fp(beta_s),
      fp(ga), fp(gc), (float)scale, B, T, H, K, BT, NT);
  k_wy_v7<<<BH * NT, v7::GEMM_THREADS, wy_sh>>>(fp(v), fp(kn), fp(gc),
      fp(beta_s), fp(X), fp(Afull), fp(w), fp(u), B, T, H, K, V, BT, NT);
  k_aqk_v7<<<BH, wth>>>(fp(qs), fp(kn), fp(gc), fp(Aqk), B, T, H, K, BT, NT);
  {
    int VC = 32, VS = V / VC;
    int cbuf = 2 * BT * VC; if (VC * K > cbuf) cbuf = VC * K;
    int rsh = (VC * K + cbuf) * sizeof(float)
            + (VC * K + 2 * BT * K + VC * BT + BT * BT + K * BT) * sizeof(bf16);
    cudaFuncSetAttribute(k_recur_seq_v7<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
    k_recur_seq_v7<true><<<BH * VS, v7::GEMM_THREADS, rsh>>>(fp(qs), fp(kn), fp(gc),
        fp(w), fp(u), fp(initial_state), fp(Aqk), fp(vcorr), fp(Sentry),
        fp(o), fp(fs), B, T, H, K, V, BT, NT);
  }

  auto dq = torch::zeros_like(q), dk = torch::zeros_like(k);
  auto dv = torch::zeros_like(v), dgr = torch::zeros_like(g);
  auto dbeta = torch::zeros_like(beta);
  auto dA_log = torch::zeros_like(A_log), ddt_bias = torch::zeros_like(dt_bias);
  auto dh0 = torch::zeros_like(initial_state);

  BwdScratch sc;
  auto s_dgc = torch::zeros({BH, T, K}, opt); sc.dgc = fp(s_dgc);
  auto s_dqs = torch::zeros({BH, T, K}, opt); sc.dqs = fp(s_dqs);
  auto s_dkn = torch::zeros({BH, T, K}, opt); sc.dkn = fp(s_dkn);
  auto s_dbeta_s = torch::zeros({BH, T}, opt); sc.dbeta_s = fp(s_dbeta_s);
  auto s_dS = torch::zeros({BH, K, V}, opt); sc.dS = fp(s_dS);
  auto s_dS_entry = torch::zeros({BH, K, V}, opt); sc.dS_entry = fp(s_dS_entry);
  auto s_dEL = torch::zeros({BH, K}, opt); sc.dEL = fp(s_dEL);
  auto s_dvcorr = torch::zeros({BH, BT, V}, opt); sc.dvcorr = fp(s_dvcorr);
  auto s_du = torch::zeros({BH, BT, V}, opt); sc.du = fp(s_du);
  auto s_dAqk = torch::zeros({BH, BT, BT}, opt); sc.dAqk = fp(s_dAqk);
  auto s_dAfull = torch::zeros({BH, BT, BT}, opt); sc.dAfull = fp(s_dAfull);
  auto s_dX = torch::zeros({BH, BT, BT}, opt); sc.dX = fp(s_dX);
  auto s_Xf = torch::zeros({BH, BT, BT}, opt); sc.Xf = fp(s_Xf);
  auto s_tmp = torch::zeros({BH, BT, BT}, opt); sc.tmp = fp(s_tmp);
  auto s_dAraw = torch::zeros({BH, BT, BT}, opt); sc.dAraw = fp(s_dAraw);
  auto s_dQE = torch::zeros({BH, BT, K}, opt); sc.dQE = fp(s_dQE);
  auto s_dP = torch::zeros({BH, BT, K}, opt); sc.dP = fp(s_dP);
  auto s_dw = torch::zeros({BH, BT, K}, opt); sc.dw = fp(s_dw);
  auto s_dkg = torch::zeros({BH, BT, K}, opt); sc.dkg = fp(s_dkg);

  k_bwd_recur_v7<<<BH, 32>>>(fp(qs), fp(kn), fp(gc), fp(v), fp(w),
      fp(Afull), fp(X), fp(Aqk), fp(vcorr), fp(Sentry), fp(beta_s),
      fp(do_), fp(dht), sc, fp(dv), fp(dh0), B, T, H, K, V, BT);
  k_bwd_pre_v7<<<BH, NT>>>(fp(q), fp(k), fp(qn), fp(kn), fp(rq), fp(rk),
      fp(beta_s), fp(ga), fp(A_log), sc.dgc, sc.dqs, sc.dkn, sc.dbeta_s,
      fp(dq), fp(dk), fp(dgr), fp(dbeta), fp(dA_log), fp(ddt_bias),
      (float)scale, B, T, H, K, BT, NT);
  cudaDeviceSynchronize();
  return {dq, dk, dv, dgr, dbeta, dA_log, ddt_bias, dh0};
}

// Report theoretical occupancy of the recur kernels at the actual K=V=128
// launch config (no HW counters needed; uses the CUDA occupancy API).
static std::vector<int64_t> v8_occupancy(int64_t K, int64_t V) {
  int VC = VCOLS8;
  int cbuf = 2 * BT16 * VC; if (VC * (int)K > cbuf) cbuf = VC * K;
  int opsz = BT16 * K + BT16 * K + BT16 * BT16 + (int)K * BT16;
  int dev = 0; cudaGetDevice(&dev);
  cudaDeviceProp prop; cudaGetDeviceProperties(&prop, dev);
  auto report = [&](const void* fn, int threads, int shmem) -> std::vector<int64_t> {
    cudaFuncSetAttribute(fn, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem);
    int blocks = 0;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, fn, threads, shmem);
    cudaFuncAttributes at; cudaFuncGetAttributes(&at, fn);
    int warps_per_block = (threads + 31) / 32;
    int active_warps = blocks * warps_per_block;
    int max_warps = prop.maxThreadsPerMultiProcessor / 32;
    return {blocks, active_warps, max_warps, at.numRegs, shmem,
            (int64_t)prop.sharedMemPerMultiprocessor};
  };
  int rsh_b = (VC * (int)K + cbuf) * sizeof(float)
            + (VC * (int)K + 1 * opsz + VC * BT16) * sizeof(bf16);
  int rsh_w = (VC * (int)K + cbuf) * sizeof(float)
            + (VC * (int)K + 2 * opsz + VC * BT16) * sizeof(bf16);
  auto bl = report((const void*)k_recur_v8<false>, v8::GEMM16_THREADS, rsh_b);
  auto ws = report((const void*)k_recur_v8_ws, WS_COMPUTE + 32, rsh_w);
  // pack: [bl: blocks,warps,maxwarps,regs,shmem,smShmem][ws: ...]
  std::vector<int64_t> out;
  out.insert(out.end(), bl.begin(), bl.end());
  out.insert(out.end(), ws.begin(), ws.end());
  return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fwd", &kda_cuda_fwd,
        "KDA chunk forward (CUDA v8: CHUNK=16, bounded gate, all-TC)",
        pybind11::arg("q"), pybind11::arg("k"), pybind11::arg("v"),
        pybind11::arg("g"), pybind11::arg("beta"), pybind11::arg("A_log"),
        pybind11::arg("dt_bias"), pybind11::arg("scale"),
        pybind11::arg("initial_state"), pybind11::arg("chunk_size"),
        pybind11::arg("lower_bound") = -5.0, pybind11::arg("pipeline") = false,
        pybind11::arg("warp_spec") = false);
  m.def("bwd", &kda_cuda_bwd, "KDA chunk backward (CUDA v8, legacy v7 path)");
  m.def("occupancy", &v8_occupancy, "theoretical occupancy of recur kernels",
        pybind11::arg("K") = 128, pybind11::arg("V") = 128);
#ifdef V8_PROFILE
  m.def("prof_reset", []() {
    unsigned long long z[16] = {0};
    cudaMemcpyToSymbol(g_prof_k1, z, sizeof(z));
    cudaMemcpyToSymbol(g_prof_k2, z, sizeof(z));
  }, "zero the in-kernel clock64 accumulators");
  m.def("prof_read", []() {
    unsigned long long k1[16], k2[16];
    cudaMemcpyFromSymbol(k1, g_prof_k1, sizeof(k1));
    cudaMemcpyFromSymbol(k2, g_prof_k2, sizeof(k2));
    std::vector<int64_t> out;
    for (int i = 0; i < 16; ++i) out.push_back((int64_t)k1[i]);
    for (int i = 0; i < 16; ++i) out.push_back((int64_t)k2[i]);
    return out;
  }, "read [k1[0..15], k2[0..15]] cycle accumulators");
#endif
}


/**
K1 k_prepare_v8 (chunk-parallel prep, 27% of fwd)                                  
                                                                                     
  ┌──────────────────────────┬──────────┐                                            
  │            op            │  share   │                                            
  ├──────────────────────────┼──────────┤
  │ l2norm + gate            │ 37.9%    │
  ├──────────────────────────┼──────────┤
  │ tri-inverse (WY, 1 warp) │ 23.7%    │
  ├──────────────────────────┼──────────┤
  │ kd/ki/qd/kr emit         │ 10.0%    │
  ├──────────────────────────┼──────────┤
  │ w/u transpose            │ 8.5%     │
  ├──────────────────────────┼──────────┤
  │ w/u gemm                 │ 6.8%     │
  ├──────────────────────────┼──────────┤
  │ L gemm                   │ 3.8%     │
  ├──────────────────────────┼──────────┤
  │ Aqk gemm                 │ 3.4%     │
  ├──────────────────────────┼──────────┤
  │ cumsum                   │ 3.2%     │
  ├──────────────────────────┼──────────┤
  │ L mask / Af / Aqk store  │ <1% each │
  └──────────────────────────┴──────────┘

  K2 k_recur_v8_ws (serial recurrence, 71% of fwd)

  ┌─────────────────────────────┬───────┐
  │             op              │ share │
  ├─────────────────────────────┼───────┤
  │ S update (S=S·exp(gl)+newS) │ 39.7% │
  ├─────────────────────────────┼───────┤
  │ vcorr = u − wS              │ 16.8% │
  ├─────────────────────────────┼───────┤
  │ St → bf16 cast              │ 15.0% │
  ├─────────────────────────────┼───────┤
  │ gemm2/3 (qS, AV)            │ 10.6% │
  ├─────────────────────────────┼───────┤
  │ gemm1 (wS)                  │ 9.8%  │
  ├─────────────────────────────┼───────┤
  │ gemm4 (newS)                │ 4.9%  │
  ├─────────────────────────────┼───────┤
  │ o write                     │ 2.3%  │
  ├─────────────────────────────┼───────┤
  │ stall: wait-ops             │ 0.9%  │
  └─────────────────────────────┴───────┘

  Aggregated K2: gemms = 25%, elementwise = 74%, load-stall = 0.9%.

  What this says

  - The gemms are NOT the bottleneck. In K2 the 4 tensor-core gemms are only 25% of
  cycles; the tiles are tiny (16×32×128 etc.) so they finish fast. 74% is elementwise
   + the barriers between substeps — S update (the fp32 S·exp(gl)+newS over VC·K with
   an expf per element) alone is 40%, plus the St→bf16 recast (15%) and vcorr (17%).
  - The load-warp prefetch is working — stall:wait-ops is only 0.9%, confirming the
  cp.async overlap fully hides operand loading. The cost is purely intra-chunk
  compute/barrier latency on the serial S chain.
  - K1 is dominated by non-gemm work too: l2norm+gate (38%) and the 1-warp triangular
   inverse (24%) together are 62%; all 4 gemms combined are only ~14%.
*/

/**

      Long-context bench done (B=1, H=64, K=V=128):                                      
                                                                                     
  ┌──────┬───────┬─────────────┬────────┬───────────┬────────────────┐               
  │  T   │  fla  │ v8 baseline │ v8+ws  │ ws vs fla │ ws vs baseline │               
  ├──────┼───────┼─────────────┼────────┼───────────┼────────────────┤               
  │ 16K  │ 10.50 │ 50.49       │ 25.76  │ 0.41x     │ 1.96x          │               
  ├──────┼───────┼─────────────┼────────┼───────────┼────────────────┤               
  │ 64K  │ 41.93 │ 197.15      │ 101.94 │ 0.41x     │ 1.93x          │               
  ├──────┼───────┼─────────────┼────────┼───────────┼────────────────┤
  │ 128K │ 83.17 │ 401.80      │ 203.57 │ 0.41x     │ 1.97x          │
  └──────┴───────┴─────────────┴────────┴───────────┴────────────────┘                                                              
  ┌──────┬───────┬────────┬───────────┬────────────────┐
  │  T   │  fla  │ v8+ws  │ ws vs fla │ ws vs baseline │                             
  ├──────┼───────┼────────┼───────────┼────────────────┤    
  │ 16K  │ 10.50 │ 25.76  │ 0.41x     │ 1.96x          │
  ├──────┼───────┼────────┼───────────┼────────────────┤
  │ 64K  │ 41.93 │ 101.94 │ 0.41x     │ 1.93x          │
  ├──────┼───────┼────────┼───────────┼────────────────┤
  │ 128K │ 83.17 │ 203.57 │ 0.41x     │ 1.97x          │
  └──────┴───────┴────────┴───────────┴────────────────┘
*/

v8+ws — 相较 v2 提升 16.2x / 7.7x

kda_chunk_cuda_v8.cu（load_cuda_v8，check_v8.py）。受 FlashKDA 启发的前向重设计。按重要性排序的优化：

1. 结构性：CHUNK=16 + bounded 门控 → 所有矩阵乘在张量核心上

切换至 BT=16 和 bounded 门控ga = lower_bound · sigmoid(exp(A_log)·(g_raw+dt_bias)) ∈ [lb, 0)。BT=16 且 lb=−5 时，块内最差累积和为 15·(−5) = −75，因此exp(75) ≈ 3.7e32 < bf16 最大值 3.4e38。每个矩阵乘 — 包括 v2…v7 中必须保持标量fp32 的 L 和 Aqk — 现在均无溢出风险，可在 bf16 张量核心上运行，使用新的 16-tile CuTe 辅助函数 cute_gemm16 （cute_gemm_v8.cuh，MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>，AtomLayout<1,2,1>，64 线程，fp32 累加）。这是最大的杠杆 — 它是后续一切的前提。

2. opt2 — 融合单次准备核函数（K1）

k_prepare_v8，grid=BH·NT，每个 (chunk,head) 一个 CTA。将原k_pre + k_aqk + k_wy 融合为单核：l2norm → bounded 门控 → cumsum → 发射kd/ki/qd/v_bf → L=kd@ki^T（张量核心）→ 三角逆（1 warp，Neumann）→ Af →Aqk=qd@ki^T（张量核心）→ w=Af@kd, u=Af@v（张量核心，直写全局）。qd/kd/Aqk/Af/w/u不再在核函数间往返全局内存。之所以有效，正因重型矩阵乘已在张量核心上，而非标量。

3. opt3 — V 拆分 + 移除前向写入（K2）

k_recur_v8，grid=BH·VS（VCOLS=32），循环层有 BH·VS 个 block（取代 v2 的 BH，这是唯一有效的并行化杠杆，v3/v7 中也曾出现）。S^T[VC,K] 在共享内存中跨串行 chunk 循环驻留。前向路径不再写入每 chunk 的Sentry/vcorr 快照（这些仅反向需要）。

4. opt1 — warp 专用化加载 warp + cp.async（胜出的流水线）

k_recur_v8_ws：blockDim = 96 计算（3 warp）+ 32 专用加载 warp。加载 warp 使用 cp.async.cg（128-bit）预取 chunk n+1 的 St 无关操作数（w/qd/Aqk/kr），同时计算 warp 运行 chunk n 的 4 个矩阵乘 — 真正的跨 warp计算/加载重叠，FlashKDA 的技术。通过 PTX 命名 barrier（barrier.sync/arrive）同步，加载 warp 从不会被计算专用的块内 barrier 阻塞。K1 将操作数写为 bf16 且预转置（kr 为 [K,16]），因此加载 warp 可直接复制，无需转换/转置。

• 同 warp 双缓冲（pipeline=True）反而回退（占用率下降，无真实重叠）；专用加载 warp 胜出因为重叠是真实的。这是首个突破占用率天花板、通过隐藏串行 S 延迟来获得收益的杠杆。

v8 性能分析（clock64，-DV8_PROFILE）

• 核函数拆分 @128k：K2 71%，K1 27%，输入转换 1.8%。
• K2 各操作：矩阵乘仅 ~35%，逐元素 + barrier ~64%，加载停顿 0.9% （cp.async 完全隐藏加载）。瓶颈 = 串行 S 链上的块内逐元素/barrier 延迟， 而非矩阵乘。

v8 已应用的修复（expf 外提）

S 更新 St=St·expf(gl[d])+C4 对每块 VC·K=4096 个元素重复计算 expf， 但 EL 仅依赖于 d（128 个独立值）→ 32 倍冗余 SFU 超越函数。 外提至 __shared__ EL_sh[256]（128 次 expf + 仅 FMA 循环）。S 更新区域 39.7%→24.4%；墙钟降低约 3%。

---

6 优化 4

// KDA chunk reference (CUDA, bf16) — v8 CHUNK=16 + bounded-gate redesign.
//
// FORWARD-ONLY. This is a structural redesign over v7 that adopts FlashKDA's
// two key choices so that the three FlashKDA-inspired optimizations actually
// pay off (in the BT=64 + softplus build they were neutral/regressive):
//
//   * CHUNK = 16 (BT=16) and the BOUNDED gate
//        ga = lower_bound * sigmoid(exp(A_log) * (g_raw + dt_bias))   in [lb, 0)
//     With BT=16 and lb=-5 the worst within-chunk cumsum span is 15*(-5)=-75,
//     so exp(-gc) <= e^75 ~ 3.7e32 < bf16 max 3.4e38 => EVERY matmul (incl. the
//     L matrix and Aqk, which v7 had to keep scalar fp32 because ki=kn*exp(-gc)
//     overflowed) is now overflow-safe and goes to bf16 tensor cores.
//
// Two-kernel structure mirroring FlashKDA (forward only):
//   K1  k_prepare_v8 (opt2 = fusion): grid=BH*NT, 1 CTA per (chunk,head), 16
//       rows. l2norm + bounded gate + cumsum, then on TC (cute_gemm16):
//         L = kd@ki^T (strict-lower, *(-beta)), triangular inverse (I+L)^-1,
//         Aqk = qd@ki^T (lower-tri), w = Af@kd, u = Af@v.
//       kd=kn*e^{gc}, ki=kn*e^{-gc}, qd=qn*scale*e^{gc}. Stores w,u,qd,Aqk,kr,
//       glast to a per-chunk workspace (former k_pre+k_aqk+k_wy fused into one).
//   K2  k_recur_v8 (opt1 + opt3): grid=BH*VS, V-split (VCOLS=32), S^T[VC,K]
//       resident in shared, serial chunk loop. Per chunk on TC:
//         vcorr = u - w@S^T ; o = qd@S^T + Aqk@vcorr ; S^T = S^T*e^{glast} + vcorr^T@kr^T.
//       opt3: forward skips per-chunk Sentry/vcorr global writes (only o and
//       final_state written).  opt1: with BT=16 the recur kernel is ~50KB (not
//       ~100KB), so an occupancy-safe double-buffer of the St-INDEPENDENT
//       operands (w/qd/Aqk/kr) is feasible (PIPELINE template flag) — these are
//       prepared in K1 and don't depend on the running state, so chunk t+1 can
//       be staged while chunk t computes.
//
// Validation: chunk_kda rejects chunk_size=16, so v8 (BT=16) is validated vs
//   chunk_kda(safe_gate=True, lower_bound=-5, chunk_size=64) — the forward o /
//   final_state are mathematically chunk-size-independent — and cross-checked vs
//   fused_recurrent_kda(lower_bound=-5).
//
// Assumes HV==H, fixed length, K==V multiple of 16.
//
// NOTE: the BACKWARD path below is left as the v7 (softplus / BT in {32,64})
// kernels verbatim and is intentionally NOT updated or validated — v8 scope is
// forward-only. kda_cuda_bwd still launches the old *_v7 kernels so the module
// keeps compiling, but it does not match v8's CHUNK=16 + bounded-gate forward.

#include <cuda_runtime.h>
#include <torch/extension.h>

#include <vector>

#include "cute_gemm_v7.cuh"
#include "cute_gemm_v8.cuh"

namespace {

using v7::bf16;
constexpr float L2_EPS = 1e-6f;

__device__ inline float softplusf(float x) {
  if (x > 20.0f) return x;
  if (x < -20.0f) return expf(x);
  return log1pf(expf(x));
}
__device__ inline float sigmoidf(float x) { return 1.0f / (1.0f + expf(-x)); }
// store a float into an output slot of type OT (float passthrough / bf16 cast).
__device__ __forceinline__ void ot_store(float* p, float v) { *p = v; }
__device__ __forceinline__ void ot_store(bf16* p, float v) { *p = __float2bfloat16(v); }
__device__ inline float warp_sum(float v) {
  for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o);
  return v;
}
// Named barrier: synchronize exactly `cnt` threads on logical barrier `id`
// (id 1..15). Lets a warp-specialized kernel sync only its compute warps
// (cnt=COMPUTE_THREADS) while the dedicated load warp stays off that barrier.
__device__ __forceinline__ void bar_sync(int id, int cnt) {
  asm volatile("barrier.sync %0, %1;" :: "r"(id), "r"(cnt) : "memory");
}
// Arrive (no wait) on named barrier `id` expecting `cnt` participants.
__device__ __forceinline__ void bar_arrive(int id, int cnt) {
  asm volatile("barrier.arrive %0, %1;" :: "r"(id), "r"(cnt) : "memory");
}
// 16-byte cp.async global->shared (8 bf16). dst must be a shared address.
__device__ __forceinline__ void cp_async16(bf16* dst, const bf16* src) {
  unsigned s = (unsigned)__cvta_generic_to_shared(dst);
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: "r"(s), "l"(src));
}
__device__ __forceinline__ void cp_async_commit() {
  asm volatile("cp.async.commit_group;\n");
}
__device__ __forceinline__ void cp_async_wait_all() {
  asm volatile("cp.async.wait_group 0;\n");
}

// ---- optional in-kernel per-region timing (clock64). OFF by default: when
// V8_PROFILE is not defined every macro below expands to nothing, so the
// production kernels are byte-identical. ON: thread 0 accumulates cycle deltas
// per region into a __device__ global, read back via the v8 pybind. ----
#ifdef V8_PROFILE
__device__ unsigned long long g_prof_k1[16];
__device__ unsigned long long g_prof_k2[16];
#define PROF_DECL long long _pc = 0;
#define PROF_START() do { if (threadIdx.x == 0) _pc = clock64(); } while (0)
#define PROF_ADD(arr, i) do { if (threadIdx.x == 0) { \
    long long _n = clock64(); atomicAdd(&arr[i], (unsigned long long)(_n - _pc)); _pc = _n; } } while (0)
#else
#define PROF_DECL
#define PROF_START()
#define PROF_ADD(arr, i)
#endif

__device__ inline int iTK(int bh, int t, int d, int T, int K) { return bh * T * K + t * K + d; }
__device__ inline int iTV(int bh, int t, int vd, int T, int V) { return bh * T * V + t * V + vd; }
__device__ inline int inTK(int b, int t, int h, int d, int T, int H, int K) {
  return ((b * T + t) * H + h) * K + d;
}
__device__ inline int inTV(int b, int t, int h, int vd, int T, int H, int V) {
  return ((b * T + t) * H + h) * V + vd;
}

// ================= FORWARD =================



// ================= v8 FORWARD (CHUNK=16, bounded gate, all-TC) =================
//
// Per-chunk workspace produced by k_prepare_v8 (layout [BH, NT, ...]):
//   w8   [BH,NT,16,K]   = Af @ kd            (TC)
//   u8   [BH,NT,16,V]   = Af @ v             (TC)
//   qd8  [BH,NT,16,K]   = qn*scale*exp(gc)   (q_decayed, ready for q@S)
//   kr8  [BH,NT,16,K]   = kn*exp(glast-gc)   (k_restored, for S update)
//   Aqk8 [BH,NT,16,16]  = qd@ki^T lower-tri  (TC)
//   gl8  [BH,NT,K]      = gc_last per dim    (decay EL = exp(gl8))
constexpr int BT16 = 16;

extern __shared__ char prep_smem[];
// OT = workspace output dtype for the per-chunk operands (w8/u8/qd8/kr8/Aqk8).
// float = baseline/pipeline K2; bf16 = warp_spec K2 (lets it cp.async-load
// directly with no fp32->bf16 conversion). gl8 stays fp32 either way.
template <typename OT>
__global__ void k_prepare_v8(
    const float* q_raw, const float* k_raw, const float* g_raw,
    const float* beta_raw, const float* A_log, const float* dt_bias,
    const float* v_in,
    OT* w8, OT* u8, OT* qd8, OT* kr8, OT* Aqk8, float* gl8,
    float scale, float lower_bound, int B, int T, int H, int K, int V, int NT) {
  int blk = blockIdx.x;
  if (blk >= B * H * NT) return;
  int bh = blk / NT, n = blk % NT;
  int b = bh / H, h = bh % H;
  int base = n * BT16, tid = threadIdx.x;
  int warp = tid >> 5, lane = tid & 31, nwarps = blockDim.x >> 5;
  float aexp = expf(A_log[h]);
  PROF_DECL PROF_START();

  // shared layout (fp32 first for 16B align, then bf16):
  //   gc[16,K] | qn[16,K] | kn[16,K] | L/Af fp32[16,16]
  //   | kd_bf[16,K] | ki_bf[16,K] | v_bf[16,V] | qd_bf[16,K] | Af_bf[16,16]
  float* gc_s = (float*)prep_smem;                 // [16,K] (ga then cumsum)
  float* qn_s = gc_s + BT16 * K;                   // [16,K]
  float* kn_s = qn_s + BT16 * K;                   // [16,K]
  float* Lf = kn_s + BT16 * K;                     // [16,16] fp32 (L then Af)
  bf16* kd_bf = (bf16*)(Lf + BT16 * BT16);         // [16,K]
  bf16* ki_bf = kd_bf + BT16 * K;                  // [16,K]
  bf16* v_bf = ki_bf + BT16 * K;                   // [16,V]
  bf16* qd_bf = v_bf + BT16 * V;                   // [16,K]
  bf16* Af_bf = qd_bf + BT16 * K;                  // [16,16]
  __shared__ float beta_sh[BT16];
  __shared__ float gl_sh[256];  // glast per dim (K<=256)

  // ---- l2norm q,k + bounded gate ; cumsum done below ----
  // one warp per row of the 16-row chunk.
  for (int c = warp; c < BT16; c += nwarps) {
    int t = base + c;
    float sq = 0.f, sk = 0.f;
    for (int d = lane; d < K; d += 32) {
      float a = q_raw[inTK(b, t, h, d, T, H, K)]; sq += a * a;
      float cc = k_raw[inTK(b, t, h, d, T, H, K)]; sk += cc * cc;
    }
    sq = warp_sum(sq); sk = warp_sum(sk);
    float r_q = rsqrtf(sq + L2_EPS), r_k = rsqrtf(sk + L2_EPS);
    for (int d = lane; d < K; d += 32) {
      qn_s[c * K + d] = q_raw[inTK(b, t, h, d, T, H, K)] * r_q;
      kn_s[c * K + d] = k_raw[inTK(b, t, h, d, T, H, K)] * r_k;
      float x = g_raw[inTK(b, t, h, d, T, H, K)] + dt_bias[h * K + d];
      gc_s[c * K + d] = lower_bound * sigmoidf(aexp * x);  // bounded gate in [lb,0)
      v_bf[c * V + d] = __float2bfloat16(v_in[inTV(b, t, h, d, T, H, V)]);
    }
    if (lane == 0) beta_sh[c] = sigmoidf(beta_raw[(b * T + t) * H + h]);
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 0);  // l2norm + gate

  // cumsum gc over the 16 rows (per dim d); one thread per dim handles column.
  for (int d = tid; d < K; d += blockDim.x) {
    float acc = 0.f;
    for (int c = 0; c < BT16; ++c) { acc += gc_s[c * K + d]; gc_s[c * K + d] = acc; }
    gl_sh[d] = acc;                 // glast for this dim
    gl8[(bh * NT + n) * K + d] = acc;
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 1);  // cumsum

  // kd = kn*exp(gc), ki = kn*exp(-gc), qd = qn*scale*exp(gc), kr = kn*exp(glast-gc).
  // kr8 is stored TRANSPOSED as [K,16] (krT[d,c]) so K2 loads it with a straight
  // copy / cp.async (no transpose needed at consume time).
  for (int idx = tid; idx < BT16 * K; idx += blockDim.x) {
    int c = idx / K, d = idx % K;
    float gcv = gc_s[idx];
    float kn = kn_s[idx], qn = qn_s[idx];
    float egc = expf(gcv), eneg = expf(-gcv);
    float qd = qn * scale * egc;
    kd_bf[idx] = __float2bfloat16(kn * egc);
    ki_bf[idx] = __float2bfloat16(kn * eneg);
    qd_bf[idx] = __float2bfloat16(qd);
    ot_store(&qd8[(bh * NT + n) * BT16 * K + idx], qd);            // q_decayed -> global
    ot_store(&kr8[(bh * NT + n) * BT16 * K + d * BT16 + c], kn * expf(gl_sh[d] - gcv));// krT[d,c]
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 2);  // kd/ki/qd/kr emit

  // ---- L = kd @ ki^T  (TC) ; strict-lower, store Araw = -beta*L into Lf ----
  v8::cute_gemm16_d(kd_bf, ki_bf, Lf, BT16, BT16, K);
  __syncthreads();
  PROF_ADD(g_prof_k1, 3);  // L gemm
  for (int idx = tid; idx < BT16 * BT16; idx += blockDim.x) {
    int c = idx / BT16, i = idx % BT16;
    Lf[idx] = (c > i) ? (-beta_sh[c] * Lf[idx]) : 0.f;
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 4);  // L mask

  // ---- triangular inverse: X = (I - Araw)^{-1} - I via forward substitution ----
  // (identical recurrence to v7/v2: X[c,i] = Araw[c,i] + sum_{i<m<c} Araw[c,m]*X[m,i],
  //  with Araw = -beta*tril(L) already in Lf; result is (I+L)^{-1} off-diagonal.)
  if (tid < 32) {
    int ln = tid;
    for (int c = 1; c < BT16; ++c) {
      int i = ln;
      float a = 0.f;
      if (i < c) {
        a = Lf[c * BT16 + i];
        for (int mm = i + 1; mm < c; ++mm) a += Lf[c * BT16 + mm] * Lf[mm * BT16 + i];
      }
      __syncwarp();
      if (i < c) Lf[c * BT16 + i] = a;
      __syncwarp();
    }
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 5);  // triangular inverse
  // Af = (X + I) * beta_s[i]  (fold beta into columns), and bf16 copy.
  for (int idx = tid; idx < BT16 * BT16; idx += blockDim.x) {
    int c = idx / BT16, i = idx % BT16;
    float xci = Lf[idx] + (c == i ? 1.0f : 0.0f);
    Af_bf[idx] = __float2bfloat16(xci * beta_sh[i]);
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 6);  // Af build

  // ---- Aqk = qd @ ki^T  (TC) ; lower-triangular (j<=c) ----
  // reuse Lf as fp32 output scratch for Aqk.
  v8::cute_gemm16_d(qd_bf, ki_bf, Lf, BT16, BT16, K);
  __syncthreads();
  PROF_ADD(g_prof_k1, 7);  // Aqk gemm
  for (int idx = tid; idx < BT16 * BT16; idx += blockDim.x) {
    int c = idx / BT16, j = idx % BT16;
    ot_store(&Aqk8[(bh * NT + n) * BT16 * BT16 + idx], (j <= c) ? Lf[idx] : 0.f);
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 8);  // Aqk store

  // NOTE: the __syncthreads that used to sit here (between the Aqk8 store and
  // the w-staging below) is REMOVED: the Aqk store only reads Lf + writes
  // global, while w/u staging overwrites ki_bf/qd_bf, whose last reader was the
  // Aqk gemm — already separated by the sync right after that gemm. (answer to
  // "算aqk和w、v也没有依赖，不用__syncthreads")
  // ---- w = Af @ kd (->w8) ; u = Af @ v (->u8) ----
  // w[16,K]=Af[16,16]@kd[16,K]; TN form C=A@B^T needs B[K,16]=kd^T. Give w and
  // u SEPARATE transpose buffers (kd^T in ki_bf, v^T in qd_bf) so the two
  // cg-gemms need NO sync between them (different inputs + different outputs).
  // (answer to "算w和u可以并行，不需要中间__syncthreads")
  bf16* tBw = ki_bf;  // kd^T [K,16]
  bf16* tBu = qd_bf;  // v^T  [V,16]
  for (int idx = tid; idx < K * BT16; idx += blockDim.x) {
    int d = idx / BT16, c = idx % BT16;
    tBw[d * BT16 + c] = kd_bf[c * K + d];
  }
  for (int idx = tid; idx < V * BT16; idx += blockDim.x) {
    int vd = idx / BT16, c = idx % BT16;
    tBu[vd * BT16 + c] = v_bf[c * V + vd];
  }
  __syncthreads();
  PROF_ADD(g_prof_k1, 9);  // w/u transpose
  if constexpr (cute::is_same<OT, bf16>::value) {
    v8::cute_gemm16_cg_bf_d(Af_bf, tBw, w8 + (bh * NT + n) * BT16 * K, K, BT16, K, BT16);
    v8::cute_gemm16_cg_bf_d(Af_bf, tBu, u8 + (bh * NT + n) * BT16 * V, V, BT16, V, BT16);
  } else {
    v8::cute_gemm16_cg_d(Af_bf, tBw, w8 + (bh * NT + n) * BT16 * K, K, BT16, K, BT16);
    v8::cute_gemm16_cg_d(Af_bf, tBu, u8 + (bh * NT + n) * BT16 * V, V, BT16, V, BT16);
  }
  PROF_ADD(g_prof_k1, 10);  // w/u gemm
}



// ---- warp-specialized recurrence (real load/compute overlap) ----
// blockDim = WS_COMPUTE (96 = 3 compute warps) + 32 (1 dedicated load warp).
// Operands now arrive as BF16 in global (K1 emits bf16 for this path) and kr8
// is pre-transposed [K,16], so the load warp uses 128-bit cp.async straight
// into shared with NO conversion/transpose. Double-buffered: load warp stages
// chunk n+1 while the compute warps run chunk n's 4 gemms (FlashKDA load-warp).
// The 4 cute gemms use exactly 64 threads, so they are guarded to tid<64; the
// extra compute warp (tid 64..95) only accelerates the elementwise loops.
//   vcorr stays in shared (CuTe helper hides fragment layout; vcorr is a
//   B-operand in AV but A-operand in newS -> incompatible frags).
constexpr int WS_COMPUTE = 128;                 // 4 compute warps (sweet spot: keeps 3 blk/SM)
constexpr int WS_GEMM = v8::GEMM16_THREADS;     // 64 threads do the MMA
// barrier ids: 1 = compute-only; 2/3 = operand-ready slot0/1; 4/5 = free slot0/1
extern __shared__ char recur8ws_smem[];
__global__ void k_recur_v8_ws(
    const bf16* w8, const bf16* u8, const bf16* qd8, const bf16* kr8,
    const bf16* Aqk8, const float* gl8, const float* h0,
    float* o_out, float* fs_out,
    int B, int T, int H, int K, int V, int NT) {
  int VS = V / VCOLS8;
  int blk = blockIdx.x;
  if (blk >= B * H * VS) return;
  int bh = blk / VS, vsblk = blk % VS, col0 = vsblk * VCOLS8;
  int b = bh / H, h = bh % H, tid = threadIdx.x;
  bool is_load = tid >= WS_COMPUTE;          // load warp = last 32 threads
  int ll = tid - WS_COMPUTE;                 // load-warp lane 0..31
  int WSALL = WS_COMPUTE + 32;               // participants on ready/free barriers

  int CBUF = 2 * BT16 * VCOLS8;
  if (VCOLS8 * K > CBUF) CBUF = VCOLS8 * K;
  float* St_f = (float*)recur8ws_smem;        // [VC,K]
  float* Cbuf = St_f + VCOLS8 * K;            // gemm scratch
  bf16* St_bf = (bf16*)(Cbuf + CBUF);         // [VC,K]
  // v9: operand buffer now also holds the prefetched u column-slice [16,VC] so
  // the load warp hides u's global latency (was the only un-prefetched operand).
  int OPSZ = BT16 * K + BT16 * K + BT16 * BT16 + K * BT16 + BT16 * VCOLS8;
  bf16* opbuf = St_bf + VCOLS8 * K;           // [2 * OPSZ]
  bf16* vcorrT_bf = opbuf + 2 * OPSZ;         // [VC,16]
  float* C1 = Cbuf;
  float* C2 = Cbuf;
  float* C3 = Cbuf + BT16 * VCOLS8;
  float* C4 = Cbuf;
  // EL = exp(glast) per dim, now DOUBLE-BUFFERED per slot and produced by the
  // load warp (it reads gl8 from global + does the K expf's off the compute
  // critical path). [2][K<=256], float4-aligned. Freed with the operand slot.
  __shared__ __align__(16) float EL_sh2[2 * 256];

  auto W_OF = [&](int slot) { return opbuf + slot * OPSZ; };
  auto QD_OF = [&](int slot) { return opbuf + slot * OPSZ + BT16 * K; };
  auto AQ_OF = [&](int slot) { return opbuf + slot * OPSZ + 2 * BT16 * K; };
  auto KR_OF = [&](int slot) { return opbuf + slot * OPSZ + 2 * BT16 * K + BT16 * BT16; };
  auto U_OF = [&](int slot) { return opbuf + slot * OPSZ + 2 * BT16 * K + BT16 * BT16 + K * BT16; };
  auto EL_OF = [&](int slot) { return EL_sh2 + slot * 256; };

  if (is_load) {
    // -------- dedicated load warp: cp.async chunk operands ahead of compute --
    // producer: wait FREE(slot) -> cp.async -> commit+wait -> signal READY(slot).
    for (int n = 0; n < NT; ++n) {
      int slot = n & 1;
      bar_sync(4 + slot, WSALL);               // wait slot freed (primed by compute)
      const bf16* w = w8 + (bh * NT + n) * BT16 * K;
      const bf16* qd = qd8 + (bh * NT + n) * BT16 * K;
      const bf16* Aq = Aqk8 + (bh * NT + n) * BT16 * BT16;
      const bf16* kr = kr8 + (bh * NT + n) * BT16 * K;  // pre-transposed [K,16]
      const bf16* u = u8 + (bh * NT + n) * BT16 * V;     // v9: prefetch u col-slice
      bf16* w_bf = W_OF(slot); bf16* qd_bf = QD_OF(slot);
      bf16* Aqk_bf = AQ_OF(slot); bf16* krT_bf = KR_OF(slot);
      bf16* u_bf = U_OF(slot);                           // dst [16,VC] contiguous
      // each cp.async moves 8 bf16 (16B). all these counts are multiples of 8.
      for (int v = ll; v < (BT16 * K) / 8; v += 32) cp_async16(w_bf + v * 8, w + v * 8);
      for (int v = ll; v < (BT16 * K) / 8; v += 32) cp_async16(qd_bf + v * 8, qd + v * 8);
      for (int v = ll; v < (BT16 * BT16) / 8; v += 32) cp_async16(Aqk_bf + v * 8, Aq + v * 8);
      for (int v = ll; v < (K * BT16) / 8; v += 32) cp_async16(krT_bf + v * 8, kr + v * 8);
      // u col-slice: dst row c is contiguous VC bf16; src row c at u + c*V + col0
      // (strided by V). VC=32 -> VC/8=4 segs per row.
      for (int v = ll; v < (BT16 * VCOLS8) / 8; v += 32) {
        int c = v / (VCOLS8 / 8), seg = v % (VCOLS8 / 8);
        cp_async16(u_bf + v * 8, u + c * V + col0 + seg * 8);
      }
      cp_async_commit();
      // EL = exp(glast[d]) for this chunk: load warp does the K expf's so the
      // compute warps' S-update reads EL straight from shared (transcendentals
      // off the critical path). gl8 is fp32 global; EL slot is fp32 shared.
      const float* gl = gl8 + (bh * NT + n) * K;
      float* el = EL_OF(slot);
      for (int d = ll; d < K; d += 32) el[d] = expf(gl[d]);
      cp_async_wait_all();                     // ensure landed before compute reads
      bar_arrive(2 + slot, WSALL);             // operands ready for compute
    }
    return;
  }

  // -------- compute warps (96): carry St_f, run the 4 gemms per chunk --------
  for (int idx = tid; idx < VCOLS8 * K; idx += WS_COMPUTE) {
    int vd = idx / K, d = idx % K;
    St_f[idx] = h0[bh * K * V + d * V + (col0 + vd)];
  }
  bar_sync(1, WS_COMPUTE);
  // prime FREE(slot): both buffers start empty so the load warp may fill them.
  bar_arrive(4 + 0, WSALL);
  bar_arrive(4 + 1, WSALL);
  PROF_DECL PROF_START();

  for (int n = 0; n < NT; ++n) {
    int slot = n & 1;
    bar_sync(2 + slot, WSALL);                 // wait operands staged by load warp
    PROF_ADD(g_prof_k2, 7);  // stall: wait operands
    bf16* w_bf = W_OF(slot); bf16* qd_bf = QD_OF(slot);
    bf16* Aqk_bf = AQ_OF(slot); bf16* krT_bf = KR_OF(slot);
    bf16* u_bf = U_OF(slot);                   // v9: u col-slice [16,VC] in shared
    // gl/EL is produced by the load warp into EL_OF(slot); compute reads it below.

    // St -> bf16 (vectorized: float4 read, 2x bf162 write; VC*K%8==0).
    {
      const float4* sf4 = reinterpret_cast<const float4*>(St_f);
      uint2* sb2 = reinterpret_cast<uint2*>(St_bf);   // 4 bf16 = uint2
      for (int i = tid; i < (VCOLS8 * K) / 4; i += WS_COMPUTE) {
        float4 f = sf4[i];
        __nv_bfloat162 lo = __floats2bfloat162_rn(f.x, f.y);
        __nv_bfloat162 hi = __floats2bfloat162_rn(f.z, f.w);
        uint2 packed;
        packed.x = *reinterpret_cast<unsigned*>(&lo);
        packed.y = *reinterpret_cast<unsigned*>(&hi);
        sb2[i] = packed;
      }
    }
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 0);  // St->bf16

    // gemm1: wS = w @ S^T ; vcorr = u - wS -> vcorrT[VC,16]  (gemm: 64 threads)
    if (tid < WS_GEMM) v8::cute_gemm16_d(w_bf, St_bf, C1, BT16, VCOLS8, K);
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 1);  // gemm1 wS
    for (int idx = tid; idx < BT16 * VCOLS8; idx += WS_COMPUTE) {
      int c = idx / VCOLS8, vd = idx % VCOLS8;
      float vc = __bfloat162float(u_bf[c * VCOLS8 + vd]) - C1[idx];  // u from shared
      vcorrT_bf[vd * BT16 + c] = __float2bfloat16(vc);
    }
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 2);  // vcorr

    // gemm2/3: qS = qd @ S^T ; AV = Aqk @ vcorr ; o = qS + AV.
    // These two gemms are INDEPENDENT (qS reads St_bf, AV reads vcorrT_bf), so
    // run them in parallel on two disjoint 64-thread warp-groups: group0
    // (tid 0..63) does qS, group1 (tid 64..127) does AV.
    if (tid < WS_GEMM) {
      v8::cute_gemm16_d_idx(qd_bf, St_bf, C2, BT16, VCOLS8, K, tid);
    } else if (tid < 2 * WS_GEMM) {
      v8::cute_gemm16_d_idx(Aqk_bf, vcorrT_bf, C3, BT16, VCOLS8, BT16, tid - WS_GEMM);
    }
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 3);  // gemm2/3 qS,AV (parallel)
    for (int idx = tid; idx < BT16 * VCOLS8; idx += WS_COMPUTE) {
      int c = idx / VCOLS8, vd = idx % VCOLS8;
      o_out[inTV(b, n * BT16 + c, h, col0 + vd, T, H, V)] = C2[idx] + C3[idx];
    }
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 4);  // o write

    // gemm4: newS[VC,K] = vcorrT @ krT^T ; S^T = S^T*EL + newS
    if (tid < WS_GEMM) v8::cute_gemm16_d(vcorrT_bf, krT_bf, C4, VCOLS8, K, BT16);
    bar_sync(1, WS_COMPUTE);                    // newS (C4) visible to all compute
    PROF_ADD(g_prof_k2, 5);  // gemm4 newS
    // S^T = S^T*EL + newS (vectorized float4; K%4==0 so each group stays in a row,
    // d = (4i)%K is a multiple of 4 -> EL float4-aligned). EL produced by load warp.
    {
      float4* st4 = reinterpret_cast<float4*>(St_f);
      const float4* c44 = reinterpret_cast<const float4*>(C4);
      const float4* el4 = reinterpret_cast<const float4*>(EL_OF(slot));
      int KQ = K / 4;
      for (int i = tid; i < (VCOLS8 * K) / 4; i += WS_COMPUTE) {
        int dq = i % KQ;                       // which float4 of EL within the row
        float4 s = st4[i], cc = c44[i], e = el4[dq];
        s.x = s.x * e.x + cc.x; s.y = s.y * e.y + cc.y;
        s.z = s.z * e.z + cc.z; s.w = s.w * e.w + cc.w;
        st4[i] = s;
      }
    }
    bar_sync(1, WS_COMPUTE);
    bar_arrive(4 + slot, WSALL);               // slot (krT + EL) fully consumed -> free
    PROF_ADD(g_prof_k2, 6);  // S update
  }

  for (int idx = tid; idx < VCOLS8 * K; idx += WS_COMPUTE) {
    int vd = idx / K, d = idx % K;
    fs_out[bh * K * V + d * V + (col0 + vd)] = St_f[idx];
  }
}


// ---------------- host launchers ----------------

static int fwd_block_threads(int NT) {
  int warps = NT < 32 ? NT : 32;
  if (warps < 1) warps = 1;
  return warps * 32;
}

std::vector<torch::Tensor> kda_cuda_fwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size,
    double lower_bound, bool pipeline, bool warp_spec) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  // v8 forces CHUNK=16 (bounded-gate redesign); chunk_size arg is ignored.
  const int BT = 16;
  int NT = T / BT, BH = B * H;

  auto o = torch::zeros({B, T, H, V}, opt);
  auto fs = torch::zeros({B, H, K, V}, opt);
  // per-chunk workspaces. warp_spec consumes them as bf16 (cp.async, no
  // conversion); baseline/pipeline consume them as fp32. gl8 is always fp32.
  auto wopt = warp_spec ? torch::TensorOptions().dtype(torch::kBFloat16).device(q.device()) : opt;
  auto w8 = torch::zeros({BH, NT, BT, K}, wopt);
  auto u8 = torch::zeros({BH, NT, BT, V}, wopt);
  auto qd8 = torch::zeros({BH, NT, BT, K}, wopt);
  auto kr8 = torch::zeros({BH, NT, BT, K}, wopt);
  auto Aqk8 = torch::zeros({BH, NT, BT, BT}, wopt);
  auto gl8 = torch::zeros({BH, NT, K}, opt);

  // ---- K1: fused prepare (l2norm + bounded gate + cumsum + L/inv/Aqk/w/u) ----
  int mxKV = K > V ? K : V;
  int prep_sh = (BT * K + BT * K + BT * K + BT * BT) * sizeof(float)
              + (BT * K + BT * K + BT * V + BT * K + BT * BT) * sizeof(bf16);
  int prep_threads = 128;
  if (warp_spec) {
    cudaFuncSetAttribute(k_prepare_v8<bf16>, cudaFuncAttributeMaxDynamicSharedMemorySize, prep_sh);
    k_prepare_v8<bf16><<<BH * NT, prep_threads, prep_sh>>>(
        fp(q), fp(k), fp(g), fp(beta), fp(A_log), fp(dt_bias), fp(v),
        bp(w8), bp(u8), bp(qd8), bp(kr8), bp(Aqk8), fp(gl8),
        (float)scale, (float)lower_bound, B, T, H, K, V, NT);
  } else {
    cudaFuncSetAttribute(k_prepare_v8<float>, cudaFuncAttributeMaxDynamicSharedMemorySize, prep_sh);
    k_prepare_v8<float><<<BH * NT, prep_threads, prep_sh>>>(
        fp(q), fp(k), fp(g), fp(beta), fp(A_log), fp(dt_bias), fp(v),
        fp(w8), fp(u8), fp(qd8), fp(kr8), fp(Aqk8), fp(gl8),
        (float)scale, (float)lower_bound, B, T, H, K, V, NT);
  }

  // ---- K2: V-split serial recurrence ----
  {
    int VC = VCOLS8, VS = V / VC;
    int cbuf = 2 * BT * VC; if (VC * K > cbuf) cbuf = VC * K;
    int opsz = BT * K + BT * K + BT * BT + K * BT;
    if (warp_spec) {
      // warp-specialized: dedicated load warp (cp.async) + 2 operand buffers.
      // v9: each operand buffer also holds the prefetched u col-slice [16,VC].
      int opsz_ws = opsz + BT * VC;
      int rsh = (VC * K + cbuf) * sizeof(float)
              + (VC * K + 2 * opsz_ws + VC * BT) * sizeof(bf16);
      int thr = WS_COMPUTE + 32;  // 96 compute + 1 load warp
      cudaFuncSetAttribute(k_recur_v8_ws, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
      k_recur_v8_ws<<<BH * VS, thr, rsh>>>(
          bp(w8), bp(u8), bp(qd8), bp(kr8), bp(Aqk8), fp(gl8),
          fp(initial_state), fp(o), fp(fs), B, T, H, K, V, NT);
    } else {
      int nbuf = pipeline ? 2 : 1;
      int rsh = (VC * K + cbuf) * sizeof(float)
              + (VC * K + nbuf * opsz + VC * BT) * sizeof(bf16);
      if (pipeline) {
        cudaFuncSetAttribute(k_recur_v8<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
        k_recur_v8<true><<<BH * VS, v8::GEMM16_THREADS, rsh>>>(
            fp(w8), fp(u8), fp(qd8), fp(kr8), fp(Aqk8), fp(gl8),
            fp(initial_state), fp(o), fp(fs), B, T, H, K, V, NT);
      } else {
        cudaFuncSetAttribute(k_recur_v8<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
        k_recur_v8<false><<<BH * VS, v8::GEMM16_THREADS, rsh>>>(
            fp(w8), fp(u8), fp(qd8), fp(kr8), fp(Aqk8), fp(gl8),
            fp(initial_state), fp(o), fp(fs), B, T, H, K, V, NT);
      }
    }
  }
  (void)mxKV;
  cudaDeviceSynchronize();
  return {o, fs};
}

std::vector<torch::Tensor> kda_cuda_bwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size,
    torch::Tensor do_, torch::Tensor dht) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);
  do_ = do_.contiguous().to(torch::kFloat32);
  dht = dht.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  int BT = chunk_size, NT = T / BT, BH = B * H;

  auto qn = torch::zeros({BH, T, K}, opt), kn = torch::zeros({BH, T, K}, opt);
  auto qs = torch::zeros({BH, T, K}, opt), ga = torch::zeros({BH, T, K}, opt);
  auto gc = torch::zeros({BH, T, K}, opt);
  auto rq = torch::zeros({BH, T}, opt), rk = torch::zeros({BH, T}, opt);
  auto beta_s = torch::zeros({BH, T}, opt);
  auto X = torch::zeros({BH, NT, BT, BT}, opt);
  auto Afull = torch::zeros({BH, NT, BT, BT}, opt);
  auto Aqk = torch::zeros({BH, NT, BT, BT}, opt);
  auto w = torch::zeros({BH, T, K}, opt), u = torch::zeros({BH, T, V}, opt);
  auto vcorr = torch::zeros({BH, T, V}, opt);
  auto Sentry = torch::zeros({BH, NT + 1, K, V}, opt);
  auto o = torch::zeros({B, T, H, V}, opt), fs = torch::zeros({B, H, K, V}, opt);

  int wth = fwd_block_threads(NT);
  int mxKV = K > V ? K : V;
  int wy_sh = BT * BT * sizeof(float) + BT * BT * sizeof(bf16)
            + mxKV * BT * sizeof(bf16);
  cudaFuncSetAttribute(k_wy_v7, cudaFuncAttributeMaxDynamicSharedMemorySize, wy_sh);

  k_pre_v7<<<BH, wth>>>(fp(q), fp(k), fp(g), fp(beta), fp(A_log),
      fp(dt_bias), fp(qn), fp(kn), fp(qs), fp(rq), fp(rk), fp(beta_s),
      fp(ga), fp(gc), (float)scale, B, T, H, K, BT, NT);
  k_wy_v7<<<BH * NT, v7::GEMM_THREADS, wy_sh>>>(fp(v), fp(kn), fp(gc),
      fp(beta_s), fp(X), fp(Afull), fp(w), fp(u), B, T, H, K, V, BT, NT);
  k_aqk_v7<<<BH, wth>>>(fp(qs), fp(kn), fp(gc), fp(Aqk), B, T, H, K, BT, NT);
  {
    int VC = 32, VS = V / VC;
    int cbuf = 2 * BT * VC; if (VC * K > cbuf) cbuf = VC * K;
    int rsh = (VC * K + cbuf) * sizeof(float)
            + (VC * K + 2 * BT * K + VC * BT + BT * BT + K * BT) * sizeof(bf16);
    cudaFuncSetAttribute(k_recur_seq_v7<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
    k_recur_seq_v7<true><<<BH * VS, v7::GEMM_THREADS, rsh>>>(fp(qs), fp(kn), fp(gc),
        fp(w), fp(u), fp(initial_state), fp(Aqk), fp(vcorr), fp(Sentry),
        fp(o), fp(fs), B, T, H, K, V, BT, NT);
  }

  auto dq = torch::zeros_like(q), dk = torch::zeros_like(k);
  auto dv = torch::zeros_like(v), dgr = torch::zeros_like(g);
  auto dbeta = torch::zeros_like(beta);
  auto dA_log = torch::zeros_like(A_log), ddt_bias = torch::zeros_like(dt_bias);
  auto dh0 = torch::zeros_like(initial_state);

  BwdScratch sc;
  auto s_dgc = torch::zeros({BH, T, K}, opt); sc.dgc = fp(s_dgc);
  auto s_dqs = torch::zeros({BH, T, K}, opt); sc.dqs = fp(s_dqs);
  auto s_dkn = torch::zeros({BH, T, K}, opt); sc.dkn = fp(s_dkn);
  auto s_dbeta_s = torch::zeros({BH, T}, opt); sc.dbeta_s = fp(s_dbeta_s);
  auto s_dS = torch::zeros({BH, K, V}, opt); sc.dS = fp(s_dS);
  auto s_dS_entry = torch::zeros({BH, K, V}, opt); sc.dS_entry = fp(s_dS_entry);
  auto s_dEL = torch::zeros({BH, K}, opt); sc.dEL = fp(s_dEL);
  auto s_dvcorr = torch::zeros({BH, BT, V}, opt); sc.dvcorr = fp(s_dvcorr);
  auto s_du = torch::zeros({BH, BT, V}, opt); sc.du = fp(s_du);
  auto s_dAqk = torch::zeros({BH, BT, BT}, opt); sc.dAqk = fp(s_dAqk);
  auto s_dAfull = torch::zeros({BH, BT, BT}, opt); sc.dAfull = fp(s_dAfull);
  auto s_dX = torch::zeros({BH, BT, BT}, opt); sc.dX = fp(s_dX);
  auto s_Xf = torch::zeros({BH, BT, BT}, opt); sc.Xf = fp(s_Xf);
  auto s_tmp = torch::zeros({BH, BT, BT}, opt); sc.tmp = fp(s_tmp);
  auto s_dAraw = torch::zeros({BH, BT, BT}, opt); sc.dAraw = fp(s_dAraw);
  auto s_dQE = torch::zeros({BH, BT, K}, opt); sc.dQE = fp(s_dQE);
  auto s_dP = torch::zeros({BH, BT, K}, opt); sc.dP = fp(s_dP);
  auto s_dw = torch::zeros({BH, BT, K}, opt); sc.dw = fp(s_dw);
  auto s_dkg = torch::zeros({BH, BT, K}, opt); sc.dkg = fp(s_dkg);

  k_bwd_recur_v7<<<BH, 32>>>(fp(qs), fp(kn), fp(gc), fp(v), fp(w),
      fp(Afull), fp(X), fp(Aqk), fp(vcorr), fp(Sentry), fp(beta_s),
      fp(do_), fp(dht), sc, fp(dv), fp(dh0), B, T, H, K, V, BT);
  k_bwd_pre_v7<<<BH, NT>>>(fp(q), fp(k), fp(qn), fp(kn), fp(rq), fp(rk),
      fp(beta_s), fp(ga), fp(A_log), sc.dgc, sc.dqs, sc.dkn, sc.dbeta_s,
      fp(dq), fp(dk), fp(dgr), fp(dbeta), fp(dA_log), fp(ddt_bias),
      (float)scale, B, T, H, K, BT, NT);
  cudaDeviceSynchronize();
  return {dq, dk, dv, dgr, dbeta, dA_log, ddt_bias, dh0};
}

// Report theoretical occupancy of the recur kernels at the actual K=V=128
// launch config (no HW counters needed; uses the CUDA occupancy API).
static std::vector<int64_t> v8_occupancy(int64_t K, int64_t V) {
  int VC = VCOLS8;
  int cbuf = 2 * BT16 * VC; if (VC * (int)K > cbuf) cbuf = VC * K;
  int opsz = BT16 * K + BT16 * K + BT16 * BT16 + (int)K * BT16;
  int dev = 0; cudaGetDevice(&dev);
  cudaDeviceProp prop; cudaGetDeviceProperties(&prop, dev);
  auto report = [&](const void* fn, int threads, int shmem) -> std::vector<int64_t> {
    cudaFuncSetAttribute(fn, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem);
    int blocks = 0;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, fn, threads, shmem);
    cudaFuncAttributes at; cudaFuncGetAttributes(&at, fn);
    int warps_per_block = (threads + 31) / 32;
    int active_warps = blocks * warps_per_block;
    int max_warps = prop.maxThreadsPerMultiProcessor / 32;
    return {blocks, active_warps, max_warps, at.numRegs, shmem,
            (int64_t)prop.sharedMemPerMultiprocessor};
  };
  int rsh_b = (VC * (int)K + cbuf) * sizeof(float)
            + (VC * (int)K + 1 * opsz + VC * BT16) * sizeof(bf16);
  int rsh_w = (VC * (int)K + cbuf) * sizeof(float)
            + (VC * (int)K + 2 * (opsz + BT16 * VC) + VC * BT16) * sizeof(bf16);
  auto bl = report((const void*)k_recur_v8<false>, v8::GEMM16_THREADS, rsh_b);
  auto ws = report((const void*)k_recur_v8_ws, WS_COMPUTE + 32, rsh_w);
  // pack: [bl: blocks,warps,maxwarps,regs,shmem,smShmem][ws: ...]
  std::vector<int64_t> out;
  out.insert(out.end(), bl.begin(), bl.end());
  out.insert(out.end(), ws.begin(), ws.end());
  return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fwd", &kda_cuda_fwd,
        "KDA chunk forward (CUDA v8: CHUNK=16, bounded gate, all-TC)",
        pybind11::arg("q"), pybind11::arg("k"), pybind11::arg("v"),
        pybind11::arg("g"), pybind11::arg("beta"), pybind11::arg("A_log"),
        pybind11::arg("dt_bias"), pybind11::arg("scale"),
        pybind11::arg("initial_state"), pybind11::arg("chunk_size"),
        pybind11::arg("lower_bound") = -5.0, pybind11::arg("pipeline") = false,
        pybind11::arg("warp_spec") = false);
  m.def("bwd", &kda_cuda_bwd, "KDA chunk backward (CUDA v8, legacy v7 path)");
  m.def("occupancy", &v8_occupancy, "theoretical occupancy of recur kernels",
        pybind11::arg("K") = 128, pybind11::arg("V") = 128);
#ifdef V8_PROFILE
  m.def("prof_reset", []() {
    unsigned long long z[16] = {0};
    cudaMemcpyToSymbol(g_prof_k1, z, sizeof(z));
    cudaMemcpyToSymbol(g_prof_k2, z, sizeof(z));
  }, "zero the in-kernel clock64 accumulators");
  m.def("prof_read", []() {
    unsigned long long k1[16], k2[16];
    cudaMemcpyFromSymbol(k1, g_prof_k1, sizeof(k1));
    cudaMemcpyFromSymbol(k2, g_prof_k2, sizeof(k2));
    std::vector<int64_t> out;
    for (int i = 0; i < 16; ++i) out.push_back((int64_t)k1[i]);
    for (int i = 0; i < 16; ++i) out.push_back((int64_t)k2[i]);
    return out;
  }, "read [k1[0..15], k2[0..15]] cycle accumulators");
#endif
}


/**
K1 k_prepare_v8 (chunk-parallel prep, 27% of fwd)                                  
                                                                                     
  ┌──────────────────────────┬──────────┐                                            
  │            op            │  share   │                                            
  ├──────────────────────────┼──────────┤
  │ l2norm + gate            │ 37.9%    │
  ├──────────────────────────┼──────────┤
  │ tri-inverse (WY, 1 warp) │ 23.7%    │
  ├──────────────────────────┼──────────┤
  │ kd/ki/qd/kr emit         │ 10.0%    │
  ├──────────────────────────┼──────────┤
  │ w/u transpose            │ 8.5%     │
  ├──────────────────────────┼──────────┤
  │ w/u gemm                 │ 6.8%     │
  ├──────────────────────────┼──────────┤
  │ L gemm                   │ 3.8%     │
  ├──────────────────────────┼──────────┤
  │ Aqk gemm                 │ 3.4%     │
  ├──────────────────────────┼──────────┤
  │ cumsum                   │ 3.2%     │
  ├──────────────────────────┼──────────┤
  │ L mask / Af / Aqk store  │ <1% each │
  └──────────────────────────┴──────────┘

  K2 k_recur_v8_ws (serial recurrence, 71% of fwd)

  ┌─────────────────────────────┬───────┐
  │             op              │ share │
  ├─────────────────────────────┼───────┤
  │ S update (S=S·exp(gl)+newS) │ 39.7% │
  ├─────────────────────────────┼───────┤
  │ vcorr = u − wS              │ 16.8% │
  ├─────────────────────────────┼───────┤
  │ St → bf16 cast              │ 15.0% │
  ├─────────────────────────────┼───────┤
  │ gemm2/3 (qS, AV)            │ 10.6% │
  ├─────────────────────────────┼───────┤
  │ gemm1 (wS)                  │ 9.8%  │
  ├─────────────────────────────┼───────┤
  │ gemm4 (newS)                │ 4.9%  │
  ├─────────────────────────────┼───────┤
  │ o write                     │ 2.3%  │
  ├─────────────────────────────┼───────┤
  │ stall: wait-ops             │ 0.9%  │
  └─────────────────────────────┴───────┘

  Aggregated K2: gemms = 25%, elementwise = 74%, load-stall = 0.9%.

  What this says

  - The gemms are NOT the bottleneck. In K2 the 4 tensor-core gemms are only 25% of
  cycles; the tiles are tiny (16×32×128 etc.) so they finish fast. 74% is elementwise
   + the barriers between substeps — S update (the fp32 S·exp(gl)+newS over VC·K with
   an expf per element) alone is 40%, plus the St→bf16 recast (15%) and vcorr (17%).
  - The load-warp prefetch is working — stall:wait-ops is only 0.9%, confirming the
  cp.async overlap fully hides operand loading. The cost is purely intra-chunk
  compute/barrier latency on the serial S chain.
  - K1 is dominated by non-gemm work too: l2norm+gate (38%) and the 1-warp triangular
   inverse (24%) together are 62%; all 4 gemms combined are only ~14%.
*/

/**

      Long-context bench done (B=1, H=64, K=V=128):                                      
                                                                                     
  ┌──────┬───────┬─────────────┬────────┬───────────┬────────────────┐               
  │  T   │  fla  │ v8 baseline │ v8+ws  │ ws vs fla │ ws vs baseline │               
  ├──────┼───────┼─────────────┼────────┼───────────┼────────────────┤               
  │ 16K  │ 10.50 │ 50.49       │ 25.76  │ 0.41x     │ 1.96x          │               
  ├──────┼───────┼─────────────┼────────┼───────────┼────────────────┤               
  │ 64K  │ 41.93 │ 197.15      │ 101.94 │ 0.41x     │ 1.93x          │               
  ├──────┼───────┼─────────────┼────────┼───────────┼────────────────┤
  │ 128K │ 83.17 │ 401.80      │ 203.57 │ 0.41x     │ 1.97x          │
  └──────┴───────┴─────────────┴────────┴───────────┴────────────────┘                                                              
  ┌──────┬───────┬────────┬───────────┬────────────────┐
  │  T   │  fla  │ v8+ws  │ ws vs fla │ ws vs baseline │                             
  ├──────┼───────┼────────┼───────────┼────────────────┤    
  │ 16K  │ 10.50 │ 25.76  │ 0.41x     │ 1.96x          │
  ├──────┼───────┼────────┼───────────┼────────────────┤
  │ 64K  │ 41.93 │ 101.94 │ 0.41x     │ 1.93x          │
  ├──────┼───────┼────────┼───────────┼────────────────┤
  │ 128K │ 83.17 │ 203.57 │ 0.41x     │ 1.97x          │
  └──────┴───────┴────────┴───────────┴────────────────┘
*/

v9+ws — 相较 v2 提升 20.9x / 9.0x（相较 v8 提升 1.29x / 1.17x）

kda_chunk_cuda_v9.cu（load_cuda_v9，check_v9.py）。v8 的副本，对k_recur_v8_ws 做了两项纯性能改动。输出与 v8 按位一致（o 和 fs 差值为 0.00e+00）— 纯性能优化，无数学变更。直接由 v8 的 K2性能分析驱动（矩阵乘仅 ~35%，单 vcorr 循环占 24%）。

opt1 — 在加载 warp 中预取 u

u8 是唯一仍在关键路径上从全局直接读取的每 chunk 操作数（在 vcorr = u − wS 内部），这正是该循环开销达 24% 的原因 — 是同尺寸o write（2.4%）的 10 倍。v9 的加载 warp 也将 u 列切片 [16,VC] 通过cp.async 拉入双缓冲操作数区域（OPSZ += BT16·VCOLS8，新增 U_OF 访问器）。计算 warp 从共享内存读取 u_bf。→ vcorr 循环 24.2% → 7.2%。

opt2 — 向量化两个 [VC,K]=4096 逐元素循环（float4 / bf162）

• St→bf16：float4 加载 + 2× __floats2bfloat162_rn → uint2 存储。
• S 更新：St_f/C4/EL_sh 的 float4（K%4==0 保证每个 EL float4 在行内， EL_sh 使用 __align__(16)）。加载/存储/ALU 指令减少 4 倍。→ St→bf16 13.1%→12.6%，S 更新 24.4%→20.4%（占比变化较小，因为 K2 总时间已大幅缩小）。

v9 性能分析（clock64）

• K2 循环周期 11.79 → 8.44 Gcyc（−28%）。
• 现在矩阵乘主导：矩阵乘 54.3%，逐元素 44.3%，停顿 1.4%。
• B=1 提升更大（1.29x）于 B=8（1.17x）：B=1 为延迟受限，隐藏 u 的延迟收益最大； B=8 更受占用率限制，更大的操作数缓冲略微牺牲了占用率。

---

// v10 register-resident per-chunk recurrence helper (CHUNK=16, V-split).
//
// Keeps the WY-correction `vcorr` register-resident across the o-path and the
// state update via SM75_U32x1_MOVM_T (FlashKDA's technique): vcorr is produced
// in gemm1's C-accumulator, MOVM_T-transposed into a B-operand fragment once,
// then reused as B in both AV = Aqk@vcorr and newS^T = krT@vcorr. This removes
// the C1/vcorr/C2/C3 shared round-trips + barriers that v9 had between the 4
// gemms. S stays in shared (bf16 B-operand); only vcorr lives in registers.
//
//   recur_warp<K,VC>(...)  — one warp (lane 0..31) owns vcol-block `vb` (16 of
//   the CTA's VC value columns). Writes this vb's `o` straight to global and its
//   newS^T slice into newS_sh (transposed into [vcol,d] to match St_f's layout).
//
// MOVM_T semantics (validated by test_recur_v10.cu): MOVM_T(C-acc representing
// X[16,16]) yields a B-operand BU such that gemm(A, BU) == A @ X.

#pragma once

#include <cuda_bf16.h>
#include <cute/tensor.hpp>
#include <cute/atom/mma_atom.hpp>
#include <cute/atom/copy_atom.hpp>
#include <cute/arch/copy_sm75.hpp>

namespace v10 {

using bf16 = __nv_bfloat16;
using namespace cute;

// single-warp 16x16x16 MMA (32 threads); MOVM_T is a warp-level register op.
using MMA1 = decltype(make_tiled_mma(
    MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>{},
    Layout<Shape<_1, _1>>{}, Tile<_16, _16, _16>{}));

// C[16,16] += A[16,Kk] @ B[16,Kk]^T  (plain copy from row-major bf16 smem).
template <int Kk, class CFrag>
__device__ inline void accum_AxBt(const bf16* Ap, int lda, const bf16* Bp,
                                  int ldb, int lane, CFrag& rC) {
  Tensor sA = make_tensor(make_smem_ptr(Ap),
      make_layout(make_shape(_16{}, Int<Kk>{}), make_stride(lda, _1{})));
  Tensor sB = make_tensor(make_smem_ptr(Bp),
      make_layout(make_shape(_16{}, Int<Kk>{}), make_stride(ldb, _1{})));
  MMA1 mma; auto thr = mma.get_slice(lane);
  auto tCsA = thr.partition_A(sA);
  auto tCsB = thr.partition_B(sB);
  auto rA = thr.make_fragment_A(tCsA);
  auto rB = thr.make_fragment_B(tCsB);
  int nk = size<2>(rA);
#pragma unroll
  for (int ik = 0; ik < nk; ++ik) {
    copy(tCsA(_, _, ik), rA(_, _, ik));
    copy(tCsB(_, _, ik), rB(_, _, ik));
    gemm(mma, rC, rA(_, _, ik), rB(_, _, ik), rC);
  }
}

// C[16,16] += A[16,16] @ B-operand-frag^T  (B already a partitioned B fragment).
template <class CFrag, class BFrag>
__device__ inline void accum_AxBfrag(const bf16* Ap, int lda, int lane,
                                     const BFrag& rB, CFrag& rC) {
  Tensor sA = make_tensor(make_smem_ptr(Ap),
      make_layout(make_shape(_16{}, _16{}), make_stride(lda, _1{})));
  MMA1 mma; auto thr = mma.get_slice(lane);
  auto tCsA = thr.partition_A(sA);
  auto rA = thr.make_fragment_A(tCsA);
  copy(tCsA(_, _, 0), rA(_, _, 0));
  gemm(mma, rC, rA(_, _, 0), rB(_, _, 0), rC);
}

// One warp owns vcol-block vb (16 cols). All operands row-major bf16 in smem.
//   w_bf,qd_bf : [16,K]        Aqk_bf : [16,16]      krT_bf : [K,16]
//   St_bf      : [VC,K]        u_bf   : [16,VC]
//   o_base     : global ptr at (time row 0, this vb's col 0), row stride o_rs
//   newS_sh    : [VC,K] fp32 shared; this vb's newS^T written transposed [vcol,d]
template <int K, int VC>
__device__ inline void recur_warp(
    const bf16* w_bf, const bf16* qd_bf, const bf16* St_bf,
    const bf16* Aqk_bf, const bf16* krT_bf, const bf16* u_bf,
    float* o_base, int o_rs, float* newS_sh, int vb, int lane) {
  MMA1 mma; auto thr = mma.get_slice(lane);
  Tensor cref = make_tensor(make_smem_ptr((float*)nullptr),
      make_layout(make_shape(_16{}, _16{}), make_stride(_16{}, _1{})));

  const bf16* Sw = St_bf + (vb * 16) * K;   // B-operand rows for this vb [16,K]

  // gemm1: wS[time,vcol] = w @ Sw^T ; vcorr = u - wS  (kept in C-acc regs)
  auto rVC = thr.partition_fragment_C(cref); clear(rVC);
  accum_AxBt<K>(w_bf, K, Sw, K, lane, rVC);
  {
    Tensor su = make_tensor(make_smem_ptr(u_bf + vb * 16),  // [time,vcol] stride (VC,1)
        make_layout(make_shape(_16{}, _16{}), make_stride(Int<VC>{}, _1{})));
    auto rU = make_fragment_like<bf16>(rVC);
    copy(thr.partition_C(su), rU);
#pragma unroll
    for (int i = 0; i < size(rVC); ++i) rVC(i) = __bfloat162float(rU(i)) - rVC(i);
  }

  // MOVM_T: vcorr C-acc -> B-operand fragment (reused by AV and newS^T).
  auto rVCbf = make_fragment_like<bf16>(rVC);
#pragma unroll
  for (int i = 0; i < size(rVC); ++i) rVCbf(i) = __float2bfloat16(rVC(i));
  Tensor bref = make_tensor(make_smem_ptr((bf16*)nullptr),
      make_layout(make_shape(_16{}, _16{}), make_stride(_16{}, _1{})));
  auto rVCb = thr.partition_fragment_B(bref);
  {
    uint32_t* src = reinterpret_cast<uint32_t*>(&rVCbf(0));
    uint32_t* dst = reinterpret_cast<uint32_t*>(&rVCb(0));
#pragma unroll
    for (int i = 0; i < 4; ++i) SM75_U32x1_MOVM_T::copy(src[i], dst[i]);
  }

  // gemm2: qS[time,vcol] = qd @ Sw^T ;  gemm3: AV = Aqk @ vcorr ; o = qS + AV
  auto rO = thr.partition_fragment_C(cref); clear(rO);
  accum_AxBt<K>(qd_bf, K, Sw, K, lane, rO);
  accum_AxBfrag(Aqk_bf, 16, lane, rVCb, rO);
  {
    Tensor go = make_tensor(make_gmem_ptr(o_base),
        make_layout(make_shape(_16{}, _16{}), make_stride(o_rs, _1{})));
    copy(rO, thr.partition_C(go));
  }

  // gemm4: newS^T[d,vcol] = krT @ vcorr, per d-block; store transposed [vcol,d]
  // into newS_sh so the S-update reads it aligned with St_f[vcol,d].
  constexpr int NDB = K / 16;
#pragma unroll
  for (int db = 0; db < NDB; ++db) {
    auto rNS = thr.partition_fragment_C(cref); clear(rNS);
    accum_AxBfrag(krT_bf + db * 16 * 16, 16, lane, rVCb, rNS);
    // store [d-row, vcol] -> newS_sh[(vb*16+vcol)*K + (db*16+d-row)]
    Tensor sns = make_tensor(make_smem_ptr(newS_sh + (vb * 16) * K + db * 16),
        make_layout(make_shape(_16{}, _16{}), make_stride(_1{}, Int<K>{})));
    copy(rNS, thr.partition_C(sns));
  }
}

}  // namespace v10

// KDA chunk reference (CUDA, bf16) — v8 CHUNK=16 + bounded-gate redesign.
//
// FORWARD-ONLY. This is a structural redesign over v7 that adopts FlashKDA's
// two key choices so that the three FlashKDA-inspired optimizations actually
// pay off (in the BT=64 + softplus build they were neutral/regressive):
//
//   * CHUNK = 16 (BT=16) and the BOUNDED gate
//        ga = lower_bound * sigmoid(exp(A_log) * (g_raw + dt_bias))   in [lb, 0)
//     With BT=16 and lb=-5 the worst within-chunk cumsum span is 15*(-5)=-75,
//     so exp(-gc) <= e^75 ~ 3.7e32 < bf16 max 3.4e38 => EVERY matmul (incl. the
//     L matrix and Aqk, which v7 had to keep scalar fp32 because ki=kn*exp(-gc)
//     overflowed) is now overflow-safe and goes to bf16 tensor cores.
//
// Two-kernel structure mirroring FlashKDA (forward only):
//   K1  k_prepare_v8 (opt2 = fusion): grid=BH*NT, 1 CTA per (chunk,head), 16
//       rows. l2norm + bounded gate + cumsum, then on TC (cute_gemm16):
//         L = kd@ki^T (strict-lower, *(-beta)), triangular inverse (I+L)^-1,
//         Aqk = qd@ki^T (lower-tri), w = Af@kd, u = Af@v.
//       kd=kn*e^{gc}, ki=kn*e^{-gc}, qd=qn*scale*e^{gc}. Stores w,u,qd,Aqk,kr,
//       glast to a per-chunk workspace (former k_pre+k_aqk+k_wy fused into one).
//   K2  k_recur_v8 (opt1 + opt3): grid=BH*VS, V-split (VCOLS=32), S^T[VC,K]
//       resident in shared, serial chunk loop. Per chunk on TC:
//         vcorr = u - w@S^T ; o = qd@S^T + Aqk@vcorr ; S^T = S^T*e^{glast} + vcorr^T@kr^T.
//       opt3: forward skips per-chunk Sentry/vcorr global writes (only o and
//       final_state written).  opt1: with BT=16 the recur kernel is ~50KB (not
//       ~100KB), so an occupancy-safe double-buffer of the St-INDEPENDENT
//       operands (w/qd/Aqk/kr) is feasible (PIPELINE template flag) — these are
//       prepared in K1 and don't depend on the running state, so chunk t+1 can
//       be staged while chunk t computes.
//
// Validation: chunk_kda rejects chunk_size=16, so v8 (BT=16) is validated vs
//   chunk_kda(safe_gate=True, lower_bound=-5, chunk_size=64) — the forward o /
//   final_state are mathematically chunk-size-independent — and cross-checked vs
//   fused_recurrent_kda(lower_bound=-5).
//
// Assumes HV==H, fixed length, K==V multiple of 16.
//
// NOTE: the BACKWARD path below is left as the v7 (softplus / BT in {32,64})
// kernels verbatim and is intentionally NOT updated or validated — v8 scope is
// forward-only. kda_cuda_bwd still launches the old *_v7 kernels so the module
// keeps compiling, but it does not match v8's CHUNK=16 + bounded-gate forward.

#include <cuda_runtime.h>
#include <torch/extension.h>

#include <vector>

#include "cute_gemm_v7.cuh"
#include "cute_gemm_v8.cuh"
#include "cute_recur_v10.cuh"

namespace {

using v7::bf16;
constexpr float L2_EPS = 1e-6f;

__device__ inline float softplusf(float x) {
  if (x > 20.0f) return x;
  if (x < -20.0f) return expf(x);
  return log1pf(expf(x));
}
__device__ inline float sigmoidf(float x) { return 1.0f / (1.0f + expf(-x)); }
// store a float into an output slot of type OT (float passthrough / bf16 cast).
__device__ __forceinline__ void ot_store(float* p, float v) { *p = v; }
__device__ __forceinline__ void ot_store(bf16* p, float v) { *p = __float2bfloat16(v); }
__device__ inline float warp_sum(float v) {
  for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o);
  return v;
}
// Named barrier: synchronize exactly `cnt` threads on logical barrier `id`
// (id 1..15). Lets a warp-specialized kernel sync only its compute warps
// (cnt=COMPUTE_THREADS) while the dedicated load warp stays off that barrier.
__device__ __forceinline__ void bar_sync(int id, int cnt) {
  asm volatile("barrier.sync %0, %1;" :: "r"(id), "r"(cnt) : "memory");
}
// Arrive (no wait) on named barrier `id` expecting `cnt` participants.
__device__ __forceinline__ void bar_arrive(int id, int cnt) {
  asm volatile("barrier.arrive %0, %1;" :: "r"(id), "r"(cnt) : "memory");
}
// 16-byte cp.async global->shared (8 bf16). dst must be a shared address.
__device__ __forceinline__ void cp_async16(bf16* dst, const bf16* src) {
  unsigned s = (unsigned)__cvta_generic_to_shared(dst);
  asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" :: "r"(s), "l"(src));
}
__device__ __forceinline__ void cp_async_commit() {
  asm volatile("cp.async.commit_group;\n");
}
__device__ __forceinline__ void cp_async_wait_all() {
  asm volatile("cp.async.wait_group 0;\n");
}

// ---- optional in-kernel per-region timing (clock64). OFF by default: when
// V8_PROFILE is not defined every macro below expands to nothing, so the
// production kernels are byte-identical. ON: thread 0 accumulates cycle deltas
// per region into a __device__ global, read back via the v8 pybind. ----
#ifdef V8_PROFILE
__device__ unsigned long long g_prof_k1[16];
__device__ unsigned long long g_prof_k2[16];
#define PROF_DECL long long _pc = 0;
#define PROF_START() do { if (threadIdx.x == 0) _pc = clock64(); } while (0)
#define PROF_ADD(arr, i) do { if (threadIdx.x == 0) { \
    long long _n = clock64(); atomicAdd(&arr[i], (unsigned long long)(_n - _pc)); _pc = _n; } } while (0)
#else
#define PROF_DECL
#define PROF_START()
#define PROF_ADD(arr, i)
#endif

__device__ inline int iTK(int bh, int t, int d, int T, int K) { return bh * T * K + t * K + d; }
__device__ inline int iTV(int bh, int t, int vd, int T, int V) { return bh * T * V + t * V + vd; }
__device__ inline int inTK(int b, int t, int h, int d, int T, int H, int K) {
  return ((b * T + t) * H + h) * K + d;
}
__device__ inline int inTV(int b, int t, int h, int vd, int T, int H, int V) {
  return ((b * T + t) * H + h) * V + vd;
}

// ================= FORWARD =================

// ---- warp-specialized recurrence (real load/compute overlap) ----
// blockDim = WS_COMPUTE (96 = 3 compute warps) + 32 (1 dedicated load warp).
// Operands now arrive as BF16 in global (K1 emits bf16 for this path) and kr8
// is pre-transposed [K,16], so the load warp uses 128-bit cp.async straight
// into shared with NO conversion/transpose. Double-buffered: load warp stages
// chunk n+1 while the compute warps run chunk n's 4 gemms (FlashKDA load-warp).
// The 4 cute gemms use exactly 64 threads, so they are guarded to tid<64; the
// extra compute warp (tid 64..95) only accelerates the elementwise loops.
//   vcorr stays in shared (CuTe helper hides fragment layout; vcorr is a
//   B-operand in AV but A-operand in newS -> incompatible frags).
constexpr int WS_COMPUTE = 128;                 // 4 compute warps (sweet spot: keeps 3 blk/SM)
constexpr int WS_GEMM = v8::GEMM16_THREADS;     // 64 threads do the MMA
// barrier ids: 1 = compute-only; 2/3 = operand-ready slot0/1; 4/5 = free slot0/1
extern __shared__ char recur8ws_smem[];
__global__ void k_recur_v8_ws(
    const bf16* w8, const bf16* u8, const bf16* qd8, const bf16* kr8,
    const bf16* Aqk8, const float* gl8, const float* h0,
    float* o_out, float* fs_out,
    int B, int T, int H, int K, int V, int NT) {
  int VS = V / VCOLS8;
  int blk = blockIdx.x;
  if (blk >= B * H * VS) return;
  int bh = blk / VS, vsblk = blk % VS, col0 = vsblk * VCOLS8;
  int b = bh / H, h = bh % H, tid = threadIdx.x;
  bool is_load = tid >= WS_COMPUTE;          // load warp = last 32 threads
  int ll = tid - WS_COMPUTE;                 // load-warp lane 0..31
  int WSALL = WS_COMPUTE + 32;               // participants on ready/free barriers

  int CBUF = 2 * BT16 * VCOLS8;
  if (VCOLS8 * K > CBUF) CBUF = VCOLS8 * K;
  float* St_f = (float*)recur8ws_smem;        // [VC,K]
  float* Cbuf = St_f + VCOLS8 * K;            // gemm scratch
  bf16* St_bf = (bf16*)(Cbuf + CBUF);         // [VC,K]
  // v9: operand buffer now also holds the prefetched u column-slice [16,VC] so
  // the load warp hides u's global latency (was the only un-prefetched operand).
  int OPSZ = BT16 * K + BT16 * K + BT16 * BT16 + K * BT16 + BT16 * VCOLS8;
  bf16* opbuf = St_bf + VCOLS8 * K;           // [2 * OPSZ]
  bf16* vcorrT_bf = opbuf + 2 * OPSZ;         // [VC,16]
  float* C1 = Cbuf;
  float* C2 = Cbuf;
  float* C3 = Cbuf + BT16 * VCOLS8;
  float* C4 = Cbuf;
  // EL = exp(glast) per dim, now DOUBLE-BUFFERED per slot and produced by the
  // load warp (it reads gl8 from global + does the K expf's off the compute
  // critical path). [2][K<=256], float4-aligned. Freed with the operand slot.
  __shared__ __align__(16) float EL_sh2[2 * 256];

  auto W_OF = [&](int slot) { return opbuf + slot * OPSZ; };
  auto QD_OF = [&](int slot) { return opbuf + slot * OPSZ + BT16 * K; };
  auto AQ_OF = [&](int slot) { return opbuf + slot * OPSZ + 2 * BT16 * K; };
  auto KR_OF = [&](int slot) { return opbuf + slot * OPSZ + 2 * BT16 * K + BT16 * BT16; };
  auto U_OF = [&](int slot) { return opbuf + slot * OPSZ + 2 * BT16 * K + BT16 * BT16 + K * BT16; };
  auto EL_OF = [&](int slot) { return EL_sh2 + slot * 256; };

  if (is_load) {
    // -------- dedicated load warp: cp.async chunk operands ahead of compute --
    // producer: wait FREE(slot) -> cp.async -> commit+wait -> signal READY(slot).
    for (int n = 0; n < NT; ++n) {
      int slot = n & 1;
      bar_sync(4 + slot, WSALL);               // wait slot freed (primed by compute)
      const bf16* w = w8 + (bh * NT + n) * BT16 * K;
      const bf16* qd = qd8 + (bh * NT + n) * BT16 * K;
      const bf16* Aq = Aqk8 + (bh * NT + n) * BT16 * BT16;
      const bf16* kr = kr8 + (bh * NT + n) * BT16 * K;  // pre-transposed [K,16]
      const bf16* u = u8 + (bh * NT + n) * BT16 * V;     // v9: prefetch u col-slice
      bf16* w_bf = W_OF(slot); bf16* qd_bf = QD_OF(slot);
      bf16* Aqk_bf = AQ_OF(slot); bf16* krT_bf = KR_OF(slot);
      bf16* u_bf = U_OF(slot);                           // dst [16,VC] contiguous
      // each cp.async moves 8 bf16 (16B). all these counts are multiples of 8.
      for (int v = ll; v < (BT16 * K) / 8; v += 32) cp_async16(w_bf + v * 8, w + v * 8);
      for (int v = ll; v < (BT16 * K) / 8; v += 32) cp_async16(qd_bf + v * 8, qd + v * 8);
      for (int v = ll; v < (BT16 * BT16) / 8; v += 32) cp_async16(Aqk_bf + v * 8, Aq + v * 8);
      for (int v = ll; v < (K * BT16) / 8; v += 32) cp_async16(krT_bf + v * 8, kr + v * 8);
      // u col-slice: dst row c is contiguous VC bf16; src row c at u + c*V + col0
      // (strided by V). VC=32 -> VC/8=4 segs per row.
      for (int v = ll; v < (BT16 * VCOLS8) / 8; v += 32) {
        int c = v / (VCOLS8 / 8), seg = v % (VCOLS8 / 8);
        cp_async16(u_bf + v * 8, u + c * V + col0 + seg * 8);
      }
      cp_async_commit();
      // EL = exp(glast[d]) for this chunk: load warp does the K expf's so the
      // compute warps' S-update reads EL straight from shared (transcendentals
      // off the critical path). gl8 is fp32 global; EL slot is fp32 shared.
      const float* gl = gl8 + (bh * NT + n) * K;
      float* el = EL_OF(slot);
      for (int d = ll; d < K; d += 32) el[d] = expf(gl[d]);
      cp_async_wait_all();                     // ensure landed before compute reads
      bar_arrive(2 + slot, WSALL);             // operands ready for compute
    }
    return;
  }

  // -------- compute warps (96): carry St_f, run the 4 gemms per chunk --------
  for (int idx = tid; idx < VCOLS8 * K; idx += WS_COMPUTE) {
    int vd = idx / K, d = idx % K;
    St_f[idx] = h0[bh * K * V + d * V + (col0 + vd)];
  }
  bar_sync(1, WS_COMPUTE);
  // prime FREE(slot): both buffers start empty so the load warp may fill them.
  bar_arrive(4 + 0, WSALL);
  bar_arrive(4 + 1, WSALL);
  PROF_DECL PROF_START();

  for (int n = 0; n < NT; ++n) {
    int slot = n & 1;
    bar_sync(2 + slot, WSALL);                 // wait operands staged by load warp
    PROF_ADD(g_prof_k2, 7);  // stall: wait operands
    bf16* w_bf = W_OF(slot); bf16* qd_bf = QD_OF(slot);
    bf16* Aqk_bf = AQ_OF(slot); bf16* krT_bf = KR_OF(slot);
    bf16* u_bf = U_OF(slot);                   // v9: u col-slice [16,VC] in shared
    // gl/EL is produced by the load warp into EL_OF(slot); compute reads it below.

    // St -> bf16 (vectorized: float4 read, 2x bf162 write; VC*K%8==0).
    {
      const float4* sf4 = reinterpret_cast<const float4*>(St_f);
      uint2* sb2 = reinterpret_cast<uint2*>(St_bf);   // 4 bf16 = uint2
      for (int i = tid; i < (VCOLS8 * K) / 4; i += WS_COMPUTE) {
        float4 f = sf4[i];
        __nv_bfloat162 lo = __floats2bfloat162_rn(f.x, f.y);
        __nv_bfloat162 hi = __floats2bfloat162_rn(f.z, f.w);
        uint2 packed;
        packed.x = *reinterpret_cast<unsigned*>(&lo);
        packed.y = *reinterpret_cast<unsigned*>(&hi);
        sb2[i] = packed;
      }
    }
    bar_sync(1, WS_COMPUTE);
    PROF_ADD(g_prof_k2, 0);  // St->bf16

    // ---- register-resident vcorr chain (v10): warps 0,1 own vcol-block 0,1
    // (VCOLS8=32 = 2x16). vcorr is produced in gemm1's C-acc, MOVM_T'd once into
    // a B-operand and reused by AV (o-path) and newS^T (state) -> no C1/vcorr/
    // C2/C3 shared round-trips or barriers between the gemms. o is written
    // straight to global; newS^T is written transposed into Cbuf [VC,K] so the
    // S-update reads it aligned with St_f[vcol,d] (Cbuf == C4).
    if (tid < 2 * 32) {
      int warp = tid >> 5, lane = tid & 31, vb = warp;
      float* o_base = o_out +
          (((long)(b * T + n * BT16) * H + h) * V + col0 + vb * 16);
      if (K == 128)
        v10::recur_warp<128, VCOLS8>(w_bf, qd_bf, St_bf, Aqk_bf, krT_bf, u_bf,
                                     o_base, H * V, Cbuf, vb, lane);
      else
        v10::recur_warp<64, VCOLS8>(w_bf, qd_bf, St_bf, Aqk_bf, krT_bf, u_bf,
                                    o_base, H * V, Cbuf, vb, lane);
    }
    bar_sync(1, WS_COMPUTE);   // o written + newS^T (Cbuf) visible to all compute
    PROF_ADD(g_prof_k2, 5);    // gemms + o + newS (register-resident)
    // S^T = S^T*EL + newS (vectorized float4; K%4==0 so each group stays in a row,
    // d = (4i)%K is a multiple of 4 -> EL float4-aligned). EL produced by load warp.
    {
      float4* st4 = reinterpret_cast<float4*>(St_f);
      const float4* c44 = reinterpret_cast<const float4*>(C4);
      const float4* el4 = reinterpret_cast<const float4*>(EL_OF(slot));
      int KQ = K / 4;
      for (int i = tid; i < (VCOLS8 * K) / 4; i += WS_COMPUTE) {
        int dq = i % KQ;                       // which float4 of EL within the row
        float4 s = st4[i], cc = c44[i], e = el4[dq];
        s.x = s.x * e.x + cc.x; s.y = s.y * e.y + cc.y;
        s.z = s.z * e.z + cc.z; s.w = s.w * e.w + cc.w;
        st4[i] = s;
      }
    }
    bar_sync(1, WS_COMPUTE);
    bar_arrive(4 + slot, WSALL);               // slot (krT + EL) fully consumed -> free
    PROF_ADD(g_prof_k2, 6);  // S update
  }

  for (int idx = tid; idx < VCOLS8 * K; idx += WS_COMPUTE) {
    int vd = idx / K, d = idx % K;
    fs_out[bh * K * V + d * V + (col0 + vd)] = St_f[idx];
  }
}


}  // namespace

// ---------------- host launchers ----------------

static int fwd_block_threads(int NT) {
  int warps = NT < 32 ? NT : 32;
  if (warps < 1) warps = 1;
  return warps * 32;
}

std::vector<torch::Tensor> kda_cuda_fwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size,
    double lower_bound, bool pipeline, bool warp_spec) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  // v8 forces CHUNK=16 (bounded-gate redesign); chunk_size arg is ignored.
  const int BT = 16;
  int NT = T / BT, BH = B * H;

  auto o = torch::zeros({B, T, H, V}, opt);
  auto fs = torch::zeros({B, H, K, V}, opt);
  // per-chunk workspaces. warp_spec consumes them as bf16 (cp.async, no
  // conversion); baseline/pipeline consume them as fp32. gl8 is always fp32.
  auto wopt = warp_spec ? torch::TensorOptions().dtype(torch::kBFloat16).device(q.device()) : opt;
  auto w8 = torch::zeros({BH, NT, BT, K}, wopt);
  auto u8 = torch::zeros({BH, NT, BT, V}, wopt);
  auto qd8 = torch::zeros({BH, NT, BT, K}, wopt);
  auto kr8 = torch::zeros({BH, NT, BT, K}, wopt);
  auto Aqk8 = torch::zeros({BH, NT, BT, BT}, wopt);
  auto gl8 = torch::zeros({BH, NT, K}, opt);

  // ---- K1: fused prepare (l2norm + bounded gate + cumsum + L/inv/Aqk/w/u) ----
  int mxKV = K > V ? K : V;
  int prep_sh = (BT * K + BT * K + BT * K + BT * BT) * sizeof(float)
              + (BT * K + BT * K + BT * V + BT * K + BT * BT) * sizeof(bf16);
  int prep_threads = 128;
  if (warp_spec) {
    cudaFuncSetAttribute(k_prepare_v8<bf16>, cudaFuncAttributeMaxDynamicSharedMemorySize, prep_sh);
    k_prepare_v8<bf16><<<BH * NT, prep_threads, prep_sh>>>(
        fp(q), fp(k), fp(g), fp(beta), fp(A_log), fp(dt_bias), fp(v),
        bp(w8), bp(u8), bp(qd8), bp(kr8), bp(Aqk8), fp(gl8),
        (float)scale, (float)lower_bound, B, T, H, K, V, NT);
  } else {
    cudaFuncSetAttribute(k_prepare_v8<float>, cudaFuncAttributeMaxDynamicSharedMemorySize, prep_sh);
    k_prepare_v8<float><<<BH * NT, prep_threads, prep_sh>>>(
        fp(q), fp(k), fp(g), fp(beta), fp(A_log), fp(dt_bias), fp(v),
        fp(w8), fp(u8), fp(qd8), fp(kr8), fp(Aqk8), fp(gl8),
        (float)scale, (float)lower_bound, B, T, H, K, V, NT);
  }

  // ---- K2: V-split serial recurrence ----
  {
    int VC = VCOLS8, VS = V / VC;
    int cbuf = 2 * BT * VC; if (VC * K > cbuf) cbuf = VC * K;
    int opsz = BT * K + BT * K + BT * BT + K * BT;
    if (warp_spec) {
      // warp-specialized: dedicated load warp (cp.async) + 2 operand buffers.
      // v9: each operand buffer also holds the prefetched u col-slice [16,VC].
      int opsz_ws = opsz + BT * VC;
      int rsh = (VC * K + cbuf) * sizeof(float)
              + (VC * K + 2 * opsz_ws + VC * BT) * sizeof(bf16);
      int thr = WS_COMPUTE + 32;  // 96 compute + 1 load warp
      cudaFuncSetAttribute(k_recur_v8_ws, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
      k_recur_v8_ws<<<BH * VS, thr, rsh>>>(
          bp(w8), bp(u8), bp(qd8), bp(kr8), bp(Aqk8), fp(gl8),
          fp(initial_state), fp(o), fp(fs), B, T, H, K, V, NT);
    } else {
      int nbuf = pipeline ? 2 : 1;
      int rsh = (VC * K + cbuf) * sizeof(float)
              + (VC * K + nbuf * opsz + VC * BT) * sizeof(bf16);
      if (pipeline) {
        cudaFuncSetAttribute(k_recur_v8<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
        k_recur_v8<true><<<BH * VS, v8::GEMM16_THREADS, rsh>>>(
            fp(w8), fp(u8), fp(qd8), fp(kr8), fp(Aqk8), fp(gl8),
            fp(initial_state), fp(o), fp(fs), B, T, H, K, V, NT);
      } else {
        cudaFuncSetAttribute(k_recur_v8<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
        k_recur_v8<false><<<BH * VS, v8::GEMM16_THREADS, rsh>>>(
            fp(w8), fp(u8), fp(qd8), fp(kr8), fp(Aqk8), fp(gl8),
            fp(initial_state), fp(o), fp(fs), B, T, H, K, V, NT);
      }
    }
  }
  (void)mxKV;
  cudaDeviceSynchronize();
  return {o, fs};
}

std::vector<torch::Tensor> kda_cuda_bwd(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor beta, torch::Tensor A_log, torch::Tensor dt_bias,
    double scale, torch::Tensor initial_state, int64_t chunk_size,
    torch::Tensor do_, torch::Tensor dht) {
  auto opt = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
  q = q.contiguous().to(torch::kFloat32);
  k = k.contiguous().to(torch::kFloat32);
  v = v.contiguous().to(torch::kFloat32);
  g = g.contiguous().to(torch::kFloat32);
  beta = beta.contiguous().to(torch::kFloat32);
  A_log = A_log.contiguous().to(torch::kFloat32);
  dt_bias = dt_bias.contiguous().to(torch::kFloat32);
  initial_state = initial_state.contiguous().to(torch::kFloat32);
  do_ = do_.contiguous().to(torch::kFloat32);
  dht = dht.contiguous().to(torch::kFloat32);

  int B = q.size(0), T = q.size(1), H = q.size(2), K = q.size(3), V = v.size(3);
  int BT = chunk_size, NT = T / BT, BH = B * H;

  auto qn = torch::zeros({BH, T, K}, opt), kn = torch::zeros({BH, T, K}, opt);
  auto qs = torch::zeros({BH, T, K}, opt), ga = torch::zeros({BH, T, K}, opt);
  auto gc = torch::zeros({BH, T, K}, opt);
  auto rq = torch::zeros({BH, T}, opt), rk = torch::zeros({BH, T}, opt);
  auto beta_s = torch::zeros({BH, T}, opt);
  auto X = torch::zeros({BH, NT, BT, BT}, opt);
  auto Afull = torch::zeros({BH, NT, BT, BT}, opt);
  auto Aqk = torch::zeros({BH, NT, BT, BT}, opt);
  auto w = torch::zeros({BH, T, K}, opt), u = torch::zeros({BH, T, V}, opt);
  auto vcorr = torch::zeros({BH, T, V}, opt);
  auto Sentry = torch::zeros({BH, NT + 1, K, V}, opt);
  auto o = torch::zeros({B, T, H, V}, opt), fs = torch::zeros({B, H, K, V}, opt);

  int wth = fwd_block_threads(NT);
  int mxKV = K > V ? K : V;
  int wy_sh = BT * BT * sizeof(float) + BT * BT * sizeof(bf16)
            + mxKV * BT * sizeof(bf16);
  cudaFuncSetAttribute(k_wy_v7, cudaFuncAttributeMaxDynamicSharedMemorySize, wy_sh);

  k_pre_v7<<<BH, wth>>>(fp(q), fp(k), fp(g), fp(beta), fp(A_log),
      fp(dt_bias), fp(qn), fp(kn), fp(qs), fp(rq), fp(rk), fp(beta_s),
      fp(ga), fp(gc), (float)scale, B, T, H, K, BT, NT);
  k_wy_v7<<<BH * NT, v7::GEMM_THREADS, wy_sh>>>(fp(v), fp(kn), fp(gc),
      fp(beta_s), fp(X), fp(Afull), fp(w), fp(u), B, T, H, K, V, BT, NT);
  k_aqk_v7<<<BH, wth>>>(fp(qs), fp(kn), fp(gc), fp(Aqk), B, T, H, K, BT, NT);
  {
    int VC = 32, VS = V / VC;
    int cbuf = 2 * BT * VC; if (VC * K > cbuf) cbuf = VC * K;
    int rsh = (VC * K + cbuf) * sizeof(float)
            + (VC * K + 2 * BT * K + VC * BT + BT * BT + K * BT) * sizeof(bf16);
    cudaFuncSetAttribute(k_recur_seq_v7<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, rsh);
    k_recur_seq_v7<true><<<BH * VS, v7::GEMM_THREADS, rsh>>>(fp(qs), fp(kn), fp(gc),
        fp(w), fp(u), fp(initial_state), fp(Aqk), fp(vcorr), fp(Sentry),
        fp(o), fp(fs), B, T, H, K, V, BT, NT);
  }

  auto dq = torch::zeros_like(q), dk = torch::zeros_like(k);
  auto dv = torch::zeros_like(v), dgr = torch::zeros_like(g);
  auto dbeta = torch::zeros_like(beta);
  auto dA_log = torch::zeros_like(A_log), ddt_bias = torch::zeros_like(dt_bias);
  auto dh0 = torch::zeros_like(initial_state);

  BwdScratch sc;
  auto s_dgc = torch::zeros({BH, T, K}, opt); sc.dgc = fp(s_dgc);
  auto s_dqs = torch::zeros({BH, T, K}, opt); sc.dqs = fp(s_dqs);
  auto s_dkn = torch::zeros({BH, T, K}, opt); sc.dkn = fp(s_dkn);
  auto s_dbeta_s = torch::zeros({BH, T}, opt); sc.dbeta_s = fp(s_dbeta_s);
  auto s_dS = torch::zeros({BH, K, V}, opt); sc.dS = fp(s_dS);
  auto s_dS_entry = torch::zeros({BH, K, V}, opt); sc.dS_entry = fp(s_dS_entry);
  auto s_dEL = torch::zeros({BH, K}, opt); sc.dEL = fp(s_dEL);
  auto s_dvcorr = torch::zeros({BH, BT, V}, opt); sc.dvcorr = fp(s_dvcorr);
  auto s_du = torch::zeros({BH, BT, V}, opt); sc.du = fp(s_du);
  auto s_dAqk = torch::zeros({BH, BT, BT}, opt); sc.dAqk = fp(s_dAqk);
  auto s_dAfull = torch::zeros({BH, BT, BT}, opt); sc.dAfull = fp(s_dAfull);
  auto s_dX = torch::zeros({BH, BT, BT}, opt); sc.dX = fp(s_dX);
  auto s_Xf = torch::zeros({BH, BT, BT}, opt); sc.Xf = fp(s_Xf);
  auto s_tmp = torch::zeros({BH, BT, BT}, opt); sc.tmp = fp(s_tmp);
  auto s_dAraw = torch::zeros({BH, BT, BT}, opt); sc.dAraw = fp(s_dAraw);
  auto s_dQE = torch::zeros({BH, BT, K}, opt); sc.dQE = fp(s_dQE);
  auto s_dP = torch::zeros({BH, BT, K}, opt); sc.dP = fp(s_dP);
  auto s_dw = torch::zeros({BH, BT, K}, opt); sc.dw = fp(s_dw);
  auto s_dkg = torch::zeros({BH, BT, K}, opt); sc.dkg = fp(s_dkg);

  k_bwd_recur_v7<<<BH, 32>>>(fp(qs), fp(kn), fp(gc), fp(v), fp(w),
      fp(Afull), fp(X), fp(Aqk), fp(vcorr), fp(Sentry), fp(beta_s),
      fp(do_), fp(dht), sc, fp(dv), fp(dh0), B, T, H, K, V, BT);
  k_bwd_pre_v7<<<BH, NT>>>(fp(q), fp(k), fp(qn), fp(kn), fp(rq), fp(rk),
      fp(beta_s), fp(ga), fp(A_log), sc.dgc, sc.dqs, sc.dkn, sc.dbeta_s,
      fp(dq), fp(dk), fp(dgr), fp(dbeta), fp(dA_log), fp(ddt_bias),
      (float)scale, B, T, H, K, BT, NT);
  cudaDeviceSynchronize();
  return {dq, dk, dv, dgr, dbeta, dA_log, ddt_bias, dh0};
}

// Report theoretical occupancy of the recur kernels at the actual K=V=128
// launch config (no HW counters needed; uses the CUDA occupancy API).
static std::vector<int64_t> v8_occupancy(int64_t K, int64_t V) {
  int VC = VCOLS8;
  int cbuf = 2 * BT16 * VC; if (VC * (int)K > cbuf) cbuf = VC * K;
  int opsz = BT16 * K + BT16 * K + BT16 * BT16 + (int)K * BT16;
  int dev = 0; cudaGetDevice(&dev);
  cudaDeviceProp prop; cudaGetDeviceProperties(&prop, dev);
  auto report = [&](const void* fn, int threads, int shmem) -> std::vector<int64_t> {
    cudaFuncSetAttribute(fn, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem);
    int blocks = 0;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&blocks, fn, threads, shmem);
    cudaFuncAttributes at; cudaFuncGetAttributes(&at, fn);
    int warps_per_block = (threads + 31) / 32;
    int active_warps = blocks * warps_per_block;
    int max_warps = prop.maxThreadsPerMultiProcessor / 32;
    return {blocks, active_warps, max_warps, at.numRegs, shmem,
            (int64_t)prop.sharedMemPerMultiprocessor};
  };
  int rsh_b = (VC * (int)K + cbuf) * sizeof(float)
            + (VC * (int)K + 1 * opsz + VC * BT16) * sizeof(bf16);
  int rsh_w = (VC * (int)K + cbuf) * sizeof(float)
            + (VC * (int)K + 2 * (opsz + BT16 * VC) + VC * BT16) * sizeof(bf16);
  auto bl = report((const void*)k_recur_v8<false>, v8::GEMM16_THREADS, rsh_b);
  auto ws = report((const void*)k_recur_v8_ws, WS_COMPUTE + 32, rsh_w);
  // pack: [bl: blocks,warps,maxwarps,regs,shmem,smShmem][ws: ...]
  std::vector<int64_t> out;
  out.insert(out.end(), bl.begin(), bl.end());
  out.insert(out.end(), ws.begin(), ws.end());
  return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fwd", &kda_cuda_fwd,
        "KDA chunk forward (CUDA v8: CHUNK=16, bounded gate, all-TC)",
        pybind11::arg("q"), pybind11::arg("k"), pybind11::arg("v"),
        pybind11::arg("g"), pybind11::arg("beta"), pybind11::arg("A_log"),
        pybind11::arg("dt_bias"), pybind11::arg("scale"),
        pybind11::arg("initial_state"), pybind11::arg("chunk_size"),
        pybind11::arg("lower_bound") = -5.0, pybind11::arg("pipeline") = false,
        pybind11::arg("warp_spec") = false);
  m.def("bwd", &kda_cuda_bwd, "KDA chunk backward (CUDA v8, legacy v7 path)");
  m.def("occupancy", &v8_occupancy, "theoretical occupancy of recur kernels",
        pybind11::arg("K") = 128, pybind11::arg("V") = 128);
#ifdef V8_PROFILE
  m.def("prof_reset", []() {
    unsigned long long z[16] = {0};
    cudaMemcpyToSymbol(g_prof_k1, z, sizeof(z));
    cudaMemcpyToSymbol(g_prof_k2, z, sizeof(z));
  }, "zero the in-kernel clock64 accumulators");
  m.def("prof_read", []() {
    unsigned long long k1[16], k2[16];
    cudaMemcpyFromSymbol(k1, g_prof_k1, sizeof(k1));
    cudaMemcpyFromSymbol(k2, g_prof_k2, sizeof(k2));
    std::vector<int64_t> out;
    for (int i = 0; i < 16; ++i) out.push_back((int64_t)k1[i]);
    for (int i = 0; i < 16; ++i) out.push_back((int64_t)k2[i]);
    return out;
  }, "read [k1[0..15], k2[0..15]] cycle accumulators");
#endif
}

7 优化 5

优化5 相比 优化4 的唯一结构性改动:把每个 chunk 的 vcorr 修正项在 K2 的 4 个 gemm 之间保持寄存器常驻(FlashKDA 的 MOVM_T 技术),不再经 shared 往返。代码在 cute_recur_v10.cuh 的 recur_warp<K,VC>,由 kda_chunk_cuda_v10.cu 的 k_recur_v8_ws 调用(warp 0/1 各负责一个 16 列的 vcol-block)。

总览数据见 OPT_LOG_v8_v9.md(v10+ws:B=1 1.094ms / B=8 7.130ms = 24.5x / 10.2x over v2,1.08x / 1.10x over v9;正确性 2.48e-3 vs fla)。

---

7.1 核心改动概述

每 chunk 的 K2 递推有 4 个 gemm:

gemm1  wS   = w  @ S^T            (M=time, N=vcol)
       vcorr = u - wS
gemm2  qS   = qd @ S^T
gemm3  AV   = Aqk @ vcorr  ->  o = qS + AV          (o 路径)
gemm4  newS^T = krT @ vcorr                          (状态更新路径)
       S^T  = S^T * exp(glast) + newS^T

• v9:vcorr 写进共享内存 vcorrT_bf,gemm3 / gemm4 再各自从 shared 读回; 4 个 gemm 之间走 C1/vcorr/C2/C3 共享缓冲,每步一个 bar_sync(≈6/chunk)。
• v10:
  1. gemm1 在 C-累加器寄存器里直接产出 vcorr = u − wS,不落 shared;
  2. 用一次 SM75_U32x1_MOVM_T(warp 级 movmatrix.sync.m8n8.trans)把 vcorr 的 C-累加器转成 B-operand 片段 rVCb;
  3. rVCb 被 gemm3(AV)和 gemm4(newS^T)复用,不再读 shared。
  • 去掉 C1/vcorr/C2/C3 往返与对应 barrier → 每 chunk bar_sync(1) ≈6→3。
  • S 本身仍留在 shared(St_bf,bf16,作 gemm1/2 的 B 操作数);真正寄存器常驻 的只有 vcorr/U(与 FlashKDA 一致,“vcorr/U/S” 里的 S 并未进寄存器)。

---

7.2 代码逐段拆解(recur_warp)

7.2.1 MMA 布局

using MMA1 = decltype(make_tiled_mma(
    MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>{},
    Layout<Shape<_1, _1>>{}, Tile<_16, _16, _16>{}));

单 warp(32 线程)一个 16×16×16 MMA。CuTe 的 fragment 即寄存器: partition_fragment_C → 每线程 4 个 fp32(拼成 16×16);partition_fragment_B → 每线程的 bf16 B 操作数。thr.partition_C(smem) 给出"本线程负责 C 的哪些元素" 的映射,故 reg↔smem 用 copy() 即可对齐。

7.2.2 gemm1:vcorr 在 C-累加器里算出,不落 shared

auto rVC = thr.partition_fragment_C(cref); clear(rVC);
accum_AxBt<K>(w_bf, K, Sw, K, lane, rVC);   // rVC = w @ Sw^T = wS  (fp32 寄存器)
{
  Tensor su = make_tensor(make_smem_ptr(u_bf + vb * 16),         // u 本 block [time,vcol]
      make_layout(make_shape(_16{}, _16{}), make_stride(Int<VC>{}, _1{})));
  auto rU = make_fragment_like<bf16>(rVC);    // 同 C 布局的 bf16 寄存器
  copy(thr.partition_C(su), rU);              // 按 C 布局把 u 读进寄存器
#pragma unroll
  for (int i = 0; i < size(rVC); ++i)
    rVC(i) = __bfloat162float(rU(i)) - rVC(i); // vcorr = u - wS,仍在 rVC 寄存器
}

• accum_AxBt 做 rVC += w[16,K] @ Sw[16,K]^T(TN gemm,结果在 C 累加器)。
• u 用与 C 相同的线程布局读进 rU,逐元素 vcorr = u − wS。
• su stride (Int<VC>, _1):u 在 global 是 [time, 全部 VC 列] 行主序,本 block 取其中 16 列,故行跨度 = VC(=32),列跨度 = 1。
• bf16 处理:用 make_fragment_like<bf16> + __bfloat162float(torch 定义了 __CUDA_NO_BFLOAT16_CONVERSIONS__,不能 float(bf16))。

7.2.3 一次 MOVM_T:C-累加器 → B-operand 片段

auto rVCbf = make_fragment_like<bf16>(rVC);
#pragma unroll
for (int i = 0; i < size(rVC); ++i) rVCbf(i) = __float2bfloat16(rVC(i)); // fp32->bf16

Tensor bref = make_tensor(make_smem_ptr((bf16*)nullptr),
    make_layout(make_shape(_16{}, _16{}), make_stride(_16{}, _1{})));
auto rVCb = thr.partition_fragment_B(bref);   // 空的 B-operand 寄存器片段
{
  uint32_t* src = reinterpret_cast<uint32_t*>(&rVCbf(0));
  uint32_t* dst = reinterpret_cast<uint32_t*>(&rVCb(0));
#pragma unroll
  for (int i = 0; i < 4; ++i) SM75_U32x1_MOVM_T::copy(src[i], dst[i]);  // 4 个 8x8 象限
}

SM75_U32x1_MOVM_T = PTX movmatrix.sync.aligned.m8n8.trans.b16,warp 级寄存器内 转置:32 线程协作把一个 8×8 bf16 块原地转置(lane 间硬件交换,不经 shared)。 16×16 = 4 个 8×8 象限,每象限每线程 1 个 uint32(2 个 bf16),故循环 4 次;CuTe 的 C-累加器与 B-operand 片段寄存器迭代顺序对齐,src[i] → dst[i] 顺序拷贝即完成整 16×16 转换。

已验证语义(check_recur_v10.py):gemm(A, MOVM_T(X_in_C)) == A @ X(不是 A @ X^T)。

7.2.4 复用 ①:gemm2 + gemm3(o 路径)

auto rO = thr.partition_fragment_C(cref); clear(rO);
accum_AxBt<K>(qd_bf, K, Sw, K, lane, rO);     // gemm2: rO  = qd @ Sw^T = qS
accum_AxBfrag(Aqk_bf, 16, lane, rVCb, rO);    // gemm3: rO += Aqk @ vcorr  (复用 rVCb)
{
  Tensor go = make_tensor(make_gmem_ptr(o_base),
      make_layout(make_shape(_16{}, _16{}), make_stride(o_rs, _1{})));
  copy(rO, thr.partition_C(go));              // o = qS + AV 直接写 global
}

accum_AxBfrag 只把 A(Aqk)copy 进寄存器,B 直接用传入的 rVCb,不读 shared:

template <class CFrag, class BFrag>
__device__ inline void accum_AxBfrag(const bf16* Ap, int lda, int lane,
                                     const BFrag& rB, CFrag& rC) {
  ...
  copy(tCsA(_, _, 0), rA(_, _, 0));
  gemm(mma, rC, rA(_, _, 0), rB(_, _, 0), rC);  // rB 即 rVCb
}

o 算完直接从寄存器写 global(stride (o_rs,1),o_rs = H*V)。

7.2.5复用 ②:gemm4(状态更新路径)

constexpr int NDB = K / 16;
#pragma unroll
for (int db = 0; db < NDB; ++db) {
  auto rNS = thr.partition_fragment_C(cref); clear(rNS);
  accum_AxBfrag(krT_bf + db*16*16, 16, lane, rVCb, rNS);  // newS^T = krT @ vcorr(复用 rVCb)
  // [d-row, vcol] 转置写进 newS_sh[(vb*16+vcol)*K + (db*16+d)]
  Tensor sns = make_tensor(make_smem_ptr(newS_sh + (vb*16)*K + db*16),
      make_layout(make_shape(_16{}, _16{}), make_stride(_1{}, Int<K>{})));
  copy(rNS, thr.partition_C(sns));
}

• 同一份 rVCb 第二次复用:newS^T[d,vcol] = krT[d,time] @ vcorr[time,vcol],按 K/16 个 d-block 循环(K=128 → 8 块)。
• 转置落 shared 的技巧:rNS 逻辑布局 [d-row, vcol-col],写出用 stride (_1, Int<K>) —— 行跨度 1、列跨度 K,等于把 [d,vcol] 转置成 [vcol,d] 存进 newS_sh,使后面未改动的向量化 S-update 读 St_f[vcol,d] 时布局对齐。

7.2.6解开 优化5 死结的等价变形

v9 卡在:vcorr 在 gemm3 是 B-operand,但原 newS = vcorr^T@kr^T 里它是 A-operand,fragment 布局不兼容,无法共用一份寄存器。v10 把 gemm4 改写成 newS^T = krT@vcorr(vcorr 永远当 B),代价仅是结果转置落 shared(上面 stride 技巧), 换来 vcorr 只做一次 MOVM_T、两个 gemm 共用同一份 B 寄存器,彻底去掉 4 个 gemm 间的 shared 往返与 barrier。

---

7.3为什么"转置一下,C 就变成 B 了"

一句话

在这颗 MMA 里,C-累加器和 B-operand 这两种 fragment,对同一块 16×16, 32 个 lane "谁存哪个元素"的映射本身就互为转置。所以一块矩阵从 C 角色换成 B 角色, 数据值不用变,只需翻转一次"lane↔元素"映射,而 movmatrix.m8n8.trans 正是干 这件事的硬件指令。

7.3.1 fragment 是"线程→元素"的约定,不是矩阵

m16n8k16 下,记 g = lane>>2(0…7),t = lane&3(0…3):

C-累加器(fp32, M×N) 每线程 4 个值:

c0,c1 → (row=g,   col=2t, 2t+1)
c2,c3 → (row=g+8, col=2t, 2t+1)

即 g→行,t→列。

B-operand(bf16, K×N) 每线程的值:

b0,b1 → (row=2t,   col=g)     (row 为 K 维)
b2,b3 → (row=2t+8, col=g)

即 g→列,t→行。

同一个 (g,t) 线程:作 C 用持有 (行=g, 列=2t),作 B 用持有 (行=2t, 列=g)—— 行列角色对调,正是一个转置。

7.3.2 直接 reinterpret 是错的

不搬运直接把 C 寄存器当 B:某线程物理上存 V = vcorr[g, 2t],但 B 布局会解释成 B[2t, g],于是 MMA 误以为 B[2t,g] = vcorr[g,2t](当成转置),且值还在错误线程 上。必须在 warp 内跨 lane 做 8×8 转置才能修正(16×16 = 4 象限,每象限每线程 1 uint32)。

7.3。3 movmatrix 就是那次跨 lane 的 8×8 转置

movmatrix.sync.aligned.m8n8.trans.b16:整 warp 把 8×8 bf16 原地转置(数据在 lane 寄存器间硬件交换)。做完后,原 C-位置 (i,j) 的值落到 B 布局读成 (i,j) 的寄存器 槽,刚好补偿第 1 节 C 与 B 的角色对调。

7.3.4 为什么"转置"反而得到 vcorr 本身(未被转置)

关键:C 布局与 B 布局彼此就差一个转置。对 C 里数据再转一次(movmatrix),两个 转置抵消 —— 净效果是 同一个 vcorr,换了一身 B 角色的寄存器排布。矩阵内容不变, 变的只是排布。故 gemm(A, MOVM_T(X_in_C)) == A @ X。

类比

C-fragment 与 B-fragment 是同一份数据的两种"装订方式"(一种按行、一种按列)。 内容一样,但机器只认特定角色的装订。movmatrix 就是"按行拆开、按列重订"的机器; 而因这两种装订本就互为转置,重订一次正好让数据以正确姿势进入下一个 gemm。

---

8 测试结果

测试环境: H20-3e, FP32, K=V=64, B=2, H=4

Forward 耗时 (ms):

Shape	Triton	v1 (1thr/bh)	v1 (chunk-par)	v2 (warp/blk)
T=128 (NT=2)	0.97	192.6	61.7	2.77
T=512 (NT=8)	0.97	650	197	6.58
T=2048 (NT=32)	0.95	2491	739	21.8

性能提升:

• 比 v1 (1thr/bh) 快 70–115×
• 比 v1 (chunk-par) 快 22–34×
• 相比 Triton: T=128 时差距仅 2.8×

结论

1. 两点判断正确且有效:

   • Warp 切分特征维 → 合并访存 + shuffle 归约，解决了 k_pre/k_wy 的并行度瓶颈
   • Block 内按 K/V 维并行、chunk 串行 → 解决了 recurrence 中 thread-0 串行的大瓶颈

2. 当前瓶颈: Forward 仍受跨 chunk 串行递推限制 (需串行处理 NT 个 chunk)，因此 T 越大，与 Triton 的相对差距略有增大。

优化5 性能验证 (H20, T=1024, H=64, K=V=128)

单次 Forward 延迟

Config	B=1 (ms)	B=8 (ms)	vs v2 (B=1 / B=8)	vs fla (B=1 / B=8)
v2 (baseline)	26.76	73.07	1.00x / 1.00x	0.04x / 0.07x
fla (chunk cs=64)	1.009	5.058	26.5x / 14.4x	1.00x / 1.00x
v8+ws	1.651	9.488	16.2x / 7.7x	0.61x / 0.53x
v9+ws	1.280	8.080	20.9x / 9.0x	0.79x / 0.63x
v10+ws	1.094	7.130	24.5x / 10.2x	0.92x / 0.71x

逐级提升

• v9+ws vs v8+ws: 1.29× (B=1) / 1.17× (B=8)
• v10+ws vs v9+ws: 1.08× (B=1) / 1.10× (B=8)

关键结论

1. v10+ws 逼近 fla

   • B=1: 达 fla 的 0.92×
   • B=8: 达 fla 的 0.71×
   • vs v2 (baseline): 24.5× (B=1) / 10.2× (B=8) 加速

2. vcorr 寄存器常驻贡献

   • 移除 C1/vcorr/C2/C3 shared 往返 + barriers
   • 每 chunk 6 次 bar_sync → 3 次
   • 额外 8–10% 性能提升

3. 剩余差距 (vs fla)

   • fla 采用全融合 2-kernel + TMA + 更短串行链
   • v10 仍保留独立 K1/K2 以及 4 次 gemm 的 barrier 开销

正确性

• v10+ws vs fla: 2.48e-3 (K=128) / 2.63e-3 (K=64)
• 远低于 2e-2 容差 ✅

硬件配置

• GPU: H20
• 参数: T=1024, H=64, K=V=128
• fla 配置: chunk_kda(safe_gate=True, lower_bound=-5, chunk_size=64)
• v2: softplus gate + BT∈{32,64}, 全 scalar fp32