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

@@ -33,6 +33,85 @@ unsafe extern "C" {
         head_dim: i32, max_blocks_per_seq: i32,
         scale: f32, stream: *mut c_void,
     );
+    fn launch_reshape_and_cache_bf16(
+        k_src: *const c_void, v_src: *const c_void,
+        k_pool: *mut c_void, v_pool: *mut c_void,
+        block_ids: *const c_void,
+        num_tokens: i32, num_heads: i32,
+        head_dim: i32, start_pos: i32, block_size: i32,
+        stream: *mut c_void,
+    );
+    fn launch_reshape_and_cache_batched_bf16(
+        k_src: *const c_void, v_src: *const c_void,
+        k_pool: *mut c_void, v_pool: *mut c_void,
+        block_tables: *const c_void, kv_lens: *const c_void,
+        batch: i32, num_heads: i32,
+        head_dim: i32, block_size: i32, max_blocks_per_seq: i32,
+        stream: *mut c_void,
+    );
+}
+
+/// Scatter `[num_kv_heads, num_tokens, head_dim]` BF16 K/V into a paged
+/// pool for a single sequence whose block table lives at `block_ids_gpu`
+/// (int32, on device).
+///
+/// `k_pool_ptr`/`v_pool_ptr` point to one layer's pool, of logical shape
+/// `[num_blocks_total, num_kv_heads, block_size, head_dim]`.
+///
+/// All pointers must be on the same GPU as the launching context.
+///
+/// # Safety
+/// Pointers must be valid GPU pointers with the documented layouts.
+/// `block_ids_gpu` must contain at least `(start_pos + num_tokens + block_size - 1) / block_size`
+/// valid physical block ids.
+pub unsafe fn reshape_and_cache_bf16(
+    k_src: *const c_void, v_src: *const c_void,
+    k_pool_ptr: *mut c_void, v_pool_ptr: *mut c_void,
+    block_ids_gpu: *const i32,
+    num_tokens: usize, num_heads: usize,
+    head_dim: usize, start_pos: usize, block_size: usize,
+    stream: *mut c_void,
+) {
+    unsafe {
+        launch_reshape_and_cache_bf16(
+            k_src, v_src,
+            k_pool_ptr, v_pool_ptr,
+            block_ids_gpu as *const c_void,
+            num_tokens as i32, num_heads as i32,
+            head_dim as i32, start_pos as i32, block_size as i32,
+            stream,
+        );
+    }
+}
+
+/// Batched scatter for the multi-sequence decode step. Reads
+/// `block_tables` (`[batch, max_blocks_per_seq]` int32 — same buffer the
+/// paged-attention kernel reads) and `kv_lens` (`[batch]` int32, current
+/// seq_len + 1 — i.e., the index of the just-written token + 1) so the
+/// caller doesn't need a separate per-step upload of block ids.
+///
+/// # Safety
+/// All pointers must be on the same GPU. `block_tables` and `kv_lens` must
+/// already be synced to the device for the active batch.
+pub unsafe fn reshape_and_cache_batched_bf16(
+    k_src: *const c_void, v_src: *const c_void,
+    k_pool_ptr: *mut c_void, v_pool_ptr: *mut c_void,
+    block_tables_gpu: *const i32, kv_lens_gpu: *const i32,
+    batch: usize, num_heads: usize,
+    head_dim: usize, block_size: usize, max_blocks_per_seq: usize,
+    stream: *mut c_void,
+) {
+    unsafe {
+        launch_reshape_and_cache_batched_bf16(
+            k_src, v_src,
+            k_pool_ptr, v_pool_ptr,
+            block_tables_gpu as *const c_void,
+            kv_lens_gpu as *const c_void,
+            batch as i32, num_heads as i32,
+            head_dim as i32, block_size as i32, max_blocks_per_seq as i32,
+            stream,
+        );
+    }
 }
 
 fn apply_causal_mask(scores: &Tensor, offset: usize) {