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/argmax.rs

@@ -0,0 +1,65 @@
+use std::ffi::c_void;
+use xserv_tensor::{DType, Device, Tensor};
+
+unsafe extern "C" {
+    fn launch_argmax_bf16(logits: *const c_void, out_idx: *mut c_void,
+                          rows: i32, cols: i32, stream: *mut c_void);
+}
+
+/// GPU argmax over the last dim of a [rows, cols] BF16 tensor.
+///
+/// Returns a host `Vec<u32>` of length `rows`. Internally:
+///   - launches one kernel that writes [rows] i32 indices on device
+///   - D2H copies just `rows * 4` bytes (vs `rows * cols * 2` for the
+///     "copy logits to CPU then argmax" path it replaces)
+///
+/// This is the greedy-decode hot path: avoids touching the full
+/// [B, vocab] logits buffer on the host every step.
+pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
+    assert_eq!(logits.ndim(), 2, "argmax expects a 2D [rows, cols] tensor");
+    assert_eq!(logits.dtype(), DType::BF16, "argmax kernel is BF16-only");
+    assert!(logits.is_contiguous(), "argmax requires contiguous input");
+    assert!(matches!(logits.device(), Device::Cuda(_)), "argmax requires GPU input");
+
+    let rows = logits.shape()[0];
+    let cols = logits.shape()[1];
+    assert!(rows <= i32::MAX as usize);
+    assert!(cols <= i32::MAX as usize);
+
+    // Output buffer: rows * i32. Pooled allocator so this is essentially free
+    // after the first call.
+    let bytes = rows * std::mem::size_of::<i32>();
+    let mut out = xserv_cuda::allocator::cached_alloc(bytes).expect("argmax out alloc");
+
+    unsafe {
+        launch_argmax_bf16(
+            logits.data_ptr() as *const c_void,
+            out.as_mut_ptr() as *mut c_void,
+            rows as i32, cols as i32,
+            std::ptr::null_mut(),
+        );
+    }
+
+    let mut host_bytes = vec![0u8; bytes];
+    out.copy_to_host(&mut host_bytes).expect("argmax D2H");
+    drop(out); // returned to pool
+
+    let host_i32: &[i32] = unsafe {
+        std::slice::from_raw_parts(host_bytes.as_ptr() as *const i32, rows)
+    };
+    host_i32.iter().map(|&v| v as u32).collect()
+}
+
+/// Convenience: argmax of a single row [1, cols] (or [cols] reshaped to [1, cols]).
+pub fn argmax_bf16_single(logits: &Tensor) -> u32 {
+    let cols = *logits.shape().last().unwrap();
+    let rows = logits.numel() / cols;
+    assert_eq!(rows, 1, "argmax_bf16_single requires a single row");
+    let view = if logits.ndim() == 2 {
+        logits.clone()
+    } else {
+        logits.reshape(&[1, cols])
+    };
+    argmax_bf16_to_host(&view)[0]
+}
+