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:
@@ -152,6 +152,18 @@ unsafe extern "C" {
|
||||
pos0: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
// RoPE with a per-row absolute position (batched KV-cache decode, M2b): row
|
||||
// `tok`'s position is `positions[tok]`. Forward only.
|
||||
pub fn launch_rope_pos_f32(
|
||||
x: *const f32,
|
||||
positions: *const i32,
|
||||
y: *mut f32,
|
||||
tokens: i32,
|
||||
heads: i32,
|
||||
head_dim: i32,
|
||||
theta: f32,
|
||||
s: CudaStream,
|
||||
);
|
||||
// Per-row scale: y[r,c] = x[r,c] * s[r] (GRPO policy-gradient backward).
|
||||
pub fn launch_scale_rows_f32(
|
||||
x: *const f32,
|
||||
|
||||
Reference in New Issue
Block a user