gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gqa: real grouped-query attention (repeat_kv op + both SDPA paths + wiring + tests)

Browse Source 
- 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:
Gahow Wang
2026-06-18 01:37:16 +08:00
parent 62b1cb5dc7
commit 830d06ad01
15 changed files with 712 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
crates/xtrain-model/src/model.rs
View File

@@ -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);