phase 5: naive multi-head attention
- Batched GEMM via cublasGemmStridedBatchedEx - Causal mask CUDA kernel (F32 + BF16) - Element-wise scale CUDA kernel (F32 + BF16) - attention() composing: batched_matmul + scale + causal_mask + softmax - Fixed to_device/contiguous infinite recursion (GPU contiguous via CPU round-trip) - 5 attention tests passing (max_err < 3e-7 F32) - Total: 61 tests passing across all crates Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -46,6 +46,19 @@ unsafe extern "C" {
|
||||
compute_type: i32,
|
||||
algo: i32,
|
||||
) -> i32;
|
||||
fn cublasGemmStridedBatchedEx(
|
||||
handle: CublasHandle,
|
||||
transa: i32, transb: i32,
|
||||
m: i32, n: i32, k: i32,
|
||||
alpha: *const c_void,
|
||||
a: *const c_void, a_type: i32, lda: i32, stride_a: i64,
|
||||
b: *const c_void, b_type: i32, ldb: i32, stride_b: i64,
|
||||
beta: *const c_void,
|
||||
c: *mut c_void, c_type: i32, ldc: i32, stride_c: i64,
|
||||
batch_count: i32,
|
||||
compute_type: i32,
|
||||
algo: i32,
|
||||
) -> i32;
|
||||
}
|
||||
|
||||
pub struct CublasContext {
|
||||
@@ -149,3 +162,68 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
|
||||
|
||||
c
|
||||
}
|
||||
|
||||
/// Batched matrix multiplication via cuBLAS: C[b] = A[b] @ B[b]
|
||||
/// a: [..., M, K], b: [..., K, N] → [..., M, N]
|
||||
/// Leading dimensions must match and tensors must be contiguous.
|
||||
pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
|
||||
assert!(a.ndim() >= 2 && b.ndim() >= 2);
|
||||
assert_eq!(a.ndim(), b.ndim());
|
||||
assert!(a.is_contiguous() && b.is_contiguous());
|
||||
assert!(matches!(a.device(), Device::Cuda(_)));
|
||||
assert_eq!(a.dtype(), b.dtype());
|
||||
|
||||
let ndim = a.ndim();
|
||||
let m = a.shape()[ndim - 2];
|
||||
let k = a.shape()[ndim - 1];
|
||||
let n = b.shape()[ndim - 1];
|
||||
assert_eq!(b.shape()[ndim - 2], k, "inner dimension mismatch");
|
||||
|
||||
// Compute batch count from leading dimensions
|
||||
let batch: usize = a.shape()[..ndim - 2].iter().product();
|
||||
assert_eq!(
|
||||
b.shape()[..ndim - 2].iter().product::<usize>(),
|
||||
batch,
|
||||
"batch dimensions mismatch"
|
||||
);
|
||||
|
||||
let mut out_shape: Vec<usize> = a.shape()[..ndim - 2].to_vec();
|
||||
out_shape.push(m);
|
||||
out_shape.push(n);
|
||||
let c = Tensor::zeros(&out_shape, a.dtype(), a.device());
|
||||
|
||||
let dtype = a.dtype();
|
||||
let (a_type, b_type, c_type) = match dtype {
|
||||
DType::F32 => (CUDA_R_32F, CUDA_R_32F, CUDA_R_32F),
|
||||
DType::BF16 => (CUDA_R_16BF, CUDA_R_16BF, CUDA_R_16BF),
|
||||
_ => panic!("unsupported dtype for batched matmul"),
|
||||
};
|
||||
|
||||
let alpha = 1.0f32;
|
||||
let beta = 0.0f32;
|
||||
// cuBLAS strides are in elements (not bytes)
|
||||
let stride_a = (m * k) as i64;
|
||||
let stride_b = (k * n) as i64;
|
||||
let stride_c = (m * n) as i64;
|
||||
|
||||
let ctx = CublasContext::new().unwrap();
|
||||
unsafe {
|
||||
cublasSetStream_v2(ctx.handle, std::ptr::null_mut());
|
||||
// Row-major trick: C = A @ B ⟺ C^T = B^T @ A^T (col-major)
|
||||
error::check(cublasGemmStridedBatchedEx(
|
||||
ctx.handle,
|
||||
CUBLAS_OP_N, CUBLAS_OP_N,
|
||||
n as i32, m as i32, k as i32,
|
||||
&alpha as *const f32 as *const c_void,
|
||||
b.data_ptr() as _, b_type, n as i32, stride_b,
|
||||
a.data_ptr() as _, a_type, k as i32, stride_a,
|
||||
&beta as *const f32 as *const c_void,
|
||||
c.data_ptr() as *mut c_void, c_type, n as i32, stride_c,
|
||||
batch as i32,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
-1,
|
||||
)).expect("cuBLAS batched GEMM failed");
|
||||
}
|
||||
xserv_cuda::device::synchronize().unwrap();
|
||||
c
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user