From c88e2ab88cac096ccf3a139654848f040667f312 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Tue, 30 Jun 2026 12:00:03 +0800 Subject: [PATCH] =?UTF-8?q?post-train:=20M2=20=E2=80=94=20decode=20primiti?= =?UTF-8?q?ves=20(rope=5Fat=20+=20decode=5Fattention)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/xtrain-cuda/src/ffi.rs | 13 +++ crates/xtrain-tensor/src/tensor.rs | 102 +++++++++++++++++++++ crates/xtrain-tensor/tests/integration.rs | 103 ++++++++++++++++++++++ csrc/ops/nn.cu | 27 ++++++ 4 files changed, 245 insertions(+) diff --git a/crates/xtrain-cuda/src/ffi.rs b/crates/xtrain-cuda/src/ffi.rs index d757e98..170f3e1 100644 --- a/crates/xtrain-cuda/src/ffi.rs +++ b/crates/xtrain-cuda/src/ffi.rs @@ -139,6 +139,19 @@ unsafe extern "C" { period: i32, s: CudaStream, ); + // RoPE at an absolute position offset (KV-cache decode, forward only): row + // `tok`'s position is `pos0 + tok` (no modulo). For a single decode token + // (tokens == 1) the one row sits at absolute position `pos0`. + pub fn launch_rope_at_f32( + x: *const f32, + y: *mut f32, + tokens: i32, + heads: i32, + head_dim: i32, + theta: f32, + pos0: i32, + s: CudaStream, + ); pub fn launch_rope_dx_f32( dy: *const f32, dx: *mut f32, diff --git a/crates/xtrain-tensor/src/tensor.rs b/crates/xtrain-tensor/src/tensor.rs index 642c9f9..7c305fe 100644 --- a/crates/xtrain-tensor/src/tensor.rs +++ b/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)`. diff --git a/crates/xtrain-tensor/tests/integration.rs b/crates/xtrain-tensor/tests/integration.rs index 4f23916..6aa9d72 100644 --- a/crates/xtrain-tensor/tests/integration.rs +++ b/crates/xtrain-tensor/tests/integration.rs @@ -56,3 +56,106 @@ fn elementwise_scale_kernel() { r.len() ); } + +/// (c) `rope_at` (KV-cache decode RoPE at an absolute position) is bit-identical +/// to the full-sequence `rope`'s corresponding row. This is the invariant the +/// decode KV-cache relies on: a single new token RoPE'd at position `t` must equal +/// what the full-sequence forward would have produced at row `t` (so cached +/// post-RoPE K matches the full-recompute path → token-identical decode). +#[test] +fn rope_at_matches_full_rope_row() { + assert!( + device::device_count().expect("device count") > 0, + "no CUDA device" + ); + device::set_device(0).unwrap(); + + let (n, heads, hd) = (7usize, 3usize, 8usize); + let theta = 10000.0f32; + // Deterministic pseudo-random fill in [-1, 1). + let host: Vec = (0..n * heads * hd) + .map(|i| ((i * 37 % 101) as f32 / 50.0) - 1.0) + .collect(); + + // Full-sequence rope (period = n → row r gets position r). + let full = Tensor::from_slice(&host, &[n, heads, hd]).to_device(Device::Cuda(0)); + let roped_full = full + .rope(theta, n) + .to_device(Device::Cpu) + .as_slice::() + .to_vec(); + + let row_len = heads * hd; + for t in 0..n { + let row = &host[t * row_len..(t + 1) * row_len]; + let roped_row = Tensor::from_slice(row, &[1, heads, hd]) + .to_device(Device::Cuda(0)) + .rope_at(theta, t) + .to_device(Device::Cpu) + .as_slice::() + .to_vec(); + let expect = &roped_full[t * row_len..(t + 1) * row_len]; + assert_eq!( + roped_row.as_slice(), + expect, + "rope_at(pos0={t}) != full rope row {t}" + ); + } + println!("rope_at OK: bit-identical to full rope across {n} positions"); +} + +/// (d) `decode_attention` (single query vs cached K/V, no mask) equals the LAST +/// query row of the full causal `attention`. This is the core decode-engine +/// invariant: the incremental path must reproduce what the full-recompute forward +/// computes for the final position, so KV-cache greedy decode is token-identical. +/// Tolerance is fp rounding (different softmax kernel + reduction order), not bits. +#[test] +fn decode_attention_matches_full_attention_last_row() { + assert!( + device::device_count().expect("device count") > 0, + "no CUDA device" + ); + device::set_device(0).unwrap(); + + let (bh, t, hd) = (6usize, 5usize, 8usize); + let scale = 1.0 / (hd as f32).sqrt(); + let n = bh * t * hd; + let qh: Vec = (0..n).map(|i| ((i * 31 % 97) as f32 / 48.0) - 1.0).collect(); + let kh: Vec = (0..n).map(|i| ((i * 53 % 89) as f32 / 44.0) - 1.0).collect(); + let vh: Vec = (0..n).map(|i| ((i * 17 % 83) as f32 / 41.0) - 1.0).collect(); + let q = Tensor::from_slice(&qh, &[bh, t, hd]).to_device(Device::Cuda(0)); + let k = Tensor::from_slice(&kh, &[bh, t, hd]).to_device(Device::Cuda(0)); + let v = Tensor::from_slice(&vh, &[bh, t, hd]).to_device(Device::Cuda(0)); + + // Reference: full causal attention, take each head's last query row. + let (full, _) = q.attention(&k, &v, scale); + let full_h = full.to_device(Device::Cpu).as_slice::().to_vec(); + + // Decode: build Q_last [bh,1,hd] from each head's last row, attend to all K/V. + let mut ql = vec![0f32; bh * hd]; + for b in 0..bh { + let src = (b * t + (t - 1)) * hd; + ql[b * hd..(b + 1) * hd].copy_from_slice(&qh[src..src + hd]); + } + let q_last = Tensor::from_slice(&ql, &[bh, 1, hd]).to_device(Device::Cuda(0)); + let dec = q_last + .decode_attention(&k, &v, scale) + .to_device(Device::Cpu) + .as_slice::() + .to_vec(); + assert_eq!(dec.len(), bh * hd, "decode out shape"); + + let mut max_abs = 0f32; + for b in 0..bh { + for d in 0..hd { + let got = dec[b * hd + d]; + let exp = full_h[(b * t + (t - 1)) * hd + d]; + max_abs = max_abs.max((got - exp).abs()); + } + } + assert!( + max_abs < 1e-4, + "decode_attention vs full last-row max abs diff {max_abs} exceeds 1e-4" + ); + println!("decode_attention OK: matches full causal last row (bh={bh}, t={t}, max|Δ|={max_abs:.2e})"); +} diff --git a/csrc/ops/nn.cu b/csrc/ops/nn.cu index b74947f..37399ef 100644 --- a/csrc/ops/nn.cu +++ b/csrc/ops/nn.cu @@ -242,6 +242,33 @@ void launch_rope_f32(const float* x, float* y, int tokens, int heads, rope_k<<>>(x, y, heads, head_dim, theta, period); } +// RoPE at an absolute position offset (KV-cache decode-time, forward only). Same +// rotate_half as rope_k, but row `tok`'s position is `pos0 + tok` (no modulo) — +// a single new decode token sits at absolute position pos0. The training rope_k +// (position = tok % period) is left untouched, so this adds no training-path risk. +__global__ void rope_at_k(const float* x, float* y, int heads, int head_dim, + float theta, int pos0) { + int tok = blockIdx.x; + int head = blockIdx.y; + int half = head_dim / 2; + int i = threadIdx.x; + if (i >= half) return; + int pos = pos0 + tok; + float freq = powf(theta, -(float)(2 * i) / (float)head_dim); + float angle = (float)pos * freq; + float c = cosf(angle), sn = sinf(angle); + int base = (tok * heads + head) * head_dim; + float x0 = x[base + i], x1 = x[base + i + half]; + y[base + i] = x0 * c - x1 * sn; + y[base + i + half] = x1 * c + x0 * sn; +} +void launch_rope_at_f32(const float* x, float* y, int tokens, int heads, + int head_dim, float theta, int pos0, void* s) { + dim3 grid(tokens, heads); + int blk = head_dim / 2; + rope_at_k<<>>(x, y, heads, head_dim, theta, pos0); +} + __global__ void rope_dx_k(const float* dy, float* dx, int heads, int head_dim, float theta, int period) { int tok = blockIdx.x;