metrale_gpu_runtime/buffers/
sizes_q12.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Scratch footprint of the kernel-batched prefill staging, shared by the arena sizing and the admission check.
4//!
5//! Owner: gpu-runtime.
6//! Invariants: none beyond the types; the functions are pure.
7
8/// 2026-09-25: Streams `scratch` is provisioned for in the kernel-batched
9/// prefill staging. A batch whose footprint does not fit is refused by
10/// `check_kernel_batched_eligible` (metrale-model-engine) and runs per
11/// stream, so this sets how often the batched path is available, not safety.
12pub const Q12_SIZING_STREAMS: usize = 8;
13
14/// 2026-09-25: Scratch bytes of the kernel-batched prefill staging for `n`
15/// streams of `chunk_len` tokens each. Both `BufferSizes::from_config` and
16/// `check_kernel_batched_eligible` call it (or the varlen form), so they agree
17/// on whether a batch fits. The terms follow `prefill_b/batch_kernel.rs` (MoE
18/// top-K area, `n` per-stream metadata blocks), `prefill_b/stage_batched.rs`
19/// (positions, three streams under MRoPE, slots, the block-table and seq_len
20/// pointer arrays, `cu_seqlens`, `kv_lens`) and the `h_state_ptrs` table.
21pub fn q12_batched_scratch_bytes(n: usize, chunk_len: usize, top_k: usize, mrope: bool) -> usize {
22    q12_batched_scratch_bytes_varlen(n, n * chunk_len, chunk_len, top_k, mrope)
23}
24
25/// 2026-09-25: [`q12_batched_scratch_bytes`] for a ragged batch: `total_tokens`
26/// is the packed total and `max_chunk_len` sizes the per-stream metadata
27/// slots, so a varlen batch is not charged the longest stream's length for
28/// every stream.
29pub fn q12_batched_scratch_bytes_varlen(
30    n: usize,
31    total_tokens: usize,
32    max_chunk_len: usize,
33    top_k: usize,
34    mrope: bool,
35) -> usize {
36    let moe = ((total_tokens * top_k * 4 * 2) + 63) & !63;
37    let per_stream_meta = ((max_chunk_len * 16) + 64).max(4096);
38    let pos = (total_tokens * 4 + 7) & !7;
39    let pos_streams = if mrope { 3 } else { 1 };
40    let slot = (total_tokens * 8 + 7) & !7;
41    let ptrs = ((n * std::mem::size_of::<u64>()) + 7) & !7;
42    let cu_seqlens = (((n + 1) * 4) + 7) & !7;
43    let kv_lens = ((n * 4) + 7) & !7;
44    let stage_meta = pos_streams * pos + slot + 2 * ptrs + cu_seqlens + kv_lens;
45    let h_state_ptrs = n * std::mem::size_of::<u64>();
46    moe + n * per_stream_meta + stage_meta + h_state_ptrs
47}
metrale_gpu_runtime/buffers/sizes_q12.rs

metrale_gpu_runtime/buffers/
sizes_q12.rs
