hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
# 注：把 padding 位的状态清理掉，避免线性递推状态被无效 token 污染。

batch_size, seq_len, _ = hidden_states.shape
# 注：后续 reshape / decode 分支判断都会用到 seq_len。

use_precomputed_states = (
    cache_params is not None
    and cache_params.has_previous_state
    and seq_len == 1
    and cache_position is not None
)
# 注：只有 decode 单步（seq_len==1）且已有历史状态时，才走 recurrent 路径；
#     否则走 chunk 路径（prefill/训练）。

if cache_params is not None:
    conv_state = cache_params.conv_states[self.layer_idx]
    recurrent_state = cache_params.recurrent_states[self.layer_idx]
# 注：从动态缓存读取卷积状态与递推状态，供增量推理复用。

projected_states_qkvz = self.in_proj_qkvz(hidden_states)
projected_states_ba = self.in_proj_ba(hidden_states)
# 注：两路投影：
#   - qkvz 路：产出 query/key/value/z
#   - ba 路：产出 b/a（后续分别变成 beta 和 g 的输入）

query, key, value, z, b, a = self.fix_query_key_value_ordering(projected_states_qkvz, projected_states_ba)
# 注：在 fix_query_key_value_ordering 内部会先按 head 分组，再 split：
#   - qkvz 按 [q, k, v, z] 拆分
#   - ba   按 [b, a] 拆分

def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba):
    """
    Derives `query`, `key` and `value` tensors from `mixed_qkvz` and `mixed_ba`.
    """
    new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + (
        self.num_k_heads,
        2 * self.head_k_dim + 2 * self.head_v_dim * self.num_v_heads // self.num_k_heads,
    )
    new_tensor_shape_ba = mixed_ba.size()[:-1] + (self.num_k_heads, 2 * self.num_v_heads // self.num_k_heads)

    mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz)
    #   mixed_qkvz
    #      -> view -> [B, L, num_k_heads, 2*head_k_dim + 2*head_v_dim*(num_v_heads/num_k_heads)]
    mixed_ba = mixed_ba.view(*new_tensor_shape_ba)
    #   mixed_ba: [B, L, 2*num_v_heads]
    #      -> view -> [B, L, num_k_heads, 2*(num_v_heads/num_k_heads)]
    
    split_arg_list_qkvz = [
        self.head_k_dim,
        self.head_k_dim,
        (self.num_v_heads // self.num_k_heads * self.head_v_dim),
        (self.num_v_heads // self.num_k_heads * self.head_v_dim),
    ]
    split_arg_list_ba = [self.num_v_heads // self.num_k_heads, self.num_v_heads // self.num_k_heads]
    query, key, value, z = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=3)
    # 在 dim=3 上 split 成 q/k/v/z
    #   query: [B, L, num_k_heads, head_k_dim]
    #   key:   [B, L, num_k_heads, head_k_dim]
    #   value: [B, L, num_k_heads, (num_v_heads/num_k_heads)*head_v_dim]
    #   z:     [B, L, num_k_heads, (num_v_heads/num_k_heads)*head_v_dim]
    b, a = torch.split(mixed_ba, split_arg_list_ba, dim=3)
    # 在 dim=3 上 split 成 b/a
    #   b: [B, L, num_k_heads, (num_v_heads/num_k_heads)]
    #   a: [B, L, num_k_heads, (num_v_heads/num_k_heads)]

    value = value.reshape(value.size(0), value.size(1), -1, self.head_v_dim)
    #   value: [B, L, num_k_heads, (num_v_heads/num_k_heads)*head_v_dim]
    #      -> reshape -> [B, L, num_v_heads, head_v_dim]
    z = z.reshape(z.size(0), z.size(1), -1, self.head_v_dim)
    #   z: [B, L, num_k_heads, (num_v_heads/num_k_heads)*head_v_dim]
    #      -> reshape -> [B, L, num_v_heads, head_v_dim]
    b = b.reshape(b.size(0), b.size(1), self.num_v_heads)
    #   b: [B, L, num_k_heads, (num_v_heads/num_k_heads)]
    #      -> reshape -> [B, L, num_v_heads]
    a = a.reshape(a.size(0), a.size(1), self.num_v_heads)
    #   a: [B, L, num_k_heads, (num_v_heads/num_k_heads)]
    #      -> reshape -> [B, L, num_v_heads]
    return query, key, value, z, b, a

query, key, value = (x.reshape(x.shape[0], x.shape[1], -1) for x in (query, key, value))
# 注：先把头维摊平，便于拼接做一维深度卷积。

mixed_qkv = torch.cat((query, key, value), dim=-1)
mixed_qkv = mixed_qkv.transpose(1, 2)
# 注：卷积输入形状转成 [B, C, L]，C=2*key_dim+value_dim。

if use_precomputed_states:
    mixed_qkv = self.causal_conv1d_update(
        mixed_qkv,
        conv_state,
        self.conv1d.weight.squeeze(1),
        self.conv1d.bias,
        self.activation,
    )
# 注：decode 增量卷积，只更新最后一步，避免重复算历史。

else:
    if cache_params is not None:
        conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
        ## 这部分的参数格式是F.pad(input, (pad_left, pad_right))；
        ## 如果大于kernel_size，pad_left就会是负数，表示裁剪掉前面多余的部分；
        ## 如果小于kernel_size，就会在前面补零。
        cache_params.conv_states[self.layer_idx] = conv_state
    if self.causal_conv1d_fn is not None:
        mixed_qkv = self.causal_conv1d_fn(
            x=mixed_qkv,
            weight=self.conv1d.weight.squeeze(1),
            bias=self.conv1d.bias,
            activation=self.activation,
            seq_idx=None,
        )
    else:
        mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :seq_len])
# 注：prefill/训练路径：
#   - 有快实现就用 fused kernel；
#   - 否则回退到 torch conv1d + silu；
#   - 若开 cache，同步初始化 conv_state。

mixed_qkv = mixed_qkv.transpose(1, 2)
query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
# 注：卷积后再拆回 Q/K/V 三路。

query = query.reshape(query.shape[0], query.shape[1], -1, self.head_k_dim)
key = key.reshape(key.shape[0], key.shape[1], -1, self.head_k_dim)
value = value.reshape(value.shape[0], value.shape[1], -1, self.head_v_dim)
# 注：恢复为多头形状 [B, L, H, d]。

beta = b.sigmoid()
# 注：写入强度 beta∈(0,1)。

g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
# 注：得到 log-space 的衰减参数；后续在 rule 内部 exp(g)->alpha∈(0,1)。

if self.num_v_heads // self.num_k_heads > 1:
    query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
    key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
# 注：当 v-head 多于 k-head 时，把 K/Q 复制到对齐的头数，保证后续逐头计算可广播。

def torch_recurrent_gated_delta_rule(query, key, value, g, beta, initial_state, ...):

    initial_dtype = query.dtype # 备份数据类型
    if use_qk_l2norm_in_kernel:
        query = l2norm(query, dim=-1, eps=1e-6)
        key = l2norm(key, dim=-1, eps=1e-6)
    # QK 归一化（可选），有助于稳定训练和推理。
    query, key, value, beta, g = [
        x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
    ]
    # 这一步很重要，它把输入的形状从 [B, L, H, d] 转成 [B, H, L, d]，并且转换为 float32 以提高数值稳定性。对于SSM等长序列模型，记忆状态会经历累加，这种精度是必要的。

    batch_size, num_heads, sequence_length, k_head_dim = key.shape
    v_head_dim = value.shape[-1]
    scale = 1 / (query.shape[-1] ** 0.5)
    query = query * scale

    # 声明（创建）空变量存储输出和递推状态
    core_attn_out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim).to(value)
    last_recurrent_state = (
        torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
        if initial_state is None
        else initial_state.to(value)
    )

    # sequence_length 一般为1，但也可以是小批量的多步（例如 decode 时一次更新多个 token）。
    for i in range(sequence_length):
        q_t = query[:, :, i]   # q_t
        k_t = key[:, :, i]     # k_t
        v_t = value[:, :, i]   # v_t
        g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1)  # alpha_t
        beta_t = beta[:, :, i].unsqueeze(-1)                # beta_t

        # 对应公式: S_{t-1} * alpha_t
        last_recurrent_state = last_recurrent_state * g_t

        # 计算 kv_mem = S_{t-1} * k_t （公式中的矩阵乘法）
        kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2)

        # 核心Delta操作: delta = (v_t - kv_mem) * beta_t
        # 这等效于先计算 beta_t * v_t，再减去 beta_t * (S_{t-1} * k_t)
        delta = (v_t - kv_mem) * beta_t

        # Rank-1 更新: S_t = S_{t-1} + k_t ⊗ delta
        # 这行代码等价于完成公式的最终合并：
        # S_t = alpha_t * S_{t-1} - alpha_t * beta_t * S_{t-1} * (k_t k_t^T) + beta_t * v_t k_t^T
        last_recurrent_state = last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2)

        # 输出: o_t = S_t * q_t
        core_attn_out[:, :, i] = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2)

    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
    # 最后把输出转回 [B, L, H, d_v] 的形状，并恢复原始数据类型。

def torch_chunk_gated_delta_rule(
    query, key, value, g, beta,
    chunk_size=64, initial_state=None,
    output_final_state=False, use_qk_l2norm_in_kernel=False,
):
    # ── 0. 预处理：QK 归一化 + 转置 + float32（与单步递推完全一致）──
    initial_dtype = query.dtype
    if use_qk_l2norm_in_kernel:
        query = l2norm(query, dim=-1, eps=1e-6)
        key = l2norm(key, dim=-1, eps=1e-6)
    query, key, value, beta, g = [
        x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
    ]
    # 形状由 [B, L, H, d] 转为 [B, H, L, d]，便于后续按头做矩阵乘法。

    # ── 1. 对序列长度做右侧补零，使其整除 chunk_size ──
    batch_size, num_heads, sequence_length, k_head_dim = key.shape
    v_head_dim = value.shape[-1]
    pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
    # pad_size 保证 total_sequence_length = sequence_length + pad_size 能被 chunk_size 整除。
    # 若 sequence_length 已经整除，则 pad_size = 0，不做任何填充。
    query  = F.pad(query,  (0, 0, 0, pad_size))   # [B, H, L+pad, d_k]
    key    = F.pad(key,    (0, 0, 0, pad_size))
    value  = F.pad(value,  (0, 0, 0, pad_size))
    beta   = F.pad(beta,   (0, pad_size))           # [B, H, L+pad]
    g      = F.pad(g,      (0, pad_size))           # [B, H, L+pad]
    total_sequence_length = sequence_length + pad_size

    scale = 1 / (query.shape[-1] ** 0.5)
    query = query * scale                           # 缩放同单步递推

    # ── 2. 预计算带 beta 的 v 和 k，对应公式中 beta_t * v_t 和 beta_t * k_t ──
    v_beta = value * beta.unsqueeze(-1)             # [B, H, L+pad, d_v]，β_t v_t
    k_beta = key   * beta.unsqueeze(-1)             # [B, H, L+pad, d_k]，β_t k_t

    # ── 3. Reshape：将时间维切分为 (num_chunks, chunk_size) ──
    query, key, value, k_beta, v_beta = [
        x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1])
        for x in (query, key, value, k_beta, v_beta)
    ]
    # 形状变为 [B, H, num_chunks, chunk_size, d]，方便按 chunk 索引。
    g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
    # g: [B, H, num_chunks, chunk_size]

    # ── 4. 块内衰减掩码 (decay_mask) ──
    # 对每个 chunk 内的 g 做累积求和，得到从 chunk 起点到各位置的累积 log-decay。
    g = g.cumsum(dim=-1)
    # g[..., t] 现在表示 sum_{s=0}^{t} log_alpha_s（块内前缀和）。

    # decay_mask[..., t, s] = exp(g[t] - g[s])，表示从位置 s 到位置 t 的衰减系数。
    # .tril() 保证只有 t >= s 的位置有值（因果性），其余置 0。
    decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
    # decay_mask: [B, H, num_chunks, chunk_size, chunk_size]
    # 对角线 (t==s) 处值为 exp(0)=1，即当前 token 对自身无衰减。

    # ── 5. 块内 Delta Rule 并行化：三角递推求解变换矩阵 attn ──
    # 目标：高效并行地在块内实现 delta rule 的"定向擦写"效果。
    # 
    # 原始 delta 操作写成矩阵形式，块内第 t 步的状态依赖之前所有步，
    # 天然是下三角结构。以下用一次矩阵乘法 + 一个 O(C^2) 的迭代修正来完成。
    #
    # 首先构造初始下三角注意力矩阵（严格下三角，diagonal=0 被 mask 掉）：

    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0)
    attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
    # k_beta @ key^T: [B, H, num_chunks, chunk_size, chunk_size]
    # 每个 [t, s] 元素 = -beta_t * (k_t · k_s) * decay(s->t)，表示位置 s 对 t 的"抹除"贡献。
    # masked_fill(mask, 0)：上三角（含对角线）置 0，保证因果性。

    # 迭代修正：处理链式依赖（S_t 依赖 S_{t-1} 依赖 S_{t-2} ...）
    # 每次迭代把"间接路径"（跨 2 步、3 步...的衰减擦写）加到矩阵里，
    # 直到所有路径都被统计到（最多 chunk_size-1 次迭代）。
    for i in range(1, chunk_size):
        row = attn[..., i, :i].clone()       # 第 i 行（i 对前 i 个 token 的直接影响）
        sub = attn[..., :i, :i].clone()      # 左上角子矩阵（前 i 个 token 之间的影响）
        attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
        # row + row @ sub：把"i->j 的直接路径"与"j->k 的间接路径"合并，
        # 类似并行前缀求和，完成下三角系统的递归展开。

    # 加上单位矩阵：对角线代表每个 token 对自身的恒等变换（无擦写时直接写入）
    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
    # attn 此时是块内完整的"Delta变换矩阵"，形状 [B, H, num_chunks, C, C]。

    # ── 6. 用变换矩阵更新 value 和 k_cumdecay ──
    value = attn @ v_beta
    # 将每个位置的 v_beta（beta_t * v_t）经过 delta 变换矩阵聚合，
    # 得到块内已完成 delta rule 修正的有效 value。
    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
    # k_cumdecay[..., t, :] 表示"t 时刻的 key 方向（含衰减与 delta 修正）"，

    # 用于块间传播：告诉下一个 chunk，当前 chunk 末尾对历史状态 S 的影响方向。

    # ── 7. 初始化块间递推状态与输出缓冲 ──
    last_recurrent_state = (
        torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
        if initial_state is None
        else initial_state.to(value)
    )
    # last_recurrent_state: [B, H, d_k, d_v]，跨 chunk 传递的 S 矩阵，与单步递推中完全相同。
    core_attn_out = torch.zeros_like(value)
    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)
    # 重新定义 mask：此处为严格上三角（diagonal=1），用于块内 Q@K^T 的因果掩码。

    # ── 8. 块间循环：对每个 chunk 计算输出并更新跨块状态 ──
    for i in range(0, total_sequence_length // chunk_size):
        q_i = query[:, :, i]       # [B, H, C, d_k]，当前 chunk 的 query
        k_i = key[:, :, i]         # [B, H, C, d_k]
        v_i = value[:, :, i]       # [B, H, C, d_v]，已经过 delta 变换的 value

        # 块内注意力（上一 chunk 内的因果 Q@K^T，加上衰减掩码）
        attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
        # attn[..., t, s] = (q_t · k_s) * decay(s->t)，因果下三角，[B, H, C, C]

        # 历史状态对当前 chunk 每个 token 的贡献（"跨块 key 对历史 S 的读取"）
        v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state
        # v_prime[..., t, :] = k_cumdecay[t] @ S_{prev}
        # 表示：如果当前 token 的 key 方向在历史 S 中已有存储，则 v_prime 是需要"擦除"的部分。

        # 修正后的 value：从 delta 变换后的 v 中减去历史状态的贡献
        v_new = v_i - v_prime
        # v_new 是最终"净写入量"，即在 delta rule 下，实际新增到 S 中的信息。

        # 历史状态直接经过衰减后对当前 Q 的贡献（跨块远程记忆检索）
        attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
        # g[:, :, i, :] 是当前 chunk 每个位置的累积 log-decay（从 chunk 起点到该位置）；
        # exp() 后乘以 q_i，再与 S_{prev} 做矩阵乘，得到历史记忆对当前查询的响应。
        # 注意：这里 q 已经乘上了块内累积衰减，反映了"历史 S 传入本 chunk 后继续衰减"。

        # 最终输出 = 跨块远程记忆 + 块内局部注意力（作用于修正后的 value）
        core_attn_out[:, :, i] = attn_inter + attn @ v_new

        # 更新跨块状态 S（递推到下一个 chunk）：
        # S_{next} = S_{prev} * exp(g_last)                        # 全局衰减（用 chunk 末尾的累积 decay）
        #           + K_i^T @ v_new（按各位置的"剩余衰减"加权）     # 块内净写入
        last_recurrent_state = (
            last_recurrent_state * g[:, :, i, -1, None, None].exp()
            # g[..., -1] 是当前 chunk 末尾的累积 log-decay，exp() 后即 alpha_{chunk_end}，
            # 表示整个 chunk 对历史状态的全局遗忘。
            + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
            # (g_last - g_t) 是从位置 t 到 chunk 末的"剩余衰减"；
            # k_i 乘以此系数后转置，再与 v_new 做外积累加，
            # 等效于把块内每个 token 的写入贡献按其到 chunk 末的衰减正确加权后汇总进 S。
        )

    # ── 9. 后处理：截断 pad、转回原始形状与数据类型 ──
    if not output_final_state:
        last_recurrent_state = None
    core_attn_out = core_attn_out.reshape(
        core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1]
    )
    core_attn_out = core_attn_out[:, :, :sequence_length]   # 去掉右侧补零
    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
    # 转回 [B, L, H, d_v]，恢复原始精度。
    return core_attn_out, last_recurrent_state

# 更新缓存中的递推状态，供下一个 decode 步骤复用
if cache_params is not None:
    cache_params.recurrent_states[self.layer_idx] = last_recurrent_state
# 注：只在有 cache 时写回；训练/prefill 无 cache 则跳过。

z_shape_og = z.shape
# 注：保存 z 的原始形状 [B, L, H, d_v]，后续 reshape 回来时需要。

# 将 core_attn_out 和 z 都压平到二维，满足 RMSNormGated 的输入要求
core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
z = z.reshape(-1, z.shape[-1])
# 注：RMSNormGated 期望输入为 [B*L*H, d_v]，故先把前三维合并。

core_attn_out = self.norm(core_attn_out, z)
# 注：self.norm 是 RMSNormGated，公式为：
#   output = RMSNorm(core_attn_out) * sigmoid(z)
# 其中：
#   - RMSNorm(core_attn_out) 对 attention 输出做归一化，稳定数值范围；
#   - sigmoid(z) 作为输出门控，z 是之前从 qkvz 投影中拆出的门控分支；
#   - 两者逐元素相乘，实现对输出的动态缩放（类似 GLU 结构）。

core_attn_out = core_attn_out.reshape(z_shape_og)
# 注：恢复为 [B, L, H, d_v] 的多头形状。

core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1)
# 注：把头维摊平，变为 [B, L, H*d_v]，准备送入线性投影层。

output = self.out_proj(core_attn_out)
# 注：out_proj 是标准线性层，将 [B, L, H*d_v] 映射回模型维度 [B, L, D]，
# 与标准 Attention 的输出投影完全对称。