kernels: reshape_and_cache, GPU argmax, single-launch GEMV

Three new CUDA kernels and one rewrite:

- reshape_and_cache: scatter K/V into paged pool in a single kernel per
  layer, replacing the Rust-side per-token per-head cudaMemcpy loop.
  Includes both single-sequence (prefill) and batched (decode) variants.

- argmax: GPU-side BF16 argmax with warp-shuffle reduction. Greedy
  decode now only D2H-transfers B×4 bytes (token ids) instead of the
  full [B, vocab] logits tensor.

- GEMV rewrite: fused zero-init inside the K-split kernel eliminates
  the cudaMemsetAsync call, reducing launches from 3 to 2 per GEMV.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

crates/xserv-kernels/src/gemm.rs

@@ -1,8 +1,16 @@
 use std::cell::RefCell;
 use std::ffi::c_void;
 use xserv_cuda::error::{self, Result};
+use xserv_cuda::GpuBuffer;
 use xserv_tensor::{DType, Device, Tensor};
 
+const CUBLAS_WORKSPACE_BYTES: usize = 32 * 1024 * 1024;
+
+// GEMV: single-kernel, no FP32 temp buffer needed
+unsafe extern "C" {
+    fn launch_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void);
+}
+
 #[derive(Debug, Clone, Copy)]
 pub enum GemmBackend {
     Naive,
@@ -16,7 +24,6 @@ unsafe extern "C" {
     fn launch_gemm_naive_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
     fn launch_gemm_tiled_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
     fn launch_gemm_tiled_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
-    fn launch_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void);
 }
 
 // --- FFI: cuBLAS ---
@@ -36,6 +43,7 @@ unsafe extern "C" {
     fn cublasCreate_v2(handle: *mut CublasHandle) -> i32;
     fn cublasDestroy_v2(handle: CublasHandle) -> i32;
     fn cublasSetStream_v2(handle: CublasHandle, stream: *mut c_void) -> i32;
+    fn cublasSetWorkspace_v2(handle: CublasHandle, workspace: *mut c_void, size: usize) -> i32;
     fn cublasGemmEx(
         handle: CublasHandle,
         transa: i32, transb: i32,
@@ -65,13 +73,25 @@ unsafe extern "C" {
 
 pub struct CublasContext {
     handle: CublasHandle,
+    /// Dedicated 32 MiB workspace owned by this handle. Held to keep the GPU
+    /// buffer alive for the lifetime of the handle; cuBLAS reads/writes into
+    /// it during GEMM. Dropped after `cublasDestroy_v2` so cuBLAS can't touch
+    /// freed memory.
+    _workspace: Option<GpuBuffer>,
 }
 
 impl CublasContext {
     pub fn new() -> Result<Self> {
         let mut handle = std::ptr::null_mut();
         error::check(unsafe { cublasCreate_v2(&mut handle) })?;
-        Ok(Self { handle })
+        // Attach a per-handle workspace. cublasSetWorkspace requires the
+        // pointer to remain valid until destroy or until a new workspace is
+        // set, so we keep the GpuBuffer in this struct.
+        let mut workspace = GpuBuffer::alloc(CUBLAS_WORKSPACE_BYTES)?;
+        error::check(unsafe {
+            cublasSetWorkspace_v2(handle, workspace.as_mut_ptr() as *mut c_void, CUBLAS_WORKSPACE_BYTES)
+        })?;
+        Ok(Self { handle, _workspace: Some(workspace) })
     }
 }
 
@@ -80,6 +100,7 @@ impl Drop for CublasContext {
         if !self.handle.is_null() {
             unsafe { cublasDestroy_v2(self.handle) };
         }
+        // _workspace drops here, after cublasDestroy_v2 has released the handle.
     }
 }
 
@@ -152,7 +173,6 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
             }
         }
         GemmBackend::CuBlas => {
-            // Fast path: custom GEMV for M=1 BF16 (bandwidth-optimal decode)
             if m == 1 && dtype == DType::BF16 {
                 let mut fp32_buf = xserv_cuda::allocator::cached_alloc(n * 4).unwrap();
                 unsafe {
@@ -163,11 +183,7 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
                         null_stream,
                     );
                 }
-                // fp32_buf returned to caching allocator pool on drop
             } else {
-                // cuBLAS uses column-major, but we have row-major tensors.
-                // Trick: compute C^T = B^T @ A^T, which gives us C in row-major.
-                // cuBLAS sees our row-major data as column-major transposed.
                 let alpha = 1.0f32;
                 let beta = 0.0f32;
 
@@ -179,19 +195,17 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
 
                 with_cublas(|handle| unsafe {
                     cublasSetStream_v2(handle, null_stream);
-                    // Row-major trick: swap A/B and transpose flags
-                    // C(row-major) = A @ B  <=>  C^T(col-major) = B^T @ A^T
                     error::check(cublasGemmEx(
                         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_ptr, b_type, n as i32,  // B as col-major = B^T
-                        a_ptr, a_type, k as i32,  // A as col-major = A^T
+                        b_ptr, b_type, n as i32,
+                        a_ptr, a_type, k as i32,
                         &beta as *const f32 as *const c_void,
-                        c_ptr, c_type, n as i32,  // C as col-major = C^T
+                        c_ptr, c_type, n as i32,
                         CUBLAS_COMPUTE_32F,
-                        -1, // default algo
+                        -1,
                     )).expect("cuBLAS GEMM failed");
                 });
             }