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:
@@ -37,6 +37,7 @@ fn main() {
|
||||
.file("../../csrc/ops/optim.cu")
|
||||
.file("../../csrc/ops/attention.cu")
|
||||
.file("../../csrc/ops/flash_attention.cu")
|
||||
.file("../../csrc/ops/repeat_kv.cu")
|
||||
.file("../../csrc/ops/cast.cu")
|
||||
.file("../../csrc/ops/dropout.cu")
|
||||
.compile("xtrain_cuda_kernels");
|
||||
|
||||
@@ -296,6 +296,37 @@ unsafe extern "C" {
|
||||
);
|
||||
}
|
||||
|
||||
// GQA repeat_kv head broadcast (csrc/ops/repeat_kv.cu, Phase T15). Expands a K/V
|
||||
// tensor from [batch·num_kv, S, hd] to the full [batch·nh, S, hd] so the SDPA
|
||||
// (composed or flash, both untouched) sees a full set of heads. Forward gathers
|
||||
// (out head qh ← kv head qh/group, group = nh/num_kv); backward sums the `group`
|
||||
// query heads sharing each kv head (deterministic, no atomics). All F32.
|
||||
#[cfg(not(no_cuda))]
|
||||
unsafe extern "C" {
|
||||
// Forward: out[b·nh+qh] = in[b·num_kv + qh/group], per [S,hd] head block.
|
||||
pub fn launch_repeat_kv_fwd_f32(
|
||||
input: *const f32,
|
||||
out: *mut f32,
|
||||
batch: i32,
|
||||
nh: i32,
|
||||
num_kv: i32,
|
||||
seq: i32,
|
||||
hd: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
// Backward: din[b·num_kv+kvh] = Σ_{r<group} dout[b·nh + kvh·group + r].
|
||||
pub fn launch_repeat_kv_bwd_f32(
|
||||
dout: *const f32,
|
||||
din: *mut f32,
|
||||
batch: i32,
|
||||
nh: i32,
|
||||
num_kv: i32,
|
||||
seq: i32,
|
||||
hd: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
}
|
||||
|
||||
// GPU-side optimizer kernels (csrc/ops/optim.cu): AdamW step (m/v on device) and
|
||||
// the global grad-norm reduction + in-place rescale (Phase T7).
|
||||
#[cfg(not(no_cuda))]
|
||||
|
||||
Reference in New Issue
Block a user