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:
2026-05-21 21:17:23 +08:00
parent c8e8153702
commit 6035ffdc0b
10 changed files with 550 additions and 12 deletions

View File

@@ -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
}