Gahow Wang
326a6fadfe
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>
1347 lines
49 KiB
Rust
1347 lines
49 KiB
Rust
//! The `Tensor` type: shape/strides/dtype over reference-counted [`Storage`],
|
||
//! with host↔device transfer and one elementwise op (`scale`) wired end-to-end
|
||
//! through a CUDA kernel.
|
||
|
||
use crate::dtype::{DType, TensorDType};
|
||
use crate::shape::{self, Dims};
|
||
use crate::storage::{Device, Storage};
|
||
|
||
/// Multi-dimensional array backed by CPU or GPU storage.
|
||
///
|
||
/// Strides are in elements (row-major). T2 tensors created here are always
|
||
/// contiguous; the `strides`/`offset` fields exist so later phases can add
|
||
/// zero-copy views without changing this type's shape.
|
||
#[derive(Clone)]
|
||
pub struct Tensor {
|
||
storage: Storage,
|
||
shape: Dims,
|
||
strides: Dims,
|
||
offset: usize,
|
||
dtype: DType,
|
||
}
|
||
|
||
impl Tensor {
|
||
// --- Creation ---
|
||
|
||
/// Build a contiguous CPU tensor from a typed host slice.
|
||
pub fn from_slice<T: TensorDType>(data: &[T], shape: &[usize]) -> Self {
|
||
let numel = shape::num_elements(shape);
|
||
assert_eq!(
|
||
data.len(),
|
||
numel,
|
||
"data length {} != shape numel {numel}",
|
||
data.len()
|
||
);
|
||
let bytes = unsafe {
|
||
std::slice::from_raw_parts(data.as_ptr() as *const u8, numel * T::DTYPE.size_bytes())
|
||
};
|
||
Self {
|
||
storage: Storage::cpu(bytes.to_vec()),
|
||
shape: Dims::from_slice(shape),
|
||
strides: shape::contiguous_strides(shape),
|
||
offset: 0,
|
||
dtype: T::DTYPE,
|
||
}
|
||
}
|
||
|
||
/// Zero-filled contiguous tensor on the given device.
|
||
pub fn zeros(shape: &[usize], dtype: DType, device: Device) -> Self {
|
||
let len_bytes = shape::num_elements(shape) * dtype.size_bytes();
|
||
let storage = Storage::zeros(len_bytes, device).expect("zeros alloc failed");
|
||
Self {
|
||
storage,
|
||
shape: Dims::from_slice(shape),
|
||
strides: shape::contiguous_strides(shape),
|
||
offset: 0,
|
||
dtype,
|
||
}
|
||
}
|
||
|
||
// --- Properties ---
|
||
|
||
pub fn shape(&self) -> &[usize] {
|
||
&self.shape
|
||
}
|
||
pub fn strides(&self) -> &[usize] {
|
||
&self.strides
|
||
}
|
||
pub fn dtype(&self) -> DType {
|
||
self.dtype
|
||
}
|
||
pub fn ndim(&self) -> usize {
|
||
self.shape.len()
|
||
}
|
||
pub fn numel(&self) -> usize {
|
||
shape::num_elements(&self.shape)
|
||
}
|
||
pub fn offset(&self) -> usize {
|
||
self.offset
|
||
}
|
||
pub fn device(&self) -> Device {
|
||
self.storage.device()
|
||
}
|
||
pub fn is_contiguous(&self) -> bool {
|
||
shape::is_contiguous(&self.shape, &self.strides)
|
||
}
|
||
pub fn storage(&self) -> &Storage {
|
||
&self.storage
|
||
}
|
||
|
||
// --- Device transfer ---
|
||
|
||
/// Move (copy) the tensor to `device`. Returns a cheap clone if already there.
|
||
pub fn to_device(&self, device: Device) -> Self {
|
||
if self.device() == device {
|
||
return self.clone();
|
||
}
|
||
let storage = self
|
||
.storage
|
||
.to_device(device)
|
||
.expect("device transfer failed");
|
||
Self {
|
||
storage,
|
||
shape: self.shape.clone(),
|
||
strides: self.strides.clone(),
|
||
offset: self.offset,
|
||
dtype: self.dtype,
|
||
}
|
||
}
|
||
|
||
// --- dtype cast (Phase T12, bf16 mixed precision) ---
|
||
|
||
/// Cast between F32 and BF16 (the AMP bridge: fp32 master ↔ bf16 compute).
|
||
/// Same dtype returns a cheap clone. Requires a contiguous CUDA tensor.
|
||
/// I32 is not castable here (only used for token-id targets).
|
||
#[cfg(not(no_cuda))]
|
||
pub fn to_dtype(&self, target: DType) -> Self {
|
||
if self.dtype == target {
|
||
return self.clone();
|
||
}
|
||
assert!(
|
||
matches!(self.device(), Device::Cuda(_)),
|
||
"to_dtype requires a CUDA tensor"
|
||
);
|
||
assert!(self.is_contiguous(), "to_dtype requires contiguous tensor");
|
||
let n = self.numel() as i32;
|
||
let out = Tensor::zeros(&self.shape, target, self.device());
|
||
match (self.dtype, target) {
|
||
(DType::F32, DType::BF16) => unsafe {
|
||
xtrain_cuda::ffi::launch_cast_f32_to_bf16(
|
||
self.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut std::ffi::c_void,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
(DType::BF16, DType::F32) => unsafe {
|
||
xtrain_cuda::ffi::launch_cast_bf16_to_f32(
|
||
self.data_ptr() as *const std::ffi::c_void,
|
||
out.data_ptr() as *mut f32,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
(a, b) => panic!("unsupported cast {a} -> {b}"),
|
||
}
|
||
out
|
||
}
|
||
|
||
// --- Host data access (CPU only) ---
|
||
|
||
/// Typed read-only view of the data. Requires a contiguous CPU tensor.
|
||
pub fn as_slice<T: TensorDType>(&self) -> &[T] {
|
||
assert_eq!(T::DTYPE, self.dtype, "dtype mismatch");
|
||
assert_eq!(self.device(), Device::Cpu, "as_slice requires CPU tensor");
|
||
assert!(self.is_contiguous(), "as_slice requires contiguous tensor");
|
||
let bytes = self.storage.as_cpu_bytes();
|
||
let start = self.offset * self.dtype.size_bytes();
|
||
unsafe { std::slice::from_raw_parts(bytes[start..].as_ptr() as *const T, self.numel()) }
|
||
}
|
||
|
||
/// Raw element pointer at the tensor's offset (for kernel launches).
|
||
pub fn data_ptr(&self) -> *const u8 {
|
||
let byte_off = self.offset * self.dtype.size_bytes();
|
||
match self.device() {
|
||
Device::Cpu => unsafe { self.storage.as_cpu_bytes().as_ptr().add(byte_off) },
|
||
Device::Cuda(_) => unsafe { self.storage.gpu_buffer().as_ptr().add(byte_off) },
|
||
}
|
||
}
|
||
|
||
// --- Elementwise op (the T2 end-to-end kernel) ---
|
||
|
||
/// Out-of-place elementwise scale: returns a new tensor `out[i] = self[i] * alpha`.
|
||
///
|
||
/// Runs the `scale_f32` CUDA kernel. Requires a contiguous F32 tensor on the
|
||
/// GPU. Available only when CUDA was compiled in (`not(no_cuda)`).
|
||
#[cfg(not(no_cuda))]
|
||
pub fn scale(&self, alpha: f32) -> Self {
|
||
assert!(
|
||
matches!(self.dtype, DType::F32 | DType::BF16),
|
||
"scale supports F32/BF16"
|
||
);
|
||
assert!(self.is_contiguous(), "scale requires contiguous tensor");
|
||
assert!(
|
||
matches!(self.device(), Device::Cuda(_)),
|
||
"scale requires a CUDA tensor"
|
||
);
|
||
|
||
let out = Tensor::zeros(&self.shape, self.dtype, self.device());
|
||
let n = self.numel() as i32;
|
||
match self.dtype {
|
||
DType::F32 => unsafe {
|
||
xtrain_cuda::ffi::launch_scale_f32(
|
||
self.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut f32,
|
||
alpha,
|
||
n,
|
||
std::ptr::null_mut(), // default stream
|
||
);
|
||
},
|
||
DType::BF16 => unsafe {
|
||
xtrain_cuda::ffi::launch_scale_bf16(
|
||
self.data_ptr() as *const std::ffi::c_void,
|
||
out.data_ptr() as *mut std::ffi::c_void,
|
||
alpha,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
_ => unreachable!(),
|
||
}
|
||
out
|
||
}
|
||
|
||
// --- GEMM (the T3 kernels) ---
|
||
|
||
/// Matrix multiply: `C = self @ other`. `self`:[M,K], `other`:[K,N] → [M,N].
|
||
///
|
||
/// Routes through cuBLAS `Sgemm` (Phase T7). fp32, so it is the same GEMM as
|
||
/// the T3 tiled kernel up to rounding order. Requires contiguous F32 tensors
|
||
/// on the same GPU. Available only when CUDA is compiled in.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn matmul(&self, other: &Tensor) -> Self {
|
||
assert_eq!(self.dtype, other.dtype, "matmul dtype mismatch");
|
||
assert!(
|
||
matches!(self.dtype, DType::F32 | DType::BF16),
|
||
"matmul supports F32/BF16"
|
||
);
|
||
assert_eq!(self.ndim(), 2, "matmul requires 2D lhs");
|
||
assert_eq!(other.ndim(), 2, "matmul requires 2D rhs");
|
||
assert_eq!(
|
||
self.shape[1], other.shape[0],
|
||
"inner dimension mismatch: {:?} @ {:?}",
|
||
self.shape, other.shape
|
||
);
|
||
assert!(
|
||
self.is_contiguous() && other.is_contiguous(),
|
||
"matmul requires contiguous tensors"
|
||
);
|
||
assert_eq!(self.device(), other.device(), "matmul device mismatch");
|
||
assert!(
|
||
matches!(self.device(), Device::Cuda(_)),
|
||
"matmul requires CUDA tensors"
|
||
);
|
||
|
||
let m = self.shape[0];
|
||
let k = self.shape[1];
|
||
let n = other.shape[1];
|
||
let out = Tensor::zeros(&[m, n], self.dtype, self.device());
|
||
match self.dtype {
|
||
// fp32 path — unchanged (bit-identical to T7/T10/T11).
|
||
DType::F32 => xtrain_cuda::cublas::sgemm(
|
||
false,
|
||
false,
|
||
m,
|
||
n,
|
||
k,
|
||
1.0,
|
||
self.data_ptr() as *const f32,
|
||
other.data_ptr() as *const f32,
|
||
0.0,
|
||
out.data_ptr() as *mut f32,
|
||
),
|
||
// bf16 path — GemmEx, bf16 in/out, fp32 accumulation.
|
||
DType::BF16 => xtrain_cuda::cublas::gemm_ex(
|
||
false,
|
||
false,
|
||
m,
|
||
n,
|
||
k,
|
||
1.0,
|
||
self.data_ptr() as *const std::ffi::c_void,
|
||
other.data_ptr() as *const std::ffi::c_void,
|
||
0.0,
|
||
out.data_ptr() as *mut std::ffi::c_void,
|
||
),
|
||
_ => unreachable!(),
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Out-of-place 2D transpose: returns a new contiguous tensor `out[j,i] =
|
||
/// self[i,j]`. Requires a contiguous F32 CUDA tensor.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn transpose_2d(&self) -> Self {
|
||
assert_eq!(self.ndim(), 2, "transpose_2d requires 2D tensor");
|
||
assert!(self.is_contiguous(), "transpose requires contiguous tensor");
|
||
if self.dtype == DType::BF16 {
|
||
return self
|
||
.to_dtype(DType::F32)
|
||
.transpose_2d()
|
||
.to_dtype(DType::BF16);
|
||
}
|
||
assert_eq!(self.dtype, DType::F32, "transpose supports F32/BF16");
|
||
assert!(
|
||
matches!(self.device(), Device::Cuda(_)),
|
||
"transpose requires a CUDA tensor"
|
||
);
|
||
|
||
let rows = self.shape[0];
|
||
let cols = self.shape[1];
|
||
let out = Tensor::zeros(&[cols, rows], DType::F32, self.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_transpose_f32(
|
||
self.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut f32,
|
||
rows as i32,
|
||
cols as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Backward of `C = A @ B` given the upstream gradient `dC` (shape [M,N]).
|
||
/// Returns `(dA, dB)` where `dA = dC @ Bᵀ` ([M,K]) and `dB = Aᵀ @ dC`
|
||
/// ([K,N]). All tensors contiguous F32 on the same GPU.
|
||
///
|
||
/// Phase T7: cuBLAS applies the transposes internally via its op flags, so we
|
||
/// avoid the two transpose kernels (and their allocations) the T3 version ran.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn matmul_backward(a: &Tensor, b: &Tensor, dc: &Tensor) -> (Tensor, Tensor) {
|
||
assert_eq!(a.ndim(), 2, "matmul_backward requires 2D A");
|
||
assert_eq!(b.ndim(), 2, "matmul_backward requires 2D B");
|
||
assert_eq!(dc.ndim(), 2, "matmul_backward requires 2D dC");
|
||
assert_eq!(a.shape[1], b.shape[0], "A/B inner dim mismatch");
|
||
assert_eq!(dc.shape[0], a.shape[0], "dC rows != A rows (M)");
|
||
assert_eq!(dc.shape[1], b.shape[1], "dC cols != B cols (N)");
|
||
|
||
assert_eq!(a.dtype, b.dtype, "matmul_backward dtype mismatch");
|
||
assert_eq!(a.dtype, dc.dtype, "matmul_backward dtype mismatch");
|
||
let (m, k, n) = (a.shape[0], a.shape[1], b.shape[1]);
|
||
let dt = a.dtype;
|
||
// dA[M,K] = dC[M,N] · Bᵀ (B stored [K,N], transposed by cuBLAS)
|
||
let da = Tensor::zeros(&[m, k], dt, a.device());
|
||
// dB[K,N] = Aᵀ · dC[M,N] (A stored [M,K], transposed by cuBLAS)
|
||
let db = Tensor::zeros(&[k, n], dt, a.device());
|
||
match dt {
|
||
DType::F32 => {
|
||
xtrain_cuda::cublas::sgemm(
|
||
false,
|
||
true,
|
||
m,
|
||
k,
|
||
n,
|
||
1.0,
|
||
dc.data_ptr() as *const f32,
|
||
b.data_ptr() as *const f32,
|
||
0.0,
|
||
da.data_ptr() as *mut f32,
|
||
);
|
||
xtrain_cuda::cublas::sgemm(
|
||
true,
|
||
false,
|
||
k,
|
||
n,
|
||
m,
|
||
1.0,
|
||
a.data_ptr() as *const f32,
|
||
dc.data_ptr() as *const f32,
|
||
0.0,
|
||
db.data_ptr() as *mut f32,
|
||
);
|
||
}
|
||
DType::BF16 => {
|
||
xtrain_cuda::cublas::gemm_ex(
|
||
false,
|
||
true,
|
||
m,
|
||
k,
|
||
n,
|
||
1.0,
|
||
dc.data_ptr() as *const std::ffi::c_void,
|
||
b.data_ptr() as *const std::ffi::c_void,
|
||
0.0,
|
||
da.data_ptr() as *mut std::ffi::c_void,
|
||
);
|
||
xtrain_cuda::cublas::gemm_ex(
|
||
true,
|
||
false,
|
||
k,
|
||
n,
|
||
m,
|
||
1.0,
|
||
a.data_ptr() as *const std::ffi::c_void,
|
||
dc.data_ptr() as *const std::ffi::c_void,
|
||
0.0,
|
||
db.data_ptr() as *mut std::ffi::c_void,
|
||
);
|
||
}
|
||
_ => panic!("matmul_backward supports F32/BF16"),
|
||
}
|
||
(da, db)
|
||
}
|
||
|
||
// --- Transformer / autograd op primitives (the T4 kernels) ---
|
||
//
|
||
// Each is a thin, contiguous-F32-on-GPU wrapper over a kernel in
|
||
// csrc/ops/nn.cu. The autograd `Var` layer (xtrain-autodiff) builds nodes on
|
||
// top of these; the analytic backwards are derived in docs/03-autograd-engine.md.
|
||
|
||
/// Elementwise `out = self + other` (same shape).
|
||
#[cfg(not(no_cuda))]
|
||
pub fn add(&self, other: &Tensor) -> Self {
|
||
self.check_binary(other, "add");
|
||
let out = Tensor::zeros(&self.shape, self.dtype, self.device());
|
||
let n = self.numel() as i32;
|
||
match self.dtype {
|
||
DType::F32 => unsafe {
|
||
xtrain_cuda::ffi::launch_add_f32(
|
||
self.data_ptr() as *const f32,
|
||
other.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut f32,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
DType::BF16 => unsafe {
|
||
xtrain_cuda::ffi::launch_add_bf16(
|
||
self.data_ptr() as *const std::ffi::c_void,
|
||
other.data_ptr() as *const std::ffi::c_void,
|
||
out.data_ptr() as *mut std::ffi::c_void,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
_ => unreachable!(),
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Elementwise `out = self * other` (same shape, Hadamard product).
|
||
#[cfg(not(no_cuda))]
|
||
pub fn mul(&self, other: &Tensor) -> Self {
|
||
self.check_binary(other, "mul");
|
||
let out = Tensor::zeros(&self.shape, self.dtype, self.device());
|
||
let n = self.numel() as i32;
|
||
match self.dtype {
|
||
DType::F32 => unsafe {
|
||
xtrain_cuda::ffi::launch_mul_f32(
|
||
self.data_ptr() as *const f32,
|
||
other.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut f32,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
DType::BF16 => unsafe {
|
||
xtrain_cuda::ffi::launch_mul_bf16(
|
||
self.data_ptr() as *const std::ffi::c_void,
|
||
other.data_ptr() as *const std::ffi::c_void,
|
||
out.data_ptr() as *mut std::ffi::c_void,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
_ => unreachable!(),
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Broadcast bias add: `out[r,c] = self[r,c] + bias[c]`.
|
||
/// `self`:[rows,cols], `bias`:[cols].
|
||
#[cfg(not(no_cuda))]
|
||
pub fn add_bias(&self, bias: &Tensor) -> Self {
|
||
assert_eq!(self.ndim(), 2, "add_bias requires 2D input");
|
||
assert_eq!(bias.ndim(), 1, "bias must be 1D");
|
||
assert_eq!(self.shape[1], bias.shape[0], "bias len != cols");
|
||
assert_eq!(self.dtype, bias.dtype, "add_bias dtype mismatch");
|
||
let (rows, cols) = (self.shape[0], self.shape[1]);
|
||
let out = Tensor::zeros(&self.shape, self.dtype, self.device());
|
||
match self.dtype {
|
||
DType::F32 => unsafe {
|
||
xtrain_cuda::ffi::launch_add_bias_f32(
|
||
self.data_ptr() as *const f32,
|
||
bias.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut f32,
|
||
rows as i32,
|
||
cols as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
DType::BF16 => unsafe {
|
||
xtrain_cuda::ffi::launch_add_bias_bf16(
|
||
self.data_ptr() as *const std::ffi::c_void,
|
||
bias.data_ptr() as *const std::ffi::c_void,
|
||
out.data_ptr() as *mut std::ffi::c_void,
|
||
rows as i32,
|
||
cols as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
_ => panic!("add_bias supports F32/BF16"),
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Column-sum over rows: `out[c] = sum_r self[r,c]`. This is the bias
|
||
/// backward (sum the upstream grad over the broadcast dim). `self`:[rows,cols]
|
||
/// → [cols].
|
||
#[cfg(not(no_cuda))]
|
||
pub fn sum_rows(&self) -> Self {
|
||
assert_eq!(self.ndim(), 2, "sum_rows requires 2D input");
|
||
let (rows, cols) = (self.shape[0], self.shape[1]);
|
||
let out = Tensor::zeros(&[cols], self.dtype, self.device());
|
||
match self.dtype {
|
||
DType::F32 => unsafe {
|
||
xtrain_cuda::ffi::launch_sum_rows_f32(
|
||
self.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut f32,
|
||
rows as i32,
|
||
cols as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
DType::BF16 => unsafe {
|
||
xtrain_cuda::ffi::launch_sum_rows_bf16(
|
||
self.data_ptr() as *const std::ffi::c_void,
|
||
out.data_ptr() as *mut std::ffi::c_void,
|
||
rows as i32,
|
||
cols as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
_ => panic!("sum_rows supports F32/BF16"),
|
||
}
|
||
out
|
||
}
|
||
|
||
/// RMSNorm forward: `y[r,c] = x[r,c] * inv_rms[r] * gamma[c]` with
|
||
/// `inv_rms = rsqrt(mean(x²) + eps)`. `self`:[rows,cols], `gamma`:[cols].
|
||
/// Returns `(y, inv_rms)`; `inv_rms`:[rows] is cached for backward.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn rms_norm(&self, gamma: &Tensor, eps: f32) -> (Tensor, Tensor) {
|
||
assert_eq!(self.ndim(), 2, "rms_norm requires 2D input");
|
||
assert_eq!(gamma.ndim(), 1, "gamma must be 1D");
|
||
assert_eq!(self.shape[1], gamma.shape[0], "gamma len != cols");
|
||
// bf16: compute the reduction in fp32 (standard AMP), downcast y back to
|
||
// bf16. inv_rms stays fp32 (the cache the fp32 backward kernel consumes).
|
||
if self.dtype == DType::BF16 {
|
||
let (y, inv_rms) = self
|
||
.to_dtype(DType::F32)
|
||
.rms_norm(&gamma.to_dtype(DType::F32), eps);
|
||
return (y.to_dtype(DType::BF16), inv_rms);
|
||
}
|
||
let (rows, cols) = (self.shape[0], self.shape[1]);
|
||
let y = Tensor::zeros(&self.shape, DType::F32, self.device());
|
||
let inv_rms = Tensor::zeros(&[rows], DType::F32, self.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_rms_norm_f32(
|
||
self.data_ptr() as *const f32,
|
||
gamma.data_ptr() as *const f32,
|
||
y.data_ptr() as *mut f32,
|
||
inv_rms.data_ptr() as *mut f32,
|
||
rows as i32,
|
||
cols as i32,
|
||
eps,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
(y, inv_rms)
|
||
}
|
||
|
||
/// RMSNorm backward. Inputs are the forward `x`, `gamma`, upstream `dy`, and
|
||
/// the cached `inv_rms`. Returns `(dx, dgamma)`.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn rms_norm_backward(
|
||
x: &Tensor,
|
||
gamma: &Tensor,
|
||
dy: &Tensor,
|
||
inv_rms: &Tensor,
|
||
) -> (Tensor, Tensor) {
|
||
// bf16: upcast (x, gamma, dy) to fp32, run the fp32 backward, downcast the
|
||
// grads back to bf16 (inv_rms is already the fp32 cache).
|
||
if x.dtype == DType::BF16 {
|
||
let (dx, dgamma) = Tensor::rms_norm_backward(
|
||
&x.to_dtype(DType::F32),
|
||
&gamma.to_dtype(DType::F32),
|
||
&dy.to_dtype(DType::F32),
|
||
inv_rms,
|
||
);
|
||
return (dx.to_dtype(DType::BF16), dgamma.to_dtype(DType::BF16));
|
||
}
|
||
let (rows, cols) = (x.shape[0], x.shape[1]);
|
||
let dx = Tensor::zeros(&[rows, cols], DType::F32, x.device());
|
||
let dgamma = Tensor::zeros(&[cols], DType::F32, x.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_rms_norm_dx_f32(
|
||
x.data_ptr() as *const f32,
|
||
gamma.data_ptr() as *const f32,
|
||
dy.data_ptr() as *const f32,
|
||
inv_rms.data_ptr() as *const f32,
|
||
dx.data_ptr() as *mut f32,
|
||
rows as i32,
|
||
cols as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
xtrain_cuda::ffi::launch_rms_norm_dgamma_f32(
|
||
x.data_ptr() as *const f32,
|
||
dy.data_ptr() as *const f32,
|
||
inv_rms.data_ptr() as *const f32,
|
||
dgamma.data_ptr() as *mut f32,
|
||
rows as i32,
|
||
cols as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
(dx, dgamma)
|
||
}
|
||
|
||
/// SiLU forward: `y = x * sigmoid(x)`, elementwise.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn silu(&self) -> Self {
|
||
assert!(
|
||
matches!(self.dtype, DType::F32 | DType::BF16),
|
||
"silu supports F32/BF16"
|
||
);
|
||
let out = Tensor::zeros(&self.shape, self.dtype, self.device());
|
||
let n = self.numel() as i32;
|
||
match self.dtype {
|
||
DType::F32 => unsafe {
|
||
xtrain_cuda::ffi::launch_silu_f32(
|
||
self.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut f32,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
DType::BF16 => unsafe {
|
||
xtrain_cuda::ffi::launch_silu_bf16(
|
||
self.data_ptr() as *const std::ffi::c_void,
|
||
out.data_ptr() as *mut std::ffi::c_void,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
_ => unreachable!(),
|
||
}
|
||
out
|
||
}
|
||
|
||
/// SiLU backward: `dx = dy * (sig + x*sig*(1-sig))`, `sig = sigmoid(x)`.
|
||
/// Inputs are the forward `x` and upstream `dy`.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn silu_backward(x: &Tensor, dy: &Tensor) -> Self {
|
||
let dx = Tensor::zeros(&x.shape, x.dtype, x.device());
|
||
let n = x.numel() as i32;
|
||
match x.dtype {
|
||
DType::F32 => unsafe {
|
||
xtrain_cuda::ffi::launch_silu_dx_f32(
|
||
x.data_ptr() as *const f32,
|
||
dy.data_ptr() as *const f32,
|
||
dx.data_ptr() as *mut f32,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
DType::BF16 => unsafe {
|
||
xtrain_cuda::ffi::launch_silu_dx_bf16(
|
||
x.data_ptr() as *const std::ffi::c_void,
|
||
dy.data_ptr() as *const std::ffi::c_void,
|
||
dx.data_ptr() as *mut std::ffi::c_void,
|
||
n,
|
||
std::ptr::null_mut(),
|
||
);
|
||
},
|
||
_ => panic!("silu_backward supports F32/BF16"),
|
||
}
|
||
dx
|
||
}
|
||
|
||
/// 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, 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"
|
||
);
|
||
if self.dtype == DType::BF16 {
|
||
return self
|
||
.to_dtype(DType::F32)
|
||
.rope(theta, period)
|
||
.to_dtype(DType::BF16);
|
||
}
|
||
let out = Tensor::zeros(&self.shape, DType::F32, self.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_rope_f32(
|
||
self.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut f32,
|
||
tokens as i32,
|
||
heads as i32,
|
||
head_dim as i32,
|
||
theta,
|
||
period as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
out
|
||
}
|
||
|
||
/// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
|
||
/// orthogonal map, so it needs no cached forward values, only `theta`/`period`.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn rope_backward(dy: &Tensor, theta: f32, period: usize) -> Self {
|
||
if dy.dtype == DType::BF16 {
|
||
return Tensor::rope_backward(&dy.to_dtype(DType::F32), theta, period)
|
||
.to_dtype(DType::BF16);
|
||
}
|
||
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 {
|
||
xtrain_cuda::ffi::launch_rope_dx_f32(
|
||
dy.data_ptr() as *const f32,
|
||
dx.data_ptr() as *mut f32,
|
||
tokens as i32,
|
||
heads as i32,
|
||
head_dim as i32,
|
||
theta,
|
||
period as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
dx
|
||
}
|
||
|
||
/// Row-wise safe softmax over the last dim. `self`:[rows,cols].
|
||
#[cfg(not(no_cuda))]
|
||
pub fn softmax(&self) -> Self {
|
||
assert_eq!(self.ndim(), 2, "softmax requires 2D input");
|
||
if self.dtype == DType::BF16 {
|
||
return self.to_dtype(DType::F32).softmax().to_dtype(DType::BF16);
|
||
}
|
||
let (rows, cols) = (self.shape[0], self.shape[1]);
|
||
let out = Tensor::zeros(&self.shape, DType::F32, self.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_softmax_f32(
|
||
self.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut f32,
|
||
rows as i32,
|
||
cols as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Softmax backward (Jacobian): `dx[r,c] = y[r,c]*(dy[r,c] - sum_c'(dy*y))`.
|
||
/// Inputs are the forward output `y` and upstream `dy`.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn softmax_backward(y: &Tensor, dy: &Tensor) -> Self {
|
||
if y.dtype == DType::BF16 {
|
||
return Tensor::softmax_backward(&y.to_dtype(DType::F32), &dy.to_dtype(DType::F32))
|
||
.to_dtype(DType::BF16);
|
||
}
|
||
let (rows, cols) = (y.shape[0], y.shape[1]);
|
||
let dx = Tensor::zeros(&y.shape, DType::F32, y.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_softmax_dx_f32(
|
||
y.data_ptr() as *const f32,
|
||
dy.data_ptr() as *const f32,
|
||
dx.data_ptr() as *mut f32,
|
||
rows as i32,
|
||
cols as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
dx
|
||
}
|
||
|
||
/// Cross-entropy forward over logits `self`:[rows,cols] with one I32 target
|
||
/// per row. Returns `(probs, loss)` where `probs`:[rows,cols] is the softmax
|
||
/// (cached for backward) and `loss`:[rows] is the per-row negative log-likelihood.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn cross_entropy(&self, target: &Tensor) -> (Tensor, Tensor) {
|
||
assert_eq!(self.ndim(), 2, "cross_entropy requires 2D logits");
|
||
assert_eq!(target.dtype, DType::I32, "target must be I32");
|
||
assert_eq!(target.numel(), self.shape[0], "one target per row");
|
||
// CE math (log-sum-exp) is fp32 (probs/loss cached fp32). The model casts
|
||
// logits→fp32 before CE; this guard keeps the op robust to bf16 logits.
|
||
if self.dtype == DType::BF16 {
|
||
return self.to_dtype(DType::F32).cross_entropy(target);
|
||
}
|
||
let (rows, cols) = (self.shape[0], self.shape[1]);
|
||
let probs = Tensor::zeros(&self.shape, DType::F32, self.device());
|
||
let loss = Tensor::zeros(&[rows], DType::F32, self.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_cross_entropy_fwd_f32(
|
||
self.data_ptr() as *const f32,
|
||
target.data_ptr() as *const i32,
|
||
probs.data_ptr() as *mut f32,
|
||
loss.data_ptr() as *mut f32,
|
||
rows as i32,
|
||
cols as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
(probs, loss)
|
||
}
|
||
|
||
/// Cross-entropy backward: `dx = scale * (probs - onehot(target))`. With
|
||
/// `scale = upstream / rows`, this is the gradient of the mean per-row loss.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn cross_entropy_backward(probs: &Tensor, target: &Tensor, scale: f32) -> Self {
|
||
let (rows, cols) = (probs.shape[0], probs.shape[1]);
|
||
let dx = Tensor::zeros(&probs.shape, DType::F32, probs.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_cross_entropy_dx_f32(
|
||
probs.data_ptr() as *const f32,
|
||
target.data_ptr() as *const i32,
|
||
dx.data_ptr() as *mut f32,
|
||
rows as i32,
|
||
cols as i32,
|
||
scale,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
dx
|
||
}
|
||
|
||
// --- Structural / model ops (the T5 kernels) ---
|
||
|
||
/// Reshape to `new_shape` (must keep `numel`). Pure metadata change on a
|
||
/// contiguous tensor — no data movement, shares the same storage. The
|
||
/// multi-head layout `[seq, n_heads*head_dim] <-> [seq, n_heads, head_dim]`
|
||
/// is exactly this.
|
||
pub fn reshape(&self, new_shape: &[usize]) -> Self {
|
||
assert!(self.is_contiguous(), "reshape requires a contiguous tensor");
|
||
assert_eq!(
|
||
shape::num_elements(new_shape),
|
||
self.numel(),
|
||
"reshape numel mismatch: {:?} -> {:?}",
|
||
self.shape.as_slice(),
|
||
new_shape
|
||
);
|
||
Self {
|
||
storage: self.storage.clone(),
|
||
shape: Dims::from_slice(new_shape),
|
||
strides: shape::contiguous_strides(new_shape),
|
||
offset: self.offset,
|
||
dtype: self.dtype,
|
||
}
|
||
}
|
||
|
||
/// Embedding gather: `out[s,:] = self[ids[s], :]`. `self`:[vocab,dim] table,
|
||
/// `ids`:[seq] I32 → out:[seq,dim].
|
||
#[cfg(not(no_cuda))]
|
||
pub fn embedding(&self, ids: &Tensor) -> Self {
|
||
assert_eq!(self.dtype, DType::F32, "embedding table must be F32");
|
||
assert_eq!(self.ndim(), 2, "embedding table must be [vocab,dim]");
|
||
assert_eq!(ids.dtype, DType::I32, "embedding ids must be I32");
|
||
assert_eq!(ids.ndim(), 1, "embedding ids must be 1D");
|
||
let (seq, dim) = (ids.shape[0], self.shape[1]);
|
||
let out = Tensor::zeros(&[seq, dim], DType::F32, self.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_embedding_fwd_f32(
|
||
self.data_ptr() as *const f32,
|
||
ids.data_ptr() as *const i32,
|
||
out.data_ptr() as *mut f32,
|
||
seq as i32,
|
||
dim as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Embedding backward (scatter-add): `dtable[ids[s],:] += dout[s,:]`, where
|
||
/// `dout`:[seq,dim], `ids`:[seq] I32. `vocab` sizes the output table.
|
||
#[cfg(not(no_cuda))]
|
||
pub fn embedding_backward(dout: &Tensor, ids: &Tensor, vocab: usize) -> Self {
|
||
let (seq, dim) = (dout.shape[0], dout.shape[1]);
|
||
let dtable = Tensor::zeros(&[vocab, dim], DType::F32, dout.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_embedding_bwd_f32(
|
||
dout.data_ptr() as *const f32,
|
||
ids.data_ptr() as *const i32,
|
||
dtable.data_ptr() as *mut f32,
|
||
seq as i32,
|
||
dim as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
dtable
|
||
}
|
||
|
||
/// 3D axis-(0,1) transpose: `self`:[a,b,c] → [b,a,c], `out[j,i,k]=self[i,j,k]`.
|
||
/// Lays out multi-head attention (`[seq,heads,hd] <-> [heads,seq,hd]`). Its
|
||
/// own backward is the same op (swap a,b).
|
||
#[cfg(not(no_cuda))]
|
||
pub fn transpose_3d01(&self) -> Self {
|
||
assert_eq!(self.ndim(), 3, "transpose_3d01 requires a 3D tensor");
|
||
assert!(self.is_contiguous(), "transpose_3d01 requires contiguous");
|
||
if self.dtype == DType::BF16 {
|
||
return self
|
||
.to_dtype(DType::F32)
|
||
.transpose_3d01()
|
||
.to_dtype(DType::BF16);
|
||
}
|
||
assert_eq!(self.dtype, DType::F32, "transpose_3d01 supports F32/BF16");
|
||
let (a, b, c) = (self.shape[0], self.shape[1], self.shape[2]);
|
||
let out = Tensor::zeros(&[b, a, c], DType::F32, self.device());
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_transpose_3d01_f32(
|
||
self.data_ptr() as *const f32,
|
||
out.data_ptr() as *mut f32,
|
||
a as i32,
|
||
b as i32,
|
||
c as i32,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
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");
|
||
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;
|
||
|
||
// scores[bh,seq,seq] = Q[bh,seq,hd] · Kᵀ[bh,hd,seq] (GEMM in self dtype)
|
||
let scores = Tensor::zeros(&[bh, seq, seq], dt, dev);
|
||
strided_batched_gemm(
|
||
dt,
|
||
false,
|
||
true,
|
||
seq,
|
||
seq,
|
||
hd,
|
||
self.data_ptr(),
|
||
seq * hd,
|
||
k.data_ptr(),
|
||
seq * hd,
|
||
scores.data_ptr(),
|
||
seq * seq,
|
||
bh,
|
||
);
|
||
// probs = softmax(causal(scores · scale)). Softmax math is fp32 (stable);
|
||
// for bf16 we upcast scores → f32 → kernel → downcast probs back to bf16
|
||
// (so the cached probs activation is half-size). One block per [bh·seq] row.
|
||
let scores_f32 = scores.to_dtype(DType::F32);
|
||
let probs_f32 = Tensor::zeros(&[bh, seq, seq], DType::F32, dev);
|
||
unsafe {
|
||
xtrain_cuda::ffi::launch_softmax_causal_f32(
|
||
scores_f32.data_ptr() as *const f32,
|
||
probs_f32.data_ptr() as *mut f32,
|
||
(bh * seq) as i32,
|
||
seq as i32,
|
||
scale,
|
||
std::ptr::null_mut(),
|
||
);
|
||
}
|
||
let probs = probs_f32.to_dtype(dt);
|
||
// out[bh,seq,hd] = probs[bh,seq,seq] · V[bh,seq,hd]
|
||
let out = Tensor::zeros(&[bh, seq, hd], dt, dev);
|
||
strided_batched_gemm(
|
||
dt,
|
||
false,
|
||
false,
|
||
seq,
|
||
hd,
|
||
seq,
|
||
probs.data_ptr(),
|
||
seq * seq,
|
||
v.data_ptr(),
|
||
seq * hd,
|
||
out.data_ptr(),
|
||
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();
|
||
let dt = q.dtype;
|
||
|
||
// dP[bh,seq,seq] = dOut[bh,seq,hd] · Vᵀ[bh,hd,seq]
|
||
let dp = Tensor::zeros(&[bh, seq, seq], dt, dev);
|
||
strided_batched_gemm(
|
||
dt,
|
||
false,
|
||
true,
|
||
seq,
|
||
seq,
|
||
hd,
|
||
dout.data_ptr(),
|
||
seq * hd,
|
||
v.data_ptr(),
|
||
seq * hd,
|
||
dp.data_ptr(),
|
||
seq * seq,
|
||
bh,
|
||
);
|
||
// dV[bh,seq,hd] = Pᵀ[bh,seq,seq] · dOut[bh,seq,hd]
|
||
let dv = Tensor::zeros(&[bh, seq, hd], dt, dev);
|
||
strided_batched_gemm(
|
||
dt,
|
||
true,
|
||
false,
|
||
seq,
|
||
hd,
|
||
seq,
|
||
probs.data_ptr(),
|
||
seq * seq,
|
||
dout.data_ptr(),
|
||
seq * hd,
|
||
dv.data_ptr(),
|
||
seq * hd,
|
||
bh,
|
||
);
|
||
// dScores = softmax Jacobian (per row) applied to dP, then ×scale.
|
||
// softmax_backward + scale are dtype-aware (fp32 math inside for bf16).
|
||
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], dt, dev);
|
||
strided_batched_gemm(
|
||
dt,
|
||
false,
|
||
false,
|
||
seq,
|
||
hd,
|
||
seq,
|
||
dscores.data_ptr(),
|
||
seq * seq,
|
||
k.data_ptr(),
|
||
seq * hd,
|
||
dq.data_ptr(),
|
||
seq * hd,
|
||
bh,
|
||
);
|
||
// dK[bh,seq,hd] = dScoresᵀ[bh,seq,seq] · Q[bh,seq,hd]
|
||
let dk = Tensor::zeros(&[bh, seq, hd], dt, dev);
|
||
strided_batched_gemm(
|
||
dt,
|
||
true,
|
||
false,
|
||
seq,
|
||
hd,
|
||
seq,
|
||
dscores.data_ptr(),
|
||
seq * seq,
|
||
q.data_ptr(),
|
||
seq * hd,
|
||
dk.data_ptr(),
|
||
seq * hd,
|
||
bh,
|
||
);
|
||
(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).
|
||
#[cfg(not(no_cuda))]
|
||
pub fn transpose_4d12(&self) -> Self {
|
||
assert_eq!(self.ndim(), 4, "transpose_4d12 requires a 4D tensor");
|
||
assert!(self.is_contiguous(), "transpose_4d12 requires contiguous");
|
||
if self.dtype == DType::BF16 {
|
||
return self
|
||
.to_dtype(DType::F32)
|
||
.transpose_4d12()
|
||
.to_dtype(DType::BF16);
|
||
}
|
||
assert_eq!(self.dtype, DType::F32, "transpose_4d12 supports F32/BF16");
|
||
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) {
|
||
assert!(
|
||
matches!(self.dtype, DType::F32 | DType::BF16),
|
||
"{op} supports F32/BF16"
|
||
);
|
||
assert_eq!(self.dtype, other.dtype, "{op} dtype mismatch");
|
||
assert_eq!(self.shape(), other.shape(), "{op} shape mismatch");
|
||
assert_eq!(self.device(), other.device(), "{op} device mismatch");
|
||
assert!(
|
||
self.is_contiguous() && other.is_contiguous(),
|
||
"{op} requires contiguous tensors"
|
||
);
|
||
}
|
||
}
|
||
|
||
/// Dispatch a strided-batched GEMM on `dt`: fp32 → `sgemm_strided_batched`,
|
||
/// bf16 → `gemm_ex_strided_batched` (bf16 in/out, fp32 accum). Pointers are the
|
||
/// raw `data_ptr()` bytes of contiguous same-dtype tensors. `alpha=1, beta=0`.
|
||
/// The fp32 path is bit-identical to the inlined T10 call it replaces.
|
||
#[cfg(not(no_cuda))]
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn strided_batched_gemm(
|
||
dt: DType,
|
||
trans_a: bool,
|
||
trans_b: bool,
|
||
m: usize,
|
||
n: usize,
|
||
k: usize,
|
||
a: *const u8,
|
||
stride_a: usize,
|
||
b: *const u8,
|
||
stride_b: usize,
|
||
c: *const u8,
|
||
stride_c: usize,
|
||
batch: usize,
|
||
) {
|
||
match dt {
|
||
DType::F32 => xtrain_cuda::cublas::sgemm_strided_batched(
|
||
trans_a,
|
||
trans_b,
|
||
m,
|
||
n,
|
||
k,
|
||
1.0,
|
||
a as *const f32,
|
||
stride_a,
|
||
b as *const f32,
|
||
stride_b,
|
||
0.0,
|
||
c as *mut f32,
|
||
stride_c,
|
||
batch,
|
||
),
|
||
DType::BF16 => xtrain_cuda::cublas::gemm_ex_strided_batched(
|
||
trans_a,
|
||
trans_b,
|
||
m,
|
||
n,
|
||
k,
|
||
1.0,
|
||
a as *const std::ffi::c_void,
|
||
stride_a,
|
||
b as *const std::ffi::c_void,
|
||
stride_b,
|
||
0.0,
|
||
c as *mut std::ffi::c_void,
|
||
stride_c,
|
||
batch,
|
||
),
|
||
_ => panic!("strided_batched_gemm supports F32/BF16"),
|
||
}
|
||
}
|
||
|
||
impl std::fmt::Debug for Tensor {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
write!(
|
||
f,
|
||
"Tensor(shape={:?}, dtype={}, device={}, contiguous={})",
|
||
self.shape.as_slice(),
|
||
self.dtype,
|
||
self.device(),
|
||
self.is_contiguous()
|
||
)
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn from_slice_shape_and_data() {
|
||
let t = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
|
||
assert_eq!(t.shape(), &[2, 3]);
|
||
assert_eq!(t.strides(), &[3, 1]);
|
||
assert_eq!(t.numel(), 6);
|
||
assert_eq!(t.device(), Device::Cpu);
|
||
assert!(t.is_contiguous());
|
||
assert_eq!(t.as_slice::<f32>(), &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
|
||
}
|
||
|
||
#[test]
|
||
fn zeros_cpu() {
|
||
let t = Tensor::zeros(&[4], DType::F32, Device::Cpu);
|
||
assert_eq!(t.as_slice::<f32>(), &[0.0, 0.0, 0.0, 0.0]);
|
||
}
|
||
}
|