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:
@@ -376,6 +376,27 @@ pub fn flash_attention(q: &Var, k: &Var, v: &Var, scale: f32) -> Var {
|
||||
)
|
||||
}
|
||||
|
||||
/// GQA repeat_kv head broadcast (Phase T15). `kv`:[batch·num_kv, seq, head_dim]
|
||||
/// (a K or V tensor) → `[batch·nh, seq, head_dim]`, each KV head broadcast to its
|
||||
/// `group = nh/num_kv` query heads (qh ← kv head qh/group, contiguous groups —
|
||||
/// matches xserv's repeat_kv). Feeds the unchanged composed/flash SDPA so GQA is
|
||||
/// "free" for both. Backward SUMS the `group` query heads sharing each KV head back
|
||||
/// onto it (the multi-group grad accumulation). `nh == num_kv` (group 1) is identity
|
||||
/// → bit-identical to the MHA path. `batch` lets the op recover num_kv from kv's bh.
|
||||
pub fn repeat_kv(kv: &Var, nh: usize, batch: usize) -> Var {
|
||||
let bh_kv = kv.value().shape()[0];
|
||||
let num_kv = bh_kv / batch;
|
||||
let out = kv.value().repeat_kv(nh, batch);
|
||||
Var::from_op(
|
||||
out,
|
||||
vec![kv.clone()],
|
||||
Box::new(move |dout, parents| {
|
||||
let din = Tensor::repeat_kv_backward(dout, num_kv, batch);
|
||||
Var::push_grad(&parents[0], din);
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
/// Cross-entropy mean loss over logits `x:[rows,cols]` with one I32 target per
|
||||
/// row. Returns a scalar [`Var`]. Backward: `dx = (probs - onehot)/rows`,
|
||||
/// scaled by the upstream scalar grad.
|
||||
|
||||
Reference in New Issue
Block a user