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);
|
||||
|
||||
268
crates/xtrain-model/tests/gqa.rs
Normal file
268
crates/xtrain-model/tests/gqa.rs
Normal file
@@ -0,0 +1,268 @@
|
||||
// T15 GQA correctness gate. Real grouped-query attention (num_kv_heads <
|
||||
// num_heads): K/V project to num_kv_heads·head_dim and are repeat_kv-broadcast to
|
||||
// the full head set before the SDPA. This test pins three things:
|
||||
//
|
||||
// 1. GQA flash == GQA composed (forward logits + loss + EVERY param grad) — the
|
||||
// repeat_kv broadcast feeds both SDPA paths unchanged, so they must agree; in
|
||||
// particular the wk/wv grads (which flow back through repeat_kv's group-sum)
|
||||
// must match. Parameterised over fp32 (tight) and bf16 (rounding band).
|
||||
// 2. group==1 (num_kv_heads == n_heads) is BIT-IDENTICAL to the pre-T15 MHA path
|
||||
// (a model with num_kv_heads explicitly == n_heads vs the default config):
|
||||
// forward logits + every grad |Δ|=0. The regression guard.
|
||||
// 3. wk/wv really shrank to [dim, kv_dim] under GQA (shape check).
|
||||
#![cfg(not(no_cuda))]
|
||||
|
||||
use xtrain_cuda::device;
|
||||
use xtrain_model::{Config, TinyTransformer, batched_ids_tensor};
|
||||
use xtrain_tensor::{DType, Device};
|
||||
|
||||
fn fill(n: usize, seed: u64, scale: f32) -> Vec<f32> {
|
||||
let mut state = seed
|
||||
.wrapping_mul(2862933555777941757)
|
||||
.wrapping_add(3037000493);
|
||||
(0..n)
|
||||
.map(|_| {
|
||||
state = state
|
||||
.wrapping_mul(6364136223846793005)
|
||||
.wrapping_add(1442695040888963407);
|
||||
(((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5) * 2.0 * scale
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn build(cfg: Config, device: Device, dtype: DType, flash: bool) -> TinyTransformer {
|
||||
let mut seed = 1u64;
|
||||
let m = TinyTransformer::new(cfg, device, |shape| {
|
||||
seed = seed.wrapping_add(1);
|
||||
let n: usize = shape.iter().product();
|
||||
if shape.len() == 1 {
|
||||
fill(n, seed, 0.02).iter().map(|v| v + 1.0).collect()
|
||||
} else {
|
||||
fill(n, seed, 0.08)
|
||||
}
|
||||
});
|
||||
m.with_compute_dtype(dtype).with_flash(flash)
|
||||
}
|
||||
|
||||
fn host(t: &xtrain_tensor::Tensor) -> Vec<f32> {
|
||||
t.to_dtype(DType::F32)
|
||||
.to_device(Device::Cpu)
|
||||
.as_slice::<f32>()
|
||||
.to_vec()
|
||||
}
|
||||
|
||||
// A real GQA config: 8 query heads, 2 kv heads → group 4. seq=40 > FA_TILE=32 so
|
||||
// the flash online-softmax tile path is exercised too.
|
||||
fn gqa_cfg() -> Config {
|
||||
let mut cfg = Config::tiny();
|
||||
cfg.vocab = 16;
|
||||
cfg.n_layers = 3;
|
||||
// tiny() is 2 heads; rebuild with 8 query / 2 kv heads keeping head_dim=16.
|
||||
Config::from_arch(cfg.vocab, 8, cfg.head_dim, cfg.n_layers, cfg.ffn_hidden).with_kv_heads(2)
|
||||
}
|
||||
|
||||
fn ids_targets(cfg: &Config, batch: usize, seq: usize) -> (Vec<Vec<i32>>, Vec<Vec<i32>>) {
|
||||
let seqs = (0..batch)
|
||||
.map(|b| {
|
||||
(0..seq)
|
||||
.map(|i| ((b * 7 + i * 3 + 1) % cfg.vocab) as i32)
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
let tgts = (0..batch)
|
||||
.map(|b| {
|
||||
(0..seq)
|
||||
.map(|i| ((b * 5 + i * 2 + 2) % cfg.vocab) as i32)
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
(seqs, tgts)
|
||||
}
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
fn run_both(
|
||||
cfg: Config,
|
||||
dtype: DType,
|
||||
) -> (Vec<f32>, f32, Vec<Vec<f32>>, Vec<f32>, f32, Vec<Vec<f32>>) {
|
||||
device::set_device(0).unwrap();
|
||||
let device = Device::Cuda(0);
|
||||
let (batch, seq) = (3usize, 40usize);
|
||||
let (seqs, tgts) = ids_targets(&cfg, batch, seq);
|
||||
let ids = batched_ids_tensor(&seqs, device);
|
||||
let tgt = batched_ids_tensor(&tgts, device);
|
||||
|
||||
let off = build(cfg, device, dtype, false);
|
||||
let off_logits = host(&off.forward_batched(&ids, batch).value());
|
||||
let off_loss = off.loss_batched(&ids, &tgt, batch);
|
||||
let off_loss_val = host(&off_loss.value())[0];
|
||||
off_loss.backward();
|
||||
let off_grads: Vec<Vec<f32>> = off
|
||||
.params()
|
||||
.iter()
|
||||
.map(|p| host(&p.grad().expect("off grad")))
|
||||
.collect();
|
||||
|
||||
let on = build(cfg, device, dtype, true);
|
||||
let on_logits = host(&on.forward_batched(&ids, batch).value());
|
||||
let on_loss = on.loss_batched(&ids, &tgt, batch);
|
||||
let on_loss_val = host(&on_loss.value())[0];
|
||||
on_loss.backward();
|
||||
let on_grads: Vec<Vec<f32>> = on
|
||||
.params()
|
||||
.iter()
|
||||
.map(|p| host(&p.grad().expect("on grad")))
|
||||
.collect();
|
||||
|
||||
(
|
||||
off_logits,
|
||||
off_loss_val,
|
||||
off_grads,
|
||||
on_logits,
|
||||
on_loss_val,
|
||||
on_grads,
|
||||
)
|
||||
}
|
||||
|
||||
// GQA flash vs composed: same SDPA math on the same repeat_kv-broadcast K/V → fp32
|
||||
// agrees to reduction-order, bf16 to its rounding band.
|
||||
#[test]
|
||||
fn gqa_flash_matches_composed_fp32() {
|
||||
assert!(device::device_count().unwrap() > 0, "no CUDA device");
|
||||
let cfg = gqa_cfg();
|
||||
assert!(cfg.num_kv_heads < cfg.n_heads, "test must be real GQA");
|
||||
let (off_l, off_loss, off_g, on_l, on_loss, on_g) = run_both(cfg, DType::F32);
|
||||
|
||||
let logit_rel = off_l
|
||||
.iter()
|
||||
.zip(&on_l)
|
||||
.map(|(a, b)| (a - b).abs() / a.abs().max(1e-4))
|
||||
.fold(0.0f32, f32::max);
|
||||
let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
|
||||
println!(
|
||||
"[GQA F32] flash on/off: loss {off_loss:.6}/{on_loss:.6} (rel {loss_rel:.2e}), \
|
||||
logits max rel {logit_rel:.2e}"
|
||||
);
|
||||
assert!(
|
||||
logit_rel < 1e-3,
|
||||
"[GQA F32] logits diverged: {logit_rel:.2e}"
|
||||
);
|
||||
assert!(loss_rel < 1e-3, "[GQA F32] loss diverged: {loss_rel:.2e}");
|
||||
|
||||
let mut worst = 0.0f32;
|
||||
for (a_g, b_g) in off_g.iter().zip(&on_g) {
|
||||
for (a, b) in a_g.iter().zip(b_g) {
|
||||
worst = worst.max((a - b).abs() / a.abs().max(1e-3));
|
||||
}
|
||||
}
|
||||
println!("[GQA F32] flash on/off grad max rel = {worst:.3e}");
|
||||
assert!(worst < 2e-2, "[GQA F32] grads diverged: {worst:.3e}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gqa_flash_matches_composed_bf16() {
|
||||
assert!(device::device_count().unwrap() > 0, "no CUDA device");
|
||||
let (off_l, off_loss, off_g, on_l, on_loss, on_g) = run_both(gqa_cfg(), DType::BF16);
|
||||
|
||||
let loss_rel = (off_loss - on_loss).abs() / off_loss.abs().max(1e-4);
|
||||
println!("[GQA BF16] flash on/off: loss {off_loss:.5}/{on_loss:.5} (rel {loss_rel:.3e})");
|
||||
assert!(loss_rel < 2e-2, "[GQA BF16] loss diverged: {loss_rel:.3e}");
|
||||
|
||||
let n = off_l.len();
|
||||
let mut rels: Vec<f32> = off_l
|
||||
.iter()
|
||||
.zip(&on_l)
|
||||
.map(|(f, b)| (b - f).abs() / f.abs().max(1.0))
|
||||
.collect();
|
||||
rels.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
||||
let mean: f32 = rels.iter().sum::<f32>() / n as f32;
|
||||
let p99 = rels[(n as f32 * 0.99) as usize];
|
||||
println!("[GQA BF16] logits: mean rel {mean:.3e}, p99 rel {p99:.3e}");
|
||||
assert!(
|
||||
mean < 1e-2,
|
||||
"[GQA BF16] logits mean rel too high: {mean:.3e}"
|
||||
);
|
||||
assert!(p99 < 5e-2, "[GQA BF16] logits p99 rel too high: {p99:.3e}");
|
||||
|
||||
let mut worst = 0.0f32;
|
||||
for (a_g, b_g) in off_g.iter().zip(&on_g) {
|
||||
let scale = a_g.iter().map(|v| v.abs()).fold(0.0f32, f32::max).max(1e-6);
|
||||
let mean_err: f32 =
|
||||
a_g.iter().zip(b_g).map(|(f, b)| (f - b).abs()).sum::<f32>() / a_g.len() as f32 / scale;
|
||||
worst = worst.max(mean_err);
|
||||
}
|
||||
println!("[GQA BF16] grads: worst per-tensor scaled-mean err = {worst:.3e}");
|
||||
assert!(worst < 3e-2, "[GQA BF16] grads diverged: {worst:.3e}");
|
||||
}
|
||||
|
||||
// REGRESSION GUARD: num_kv_heads == n_heads (group 1) must be BIT-IDENTICAL to the
|
||||
// pre-T15 MHA path. Build one model with the default config (num_kv_heads ==
|
||||
// n_heads, the untouched path: repeat_kv not even invoked) and one that explicitly
|
||||
// sets num_kv_heads = n_heads, then assert forward logits + every grad match to the
|
||||
// bit. (Same composed path, so this is exact equality, not a tolerance.)
|
||||
#[test]
|
||||
fn gqa_group1_bit_identical_to_mha() {
|
||||
assert!(device::device_count().unwrap() > 0, "no CUDA device");
|
||||
device::set_device(0).unwrap();
|
||||
let device = Device::Cuda(0);
|
||||
|
||||
let mut base = Config::tiny();
|
||||
base.vocab = 16;
|
||||
base.n_layers = 3;
|
||||
let base = Config::from_arch(base.vocab, 4, base.head_dim, base.n_layers, base.ffn_hidden);
|
||||
// `explicit` sets num_kv_heads = n_heads (already the default, but exercises the
|
||||
// with_kv_heads path); they are the same config → must produce identical output.
|
||||
let explicit = base.with_kv_heads(base.n_heads);
|
||||
assert_eq!(base.num_kv_heads, explicit.num_kv_heads);
|
||||
|
||||
let (batch, seq) = (2usize, 8usize);
|
||||
let (seqs, tgts) = ids_targets(&base, batch, seq);
|
||||
let ids = batched_ids_tensor(&seqs, device);
|
||||
let tgt = batched_ids_tensor(&tgts, device);
|
||||
|
||||
let run = |cfg: Config| -> (Vec<f32>, f32, Vec<Vec<f32>>) {
|
||||
let m = build(cfg, device, DType::F32, false);
|
||||
let logits = host(&m.forward_batched(&ids, batch).value());
|
||||
let loss = m.loss_batched(&ids, &tgt, batch);
|
||||
let loss_v = host(&loss.value())[0];
|
||||
loss.backward();
|
||||
let grads = m
|
||||
.params()
|
||||
.iter()
|
||||
.map(|p| host(&p.grad().unwrap()))
|
||||
.collect();
|
||||
(logits, loss_v, grads)
|
||||
};
|
||||
let (la, sa, ga) = run(base);
|
||||
let (lb, sb, gb) = run(explicit);
|
||||
assert_eq!(la, lb, "group-1 logits must be bit-identical to MHA");
|
||||
assert_eq!(sa, sb, "group-1 loss must be bit-identical to MHA");
|
||||
for (a, b) in ga.iter().zip(&gb) {
|
||||
assert_eq!(a, b, "group-1 grad must be bit-identical to MHA");
|
||||
}
|
||||
println!("[GQA group1] bit-identical to MHA: logits + loss + all grads |Δ|=0");
|
||||
}
|
||||
|
||||
// Under GQA, wk/wv must be [dim, kv_dim] (= num_kv_heads·head_dim), wq stays [dim,dim].
|
||||
#[test]
|
||||
fn gqa_kv_proj_shape() {
|
||||
assert!(device::device_count().unwrap() > 0, "no CUDA device");
|
||||
device::set_device(0).unwrap();
|
||||
let device = Device::Cuda(0);
|
||||
let cfg = gqa_cfg();
|
||||
let m = build(cfg, device, DType::F32, false);
|
||||
let p = m.params();
|
||||
// params order: embed, then per block [attn_norm, wq, wk, wv, q_norm, k_norm, wo, ...]
|
||||
let wq = p[1].value().shape().to_vec();
|
||||
let wk = p[2].value().shape().to_vec();
|
||||
let wv = p[3].value().shape().to_vec();
|
||||
assert_eq!(wq, vec![cfg.dim, cfg.dim], "wq must be [dim,dim]");
|
||||
assert_eq!(wk, vec![cfg.dim, cfg.kv_dim()], "wk must be [dim,kv_dim]");
|
||||
assert_eq!(wv, vec![cfg.dim, cfg.kv_dim()], "wv must be [dim,kv_dim]");
|
||||
println!(
|
||||
"[GQA shapes] wq {:?} wk {:?} wv {:?} (kv_dim {})",
|
||||
wq,
|
||||
wk,
|
||||
wv,
|
||||
cfg.kv_dim()
|
||||
);
|
||||
}
|
||||
@@ -52,6 +52,10 @@ cfg = read_cfg()
|
||||
DIM = int(cfg["dim"])
|
||||
NL = int(cfg["n_layers"])
|
||||
NH = int(cfg["n_heads"])
|
||||
# GQA (T15): num_kv_heads <= n_heads; each kv head shared by group query heads.
|
||||
# Default to NH (MHA) for fixtures dumped before the field existed.
|
||||
NKV = int(cfg.get("num_kv_heads", str(NH)))
|
||||
GROUP = NH // NKV
|
||||
HD = int(cfg["head_dim"])
|
||||
EPS = float(cfg["eps"])
|
||||
THETA = float(cfg["rope_theta"])
|
||||
@@ -114,17 +118,23 @@ for L in layers:
|
||||
# Attention
|
||||
x = rms_norm(h, L["attn_norm"])
|
||||
q = (x @ L["wq"]).reshape(B * SEQ, NH, HD)
|
||||
k = (x @ L["wk"]).reshape(B * SEQ, NH, HD)
|
||||
v = (x @ L["wv"]).reshape(B * SEQ, NH, HD)
|
||||
# GQA: K/V project to NKV heads, then repeat each kv head GROUP times to NH.
|
||||
k = (x @ L["wk"]).reshape(B * SEQ, NKV, HD)
|
||||
v = (x @ L["wv"]).reshape(B * SEQ, NKV, HD)
|
||||
# Per-head QK-norm (Qwen3-style), before RoPE.
|
||||
q = rms_norm(q, L["q_norm"])
|
||||
k = rms_norm(k, L["k_norm"])
|
||||
q = rope(q) # [B*SEQ, nh, hd]
|
||||
k = rope(k)
|
||||
# Reshape to [B, NH, SEQ, HD] so attention runs within each sequence.
|
||||
k = rope(k) # [B*SEQ, nkv, hd]
|
||||
# Reshape to [B, *, SEQ, HD]; broadcast kv heads to NH (repeat_interleave along
|
||||
# the head axis: kv head kvh → query heads [kvh*GROUP, (kvh+1)*GROUP), matching
|
||||
# xtrain repeat_kv + xserv repeat_kv).
|
||||
q = q.reshape(B, SEQ, NH, HD).transpose(1, 2) # [B, nh, seq, hd]
|
||||
k = k.reshape(B, SEQ, NH, HD).transpose(1, 2)
|
||||
v = v.reshape(B, SEQ, NH, HD).transpose(1, 2)
|
||||
k = k.reshape(B, SEQ, NKV, HD).transpose(1, 2) # [B, nkv, seq, hd]
|
||||
v = v.reshape(B, SEQ, NKV, HD).transpose(1, 2)
|
||||
if GROUP > 1:
|
||||
k = k.repeat_interleave(GROUP, dim=1) # [B, nh, seq, hd]
|
||||
v = v.repeat_interleave(GROUP, dim=1)
|
||||
scale = 1.0 / math.sqrt(HD)
|
||||
scores = (q @ k.transpose(-1, -2)) * scale + mask # [B, nh, seq, seq]
|
||||
probs = torch.softmax(scores, dim=-1)
|
||||
|
||||
@@ -58,8 +58,20 @@ fn dump_for_parity() {
|
||||
// sequence-major to [B*S]=8 ids. Per-sequence RoPE position (resets at the
|
||||
// sequence boundary) + per-sequence causal masking (no cross-sequence
|
||||
// attention) are both checked against PyTorch.
|
||||
// Default: tiny MHA (2 heads). With XTRAIN_PARITY_KV_HEADS=k set, dump a real
|
||||
// GQA config (8 query heads / k kv heads) so parity.py checks GQA at B>1 — the
|
||||
// kv-projection shapes + the repeat_kv group-sum backward against PyTorch.
|
||||
let mut cfg = Config::tiny();
|
||||
cfg.vocab = 12;
|
||||
if let Ok(kv) = std::env::var("XTRAIN_PARITY_KV_HEADS") {
|
||||
let kv: usize = kv.parse().expect("XTRAIN_PARITY_KV_HEADS");
|
||||
cfg = Config::from_arch(cfg.vocab, 8, cfg.head_dim, cfg.n_layers, cfg.ffn_hidden)
|
||||
.with_kv_heads(kv);
|
||||
println!(
|
||||
"parity: GQA config (n_heads {} kv_heads {})",
|
||||
cfg.n_heads, cfg.num_kv_heads
|
||||
);
|
||||
}
|
||||
let batch = 2usize;
|
||||
let seq = 4usize;
|
||||
let ids: Vec<i32> = vec![3, 1, 4, 1, 5, 9, 2, 6]; // [B*S], sequence-major
|
||||
@@ -92,6 +104,7 @@ fn dump_for_parity() {
|
||||
writeln!(f, "dim {}", cfg.dim).unwrap();
|
||||
writeln!(f, "n_layers {}", cfg.n_layers).unwrap();
|
||||
writeln!(f, "n_heads {}", cfg.n_heads).unwrap();
|
||||
writeln!(f, "num_kv_heads {}", cfg.num_kv_heads).unwrap();
|
||||
writeln!(f, "head_dim {}", cfg.head_dim).unwrap();
|
||||
writeln!(f, "ffn_hidden {}", cfg.ffn_hidden).unwrap();
|
||||
writeln!(f, "eps {:e}", cfg.eps).unwrap();
|
||||
|
||||
Reference in New Issue
Block a user