autograd: batch dim for ops (flatten linears, batched attention)
Add the batched-forward primitives. Linears/norms/elementwise/embedding/CE already act on flat [rows,dim], so they work unchanged on [B*S,dim]; only attention + RoPE need sequence awareness: - RoPE: kernel takes a `period` (= seq len) so position = row % period, i.e. per-sequence position on a flattened batch (period == tokens = single seq). - Fused batched causal attention: new `Tensor::attention`/`attention_backward` + ops node, running QKᵀ and PV as cublasSgemmStridedBatched over the B*nh (sequence,head) blocks (new sgemm_strided_batched binding) and a causal softmax kernel (scale + per-row causal mask inline) — the whole attention is 3 launches regardless of B*nh, no per-head/per-seq loop, no host round-trip. - transpose_4d12 ([B,S,nh,hd] <-> [B,nh,S,hd]) to lay out the batched heads. grad-checks: new batched-rope, transpose_4d12, batched-attention dQ/dK/dV all pass finite-diff (attn dK 1.5e-2, dQ 7.5e-3, dV 2.9e-4; rest tighter) alongside the existing 12. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -454,13 +454,20 @@ impl Tensor {
|
||||
dx
|
||||
}
|
||||
|
||||
/// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; the position
|
||||
/// of each token is its row index. Returns the rotated tensor.
|
||||
/// RoPE forward (rotate_half). `self`:[tokens,heads,head_dim]; each token's
|
||||
/// position is `row % period`. `period` = sequence length, so a flattened
|
||||
/// batch `[B*S,heads,head_dim]` gets per-sequence positions (pass `period=S`);
|
||||
/// pass `period=tokens` for a single sequence (position = row). Returns the
|
||||
/// rotated tensor.
|
||||
#[cfg(not(no_cuda))]
|
||||
pub fn rope(&self, theta: f32) -> Self {
|
||||
pub fn rope(&self, theta: f32, period: usize) -> Self {
|
||||
assert_eq!(self.ndim(), 3, "rope requires [tokens,heads,head_dim]");
|
||||
let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
|
||||
assert_eq!(head_dim % 2, 0, "head_dim must be even");
|
||||
assert!(
|
||||
period > 0 && tokens % period == 0,
|
||||
"tokens must be a multiple of period"
|
||||
);
|
||||
let out = Tensor::zeros(&self.shape, DType::F32, self.device());
|
||||
unsafe {
|
||||
xtrain_cuda::ffi::launch_rope_f32(
|
||||
@@ -470,6 +477,7 @@ impl Tensor {
|
||||
heads as i32,
|
||||
head_dim as i32,
|
||||
theta,
|
||||
period as i32,
|
||||
std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
@@ -477,9 +485,9 @@ impl Tensor {
|
||||
}
|
||||
|
||||
/// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
|
||||
/// orthogonal map, so it needs no cached forward values, only `theta`.
|
||||
/// orthogonal map, so it needs no cached forward values, only `theta`/`period`.
|
||||
#[cfg(not(no_cuda))]
|
||||
pub fn rope_backward(dy: &Tensor, theta: f32) -> Self {
|
||||
pub fn rope_backward(dy: &Tensor, theta: f32, period: usize) -> Self {
|
||||
let (tokens, heads, head_dim) = (dy.shape[0], dy.shape[1], dy.shape[2]);
|
||||
let dx = Tensor::zeros(&dy.shape, DType::F32, dy.device());
|
||||
unsafe {
|
||||
@@ -490,6 +498,7 @@ impl Tensor {
|
||||
heads as i32,
|
||||
head_dim as i32,
|
||||
theta,
|
||||
period as i32,
|
||||
std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
@@ -667,6 +676,202 @@ impl Tensor {
|
||||
out
|
||||
}
|
||||
|
||||
// --- Batched attention (the T10 fused op) ---
|
||||
|
||||
/// Batched causal scaled-dot-product attention. `self`=Q, `k`, `v` are each
|
||||
/// `[bh, seq, head_dim]` (bh = batch·n_heads), contiguous F32 on one GPU.
|
||||
/// Computes, per batch element, `out = softmax(causal(Q·Kᵀ / √hd)) · V`. The
|
||||
/// two GEMMs run as `cublasSgemmStridedBatched` and the softmax+scale+causal
|
||||
/// mask is one kernel, so the whole attention is 3 launches regardless of bh.
|
||||
/// Returns `(out, probs)` where `probs`:[bh,seq,seq] is cached for backward.
|
||||
#[cfg(not(no_cuda))]
|
||||
pub fn attention(&self, k: &Tensor, v: &Tensor, scale: f32) -> (Tensor, Tensor) {
|
||||
assert_eq!(self.ndim(), 3, "attention Q must be [bh,seq,head_dim]");
|
||||
assert_eq!(self.shape(), k.shape(), "Q/K shape mismatch");
|
||||
assert_eq!(self.shape(), v.shape(), "Q/V shape mismatch");
|
||||
let (bh, seq, hd) = (self.shape[0], self.shape[1], self.shape[2]);
|
||||
let dev = self.device();
|
||||
|
||||
// scores[bh,seq,seq] = Q[bh,seq,hd] · Kᵀ[bh,hd,seq]
|
||||
let scores = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
|
||||
xtrain_cuda::cublas::sgemm_strided_batched(
|
||||
false,
|
||||
true,
|
||||
seq,
|
||||
seq,
|
||||
hd,
|
||||
1.0,
|
||||
self.data_ptr() as *const f32,
|
||||
seq * hd,
|
||||
k.data_ptr() as *const f32,
|
||||
seq * hd,
|
||||
0.0,
|
||||
scores.data_ptr() as *mut f32,
|
||||
seq * seq,
|
||||
bh,
|
||||
);
|
||||
// probs = softmax(causal(scores · scale)), one block per [bh·seq] row.
|
||||
let probs = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
|
||||
unsafe {
|
||||
xtrain_cuda::ffi::launch_softmax_causal_f32(
|
||||
scores.data_ptr() as *const f32,
|
||||
probs.data_ptr() as *mut f32,
|
||||
(bh * seq) as i32,
|
||||
seq as i32,
|
||||
scale,
|
||||
std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
// out[bh,seq,hd] = probs[bh,seq,seq] · V[bh,seq,hd]
|
||||
let out = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
|
||||
xtrain_cuda::cublas::sgemm_strided_batched(
|
||||
false,
|
||||
false,
|
||||
seq,
|
||||
hd,
|
||||
seq,
|
||||
1.0,
|
||||
probs.data_ptr() as *const f32,
|
||||
seq * seq,
|
||||
v.data_ptr() as *const f32,
|
||||
seq * hd,
|
||||
0.0,
|
||||
out.data_ptr() as *mut f32,
|
||||
seq * hd,
|
||||
bh,
|
||||
);
|
||||
(out, probs)
|
||||
}
|
||||
|
||||
/// Backward of [`attention`](Self::attention). Inputs: forward `q`,`k`,`v`,
|
||||
/// the cached `probs`, the upstream `dout` (all batched `[bh,seq,*]`), and the
|
||||
/// same `scale`. Returns `(dq, dk, dv)`.
|
||||
///
|
||||
/// dP = dOut · Vᵀ ; dV = Pᵀ · dOut
|
||||
/// dScores = softmax_jacobian(P, dP) · scale (scale folded back in)
|
||||
/// dQ = dScores · K ; dK = dScoresᵀ · Q
|
||||
///
|
||||
/// Masked (future) entries of P are 0, so the softmax Jacobian zeros their
|
||||
/// gradient — the causal mask needs no special handling here.
|
||||
#[cfg(not(no_cuda))]
|
||||
pub fn attention_backward(
|
||||
q: &Tensor,
|
||||
k: &Tensor,
|
||||
v: &Tensor,
|
||||
probs: &Tensor,
|
||||
dout: &Tensor,
|
||||
scale: f32,
|
||||
) -> (Tensor, Tensor, Tensor) {
|
||||
let (bh, seq, hd) = (q.shape[0], q.shape[1], q.shape[2]);
|
||||
let dev = q.device();
|
||||
|
||||
// dP[bh,seq,seq] = dOut[bh,seq,hd] · Vᵀ[bh,hd,seq]
|
||||
let dp = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
|
||||
xtrain_cuda::cublas::sgemm_strided_batched(
|
||||
false,
|
||||
true,
|
||||
seq,
|
||||
seq,
|
||||
hd,
|
||||
1.0,
|
||||
dout.data_ptr() as *const f32,
|
||||
seq * hd,
|
||||
v.data_ptr() as *const f32,
|
||||
seq * hd,
|
||||
0.0,
|
||||
dp.data_ptr() as *mut f32,
|
||||
seq * seq,
|
||||
bh,
|
||||
);
|
||||
// dV[bh,seq,hd] = Pᵀ[bh,seq,seq] · dOut[bh,seq,hd]
|
||||
let dv = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
|
||||
xtrain_cuda::cublas::sgemm_strided_batched(
|
||||
true,
|
||||
false,
|
||||
seq,
|
||||
hd,
|
||||
seq,
|
||||
1.0,
|
||||
probs.data_ptr() as *const f32,
|
||||
seq * seq,
|
||||
dout.data_ptr() as *const f32,
|
||||
seq * hd,
|
||||
0.0,
|
||||
dv.data_ptr() as *mut f32,
|
||||
seq * hd,
|
||||
bh,
|
||||
);
|
||||
// dScores = softmax Jacobian (per row) applied to dP, then ×scale.
|
||||
// Reuse the row-wise softmax backward over the flattened [bh·seq, seq].
|
||||
let dscores = Tensor::softmax_backward(
|
||||
&probs.reshape(&[bh * seq, seq]),
|
||||
&dp.reshape(&[bh * seq, seq]),
|
||||
)
|
||||
.reshape(&[bh, seq, seq]);
|
||||
let dscores = dscores.scale(scale);
|
||||
// dQ[bh,seq,hd] = dScores[bh,seq,seq] · K[bh,seq,hd]
|
||||
let dq = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
|
||||
xtrain_cuda::cublas::sgemm_strided_batched(
|
||||
false,
|
||||
false,
|
||||
seq,
|
||||
hd,
|
||||
seq,
|
||||
1.0,
|
||||
dscores.data_ptr() as *const f32,
|
||||
seq * seq,
|
||||
k.data_ptr() as *const f32,
|
||||
seq * hd,
|
||||
0.0,
|
||||
dq.data_ptr() as *mut f32,
|
||||
seq * hd,
|
||||
bh,
|
||||
);
|
||||
// dK[bh,seq,hd] = dScoresᵀ[bh,seq,seq] · Q[bh,seq,hd]
|
||||
let dk = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
|
||||
xtrain_cuda::cublas::sgemm_strided_batched(
|
||||
true,
|
||||
false,
|
||||
seq,
|
||||
hd,
|
||||
seq,
|
||||
1.0,
|
||||
dscores.data_ptr() as *const f32,
|
||||
seq * seq,
|
||||
q.data_ptr() as *const f32,
|
||||
seq * hd,
|
||||
0.0,
|
||||
dk.data_ptr() as *mut f32,
|
||||
seq * hd,
|
||||
bh,
|
||||
);
|
||||
(dq, dk, dv)
|
||||
}
|
||||
|
||||
/// 4D axis-(1,2) transpose: `self`:[a,b,c,d] → [a,c,b,d],
|
||||
/// `out[i,k,j,l]=self[i,j,k,l]`. Lays out batched multi-head attention
|
||||
/// (`[B,S,nh,hd] <-> [B,nh,S,hd]`). Its own backward is the same op (swap b,c).
|
||||
#[cfg(not(no_cuda))]
|
||||
pub fn transpose_4d12(&self) -> Self {
|
||||
assert_eq!(self.dtype, DType::F32, "transpose_4d12 only supports F32");
|
||||
assert_eq!(self.ndim(), 4, "transpose_4d12 requires a 4D tensor");
|
||||
assert!(self.is_contiguous(), "transpose_4d12 requires contiguous");
|
||||
let (a, b, c, d) = (self.shape[0], self.shape[1], self.shape[2], self.shape[3]);
|
||||
let out = Tensor::zeros(&[a, c, b, d], DType::F32, self.device());
|
||||
unsafe {
|
||||
xtrain_cuda::ffi::launch_transpose_4d12_f32(
|
||||
self.data_ptr() as *const f32,
|
||||
out.data_ptr() as *mut f32,
|
||||
a as i32,
|
||||
b as i32,
|
||||
c as i32,
|
||||
d as i32,
|
||||
std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
// Shared validation for same-shape binary elementwise ops.
|
||||
#[cfg(not(no_cuda))]
|
||||
fn check_binary(&self, other: &Tensor, op: &str) {
|
||||
|
||||
Reference in New Issue
Block a user