self.config = config
self.hidden_size = config.hidden_size # 隐藏层维度
self.num_v_heads = config.linear_num_value_heads # v的头数
self.num_k_heads = config.linear_num_key_heads # k的头数
self.head_k_dim = config.linear_key_head_dim # 单头k的维度
self.head_v_dim = config.linear_value_head_dim # 单头v的维度
self.key_dim = self.head_k_dim * self.num_k_heads # 总的k的维度
self.value_dim = self.head_v_dim * self.num_v_heads # 总的v的维度

projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2 # qkvz投影
projection_size_ba = self.num_v_heads * 2 # ba投影

self.in_proj_qkvz = nn.Linear(self.hidden_size, projection_size_qkvz, bias=False)
self.in_proj_ba = nn.Linear(self.hidden_size, projection_size_ba, bias=False)

# [b, s, key_dim * 2 + value_dim * 2]
projected_states_qkvz = self.in_proj_qkvz(hidden_states)
# [b, s, num_v_heads * 2]
projected_states_ba = self.in_proj_ba(hidden_states)

# query, key:[b, s, key_dim]
# value, z:[b, s, value_dim]
query, key, value, z = torch.split(
    projected_states_qkvz,
    [self.key_dim, self.key_dim, self.value_dim, self.value_dim],
    dim=-1)

# b,a:[b, s, num_v_heads]
b, a = torch.split(
    projected_states_ba,
    [self.num_v_heads, self.num_v_heads],
    dim=-1)

# 拼接qkv
mixed_qkv = torch.cat((query, key, value), dim=-1) # [b,s,key_dim*2 + value_dim]
# 方便做卷积，把sequence放到最后面
mixed_qkv = mixed_qkv.transpose(1, 2) # [b, key_di=*2 + value_dim, s]
# 卷积实现+silu激活
mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :seq_len])
# conv之后再交换回来
mixed_qkv = mixed_qkv.transpose(1, 2) # [b, s, key_dim*2 + value_dim]

conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
cache_params.conv_states[self.layer_idx] = conv_state
mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :s])

query, key, value = torch.split(
    mixed_qkv,
    [
        self.key_dim,
        self.key_dim,
        self.value_dim,
    ],
    dim=-1,
)
query = query.reshape(bs, s, -1, self.head_k_dim) # [b,s,num_k_heads, head_k_dim]
key = key.reshape(bs, s, -1, self.head_k_dim) # [b,s,num_k_heads, head_k_dim]
value = value.reshape(bs, s, -1, self.head_v_dim) # [b,s,num_v_heads, head_v_dim]
# 类似于GQA,保证head数一致,此时num_v_heads=num_k_heads
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)

# 处理alpha 得到遗忘门$g$
self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads)) # [num_heads]
A = torch.empty(self.num_v_heads).uniform_(0, 16) # [num_heads]
self.A_log = nn.Parameter(torch.log(A)) # [num_heads]
g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) # [b, s, num_heads]

# 处理beta
beta = b.sigmoid()

def torch_recurrent_gated_delta_rule(
    query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False
):
    initial_dtype = query.dtype
    # 对qk做l2norm标准化
    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)
    ]


    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)
    )

    # 遍历整个序列长度
    for i in range(sequence_length):
        # 取出当前t时刻的qkvgb
        q_t = query[:, :, i]
        k_t = key[:, :, i]
        v_t = value[:, :, i]
        g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1)
        beta_t = beta[:, :, i].unsqueeze(-1)

        # 对应到下面的门控的注意力计算公式
        last_recurrent_state = last_recurrent_state * g_t # 计算aS_{t-1}
        kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2) # 计算aS_{t-1}k
        delta = (v_t - kv_mem) * beta_t # 计算 (v-aS_{t-1}k)b
        last_recurrent_state = last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2) # 计算完整的
        core_attn_out[:, :, i] = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2)

    if not output_final_state:
        last_recurrent_state = None
    # [b,s,num_heads,head_v_dim]
    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
    return core_attn_out, last_recurrent_state

def forward_native(
    self,
    x: torch.Tensor,
    residual: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
    orig_dtype = x.dtype
    if residual is not None:
        x = x + residual
        residual = x

    x = x.float()
    variance = x.pow(2).mean(dim=-1, keepdim=True)
    x = x * torch.rsqrt(variance + self.variance_epsilon)
    x = x * (1.0 + self.weight.float()) # 唯一区别，这里weight加上了1
    x = x.to(orig_dtype)
    return x if residual is None else (x, residual)