cuda: fused flash-attention kernel (fwd + flash-style bwd)
csrc/ops/flash_attention.cu: a single fused fwd kernel (one block per query row, streams KV in tiles of 32, online softmax — running max/sum + rescaled V accumulator, causal mask inlined, never materializes the [bh,S,S] scores) writing out[bh,S,hd] + the per-row logsumexp L (O(N), saved for backward). flash-style bwd: recompute scores from Q/K/V + L, collapse the softmax Jacobian with D[i]=ΣdO·O, dQ owned per row, dK/dV atomicAdd across rows. Tensor::flash_attention / flash_attention_backward wrap them (bf16 upcasts Q/K/V→f32 for the kernel, same fp32-softmax policy as composed). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -36,6 +36,7 @@ fn main() {
|
||||
.file("../../csrc/ops/model.cu")
|
||||
.file("../../csrc/ops/optim.cu")
|
||||
.file("../../csrc/ops/attention.cu")
|
||||
.file("../../csrc/ops/flash_attention.cu")
|
||||
.file("../../csrc/ops/cast.cu")
|
||||
.compile("xtrain_cuda_kernels");
|
||||
}
|
||||
|
||||
@@ -243,6 +243,59 @@ unsafe extern "C" {
|
||||
);
|
||||
}
|
||||
|
||||
// Fused flash-attention (csrc/ops/flash_attention.cu, Phase T14). A SINGLE kernel
|
||||
// each for forward/backward that streams over KV tiles with an online softmax and
|
||||
// NEVER materializes the [bh,S,S] score matrix. Q/K/V/out are [bh,S,hd] row-major
|
||||
// F32; the forward saves only the per-row logsumexp `l` ([bh*S], O(N)) for backward.
|
||||
#[cfg(not(no_cuda))]
|
||||
unsafe extern "C" {
|
||||
// Forward: o[bh,S,hd] = softmax(causal(Q·Kᵀ·scale))·V, online over KV tiles.
|
||||
// Also writes l[bh*S] = per-row logsumexp (saved for backward, not the scores).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn launch_flash_attention_fwd_f32(
|
||||
q: *const f32,
|
||||
k: *const f32,
|
||||
v: *const f32,
|
||||
o: *mut f32,
|
||||
l: *mut f32,
|
||||
bh: i32,
|
||||
seq: i32,
|
||||
hd: i32,
|
||||
scale: f32,
|
||||
s: CudaStream,
|
||||
);
|
||||
// Per-row D[i]=Σ_d dO[i,d]·O[i,d] over `rows`=bh*S rows of width `hd`. Must run
|
||||
// before the backward kernel (which takes the precomputed D, not O).
|
||||
pub fn launch_flash_attention_rowdot_f32(
|
||||
d_o: *const f32,
|
||||
o: *const f32,
|
||||
d_d: *mut f32,
|
||||
rows: i32,
|
||||
hd: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
// Backward: recomputes scores from Q/K/V + saved logsumexp `l` (NO cached probs)
|
||||
// and the precomputed `d_d` (= D), produces dq/dk/dv. dq/dk/dv must be PRE-ZEROED
|
||||
// (dk/dv are accumulated across query rows via atomicAdd).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn launch_flash_attention_bwd_f32(
|
||||
q: *const f32,
|
||||
k: *const f32,
|
||||
v: *const f32,
|
||||
d_o: *const f32,
|
||||
l: *const f32,
|
||||
d_d: *mut f32,
|
||||
dq: *mut f32,
|
||||
dk: *mut f32,
|
||||
dv: *mut f32,
|
||||
bh: i32,
|
||||
seq: i32,
|
||||
hd: i32,
|
||||
scale: f32,
|
||||
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))]
|
||||
|
||||
@@ -1092,6 +1092,119 @@ impl Tensor {
|
||||
(dq, dk, dv)
|
||||
}
|
||||
|
||||
// --- Fused flash-attention (the T14 op) ---
|
||||
|
||||
/// Fused flash-attention forward (Phase T14). `self`=Q, `k`, `v` each
|
||||
/// `[bh, seq, head_dim]`, contiguous on one GPU. Computes, per batch element,
|
||||
/// `out = softmax(causal(Q·Kᵀ·scale))·V` in a SINGLE kernel that streams over
|
||||
/// KV tiles with an online softmax — the `[bh,seq,seq]` score matrix is NEVER
|
||||
/// materialized. Returns `(out, lse)` where `lse`:[bh,seq] (F32) is the per-row
|
||||
/// logsumexp cached for backward (O(N), vs the composed path's O(N²) probs).
|
||||
///
|
||||
/// The fused kernel is fp32; for bf16 we upcast Q/K/V → f32 → kernel → downcast
|
||||
/// `out` back to bf16 (same fp32-softmax policy as the composed [`attention`]),
|
||||
/// so flash and composed produce the same softmax numerics. `lse` stays fp32.
|
||||
#[cfg(not(no_cuda))]
|
||||
pub fn flash_attention(&self, k: &Tensor, v: &Tensor, scale: f32) -> (Tensor, Tensor) {
|
||||
assert_eq!(
|
||||
self.ndim(),
|
||||
3,
|
||||
"flash_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");
|
||||
assert_eq!(self.dtype, k.dtype, "Q/K dtype mismatch");
|
||||
assert_eq!(self.dtype, v.dtype, "Q/V dtype mismatch");
|
||||
let (bh, seq, hd) = (self.shape[0], self.shape[1], self.shape[2]);
|
||||
let dev = self.device();
|
||||
let dt = self.dtype;
|
||||
|
||||
let qf = self.to_dtype(DType::F32);
|
||||
let kf = k.to_dtype(DType::F32);
|
||||
let vf = v.to_dtype(DType::F32);
|
||||
let out_f32 = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
|
||||
let lse = Tensor::zeros(&[bh, seq], DType::F32, dev);
|
||||
unsafe {
|
||||
xtrain_cuda::ffi::launch_flash_attention_fwd_f32(
|
||||
qf.data_ptr() as *const f32,
|
||||
kf.data_ptr() as *const f32,
|
||||
vf.data_ptr() as *const f32,
|
||||
out_f32.data_ptr() as *mut f32,
|
||||
lse.data_ptr() as *mut f32,
|
||||
bh as i32,
|
||||
seq as i32,
|
||||
hd as i32,
|
||||
scale,
|
||||
std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
(out_f32.to_dtype(dt), lse)
|
||||
}
|
||||
|
||||
/// Backward of [`flash_attention`](Self::flash_attention). Inputs: forward
|
||||
/// `q`,`k`,`v`, the forward output `out`, the cached `lse`:[bh,seq], the upstream
|
||||
/// `dout`, and the same `scale`. Returns `(dq, dk, dv)`.
|
||||
///
|
||||
/// flash-style: NO cached probs. Recomputes scores from Q/K/V + `lse`, uses
|
||||
/// `D[i]=Σ dOᵢ·Oᵢ` to collapse the softmax Jacobian, streams KV in tiles. dQ is
|
||||
/// owned per query row; dK/dV are accumulated across rows (atomicAdd). Same
|
||||
/// fp32 kernel; bf16 callers get fp32 grads which the autograd `cast` op casts.
|
||||
#[cfg(not(no_cuda))]
|
||||
pub fn flash_attention_backward(
|
||||
q: &Tensor,
|
||||
k: &Tensor,
|
||||
v: &Tensor,
|
||||
out: &Tensor,
|
||||
lse: &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();
|
||||
let dt = q.dtype;
|
||||
|
||||
let qf = q.to_dtype(DType::F32);
|
||||
let kf = k.to_dtype(DType::F32);
|
||||
let vf = v.to_dtype(DType::F32);
|
||||
let of = out.to_dtype(DType::F32);
|
||||
let dof = dout.to_dtype(DType::F32);
|
||||
// D[i] = Σ_d dO[i,d]·O[i,d] (one scalar per query row, O(N)).
|
||||
let d = Tensor::zeros(&[bh, seq], DType::F32, dev);
|
||||
unsafe {
|
||||
xtrain_cuda::ffi::launch_flash_attention_rowdot_f32(
|
||||
dof.data_ptr() as *const f32,
|
||||
of.data_ptr() as *const f32,
|
||||
d.data_ptr() as *mut f32,
|
||||
(bh * seq) as i32,
|
||||
hd as i32,
|
||||
std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
// dq/dk/dv pre-zeroed (Tensor::zeros memsets); dk/dv accumulate via atomicAdd.
|
||||
let dq = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
|
||||
let dk = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
|
||||
let dv = Tensor::zeros(&[bh, seq, hd], DType::F32, dev);
|
||||
unsafe {
|
||||
xtrain_cuda::ffi::launch_flash_attention_bwd_f32(
|
||||
qf.data_ptr() as *const f32,
|
||||
kf.data_ptr() as *const f32,
|
||||
vf.data_ptr() as *const f32,
|
||||
dof.data_ptr() as *const f32,
|
||||
lse.data_ptr() as *const f32,
|
||||
d.data_ptr() as *mut f32,
|
||||
dq.data_ptr() as *mut f32,
|
||||
dk.data_ptr() as *mut f32,
|
||||
dv.data_ptr() as *mut f32,
|
||||
bh as i32,
|
||||
seq as i32,
|
||||
hd as i32,
|
||||
scale,
|
||||
std::ptr::null_mut(),
|
||||
);
|
||||
}
|
||||
(dq.to_dtype(dt), dk.to_dtype(dt), dv.to_dtype(dt))
|
||||
}
|
||||
|
||||
/// 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).
|
||||
|
||||
Reference in New Issue
Block a user