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>

crates/xtrain-cuda/src/cublas.rs

@@ -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}");
+    });
+}