cuda: bf16 cuBLAS GemmEx (16BF in/out, fp32 accum) + cast kernels
Add the bf16 compute primitives for T12 mixed precision: - DType::BF16 (half::bf16 as TensorDType), 2 bytes. - cublasGemmEx / cublasGemmStridedBatchedEx FFI + CUDA_R_16BF / CUBLAS_COMPUTE_32F constants (values per xserv gemm.rs). - cublas::gemm_ex / gemm_ex_strided_batched: same row-major⟺col-major transpose algebra as sgemm, bf16 in/out, fp32 accumulation. - csrc/ops/cast.cu: f32<->bf16 cast + bf16 elementwise (add/mul/scale/ silu(+dx)/add_bias/sum_rows), each load->fp32->compute->store bf16. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -19,6 +19,7 @@
|
||||
|
||||
use crate::ffi::{self, CublasHandle};
|
||||
use std::cell::RefCell;
|
||||
use std::ffi::c_void;
|
||||
|
||||
thread_local! {
|
||||
static HANDLE: RefCell<Option<CublasHandle>> = const { RefCell::new(None) };
|
||||
@@ -159,3 +160,131 @@ pub fn sgemm_strided_batched(
|
||||
assert_eq!(status, 0, "cublasSgemmStridedBatched failed: {status}");
|
||||
});
|
||||
}
|
||||
|
||||
/// bf16 row-major GEMM `C[m,n] = opA(A)·opB(B)` via `cublasGemmEx`: bf16 in/out,
|
||||
/// **fp32 accumulation** (`CUBLAS_COMPUTE_32F`) — the standard AMP matmul (Phase
|
||||
/// T12). `a`/`b`/`c` are device pointers to row-major **bf16** matrices; the
|
||||
/// row-major⟺col-major transpose algebra is identical to [`sgemm`] (we compute
|
||||
/// the col-major `Cᵀ`). `alpha`/`beta` are fp32 host scalars (compute is fp32).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn gemm_ex(
|
||||
trans_a: bool,
|
||||
trans_b: bool,
|
||||
m: usize,
|
||||
n: usize,
|
||||
k: usize,
|
||||
alpha: f32,
|
||||
a: *const c_void,
|
||||
b: *const c_void,
|
||||
beta: f32,
|
||||
c: *mut c_void,
|
||||
) {
|
||||
let lda = if trans_a { m } else { k };
|
||||
let ldb = if trans_b { k } else { n };
|
||||
let ldc = n;
|
||||
let op_a = if trans_a {
|
||||
ffi::CUBLAS_OP_T
|
||||
} else {
|
||||
ffi::CUBLAS_OP_N
|
||||
};
|
||||
let op_b = if trans_b {
|
||||
ffi::CUBLAS_OP_T
|
||||
} else {
|
||||
ffi::CUBLAS_OP_N
|
||||
};
|
||||
let bf16 = ffi::CUDA_R_16BF;
|
||||
|
||||
with_handle(|handle| {
|
||||
let status = unsafe {
|
||||
ffi::cublasGemmEx(
|
||||
handle,
|
||||
op_b,
|
||||
op_a,
|
||||
n as i32,
|
||||
m as i32,
|
||||
k as i32,
|
||||
&alpha as *const f32 as *const c_void,
|
||||
b,
|
||||
bf16,
|
||||
ldb as i32,
|
||||
a,
|
||||
bf16,
|
||||
lda as i32,
|
||||
&beta as *const f32 as *const c_void,
|
||||
c,
|
||||
bf16,
|
||||
ldc as i32,
|
||||
ffi::CUBLAS_COMPUTE_32F,
|
||||
ffi::CUBLAS_GEMM_DEFAULT,
|
||||
)
|
||||
};
|
||||
assert_eq!(status, 0, "cublasGemmEx failed: {status}");
|
||||
});
|
||||
}
|
||||
|
||||
/// Strided-batched bf16 GEMM (Phase T12) — the [`gemm_ex`] analogue of
|
||||
/// [`sgemm_strided_batched`] for the batched attention GEMMs. bf16 in/out, fp32
|
||||
/// accumulation; strides are in ELEMENTS.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn gemm_ex_strided_batched(
|
||||
trans_a: bool,
|
||||
trans_b: bool,
|
||||
m: usize,
|
||||
n: usize,
|
||||
k: usize,
|
||||
alpha: f32,
|
||||
a: *const c_void,
|
||||
stride_a: usize,
|
||||
b: *const c_void,
|
||||
stride_b: usize,
|
||||
beta: f32,
|
||||
c: *mut c_void,
|
||||
stride_c: usize,
|
||||
batch: usize,
|
||||
) {
|
||||
let lda = if trans_a { m } else { k };
|
||||
let ldb = if trans_b { k } else { n };
|
||||
let ldc = n;
|
||||
let op_a = if trans_a {
|
||||
ffi::CUBLAS_OP_T
|
||||
} else {
|
||||
ffi::CUBLAS_OP_N
|
||||
};
|
||||
let op_b = if trans_b {
|
||||
ffi::CUBLAS_OP_T
|
||||
} else {
|
||||
ffi::CUBLAS_OP_N
|
||||
};
|
||||
let bf16 = ffi::CUDA_R_16BF;
|
||||
|
||||
with_handle(|handle| {
|
||||
let status = unsafe {
|
||||
ffi::cublasGemmStridedBatchedEx(
|
||||
handle,
|
||||
op_b,
|
||||
op_a,
|
||||
n as i32,
|
||||
m as i32,
|
||||
k as i32,
|
||||
&alpha as *const f32 as *const c_void,
|
||||
b,
|
||||
bf16,
|
||||
ldb as i32,
|
||||
stride_b as i64,
|
||||
a,
|
||||
bf16,
|
||||
lda as i32,
|
||||
stride_a as i64,
|
||||
&beta as *const f32 as *const c_void,
|
||||
c,
|
||||
bf16,
|
||||
ldc as i32,
|
||||
stride_c as i64,
|
||||
batch as i32,
|
||||
ffi::CUBLAS_COMPUTE_32F,
|
||||
ffi::CUBLAS_GEMM_DEFAULT,
|
||||
)
|
||||
};
|
||||
assert_eq!(status, 0, "cublasGemmStridedBatchedEx failed: {status}");
|
||||
});
|
||||
}
|
||||
|
||||
@@ -324,3 +324,126 @@ unsafe extern "C" {
|
||||
pub const CUBLAS_OP_N: i32 = 0;
|
||||
#[cfg(not(no_cuda))]
|
||||
pub const CUBLAS_OP_T: i32 = 1;
|
||||
|
||||
// --- bf16 mixed precision (Phase T12) ---
|
||||
//
|
||||
// cudaDataType / cublasComputeType enum values (same as xserv's gemm.rs). The
|
||||
// bf16 GEMM uses bf16 in/out with fp32 accumulation (CUBLAS_COMPUTE_32F).
|
||||
#[cfg(not(no_cuda))]
|
||||
pub const CUDA_R_32F: i32 = 0;
|
||||
#[cfg(not(no_cuda))]
|
||||
pub const CUDA_R_16BF: i32 = 14;
|
||||
#[cfg(not(no_cuda))]
|
||||
pub const CUBLAS_COMPUTE_32F: i32 = 68;
|
||||
/// CUBLAS_GEMM_DEFAULT — let cuBLAS pick the algorithm.
|
||||
#[cfg(not(no_cuda))]
|
||||
pub const CUBLAS_GEMM_DEFAULT: i32 = -1;
|
||||
|
||||
#[cfg(not(no_cuda))]
|
||||
unsafe extern "C" {
|
||||
// General GEMM with explicit in/out + compute types (bf16 path). `alpha`/
|
||||
// `beta` are fp32 host scalars (compute type is fp32). Pointers are void* so
|
||||
// the same FFI serves bf16 / fp32.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn cublasGemmEx(
|
||||
handle: CublasHandle,
|
||||
transa: i32,
|
||||
transb: i32,
|
||||
m: i32,
|
||||
n: i32,
|
||||
k: i32,
|
||||
alpha: *const std::ffi::c_void,
|
||||
a: *const std::ffi::c_void,
|
||||
a_type: i32,
|
||||
lda: i32,
|
||||
b: *const std::ffi::c_void,
|
||||
b_type: i32,
|
||||
ldb: i32,
|
||||
beta: *const std::ffi::c_void,
|
||||
c: *mut std::ffi::c_void,
|
||||
c_type: i32,
|
||||
ldc: i32,
|
||||
compute_type: i32,
|
||||
algo: i32,
|
||||
) -> i32;
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn cublasGemmStridedBatchedEx(
|
||||
handle: CublasHandle,
|
||||
transa: i32,
|
||||
transb: i32,
|
||||
m: i32,
|
||||
n: i32,
|
||||
k: i32,
|
||||
alpha: *const std::ffi::c_void,
|
||||
a: *const std::ffi::c_void,
|
||||
a_type: i32,
|
||||
lda: i32,
|
||||
stride_a: i64,
|
||||
b: *const std::ffi::c_void,
|
||||
b_type: i32,
|
||||
ldb: i32,
|
||||
stride_b: i64,
|
||||
beta: *const std::ffi::c_void,
|
||||
c: *mut std::ffi::c_void,
|
||||
c_type: i32,
|
||||
ldc: i32,
|
||||
stride_c: i64,
|
||||
batch_count: i32,
|
||||
compute_type: i32,
|
||||
algo: i32,
|
||||
) -> i32;
|
||||
}
|
||||
|
||||
// bf16 cast + elementwise kernels (csrc/ops/cast.cu). Pointers are void* (bf16
|
||||
// buffers); f32 sides are typed. The activation stream flows bf16; the math
|
||||
// accumulates in fp32 inside each kernel.
|
||||
#[cfg(not(no_cuda))]
|
||||
unsafe extern "C" {
|
||||
pub fn launch_cast_f32_to_bf16(input: *const f32, out: *mut c_void, n: i32, s: CudaStream);
|
||||
pub fn launch_cast_bf16_to_f32(input: *const c_void, out: *mut f32, n: i32, s: CudaStream);
|
||||
|
||||
pub fn launch_add_bf16(
|
||||
a: *const c_void,
|
||||
b: *const c_void,
|
||||
out: *mut c_void,
|
||||
n: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
pub fn launch_mul_bf16(
|
||||
a: *const c_void,
|
||||
b: *const c_void,
|
||||
out: *mut c_void,
|
||||
n: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
pub fn launch_scale_bf16(
|
||||
input: *const c_void,
|
||||
out: *mut c_void,
|
||||
alpha: f32,
|
||||
n: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
pub fn launch_silu_bf16(x: *const c_void, y: *mut c_void, n: i32, s: CudaStream);
|
||||
pub fn launch_silu_dx_bf16(
|
||||
x: *const c_void,
|
||||
dy: *const c_void,
|
||||
dx: *mut c_void,
|
||||
n: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
pub fn launch_add_bias_bf16(
|
||||
x: *const c_void,
|
||||
bias: *const c_void,
|
||||
out: *mut c_void,
|
||||
rows: i32,
|
||||
cols: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
pub fn launch_sum_rows_bf16(
|
||||
dout: *const c_void,
|
||||
dbias: *mut c_void,
|
||||
rows: i32,
|
||||
cols: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user