Grouped-Query & Multi-Query Attention

• Ensure proper multi-dimensional tensor broadcasting and keepdims retention.
• Enforce numerical guards (eps clamping and overflow thresholds) during exponentiation and division.
• Verify train versus eval mode behavioral distinctions (e.g. frozen running statistics and dropout bypass).

💻 Industrial Code Implementation

import numpy as np

def repeat_kv(x: np.ndarray, num_rep: int) -> np.ndarray:
    """
    将 KV 头重复展开以对齐 Query 头的数量。
    x 形状: (B, H_kv, S, D_k)
    返回: (B, H_q, S, D_k)，其中 H_q = H_kv * num_rep
    """
    if num_rep == 1:
        return x
    B, H_kv, S, D_k = x.shape
    # 增加扩展维度并广播
    x = x[:, :, np.newaxis, :, :]  # (B, H_kv, 1, S, D_k)
    x = np.repeat(x, num_rep, axis=2)  # (B, H_kv, num_rep, S, D_k)
    return x.reshape(B, H_kv * num_rep, S, D_k)

def grouped_query_attention(
    q: np.ndarray,      # (B, H_q, S, D_k)
    k: np.ndarray,      # (B, H_kv, S, D_k)
    v: np.ndarray,      # (B, H_kv, S, D_k)
    is_causal: bool = True
) -> np.ndarray:
    B, H_q, S, D_k = q.shape
    H_kv = k.shape[1]
    assert H_q % H_kv == 0, "Query 头数必须是 KV 头数的整数倍"
    num_rep = H_q // H_kv
    
    # 将 KV 广播扩展为与 Q 头数相同
    k_expanded = repeat_kv(k, num_rep)  # (B, H_q, S, D_k)
    v_expanded = repeat_kv(v, num_rep)  # (B, H_q, S, D_k)
    
    # 执行常规多头注意力点积
    scores = np.matmul(q, k_expanded.swapaxes(-1, -2)) / np.sqrt(D_k)
    if is_causal:
        mask = np.triu(np.ones((S, S), dtype=bool), k=1)
        scores = np.where(mask, -1e9, scores)
    
    scores_max = np.max(scores, axis=-1, keepdims=True)
    probs = np.exp(scores - scores_max)
    probs /= np.sum(probs, axis=-1, keepdims=True)
    return np.matmul(probs, v_expanded)

import numpy as np
B, S, D_k = 2, 4, 8
H_q = 8
H_kv = 2  # 4 个 Q 头共享 1 个 KV 头
q = np.random.randn(B, H_q, S, D_k)
k = np.random.randn(B, H_kv, S, D_k)
v = np.random.randn(B, H_kv, S, D_k)
out = grouped_query_attention(q, k, v)
assert out.shape == (B, H_q, S, D_k), "输出形状不正确"
print("✓ GQA 分组查询注意力自测通过")