bsz, q_len, _ = hidden_states.size()
        
# 計(jì)算壓縮后的Q，再還原成高維
# [B, q_len, hidden_size]
# 即[B, q_len, num_head * q_head_dim]
q = self.w_uq(self.q_a_layernorm(self.w_dq(hidden_states)))
# [B, num_head, q_len, q_head_dim]
q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)
# 包含當(dāng)前位置可用上下文的長(zhǎng)度
kv_seq_len = q.size(-2)
if past_key_value is not None:
    if self.layer_idx is None:
        raise ValueError(
            f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
            "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
            "with a layer index."
        )
    kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
# 得到當(dāng)前壓縮后的kv, c_t^{kv}
# [B, q_len, d_c]
compressed_kv = self.w_dkv(hidden_states)

# 將當(dāng)前位置之前的壓縮后的kv拼接到前面
if past_key_value is not None:
    # 得到的應(yīng)該是[B, kv_seq_len, d_c], c^{kv}
    compressed_kv = past_key_value.update(compressed_kv)
# 計(jì)算得到k^C和v^C
# [B, num_head, kv_seq_len, q_head_dim]
k = self.w_uk(compressed_kv).view(bsz, -1, self.num_heads, self.q_head_dim).transpose(1, 2)
v = self.w_uv(compressed_kv).view(bsz, -1, self.num_heads, self.q_head_dim).transpose(1, 2)

# 注意力權(quán)重
# [B, num_head, q_len, kv_seq_len]
attn_weights = (
    torch.matmul(q, k.transpose(2, 3)) * self.softmax_scale
)
...
attn_weights = nn.functional.softmax(
    attn_weights, dim=-1, dtype=torch.float32
).to(query_states.dtype)
attn_weights = nn.functional.dropout(
    attn_weights, p=self.attention_dropout, training=self.training
)
# [B, num_head, q_len, q_head_dim]
attn_output = torch.matmul(attn_weights, v)
...