⚡ KDA · Key-Dependent Attention

# kda.py - Chunk-KDA 入口函数
def chunk_kda(q, k, v, g, beta, scale=None, initial_state=None, ...):
    if scale is None:
        scale = k.shape[-1] ** -0.5

    # 可选的 L2 归一化
    if use_qk_l2norm_in_kernel:
        q = l2norm_fwd(q.contiguous(), use_tiled_kernel=True)
        k = l2norm_fwd(k.contiguous(), use_tiled_kernel=True)

    o, final_state = chunk_kda_fwd(
        q=q, k=k, v=v.contiguous(), g=g.contiguous(),
        beta=beta.contiguous(), scale=scale, ...)
    return o, final_state

# kda.py - 分块前向传播核心逻辑
def chunk_kda_fwd(q, k, v, g, beta, scale, initial_state, ...):
    chunk_size = FLA_CHUNK_SIZE  # 默认 64
    # Step 1: 累积门控和
    g = chunk_local_cumsum(g, chunk_size=chunk_size, ...)
    g = g * RCP_LN2

    # Step 2: 计算 K·K^T 和 Q·K^T 矩阵
    A, Aqk = chunk_kda_scaled_dot_kkt_fwd(q=q, k=k, gk=g, beta=beta, ...)

    # Step 3: 求解下三角矩阵逆
    A = solve_tril_kda(A=A, ...)

    # Step 4: 重计算 w 和 u
    w, u, _, kg = recompute_w_u_fwd(k=k, v=v, beta=beta, A=A, gk=g, ...)

    # Step 5: 计算 h-state 隐状态
    h, v_new, final_state = chunk_gated_delta_rule_fwd_h_kda(
        k=kg, w=w, u=u, gk=g, initial_state=initial_state, ...)

    # Step 6: 计算最终输出
    o = chunk_gla_fwd_o_gk(q=q, v=v_new, g=g, A=Aqk, h=h, scale=scale, ...)
    return o, final_state

# fused_recurrent_kda.py - 逐 token 循环递推
@triton.jit(do_not_specialize=["N", "T"])
def fused_recurrent_gated_delta_rule_fwd_kernel(q, k, v, g, beta, o, ...):
    # 初始化隐状态
    b_h = tl.zeros([BV, BK], dtype=tl.float32)
    if USE_INITIAL_STATE:
        b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)

    # 主循环：逐 token 递推
    for i_t in range(0, T):
        b_q = tl.load(p_q, ...).to(tl.float32)
        b_k = tl.load(p_k, ...).to(tl.float32)
        b_v = tl.load(p_v, ...).to(tl.float32)

        # KDA 特殊处理：向量门控
        if IS_KDA:
            b_gk = tl.load(p_gk).to(tl.float32)
            b_h *= tl.exp(b_gk[None, :])  # 门控衰减

        # Delta Rule 更新
        b_v -= tl.sum(b_h * b_k[None, :], 1)  # 残差计算
        b_v *= b_beta
        b_h += b_v[:, None] * b_k[None, :]     # 状态更新
        b_o = tl.sum(b_h * b_q[None, :], 1)    # 输出计算
        tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)

# solve_tril.py - 下三角矩阵求逆
@triton.jit(do_not_specialize=["T"])
def solve_tril_16x16_kernel_kda(A, Ai, cu_seqlens, chunk_indices, T, H, BT):
    # 加载 16x16 子块
    b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)
    b_A = -tl.where(m_A, b_A, 0)  # 取严格下三角部分

    # 迭代求解逆矩阵
    for i in range(2, min(16, T - i_t * 16)):
        b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset)
        b_a = b_a + tl.sum(b_a[:, None] * b_A, 0)
        b_A = tl.where((o_i == i)[:, None], b_a, b_A)

    b_A += m_I  # 加上单位矩阵
    tl.store(p_Ai, b_A.to(...))

# kda.py 中的自动调优配置
BT_LIST_AUTOTUNE = [32, 64, 128]  # 块大小搜索空间
NUM_WARPS_AUTOTUNE = [4, 8, 16, 32]  # Warp 数量搜索空间

@triton.autotune(
    configs=[
        triton.Config({"BT": bt}, num_warps=nw, num_stages=ns)
        for bt in BT_LIST_AUTOTUNE
        for nw in NUM_WARPS_AUTOTUNE
        for ns in [2, 3]
    ],
    key=["H", "D"],
)
@triton.jit
def kda_gate_fwd_kernel(g, A, y, g_bias, beta, threshold, T, H, D, BT, BD):
    # Kernel 实现...
    pass