程序题：分组查询注意力

本题旨在考察你对分组查询注意力 (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()