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>
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user