gqa: real grouped-query attention (repeat_kv op + both SDPA paths + wiring + tests)
- repeat_kv CUDA kernel: fwd head-block gather, bwd DETERMINISTIC group-sum (each kv head sums its group of query-head grads; no atomics) + Tensor/ops node. - Config gains num_kv_heads (default = n_heads → MHA); wk/wv project to kv_dim; attention() repeat_kv-broadcasts K/V to nh heads before the UNCHANGED composed & flash SDPA → GQA on both paths. group=1 is identity → MHA bit-identical. - --kv-heads flag on train/train_ddp/export_safetensors/greedy_sample; export writes real num_key_value_heads (xserv repeat_kv grouping aligned). - Tests: repeat_kv grad-check (group>1 grad-sum + group=1 identity); model gqa.rs (GQA flash==composed fp32/bf16, group=1 bit-identical to MHA, kv-proj shape); parity_dump+parity.py GQA path (repeat_interleave) via XTRAIN_PARITY_KV_HEADS. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -10,8 +10,15 @@ pub struct Config {
|
||||
pub dim: usize,
|
||||
/// Number of decoder blocks.
|
||||
pub n_layers: usize,
|
||||
/// Number of attention heads.
|
||||
/// Number of attention (query) heads.
|
||||
pub n_heads: usize,
|
||||
/// Number of key/value heads (Phase T15, GQA). Each KV head is shared by a
|
||||
/// group of `n_heads / num_kv_heads` query heads (repeat_kv). Must divide
|
||||
/// `n_heads`. `num_kv_heads == n_heads` (the default) = MHA, bit-identical to
|
||||
/// the pre-T15 path; `num_kv_heads < n_heads` = real grouped-query attention,
|
||||
/// shrinking the K/V projections to `num_kv_heads * head_dim` and exported as a
|
||||
/// real `num_key_value_heads`.
|
||||
pub num_kv_heads: usize,
|
||||
/// Per-head dimension (`dim / n_heads`).
|
||||
pub head_dim: usize,
|
||||
/// SwiGLU hidden width (gate/up project to this, down projects back).
|
||||
@@ -37,6 +44,7 @@ impl Config {
|
||||
dim: n_heads * head_dim,
|
||||
n_layers: 2,
|
||||
n_heads,
|
||||
num_kv_heads: n_heads, // default = MHA
|
||||
head_dim,
|
||||
ffn_hidden: 64,
|
||||
eps: 1e-5,
|
||||
@@ -62,6 +70,7 @@ impl Config {
|
||||
dim: n_heads * head_dim,
|
||||
n_layers,
|
||||
n_heads,
|
||||
num_kv_heads: n_heads, // default = MHA; set via with_kv_heads for GQA
|
||||
head_dim,
|
||||
ffn_hidden,
|
||||
eps: 1e-5,
|
||||
@@ -70,6 +79,27 @@ impl Config {
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the number of K/V heads (Phase T15, GQA). Builder-style so existing
|
||||
/// `from_arch` call sites stay MHA unless they opt in. Asserts `num_kv_heads`
|
||||
/// divides `n_heads`.
|
||||
pub fn with_kv_heads(mut self, num_kv_heads: usize) -> Self {
|
||||
assert!(num_kv_heads > 0, "num_kv_heads must be > 0");
|
||||
assert_eq!(
|
||||
self.n_heads % num_kv_heads,
|
||||
0,
|
||||
"n_heads {} not divisible by num_kv_heads {num_kv_heads}",
|
||||
self.n_heads
|
||||
);
|
||||
self.num_kv_heads = num_kv_heads;
|
||||
self
|
||||
}
|
||||
|
||||
/// KV projection width (`num_kv_heads * head_dim`). For GQA this is smaller than
|
||||
/// `dim`; for MHA it equals `dim`.
|
||||
pub fn kv_dim(&self) -> usize {
|
||||
self.num_kv_heads * self.head_dim
|
||||
}
|
||||
|
||||
/// Transformer-core parameter count: everything except the token embedding and
|
||||
/// the LM head (the two `vocab × dim` tables). This is the figure the scaling
|
||||
/// ladder is sized against — the 50257-vocab embed+lm_head adds a fixed ~25M on
|
||||
@@ -82,7 +112,8 @@ impl Config {
|
||||
pub fn num_params(&self) -> usize {
|
||||
let per_layer = 2 * self.dim // 2 rmsnorm gammas
|
||||
+ 2 * self.head_dim // q/k per-head norm gammas
|
||||
+ 3 * self.dim * self.dim // q/k/v proj
|
||||
+ self.dim * self.dim // q proj [dim,dim]
|
||||
+ 2 * self.dim * self.kv_dim() // k/v proj [dim,kv_dim] (GQA: smaller)
|
||||
+ self.dim * self.dim // out proj
|
||||
+ 2 * self.dim * self.ffn_hidden // gate/up proj
|
||||
+ self.ffn_hidden * self.dim; // down proj
|
||||
|
||||
@@ -13,8 +13,8 @@ use xtrain_tensor::{DType, Device, Tensor};
|
||||
struct Block {
|
||||
attn_norm: Var, // [dim]
|
||||
wq: Var, // [dim, dim]
|
||||
wk: Var, // [dim, dim]
|
||||
wv: Var, // [dim, dim]
|
||||
wk: Var, // [dim, kv_dim] — kv_dim = num_kv_heads·head_dim (GQA; = dim for MHA)
|
||||
wv: Var, // [dim, kv_dim]
|
||||
q_norm: Var, // [head_dim] — per-head QK-norm (Qwen3-style)
|
||||
k_norm: Var, // [head_dim]
|
||||
wo: Var, // [dim, dim]
|
||||
@@ -91,8 +91,9 @@ impl TinyTransformer {
|
||||
.map(|_| Block {
|
||||
attn_norm: mk(&[cfg.dim]),
|
||||
wq: mk(&[cfg.dim, cfg.dim]),
|
||||
wk: mk(&[cfg.dim, cfg.dim]),
|
||||
wv: mk(&[cfg.dim, cfg.dim]),
|
||||
// GQA (T15): K/V project to num_kv_heads·head_dim (= dim when MHA).
|
||||
wk: mk(&[cfg.dim, cfg.kv_dim()]),
|
||||
wv: mk(&[cfg.dim, cfg.kv_dim()]),
|
||||
q_norm: mk(&[cfg.head_dim]),
|
||||
k_norm: mk(&[cfg.head_dim]),
|
||||
wo: mk(&[cfg.dim, cfg.dim]),
|
||||
@@ -435,37 +436,47 @@ fn attention(
|
||||
wo: &Var,
|
||||
) -> Var {
|
||||
let (nh, hd) = (cfg.n_heads, cfg.head_dim);
|
||||
let num_kv = cfg.num_kv_heads; // GQA (T15): K/V have fewer heads than Q
|
||||
let total = batch * seq;
|
||||
let bh = batch * nh;
|
||||
let scale = 1.0 / (hd as f32).sqrt();
|
||||
|
||||
// Project, qk-norm + RoPE, then lay out as a batched [B*nh, seq, hd] tensor.
|
||||
// [B*S,dim] @ [dim,dim] = [B*S,dim]
|
||||
// reshape [B*S, nh, hd]
|
||||
// Project, qk-norm + RoPE, then lay out as a batched [B*heads, seq, hd] tensor.
|
||||
// `heads` = nh for Q, num_kv for K/V (GQA; equal for MHA).
|
||||
// [B*S,dim] @ [dim,heads*hd] = [B*S, heads*hd]
|
||||
// reshape [B*S, heads, hd]
|
||||
// qk-norm per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
|
||||
// rope [B*S, nh, hd] with per-sequence position (period = seq)
|
||||
// reshape [B, S, nh, hd] → transpose(1,2) → [B, nh, S, hd] → [B*nh, S, hd]
|
||||
let to_bh = |proj: Var, norm: Option<&Var>| -> Var {
|
||||
let r = ops::reshape(&proj, &[total, nh, hd]);
|
||||
// rope [B*S, heads, hd] with per-sequence position (period = seq)
|
||||
// reshape [B, S, heads, hd] → transpose(1,2) → [B, heads, S, hd] → [B*heads, S, hd]
|
||||
let to_bh = |proj: Var, heads: usize, norm: Option<&Var>| -> Var {
|
||||
let r = ops::reshape(&proj, &[total, heads, hd]);
|
||||
let r = match norm {
|
||||
// Per-head RMSNorm: flatten the (B*S,nh) head rows, norm over hd,
|
||||
// Per-head RMSNorm: flatten the (B*S,heads) head rows, norm over hd,
|
||||
// restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
|
||||
Some(gamma) => {
|
||||
let flat = ops::reshape(&r, &[total * nh, hd]);
|
||||
let flat = ops::reshape(&r, &[total * heads, hd]);
|
||||
let normed = ops::rms_norm(&flat, &norm_gamma(cdt, gamma), cfg.eps);
|
||||
let r = ops::reshape(&normed, &[total, nh, hd]);
|
||||
let r = ops::reshape(&normed, &[total, heads, hd]);
|
||||
ops::rope(&r, cfg.rope_theta, seq)
|
||||
}
|
||||
None => r,
|
||||
};
|
||||
let r = ops::reshape(&r, &[batch, seq, nh, hd]);
|
||||
let t = ops::transpose_4d12(&r); // [B, nh, S, hd]
|
||||
ops::reshape(&t, &[bh, seq, hd]) // [B*nh, S, hd]
|
||||
let r = ops::reshape(&r, &[batch, seq, heads, hd]);
|
||||
let t = ops::transpose_4d12(&r); // [B, heads, S, hd]
|
||||
ops::reshape(&t, &[batch * heads, seq, hd]) // [B*heads, S, hd]
|
||||
};
|
||||
|
||||
let q = to_bh(linear(cdt, x, wq), Some(q_norm));
|
||||
let k = to_bh(linear(cdt, x, wk), Some(k_norm));
|
||||
let v = to_bh(linear(cdt, x, wv), None);
|
||||
let q = to_bh(linear(cdt, x, wq), nh, Some(q_norm));
|
||||
// K/V are laid out with num_kv heads, then repeat_kv-broadcast to nh heads so
|
||||
// the SDPA below (composed or flash, both unchanged) sees a full head set. The
|
||||
// broadcast's backward sums each KV head's group of query-head grads (GQA). For
|
||||
// MHA (num_kv == nh) repeat_kv is identity → bit-identical to the pre-T15 path.
|
||||
let k = to_bh(linear(cdt, x, wk), num_kv, Some(k_norm));
|
||||
let v = to_bh(linear(cdt, x, wv), num_kv, None);
|
||||
let (k, v) = if num_kv == nh {
|
||||
(k, v)
|
||||
} else {
|
||||
(ops::repeat_kv(&k, nh, batch), ops::repeat_kv(&v, nh, batch))
|
||||
};
|
||||
|
||||
// Causal SDPA over all B*nh (sequence,head) blocks. `flash` (T14) picks the
|
||||
// single fused flash kernel (online softmax, no materialized [bh,S,S] scores);
|
||||
|
||||
Reference in New Issue
Block a user