TileLang 实战：KDA 从零到一–Kimi Delta Attention

前三篇分别实现了无遗忘的 chunked 线性注意力、标量衰减、以及带删除项的 Gated DeltaNet。本篇是系列收尾：把 GDN 的标量门 $\alpha_t$ 换成逐通道向量门 $\bm{a}_t \in \mathbb{R}^{d_k}$。

$$\mathbf{S}_t = \mathbf{S}_{t-1}\operatorname{Diag}(\bm{a}_t)\big(\mathbf{I} - \beta_t\bm{k}_t\bm{k}_t^{\intercal}\big) + \beta_t\bm{v}_t\bm{k}_t^{\intercal}$$

改动只有一处：$\alpha_t$ 变成了 $\operatorname{Diag}(\bm{a}_t)$。但这一处把前三篇积累的所有便利拆掉了–衰减不再能从矩阵里外提，$\Gamma$ 从 $C\times C$ 变成 $C \times C \times d_k$，而累积衰减积的通道离散度会直接把 fp16 打穿。

本篇聚焦实现。KDA 的数学推导（递推式、逐通道 WY 表示、UT 变换、下界衰减与满秩门控的动机）见《KDA 的来龙去脉》§3–§4，这里不重复；本文只做一件事：把那些公式落成能跑的 kernel，并量化每一处数值边界。

前三篇见 Chunked 线性注意力、标量衰减、Gated DeltaNet。

0. 符号与三处变化

沿用前三篇：$\beta_t$ 写入强度、$C$ 块长、$[t]$ 块序号、$r \in [1,C]$ 块内位置、$\mathbf{S}^\intercal \in \mathbb{R}^{d_k \times d_v}$ 为 kernel 存储布局。

本篇的核心量改为累积 log 衰减（沿用《来龙去脉》的记号）：

$$\bm{\gamma}_i = \sum_{s=1}^{i}\bm{g}_s \in \mathbb{R}^{d_k}, \qquad \bm{g}_s = \log\bm{a}_s < \bm{0}\ \text{（逐分量）}$$

注意 $\bm{\gamma}_i$ 是向量–这是与前两篇最本质的区别。前三篇的 $\gamma^r$ 是标量，$\Gamma_{ij} = \gamma_i/\gamma_j$ 是一张 $C \times C$ 的表；本篇 $\bm{\gamma}_i - \bm{\gamma}_j$ 是 $d_k$ 维向量，"衰减矩阵"概念上是 $C \times C \times d_k$，不可能物化。

三处结构性变化：

第三行是全篇的技术核心。

1. 衰减长在内部：问题与出路

1.1 朴素做法的代价

$M_{ci} = (\bm{k}_c \odot e^{\bm{\gamma}_c - \bm{\gamma}_i})^\intercal\bm{k}_i$ 里的指数依赖 $(c, i, d)$ 三个下标。直接算就是 $C^2/2$ 次长度 $d_k$ 的加权内积，每次都要现算 $d_k$ 个 exp。

这不是常数因子问题–它把一次 GEMM 变成了 $C^2/2$ 个独立的向量运算，完全用不上 Tensor Core。$C = 64$、$d_k = 128$ 时是 2016 次加权内积、约 26 万次 exp2。

1.2 出路：指数可分离

关键观察是指数可以按通道拆开：

$$e^{\bm{\gamma}_c - \bm{\gamma}_i} = e^{\bm{\gamma}_c} \odot e^{-\bm{\gamma}_i}$$

于是：

$$M_{ci} = \sum_{d}\underbrace{\big(k_c[d]\,e^{\gamma_c[d]}\big)}_{\widetilde{K}^{+}[c,d]}\underbrace{\big(k_i[d]\,e^{-\gamma_i[d]}\big)}_{\widetilde{K}^{-}[i,d]} = \big(\widetilde{\mathbf{K}}^{+}\widetilde{\mathbf{K}}^{-\intercal}\big)_{ci}$$

一次 GEMM 解决，前置两个 $C \times d_k$ 的逐元素加权。实测拆分与朴素计算的差异 $1.94 \times 10^{-16}$，代数上完全等价。

1.3 代价：$e^{-\bm{\gamma}_i}$ 是大于 1 的量

第二篇讲过一个教训：把比值 $\gamma_i/\gamma_j$ 因式分解成 $\gamma_i \cdot (1/\gamma_j)$ 会物化一个指数增长的量，fp16 下 $\alpha \le 0.8$ 即 NaN。这里是同一个陷阱的逐通道版本–$e^{-\gamma_i[d]} \ge 1$，且随 $i$ 与通道衰减强度指数增长。

但本篇的处境和第二篇不同：那里有替代方案（保留比值形式），这里没有。逐通道衰减无法外提，不拆就用不上 Tensor Core。所以问题从"要不要拆"变成了"怎样让拆分在数值上安全"。

实测 $e^{-\bm{\gamma}}$ 的最大值（$C = 64$、$d_k = 128$）：

门控下界在 0.9 以上时，$e^{-\bm{\gamma}}$ 连 fp16 都装得下。 这就把 KDA 论文里"下界衰减（lower-bounded decay）"这个设计从模型层面的技巧，变成了 kernel 能否用 Tensor Core 的前提条件。

2. 下界衰减：不是精度调优，是可行性前提

2.1 无下界时 fp16 全线归零

$e^{-\bm{\gamma}}$ 会溢出，另一头 $e^{\bm{\gamma}}$ 会下溢。而逐通道门让后者严重得多–总有一些通道学到很小的 $a$，它们的 $\gamma$ 累积得最快。

实测 $\min e^{\bm{\gamma}_C}$ 与 fp16 下归零的通道数（$d_k = 128$）：

$C = 128$、下界 0.5 时全部 128 个通道在 fp16 下归零–状态被彻底清空，kernel 输出恒为块内项，跨块信息完全丢失。

加下界后：

下界 0.9 是分界线。 这一个数字同时解决了两头：$e^{\bm{\gamma}}$ 不下溢、$e^{-\bm{\gamma}}$ 不溢出。

2.2 通道离散度：逐通道门特有的问题

标量门下 $\gamma^C$ 是一个数；逐通道门下它是 $d_k$ 个数，而它们的极差决定了同一个 chunk 内不同通道的数值尺度差异：

极差 $47$ 意味着同一个 fragment 里最强和最弱通道的数值相差 20 个数量级–任何浮点格式都无法同时表示。这是逐通道门独有的病：标量门只需要担心整体的尺度漂移，逐通道门还要担心通道之间的尺度撕裂。

下界 0.9 把极差压到 1.74（$e^{1.74} = 5.7$），完全可控。

2.3 sub-chunk：第二道保险

即使有下界，$C = 128$ 时 $\max e^{-\bm{\gamma}} = 1.79\times10^3$（下界 0.9）–fp16 装得下但余量不多。把块内再切成 sub-chunk，每个 sub-chunk 内部重新起算 $\bm{\gamma}$：

指数跨度只取决于 sub-chunk 长度，与总块长无关。这与第二篇"累积积按 chunk 重置"是同一个道理，只是又下降了一层：chunk 重置控制 $\bm{\gamma}$ 本身，sub-chunk 重置控制 $e^{-\bm{\gamma}}$ 的动态范围。

代价是 sub-chunk 之间需要额外的状态传递，块内变成两层循环。$C = 64$ + 下界 0.9 时 $\max e^{-\bm{\gamma}} = 55$，不需要 sub-chunk；$C = 128$ 时建议切 32 或 64。

3. 四层参考实现

沿用前三篇框架。

3.1 参考 A：逐 token 递归

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

def ref_A(Q, K, V, g, beta):
    """g: (B,N,H,DK) 逐通道 log 门控，g < 0"""
    B, N, H, DK = Q.shape
    DV = V.shape[-1]
    O = np.zeros((B, N, H, DV))
    for b in range(B):
        for h in range(H):
            S = np.zeros((DV, DK))                      # 论文方向
            for t in range(N):
                k, v, q = K[b, t, h], V[b, t, h], Q[b, t, h]
                a, bt = np.exp(g[b, t, h]), beta[b, t, h]
                # Diag(a) 在 S 与 Householder 之间
                S = S @ np.diag(a) @ (np.eye(DK) - bt * np.outer(k, k)) \
                    + bt * np.outer(v, k)
                O[b, t, h] = S @ q
    return O

$\operatorname{Diag}(\bm{a}_t)$ 的位置很关键：它夹在 $\mathbf{S}_{t-1}$ 与 Householder 之间。写成 $\operatorname{Diag}(\bm{a})\mathbf{S}$ 或挪到括号外都会算错。

3.2 参考 B：chunkwise + 逐通道 UT 变换

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

def ref_B(Q, K, V, g, beta, C):
    B, N, H, DK = Q.shape
    DV = V.shape[-1]
    O = np.zeros((B, N, H, DV))
    for b in range(B):
        for h in range(H):
            S = np.zeros((DK, DV))                      # 转置布局
            for c in range(N // C):
                sl = slice(c * C, (c + 1) * C)
                Qc, Kc, Vc = Q[b, sl, h], K[b, sl, h], V[b, sl, h]
                gc, bt = g[b, sl, h], beta[b, sl, h]

                gam = np.cumsum(gc, axis=0)             # γ_i ∈ R^{DK}
                gC  = gam[-1]

                # M[c,i] = (k_c ⊙ e^{γ_c-γ_i})·k_i,  i < c
                M = np.zeros((C, C))
                for cc in range(C):
                    for i in range(cc):
                        M[cc, i] = np.dot(Kc[cc] * np.exp(gam[cc] - gam[i]), Kc[i])

                L  = np.tril(bt[None, :] * M, -1)       # L[c,i] = β_i M[c,i]
                Tm = np.linalg.inv(np.eye(C) + L)
                Ah = bt[:, None] * Tm                   # diag(β) T

                W  = Ah @ (Kc * np.exp(gam))            # C × DK
                U  = Ah @ Vc                            # C × DV
                Vt = U - W @ S                          # 伪值 Ṽ

                # A^qk[c,j] = (q_c ⊙ e^{γ_c-γ_j})·k_j,  j ≤ c
                Aqk = np.zeros((C, C))
                for cc in range(C):
                    for j in range(cc + 1):
                        Aqk[cc, j] = np.dot(Qc[cc] * np.exp(gam[cc] - gam[j]), Kc[j])

                O[b, sl, h] = (Qc * np.exp(gam)) @ S + Aqk @ Vt
                S = np.diag(np.exp(gC)) @ S + (Kc * np.exp(gC - gam)).T @ Vt
    return O

这份参考刻意用朴素双循环算 $\mathbf{M}$ 与 $\mathbf{A}^{qk}$–慢，但与公式逐字对应，作为基准可信。§3.3 的 kernel 镜像才用 GEMM 化写法。

$L_{ci} = \beta_i M_{ci}$ 与 $\hat{\mathbf{A}} = \operatorname{diag}(\beta)\mathbf{T}$ 这一对下标必须配套。实测另一种等价写法是 $L_{ci} = \beta_c M_{ci}$ 配 $\hat{\mathbf{A}} = \mathbf{T}\operatorname{diag}(\beta)$，两者都对，混用则错。

3.3 参考 C：kernel 控制流镜像（GEMM 化）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

def ref_C(Q, K, V, g, beta, C, block_DV):
    B, N, H, DK = Q.shape
    DV = V.shape[-1]
    O = np.zeros((B, N, H, DV))
    for bv in range(DV // block_DV):                    # ← grid.x
        dv = slice(bv * block_DV, (bv + 1) * block_DV)
        for bbh in range(B * H):                        # ← grid.y
            b, h = bbh // H, bbh % H
            S = np.zeros((DK, block_DV))
            for c in range(N // C):                     # ← T.Pipelined
                sl = slice(c * C, (c + 1) * C)
                Qc, Kc = Q[b, sl, h], K[b, sl, h]
                Vc = V[b, sl, h, dv]
                gc, bt = g[b, sl, h], beta[b, sl, h]

                gam = np.cumsum(gc, axis=0)
                gC  = gam[-1]
                Ep, Em = np.exp(gam), np.exp(-gam)      # e^{γ}, e^{-γ}

                # ── 指数可分离：两次 GEMM 代替 C²/2 次加权内积 ──
                Kp, Km = Kc * Ep, Kc * Em
                M   = np.tril(Kp @ Km.T, -1)            # 严格下三角
                Aqk = np.tril((Qc * Ep) @ Km.T)         # 含对角

                L  = np.tril(bt[None, :] * M, -1)
                # 前向替换解 T @ [Kp | Vc]，再逐行乘 β（因为 Â = diag(β)T）
                RHS = np.concatenate([Kc * Ep, Vc], axis=1)
                X = np.zeros_like(RHS)
                for r in range(C):
                    X[r] = RHS[r]
                    for j in range(r):
                        X[r] -= L[r, j] * X[j]
                X = bt[:, None] * X                     # ← β 在解完之后乘
                W, U = X[:, :DK], X[:, DK:]

                Vt = U - W @ S
                O[b, sl, h, dv] = (Qc * Ep) @ S + Aqk @ Vt
                S = Ep[-1][:, None] * S + (Kc * np.exp(gC - gam)).T @ Vt
    return O

三处与参考 B 不同的实现选择：

1. $\mathbf{M}$ 与 $\mathbf{A}^{qk}$ 走 GEMM：$\widetilde{\mathbf{K}}^{+}\widetilde{\mathbf{K}}^{-\intercal}$。$\mathbf{A}^{qk}$ 含对角（$j \le c$），$\mathbf{M}$ 不含（$i < c$）。
2. $\mathbf{W}$ 与 $\mathbf{U}$ 拼成一次前向替换：两者共用同一个 $(\mathbf{I}+\mathbf{L})$，拼接后只解一遍，省一半串行开销。
3. $\beta$ 必须在解完三角系统之后乘。$\hat{\mathbf{A}} = \operatorname{diag}(\beta)\mathbf{T}$ 展开是「先 $\mathbf{T}$ 作用、再逐行乘 $\beta$」；若把 $\beta$ 提前乘进右端项，算的就是 $\mathbf{T}\operatorname{diag}(\beta)$–那是另一种配对（需搭配 $L_{ci} = \beta_c M_{ci}$），混用则错。这个错误在 $\beta$ 全部相等时完全看不出来，我第一次写就踩了：$\beta$ 随机时相对 L2 达 $1.3\times10^{-1}$。
4. 状态更新用逐行乘代替 $\operatorname{Diag}$：Ep[-1][:, None] * S 就是 $\operatorname{Diag}(e^{\bm{\gamma}_C})\mathbf{S}$，不物化对角矩阵。

3.4 一致性验证

$B=2, H=2, N=12, d_k=d_v=4, C=4$，$\bm{a} \sim \mathcal{U}(0.90, 0.999)^{d_k}$、$\beta \sim \mathcal{U}(0.1, 0.9)$，$\bm{q},\bm{k}$ 已 L2 归一化，fp64：

三个退化检验（本篇比 GDN 多一个）：

三个 $\text{block}_{DV}$（1 / 2 / 4）的相对 L2 分别为 $2.48$ / $2.57$ / $2.57 \times 10^{-16}$–切 DV 零依赖，与第一篇的结论一致。

第一个是本篇独有且最重要的：把逐通道门退化成标量门，必须精确回到第三篇的结果。这一步能抓住"指数分离时把通道维和位置维搞混"这类错误–而那类错误在通道值本来就相同时会隐身。

4. TileLang kernel

4.0 寄存器账

$d_k = d_v = 128$、$C = 64$、$\text{block}_{DV} = 32$、128 线程：

320 超限了。 逐通道门带来的 $[C, d_k]$ 量级 fragment（$e^{\pm\bm{\gamma}}$、$\mathbf{W}$）比 $C^2$ 表更吃寄存器–$C \times d_k = 64\times128$ 是 $C^2 = 64^2$ 的两倍。

三条出路：

1. $e^{-\bm{\gamma}}$ 不常驻：它只在构造 $\mathbf{M}$、$\mathbf{A}^{qk}$ 时用，用完即弃，可以放 shared memory。省 64 reg。
2. $\mathbf{W}$ 走 shared：它是 $\mathbf{U} - \mathbf{W}\mathbf{S}$ 的中间量，不参与后续 GEMM 的累加器。省 64 reg。
3. $C$ 降到 32：所有 $C$ 相关量减半。

组合 1+2 后约 192 reg，可行。下面的 kernel 采用这个方案。

4.1 声明与累积衰减

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

@tilelang.jit(out_idx=[5])
def kda_chunk(B, H, S, DK, DV, blk=64, block_DV=32,
              num_stages=2, threads=128,
              dtype=T.float16, accum_dtype=T.float32):
    C  = blk
    NS = T.ceildiv(S, C)

    @T.prim_func
    def main(
        Q:    T.Tensor([B, S, H, DK], dtype),
        K:    T.Tensor([B, S, H, DK], dtype),
        V:    T.Tensor([B, S, H, DV], dtype),
        G:    T.Tensor([B, S, H, DK], accum_dtype),   # 逐通道 log 门控，< 0
        Beta: T.Tensor([B, S, H], accum_dtype),
        O:    T.Tensor([B, S, H, DV], dtype),
    ):
        with T.Kernel(T.ceildiv(DV, block_DV), B * H, threads=threads) as (bv, bbh):
            bb, bh = bbh // H, bbh % H
            dv0 = bv * block_DV

            Q_s  = T.alloc_shared([C, DK], dtype)
            K_s  = T.alloc_shared([C, DK], dtype)
            V_s  = T.alloc_shared([C, block_DV], dtype)
            S_s  = T.alloc_shared([DK, block_DV], dtype)
            O_s  = T.alloc_shared([C, block_DV], dtype)
            # 逐通道特有：e^{±γ} 与 W 放 shared，避免寄存器超限
            Ep_s = T.alloc_shared([C, DK], accum_dtype)   # e^{γ_r}
            Em_s = T.alloc_shared([C, DK], accum_dtype)   # e^{-γ_r}
            Kp_s = T.alloc_shared([C, DK], dtype)         # k ⊙ e^{γ}
            Km_s = T.alloc_shared([C, DK], dtype)         # k ⊙ e^{-γ}
            W_s  = T.alloc_shared([C, DK], accum_dtype)
            Vt_s = T.alloc_shared([C, block_DV], dtype)

            S_f    = T.alloc_fragment([DK, block_DV], accum_dtype)  # loop-carried
            acc_o  = T.alloc_fragment([C, block_DV], accum_dtype)
            Aqk    = T.alloc_fragment([C, C], accum_dtype)
            Aq_c   = T.alloc_fragment([C, C], dtype)
            Lmat   = T.alloc_fragment([C, C], accum_dtype)          # 三角系数
            Vt     = T.alloc_fragment([C, block_DV], accum_dtype)
            RHS_v  = T.alloc_fragment([C, block_DV], accum_dtype)
            gam    = T.alloc_fragment([C, DK], accum_dtype)         # 累积 log
            be_f   = T.alloc_fragment([C], accum_dtype)

累积衰减是逐通道的前缀和–$d_k$ 个通道各自独立 cumsum，通道之间完全并行：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

T.clear(S_f)

for i_s in T.Pipelined(NS, num_stages=num_stages):
    s0 = i_s * C
    T.copy(Q[bb, s0:s0+C, bh, :], Q_s)
    T.copy(K[bb, s0:s0+C, bh, :], K_s)
    T.copy(V[bb, s0:s0+C, bh, dv0:dv0+block_DV], V_s)

    # ── 逐通道累积 log 衰减：通道维并行，位置维串行 ──
    for d in T.Parallel(DK):
        gam[0, d] = G[bb, s0, bh, d]
    for r in T.serial(1, C):
        for d in T.Parallel(DK):
            gam[r, d] = gam[r - 1, d] + G[bb, s0 + r, bh, d]

    for r, d in T.Parallel(C, DK):
        Ep_s[r, d] = T.exp2(gam[r, d] * 1.4426950408889634)   # e^{γ}
        Em_s[r, d] = T.exp2(-gam[r, d] * 1.4426950408889634)  # e^{-γ}
        Kp_s[r, d] = T.Cast(dtype, K_s[r, d] * Ep_s[r, d])
        Km_s[r, d] = T.Cast(dtype, K_s[r, d] * Em_s[r, d])
    for r in T.Parallel(C):
        be_f[r] = Beta[bb, s0 + r, bh]

对比 GDN：那里前缀和是 $C$ 步串行、每步 1 个数；这里是 $C$ 步串行、每步 $d_k$ 个通道并行。串行长度不变，但并行度从 1 涨到 128–逐通道门在这一处反而更适合 GPU。

1.4426950408889634 是 $\log_2 e$，把 $e^x$ 转成硬件 exp2（与第三篇一致）。

4.2 指数分离：两次 GEMM 取代双重循环

1
2
3
4
5
6
7
8
9
10
11
12

# ── M = strictLower(Kp @ Km^T)：衰减 KK 矩阵 ──
T.gemm(Kp_s, Km_s, Lmat, transpose_B=True, clear_accum=True)
for i, j in T.Parallel(C, C):
    Lmat[i, j] = T.if_then_else(
        j < i, be_f[j] * Lmat[i, j], 0.0)      # L[c,i] = β_i M[c,i]

# ── A^qk = tril(Qp @ Km^T)，含对角 ──
for r, d in T.Parallel(C, DK):
    Q_s[r, d] = T.Cast(dtype, Q_s[r, d] * Ep_s[r, d])   # q ⊙ e^{γ}
T.gemm(Q_s, Km_s, Aqk, transpose_B=True, clear_accum=True)
for i, j in T.Parallel(C, C):
    Aqk[i, j] = T.if_then_else(j <= i, Aqk[i, j], 0.0)

这是全篇最关键的两行 GEMM。 朴素写法需要 $C^2/2$ 次长度 $d_k$ 的加权内积（$C=64, d_k=128$ 时约 26 万次 exp）；分离后 $e^{\pm\bm{\gamma}}$ 各算一次（$C \times d_k = 8192$ 次 exp2），剩下交给 Tensor Core。

注意两个掩码的边界不同：Lmat 取 $j < i$（严格下三角，对角会重复计入 $\beta_r$），Aqk 取 $j \le i$（含对角，自己看自己零衰减）。

4.3 前向替换：W 与 U 合并求解

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

# ── 右端项：[Kp | V]，注意 β 不在这里乘 ──
for r, d in T.Parallel(C, DK):
    W_s[r, d] = K_s[r, d] * Ep_s[r, d]
for r, d in T.Parallel(C, block_DV):
    RHS_v[r, d] = V_s[r, d]

# ── 前向替换，一次解出 T@[Kp | V] ──
for r in T.serial(C):
    for j in T.serial(r):
        for d in T.Parallel(DK):
            W_s[r, d] -= Lmat[r, j] * W_s[j, d]
        for d in T.Parallel(block_DV):
            RHS_v[r, d] -= Lmat[r, j] * RHS_v[j, d]

# ── 解完之后才逐行乘 β（Â = diag(β)T）──
for r, d in T.Parallel(C, DK):
    W_s[r, d] *= be_f[r]
for r, d in T.Parallel(C, block_DV):
    RHS_v[r, d] *= be_f[r]

# ── 伪值 Ṽ = U - W S ──
T.copy(S_f, S_s)
T.gemm(W_s, S_s, Vt, clear_accum=True)      # W S
for r, d in T.Parallel(C, block_DV):
    Vt[r, d] = RHS_v[r, d] - Vt[r, d]
T.copy(Vt, Vt_s)

$\mathbf{W}$ 与 $\mathbf{U}$ 共用同一个 $(\mathbf{I}+\mathbf{L})$，拼在一次前向替换里只付一遍串行代价。这一步的并行度是 $d_k + \text{block}_{DV} = 160$，比 GDN 的 32 好得多–逐通道门在这里又一次因为多了通道维而受益。

4.4 输出与状态更新

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

# ② 块间 + ③ 块内
T.gemm(Q_s, S_s, acc_o, clear_accum=True)   # (Q ⊙ e^γ) S
T.copy(Aqk, Aq_c)
T.gemm(Aq_c, Vt_s, acc_o)                   # A^qk Ṽ

T.copy(acc_o, O_s)
T.copy(O_s, O[bb, s0:s0+C, bh, dv0:dv0+block_DV])

# ① 状态更新：Diag(e^{γ_C}) S + (K ⊙ e^{γ_C-γ_r})^T Ṽ
for d, j in T.Parallel(DK, block_DV):
    S_f[d, j] *= Ep_s[C - 1, d]             # 逐通道！不是标量
for r, d in T.Parallel(C, DK):
    Kp_s[r, d] = T.Cast(dtype,
        K_s[r, d] * Ep_s[C - 1, d] * Em_s[r, d])   # e^{γ_C-γ_r}
T.gemm(Kp_s, Vt_s, S_f, transpose_A=True)

最后一段体现了逐通道门与前三篇最直观的差别：

• GDN：S_f[i,j] *= alpha_pow_C，一个标量乘整个状态
• KDA：S_f[d,j] *= Ep_s[C-1,d]，每个 $d_k$ 行乘各自的衰减

状态的不同行按不同速率淡出–这就是"逐通道"在 kernel 层面的全部含义。

e^{\bm{\gamma}_C - \bm{\gamma}_r} 用 Ep_s[C-1,d] * Em_s[r,d] 算，即 $e^{\gamma_C}\cdot e^{-\gamma_r}$。这里复用了已有的两张表，不必重算 exp2；但它也是 $e^{-\bm{\gamma}}$ 参与的第三处，进一步说明为什么下界不可省。

4.5 七处次序约束

比 GDN 多两处：

1. 状态更新排在输出写回之后（右移语义，四篇一致）。
2. S_f *= Ep_s[C-1,:] 在 T.gemm 之前。
3. Q_s 被原地乘 $e^{\bm{\gamma}}$ 后不可再用于块内项–但本篇的块内项恰好也用 $\bm{q}\odot e^{\bm{\gamma}}$（$\mathbf{A}^{qk}$ 的定义里就带），所以不需要重载 Q。这是与 GDN 的一处反差，容易照抄出错。
4. Lmat 只取严格下三角（$j < i$），Aqk 含对角（$j \le i$）。
5. K_s 必须保持原始值：$\mathbf{Kp}$、$\mathbf{Km}$、右端项、状态更新四处都从 K_s 派生，任何一处原地修改都会污染后续。本篇的做法是始终写入独立的 Kp_s/Km_s，K_s 只读–比 GDN 的"污染两次再重载"更清晰。
6. $\beta$ 在前向替换之后乘，不能提前混进右端项。$\hat{\mathbf{A}} = \operatorname{diag}(\beta)\mathbf{T}$ 与 $\mathbf{T}\operatorname{diag}(\beta)$ 是两种不同配对，各自要搭配 $L_{ci} = \beta_i M_{ci}$ 或 $\beta_c M_{ci}$。$\beta$ 全相等时这个错误完全隐身。
7. T.clear(S_f) 在循环外，clear_accum=True 在循环内。

第 3、5、6 条都是"照抄上一篇会错"的地方，其中第 6 条我实际踩过。

5. 四篇对照

四篇的箭头量语义完全一致：$\overleftarrow{\bm{q}}$ 乘 $e^{\bm{\gamma}_r}$、$\overrightarrow{\bm{k}}$ 乘 $e^{\bm{\gamma}_C-\bm{\gamma}_r}$、$\overrightarrow{\mathbf{S}}$ 乘 $e^{\bm{\gamma}_C}$。从第一篇的"无衰减"到本篇的"逐通道"，衰减插入的三个位置从未改变，改变的只是每个位置乘的是标量还是向量。

6. 总结

1. 逐通道门只改一个符号，却拆掉了前三篇所有便利。$\alpha_t \to \operatorname{Diag}(\bm{a}_t)$ 让累积积从标量变向量，$\Gamma$ 从 $C\times C$ 表变成概念上的 $C\times C\times d_k$–不可能物化，必须融进 GEMM。
2. 衰减长在 KK 矩阵内部，这是 KDA chunkwise 的核心难点。$M_{ci} = (\bm{k}_c \odot e^{\bm{\gamma}_c-\bm{\gamma}_i})^\intercal\bm{k}_i$ 的指数依赖三个下标，无法像标量门那样外提成逐元素乘。
3. 出路是指数可分离：$e^{\bm{\gamma}_c-\bm{\gamma}_i} = e^{\bm{\gamma}_c}\odot e^{-\bm{\gamma}_i}$，于是 $\mathbf{M} = \widetilde{\mathbf{K}}^{+}\widetilde{\mathbf{K}}^{-\intercal}$ 一次 GEMM 解决（实测等价，误差 $1.94\times10^{-16}$）。朴素写法要 $C^2/2$ 次加权内积、约 26 万次 exp；分离后只需 $C d_k = 8192$ 次 exp2 加两次 GEMM。
4. 代价是必须物化 $e^{-\bm{\gamma}}$–第二篇批判过的 $1/\gamma$ 陷阱的逐通道版本。但这次没有替代方案：不拆就用不上 Tensor Core。问题从"要不要拆"变成"如何让拆分数值安全"。
5. 门控下界是 kernel 可行性的前提，不是精度调优。实测 $C=128$、无下界时 fp16 下 $e^{\bm{\gamma}_C}$ 128/128 通道全部归零，状态被彻底清空。下界 0.9 时归零通道数为 0，同时 $\max e^{-\bm{\gamma}} = 55$ 不溢出–一个数字同时管住了下溢与溢出两头。
6. 通道离散度是逐通道门独有的病。$\bm{a}\in[0.01,0.999]$、$C=128$ 时 $\bm{\gamma}_C$ 的通道极差达 47（log 域），即同一 fragment 内最强与最弱通道相差 20 个数量级，任何浮点格式都无法同时表示。下界 0.9 把极差压到 1.74。
7. sub-chunk 是第二道保险。$e^{-\bm{\gamma}}$ 的动态范围只取决于 sub-chunk 长度：$C=128$ 整块 $1.79\times10^3$，切 32 后降到 9.02。这与第二篇"累积积按 chunk 重置"同理，只是又降一层。$C=64$ + 下界 0.9 时不需要。
8. 逐通道门在两处反而更适合 GPU。累积前缀和：GDN 是 $C$ 步串行 × 宽度 1，KDA 是 $C$ 步串行 × 宽度 $d_k$，串行长度不变而并行度从 1 涨到 128。前向替换：$\mathbf{W}$ 与 $\mathbf{U}$ 共用 $(\mathbf{I}+\mathbf{L})$ 拼成一次求解，并行度 $d_k + \text{bDV} = 160$，远高于 GDN 的 32。
9. 寄存器压力换了主角。GDN 的瓶颈是三张 $C^2$ 表；KDA 是三张 $C \times d_k$（$e^{\pm\bm{\gamma}}$、$\mathbf{W}$），$64\times128$ 是 $64^2$ 的两倍。朴素分配约 320 reg 超限，把 $e^{-\bm{\gamma}}$ 与 $\mathbf{W}$ 卸载到 shared memory 后降到约 192。
10. 三处"照抄上一篇会错"：块内项的 $\mathbf{A}^{qk}$ 定义里本就带 $\bm{q}\odot e^{\bm{\gamma}}$，不需要像 GDN 那样重载 Q；K_s 应始终只读、派生量写独立 buffer，而非 GDN 的"污染再重载"；$\beta$ 必须在解完三角系统之后乘–$\operatorname{diag}(\beta)\mathbf{T}$ 与 $\mathbf{T}\operatorname{diag}(\beta)$ 需搭配不同的 $L$ 下标，写反时 $\beta$ 全相等则隐身、$\beta$ 随机则相对 L2 达 $1.3\times10^{-1}$（我第一次写就踩了这个）。
11. 三个退化检验，比 GDN 多一个。$\bm{g}$ 各通道同值 → 回到 GDN（检验指数分离是否搞混通道维与位置维）、$\beta\to0$ → 逐通道纯衰减、$\bm{g}\to\bm{0}$ → 纯 DeltaNet。第一个最重要且本篇独有。
12. 衰减的三处落点四篇未变。$\overleftarrow{\bm{q}}$、$\overrightarrow{\bm{k}}$、$\overrightarrow{\mathbf{S}}$ 从第一篇到第四篇完全一致，变的只是每处乘标量还是乘向量。这是 GDN 那套箭头记号真正的价值–它把"衰减"隔离成了一个独立于门控形态的层面。

可迁移的启示：把一个标量参数升级成向量，代价从来不在参数量。它会连带改变哪些量能被预计算、哪些表能被物化、哪些运算能进 Tensor Core。KDA 的 $\alpha_t \to \operatorname{Diag}(\bm{a}_t)$ 只多了 $d_k$ 个数，却让衰减掩码从"一张可复用的表"变成"必须融进 GEMM 的隐式结构"，并且逼出了一个模型层面的约束（门控下界）作为 kernel 可行性的前提。当一个数值技巧成为架构设计的必要条件时，它就不再是实现细节。

参考

• Kimi Linear：arXiv:2510.26692。KDA 的架构设计、下界衰减与满秩门控
• Gated Delta Networks：arXiv:2412.06464，ICLR 2025。标量门版本的 chunkwise 算法与 UT 变换
• 本系列前三篇：Chunked 线性注意力、标量衰减、Gated DeltaNet
• KDA 数学推导详解：KDA 的来龙去脉

TileLang 实战：KDA 从零到一–Gated DeltaNet

前两篇分别实现了无遗忘的 chunked 线性注意力，以及带标量衰减的版本。两者的状态更新都只做加法：新的键值对累加进状态，旧信息靠衰减系数被动遗忘。本篇补上最后一块–删除。

递推式变成：

$$\mathbf{S}_t = \mathbf{S}_{t-1}\Big(\alpha_t\big(\mathbf{I} - \beta_t \bm{k}_t\bm{k}_t^{\intercal}\big)\Big) + \beta_t\bm{v}_t\bm{k}_t^{\intercal}$$

比上一篇多出的是 $\mathbf{I} - \beta_t\bm{k}_t\bm{k}_t^{\intercal}$ 这个广义 Householder 变换。它带来的不是又一个逐元素权重，而是一个矩阵乘在状态右侧–这一改把整个 chunkwise 算法的结构改掉了：块内不再能靠一张下三角权重表解决，需要解一个 $C \times C$ 的三角系统。

上两篇见《Chunked 线性注意力》与《标量衰减》，本文沿用其符号约定与四层参考验证框架。

0. 符号约定

沿用 GDN（arXiv:2412.06464v3）§2.2、§3.1、§3.3 与附录 A 的记号，本篇新增两个：

沿用的：$\alpha_t$ 单步衰减、$\gamma^r_{[t]} = \prod_{i=1}^{r}\alpha_i$ 累积衰减积（按 chunk 重置）、$\Gamma_{ij} = \gamma_i/\gamma_j$ 衰减感知因果掩码、$C$ 块长、$[t]$ 块序号、$r \in [1,C]$ 块内位置。

一个前提：论文对 $\bm{q}, \bm{k}$ 做 L2 归一化，即 $\|\bm{k}_t\| = 1$。这不只是训练稳定性的考虑–下面 §1.2 会看到它直接决定了 Householder 变换会不会把状态越推越大。

1. 从加法到删除：delta rule 在做什么

1.1 三种更新方式的对照

把三篇的递推式并排放，差异一目了然：

标量衰减的问题在于它不分对象：$\alpha_t$ 一乘，所有历史信息按同一比例衰减。要腾出空间写入新内容，只能把无关的旧信息一起冲淡。GDN 论文的说法是，门控擅长「快速擦除」，delta rule 擅长「定向修改」，两者互补。

delta rule 的定向性来自哪里？把它拆开看：

$$\mathbf{S}_t = \mathbf{S}_{t-1} - \underbrace{(\mathbf{S}_{t-1}\bm{k}_t)}_{\bm{v}_t^{\text{old}}}\bm{k}_t^{\intercal} + \underbrace{\big(\beta_t\bm{v}_t + (1-\beta_t)\mathbf{S}_{t-1}\bm{k}_t\big)}_{\bm{v}_t^{\text{new}}}\bm{k}_t^{\intercal}$$

读法是：先把 $\bm{k}_t$ 这个键上原有的值 $\bm{v}^{\text{old}}_t = \mathbf{S}_{t-1}\bm{k}_t$ 减掉，再写入新值 $\bm{v}^{\text{new}}_t$，而新值是旧值与目标值的凸组合，$\beta_t$ 控制替换的彻底程度。$\beta_t \to 1$ 是完全覆盖，$\beta_t \to 0$ 是不动。

关键在于这个减法只作用在 $\bm{k}_t$ 方向上，与 $\bm{k}_t$ 正交的记忆分毫不动。这就是「定向」–相比标量衰减的一刀切，delta rule 只擦掉要覆盖的那一条。

1.2 为什么是 Householder，以及 L2 归一化的作用

$\mathbf{I} - \beta_t\bm{k}_t\bm{k}_t^{\intercal}$ 是广义 Householder 变换。$\|\bm{k}_t\| = 1$ 时，它的特征值只有两种取值，结构一目了然：

• 沿 $\bm{k}_t$ 方向：特征值 $1 - \beta_t$
• 与 $\bm{k}_t$ 正交的 $d_k - 1$ 个方向：特征值 $1$

于是 $\beta_t \in (0,1)$ 时全部特征值的绝对值都不超过 1，状态每步只会被压缩、不会被放大，递推因此稳定。这正是论文对 $\bm{k}$ 做 L2 归一化的深层原因–若 $\|\bm{k}_t\| \ne 1$，沿 $\bm{k}_t$ 的特征值变成 $1 - \beta_t\|\bm{k}_t\|^2$，$\beta_t\|\bm{k}_t\|^2 > 2$ 时就会翻到 $-1$ 以下，递推放大。

顺带一提，论文脚注提到可以放开到 $\beta_t \in (0,2)$ 以允许负特征值，那是为了解锁状态跟踪能力（state tracking）。本文按 $(0,1)$ 处理。

1.3 test-time SGD 视角

论文给了一个很有启发的解释：把状态 $\mathbf{S}$ 看成一个快速权重矩阵，delta rule 就是在做在线回归的一步梯度下降。目标是 $\mathcal{L}(\mathbf{S}_t) = \frac{1}{2}\|\mathbf{S}_t\bm{k}_t - \bm{v}_t\|^2$，那么：

$$\mathbf{S}_t - \beta_t\nabla\mathcal{L}(\mathbf{S}_t) = \mathbf{S}_t - \beta_t(\mathbf{S}_t\bm{k}_t - \bm{v}_t)\bm{k}_t^{\intercal} = \mathbf{S}_t(\mathbf{I} - \beta_t\bm{k}_t\bm{k}_t^{\intercal}) + \beta_t\bm{v}_t\bm{k}_t^{\intercal}$$

$\beta_t$ 就是学习率，$\alpha_t$ 就是 weight decay。 这个视角下 gated delta rule 没有任何神秘之处–它是带权重衰减的 test-time SGD。

2. Chunkwise 形式：为什么需要解三角系统

2.1 展开递推：转移矩阵不再是标量

按块展开 $r$ 步（论文式 10）：

$$\mathbf{S}_{[t]}^{r} = \mathbf{S}_{[t]}\underbrace{\prod_{i=1}^{r}\alpha^i_{[t]}\big(\mathbf{I} - \beta^i_{[t]}\bm{k}^i_{[t]}\bm{k}^{i\intercal}_{[t]}\big)}_{\mathbf{F}^r_{[t]}} + \underbrace{\sum_{i=1}^{r}\beta^i_{[t]}\bm{v}^i_{[t]}\bm{k}^{i\intercal}_{[t]}\prod_{j=i+1}^{r}\alpha^j_{[t]}\big(\mathbf{I} - \beta^j_{[t]}\bm{k}^j_{[t]}\bm{k}^{j\intercal}_{[t]}\big)}_{\mathbf{G}^r_{[t]}}$$

对比上一篇：那里的转移量是标量 $\alpha^{r-i}$，可以直接查表。这里是矩阵连乘，$\mathbf{F}$ 与 $\mathbf{G}$ 都不能靠逐元素权重表达。

衰减部分可以先提出来–$\prod\alpha_i = \gamma^r$ 是标量，与 Householder 部分可交换：

$$\mathbf{F}^r_{[t]} = \gamma^r_{[t]}\,\mathbf{P}^r_{[t]}, \qquad \mathbf{P}^r_{[t]} = \prod_{i=1}^{r}\big(\mathbf{I} - \beta^i\bm{k}^i\bm{k}^{i\intercal}\big)$$

剩下的 $\mathbf{P}$ 是纯 DeltaNet 的部分，靠 WY 表示处理–这正是论文 §2.2 的内容，下一节完整展开。

2.2 论文 §2.2：无门控 DeltaNet 的 WY 表示

在处理带门控的版本之前，先把论文 §2.2 那套无门控 DeltaNet 的 chunkwise 推导完整走一遍。GDN 的做法本质上是在这套框架上打补丁，先看清基线，后面的改动才有参照。

本节 $\alpha_t \equiv 1$（无遗忘门），递推退化为 $\mathbf{S}_t = \mathbf{S}_{t-1}(\mathbf{I} - \beta_t\bm{k}_t\bm{k}_t^{\intercal}) + \beta_t\bm{v}_t\bm{k}_t^{\intercal}$。

2.2.1 部分展开：两个连乘（论文式 3）

按块部分展开递推：

$$\mathbf{S}^r_{[t]} = \mathbf{S}_{[t]}\underbrace{\Big(\prod_{i=1}^{r}\big(\mathbf{I} - \beta^i_{[t]}\bm{k}^i_{[t]}\bm{k}^{i\intercal}_{[t]}\big)\Big)}_{\textstyle :=\mathbf{P}^r_{[t]}} + \underbrace{\sum_{i=1}^{r}\beta^i_{[t]}\bm{v}^i_{[t]}\bm{k}^{i\intercal}_{[t]}\prod_{j=i+1}^{r}\big(\mathbf{I} - \beta^j_{[t]}\bm{k}^j_{[t]}\bm{k}^{j\intercal}_{[t]}\big)}_{\textstyle :=\mathbf{H}^r_{[t]}}$$

两个部分的角色不同，值得分清：

$\mathbf{P}$ 是纯连乘，$\mathbf{H}$ 的每一项后面还挂着一截连乘尾巴–两者都不能直接算，$C = 64$ 时 $\mathbf{P}$ 要 64 个 $d_k\times d_k$ 矩阵相乘。WY 表示的作用就是把这两个连乘各自压成一次求和。

2.2.2 经典 WY：把 Householder 连乘压成秩-$C$ 更新（论文式 4）

这是 Bischof & Van Loan (1985) 的经典结果。核心事实：$C$ 个 Householder 矩阵的乘积可以写成单位矩阵减去一个秩至多 $C$ 的修正：

$$\mathbf{P}^r_{[t]} = \mathbf{I} - \sum_{i=1}^{r}\bm{w}^i_{[t]}\bm{k}^{i\intercal}_{[t]} \in \mathbb{R}^{d_k\times d_k}, \qquad \bm{w}^r_{[t]} = \beta^r_{[t]}\Big(\bm{k}^r_{[t]} - \sum_{i=1}^{r-1}\bm{w}^i_{[t]}\big(\bm{k}^{i\intercal}_{[t]}\bm{k}^r_{[t]}\big)\Big) \in \mathbb{R}^{d_k}$$

为什么成立，看一步归纳就够：

$$\mathbf{P}^{r} = \mathbf{P}^{r-1}\big(\mathbf{I} - \beta_r\bm{k}_r\bm{k}_r^{\intercal}\big) = \mathbf{P}^{r-1} - \underbrace{\beta_r\mathbf{P}^{r-1}\bm{k}_r}_{\textstyle \bm{w}_r}\bm{k}_r^{\intercal}$$

于是 $\bm{w}_r = \beta_r\mathbf{P}^{r-1}\bm{k}_r$，把 $\mathbf{P}^{r-1} = \mathbf{I} - \sum_{i<r}\bm{w}_i\bm{k}_i^{\intercal}$ 代入就得到上面那个递推。每多一个 Householder，秩只增加 1–这就是"连乘变求和"的全部内容。

2.2.3 同一个模具：$\mathbf{H}$ 的 WY 表示（论文式 5）

$\mathbf{H}$ 的推导结构与 $\mathbf{P}$ 完全平行：

$$\mathbf{H}^r_{[t]} = \sum_{i=1}^{r}\bm{u}^i_{[t]}\bm{k}^{i\intercal}_{[t]} \in \mathbb{R}^{d_v\times d_k}, \qquad \bm{u}^r_{[t]} = \beta^r_{[t]}\Big(\bm{v}^r_{[t]} - \sum_{i=1}^{r-1}\bm{u}^i_{[t]}\big(\bm{k}^{i\intercal}_{[t]}\bm{k}^r_{[t]}\big)\Big) \in \mathbb{R}^{d_v}$$

把 $\bm{w}$ 的递推和 $\bm{u}$ 的递推并排看，会发现它们是同一个式子：

$$\bm{w}_r = \beta_r\Big(\underline{\bm{k}_r} - \sum_{i<r}\bm{w}_i(\bm{k}_i^{\intercal}\bm{k}_r)\Big), \qquad \bm{u}_r = \beta_r\Big(\underline{\bm{v}_r} - \sum_{i<r}\bm{u}_i(\bm{k}_i^{\intercal}\bm{k}_r)\Big)$$

只有下划线处不同：一个是 $\bm{k}_r$、一个是 $\bm{v}_r$。系数完全一样–这个观察是后面式 (6)(7) 能共用一个 $\mathbf{T}$ 的全部原因，也是 kernel 里 $\mathbf{W}$ 与 $\mathbf{U}$ 能拼进一次前向替换的依据（第四篇 KDA 用的就是这个技巧）。

写成矩阵形式，两个连乘都消失了：

$$\mathbf{P}_{[t]} = \mathbf{I} - \mathbf{W}_{[t]}^{\intercal}\mathbf{K}_{[t]} \in \mathbb{R}^{d_k\times d_k}, \qquad \mathbf{H}_{[t]} = \mathbf{U}_{[t]}^{\intercal}\mathbf{K}_{[t]} \in \mathbb{R}^{d_v\times d_k}$$

2.2.4 UT 变换：递推也不必串行（论文式 6、7）

式 (4)(5) 虽然把连乘变成了求和，但 $\bm{w}_r$ / $\bm{u}_r$ 自身还是串行递推。Joffrain et al. (2006) 的 UT 变换把它变成一次矩阵求逆：

$$\mathbf{T}_{[t]} = \Big[\mathbf{I} + \operatorname{strictLower}\big(\operatorname{diag}(\beta_{[t]})\mathbf{K}_{[t]}\mathbf{K}_{[t]}^{\intercal}\big)\Big]^{-1}\operatorname{diag}(\beta_{[t]}) \in \mathbb{R}^{C\times C}$$

$$\mathbf{W}_{[t]} = \mathbf{T}_{[t]}\mathbf{K}_{[t]} \in \mathbb{R}^{C\times d_k}, \qquad \mathbf{U}_{[t]} = \mathbf{T}_{[t]}\mathbf{V}_{[t]} \in \mathbb{R}^{C\times d_v}$$

$\mathbf{W}$ 与 $\mathbf{U}$ 共用同一个 $\mathbf{T}$，正是因为 §2.2.3 那两个递推的系数相同。这一步的意义在于：$\mathbf{T}$ 只跟 $\mathbf{K}$ 和 $\beta$ 有关，与 $\mathbf{V}$ 无关。

2.2.5 代回式 3：可用 Tensor Core 的算法（论文式 8、9）

$$\begin{aligned} \mathbf{S}_{[t+1]} &= \mathbf{S}_{[t]}\mathbf{P}_{[t]} + \mathbf{H}_{[t]} = \mathbf{S}_{[t]} + \big(\mathbf{U}_{[t]} - \mathbf{W}_{[t]}\mathbf{S}_{[t]}^{\intercal}\big)^{\intercal}\mathbf{K}_{[t]} \in \mathbb{R}^{d_v\times d_k} \\[4pt] \mathbf{O}_{[t]} &= \mathbf{Q}_{[t]}\mathbf{S}_{[t]}^{\intercal} + \big(\mathbf{Q}_{[t]}\mathbf{K}_{[t]}^{\intercal}\odot\mathbf{M}\big)\big(\mathbf{U}_{[t]} - \mathbf{W}_{[t]}\mathbf{S}_{[t]}^{\intercal}\big) \in \mathbb{R}^{C\times d_v} \end{aligned}$$