post-train: M2b — batched KV-cache decode (G-way, token-identical)
The rollout long-pole fix deferred from M2a: decode the G samples of one prompt in lockstep (one forward per step over the group → G× fewer kernel launches). - rope_pos(x, positions[]): RoPE with a per-row absolute position (new forward- only kernel) — G rows share one decode position. Gate: == full rope for [0..n], == rope_at(P) per row for uniform P (bit-identical). - generate_cached_batch: BatchKVCache [T, G·num_kv, hd] + batched decode_step. decode_attention is already batch-agnostic (bh = G·nh); repeat_kv(nh, batch=G) broadcasts per group. No finished-mask / ragged prompts yet (perf-only / next). - Gate (tests/decode_batch.rs): all G greedy rows token-identical to the single- sequence decode (8 query / 2 kv heads → exercises repeat_kv batching). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -822,6 +822,40 @@ impl Tensor {
|
||||
out
|
||||
}
|
||||
|
||||
/// RoPE with a PER-ROW absolute position (batched KV-cache decode, M2b).
|
||||
/// `self`:[tokens,heads,head_dim]; row `t`'s position is `positions[t]` (an
|
||||
/// I32 `[tokens]` tensor). For G-way batched decode all G rows share one decode
|
||||
/// position; for ragged batches each row carries its own. Mirrors `rope_at`'s
|
||||
/// dtype handling; forward only.
|
||||
#[cfg(not(no_cuda))]
|
||||
pub fn rope_pos(&self, positions: &Tensor, theta: f32) -> Self {
|
||||
assert_eq!(self.ndim(), 3, "rope_pos 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");
|
||||
assert_eq!(positions.dtype, DType::I32, "positions must be I32");
|
||||
assert_eq!(positions.numel(), tokens, "one position per token");
|
||||
if self.dtype == DType::BF16 {
|
||||
return self
|
||||
.to_dtype(DType::F32)
|
||||
.rope_pos(positions, theta)
|
||||
.to_dtype(DType::BF16);
|
||||
}
|
||||
let out = Tensor::zeros(&self.shape, DType::F32, self.device());
|
||||
unsafe {
|
||||
xtrain_cuda::ffi::launch_rope_pos_f32(
|
||||
self.data_ptr() as *const f32,
|
||||
positions.data_ptr() as *const i32,
|
||||
out.data_ptr() as *mut f32,
|
||||
tokens as i32,
|
||||
heads as i32,
|
||||
head_dim as i32,
|
||||
theta,
|
||||
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))]
|
||||
|
||||
Reference in New Issue
Block a user