A53019. 程序题：分组查询注意力本题旨在考察你对分组查询注意力 (Grouped-Query Attention, GQA) 单步推理的理解与实现能力。题目提供了一个带有 5 个空缺的Python函数 attn_step_gqa，你需要从给定的选项中选择正确答案填入。import math
编程题
知识点
题目描述
程序题：分组查询注意力
本题旨在考察你对分组查询注意力 (Grouped-Query Attention, GQA) 单步推理的理解与实现能力。题目提供了一个带有 5 个空缺的Python函数 attn_step_gqa，你需要从给定的选项中选择正确答案填入。
import math
import torch
from typing import Tuple
def expand_kv_for_gqa(x: torch.Tensor, n_rep: int) ->torch.Tensor:
B, Kv, L, D = x.shape
H = Kv * n_rep
x = x.unsqueeze(1).repeat(1, n_rep, 1, 1, 1) # (B, n_rep,Kv, L, D)
return x.reshape(B, H, L, D)
def attn_step_gqa(
q: torch.Tensor, # (B, H, 1, D)
k_new: torch.Tensor, # (B, Kv, 1, D)
v_new: torch.Tensor, # (B, Kv, 1, D)
k_cache: torch.Tensor, # (B, Kv, Lk, D)
v_cache: torch.Tensor, # (B, Kv, Lk, D)
*, head_rep: int, # H = Kv * head_rep causal_mask: torch.Tensor # (1, 1, 1, Lk+1)；可见位置=1，
不可见=0 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
# 1) 追加 KV 缓存（在 Kv 头域）
k_all = torch.cat([k_cache, k_new], dim=2) # (B, Kv,Lk+1, D)
v_all = torch.cat([v_cache, v_new], dim=2) # (B, Kv,Lk+1, D)
# 2) 将 KV 扩成 H 头，以与 Q 对齐（GQA/MQA 共享）
kH = ____[26]____
vH = expand_kv_for_gqa(v_all, head_rep)
# 3) 注意力打分： (B,H,1,D) @ (B,H,D,L) -> (B,H,1,L)
scale = ____[27]____
logits = ____[28]____
# 4) 掩码与归一化
logits = ____[29]____
attn = torch.softmax(logits, dim=-1)
# 5) 聚合得到输出
out = ____[30]____
# 6) 返回输出与更新后的缓存（供下一步继续用）
return out, k_all, v_all
def example_usage():
"""演示 GQA 的使用方法"""
# 配置参数
B = 2 # batch size
H = 8 # 查询头数
Kv = 2 # 键值头数（GQA 中 Kv < H）
D = 64 # 头维度
Lk = 10 # 缓存序列长度
head_rep = H // Kv # 每个 KV 头的复制次数
# 创建示例张量
q = torch.randn(B, H, 1, D)
k_new = torch.randn(B, Kv, 1, D)
v_new = torch.randn(B, Kv, 1, D)
k_cache = torch.randn(B, Kv, Lk, D)
v_cache = torch.randn(B, Kv, Lk, D)
# 创建因果掩码（允许查看所有之前的位置）
causal_mask = torch.ones(1, 1, 1, Lk + 1)
# 执行 GQA 步骤
out, k_all, v_all = attn_step_gqa(
q, k_new, v_new, k_cache, v_cache,
head_rep=head_rep,
causal_mask=causal_mask
)
print(f"输入形状:")
print(f" q: {q.shape}")
print(f" k_new: {k_new.shape}")
print(f" v_new: {v_new.shape}")
print(f" k_cache: {k_cache.shape}")
print(f" v_cache: {v_cache.shape}")
print(f"\n 输出形状:")
print(f" out: {out.shape}")
print(f" k_all (updated): {k_all.shape}")
print(f" v_all (updated): {v_all.shape}")
if __name__ == "__main__":
example_usage()