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:
@@ -269,6 +269,33 @@ void launch_rope_at_f32(const float* x, float* y, int tokens, int heads,
|
||||
rope_at_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, y, heads, head_dim, theta, pos0);
|
||||
}
|
||||
|
||||
// RoPE with a PER-ROW absolute position (batched KV-cache decode, M2b): row `tok`'s
|
||||
// position is `positions[tok]` (an i32 per token). For G-way batched decode all G
|
||||
// rows share one decode position; for ragged batches each row carries its own.
|
||||
// Forward only; the training rope_k is untouched.
|
||||
__global__ void rope_pos_k(const float* x, const int* positions, float* y,
|
||||
int heads, int head_dim, float theta) {
|
||||
int tok = blockIdx.x;
|
||||
int head = blockIdx.y;
|
||||
int half = head_dim / 2;
|
||||
int i = threadIdx.x;
|
||||
if (i >= half) return;
|
||||
int pos = positions[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_pos_f32(const float* x, const int* positions, float* y,
|
||||
int tokens, int heads, int head_dim, float theta, void* s) {
|
||||
dim3 grid(tokens, heads);
|
||||
int blk = head_dim / 2;
|
||||
rope_pos_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, positions, y, heads, head_dim, theta);
|
||||
}
|
||||
|
||||
// Per-row scale: y[r,c] = x[r,c] * s[r]. One block per row. Used by the GRPO
|
||||
// (M4) policy-gradient backward, where each completion token's row of
|
||||
// (probs − onehot) is scaled by its own per-token coefficient.
|
||||
|
||||
Reference in New Issue
Block a user