post-train: M2 — decode primitives (rope_at + decode_attention)

Two forward-only Tensor primitives the KV-cache decode engine is built on,
each gated by an isolated correctness test:

- rope_at(theta, pos0): RoPE at an absolute position (pos = pos0 + row, no
  modulo) for a single decode token, vs the training rope_k (pos = row %
  period) left untouched. New forward-only CUDA kernel, no training-path risk.
  Gate: bit-identical to the full-sequence rope's corresponding row.
- decode_attention(k, v, scale): single-query × cached-K/V SDPA, composed from
  the existing strided batched GEMM + plain (non-causal) softmax — no new
  kernel. Gate: equals the full causal attention's last query row (max |Δ| 6e-8).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-tensor/src/tensor.rs

@@ -790,6 +790,38 @@ impl Tensor {
         out
     }
 
+    /// RoPE at an absolute position offset (KV-cache decode, forward only).
+    /// `self`:[tokens,heads,head_dim]; row `r`'s position is `pos0 + r` (no
+    /// modulo). For a single new decode token pass `tokens == 1` → the one row is
+    /// rotated at absolute position `pos0`. Mirrors [`rope`](Self::rope)'s dtype
+    /// handling (bf16 → f32 → bf16); no backward (inference path).
+    #[cfg(not(no_cuda))]
+    pub fn rope_at(&self, theta: f32, pos0: usize) -> Self {
+        assert_eq!(self.ndim(), 3, "rope_at requires [tokens,heads,head_dim]");
+        let (tokens, heads, head_dim) = (self.shape[0], self.shape[1], self.shape[2]);
+        assert_eq!(head_dim % 2, 0, "head_dim must be even");
+        if self.dtype == DType::BF16 {
+            return self
+                .to_dtype(DType::F32)
+                .rope_at(theta, pos0)
+                .to_dtype(DType::BF16);
+        }
+        let out = Tensor::zeros(&self.shape, DType::F32, self.device());
+        unsafe {
+            xtrain_cuda::ffi::launch_rope_at_f32(
+                self.data_ptr() as *const f32,
+                out.data_ptr() as *mut f32,
+                tokens as i32,
+                heads as i32,
+                head_dim as i32,
+                theta,
+                pos0 as i32,
+                std::ptr::null_mut(),
+            );
+        }
+        out
+    }
+
     /// RoPE backward: apply the inverse (transpose) rotation to `dy`. RoPE is an
     /// orthogonal map, so it needs no cached forward values, only `theta`/`period`.
     #[cfg(not(no_cuda))]
@@ -1076,6 +1108,76 @@ impl Tensor {
         (out, probs)
     }
 
+    /// Decode-time (incremental) attention: a SINGLE query position against a
+    /// cached K/V of length `t` (KV-cache decode, forward only). `self` = Q
+    /// `[bh,1,head_dim]`; `k`,`v` = `[bh,t,head_dim]`, already repeat_kv-expanded
+    /// to `bh` heads. Returns out `[bh,head_dim]` (= `[bh,1,head_dim]` flattened).
+    ///
+    /// No causal mask is needed — the one query sits at the end, so every cached
+    /// key (positions `0..t`) is visible. This is exactly the LAST query row of the
+    /// full causal [`attention`](Self::attention), so KV-cache greedy decode is
+    /// token-identical to full recompute. Softmax is computed in f32 (matching the
+    /// causal path) with `scale` folded in before the exponentials.
+    #[cfg(not(no_cuda))]
+    pub fn decode_attention(&self, k: &Tensor, v: &Tensor, scale: f32) -> Self {
+        assert_eq!(self.ndim(), 3, "decode_attention Q must be [bh,1,head_dim]");
+        assert_eq!(self.shape[1], 1, "decode_attention Q seq must be 1");
+        assert_eq!(k.ndim(), 3, "decode_attention K must be [bh,t,head_dim]");
+        assert_eq!(k.shape(), v.shape(), "K/V shape mismatch");
+        assert_eq!(self.dtype, k.dtype, "Q/K dtype mismatch");
+        assert_eq!(self.dtype, v.dtype, "Q/V dtype mismatch");
+        let (bh, hd) = (self.shape[0], self.shape[2]);
+        assert_eq!(k.shape[0], bh, "Q/K batch-head mismatch");
+        assert_eq!(k.shape[2], hd, "Q/K head_dim mismatch");
+        let t = k.shape[1]; // cached length
+        let dt = self.dtype;
+        let dev = self.device();
+
+        // scores[bh,1,t] = Q[bh,1,hd] · Kᵀ[bh,hd,t]  (per-head batched GEMM).
+        // [bh,1,t] is stored identically to [bh,t]; allocate 2D so the rowwise
+        // softmax can run without a reshape.
+        let scores = Tensor::zeros(&[bh, t], dt, dev);
+        strided_batched_gemm(
+            dt,
+            false,
+            true,
+            1,
+            t,
+            hd,
+            self.data_ptr(),
+            hd,
+            k.data_ptr(),
+            t * hd,
+            scores.data_ptr(),
+            t,
+            bh,
+        );
+        // probs = softmax(scale · scores) over the t keys (f32, like the causal path).
+        let probs = scores
+            .to_dtype(DType::F32)
+            .scale(scale)
+            .softmax()
+            .to_dtype(dt);
+        // out[bh,1,hd] = probs[bh,1,t] · V[bh,t,hd].
+        let out = Tensor::zeros(&[bh, hd], dt, dev);
+        strided_batched_gemm(
+            dt,
+            false,
+            false,
+            1,
+            hd,
+            t,
+            probs.data_ptr(),
+            t,
+            v.data_ptr(),
+            t * hd,
+            out.data_ptr(),
+            hd,
+            bh,
+        );
+        out
+    }
+
     /// Backward of [`attention`](Self::attention). Inputs: forward `q`,`k`,`v`,
     /// the cached `probs`, the upstream `dout` (all batched `[bh,seq,*]`), and the
     /// same `scale`. Returns `(dq, dk, dv)`.