spark_runtime/buffers/
sizes_q12.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Q12 kernel-batched prefill scratch sizing. Split from `sizes.rs` (500-LoC cap).
4
5/// Streams assumed when provisioning scratch for the Q12 kernel-batched
6/// prefill path. The per-stream metadata region scales with N, so the
7/// scratch buffer must be sized for a realistic max concurrent batched
8/// streams. Beyond this, `check_kernel_batched_eligible` falls the dispatch
9/// back to the per-stream path (which respects the same arena cap), so this
10/// bound only governs how often the fast path is available — never safety.
11pub const Q12_SIZING_STREAMS: usize = 8;
12
13/// Exact scratch footprint (bytes) of the Q12 kernel-batched prefill staging
14/// for `n` streams of `chunk_len` tokens each. SSOT for both scratch sizing
15/// (`BufferSizes::from_config`) and the pre-flight eligibility check
16/// (`check_kernel_batched_eligible`), so the two can never disagree about
17/// whether a batch fits. Mirrors the staging layout in `batch_kernel.rs`
18/// (MoE topk area + N per-stream meta blocks) and `stage_batched.rs`
19/// (stacked positions ×(3 if MRoPE) + slots + block/seq_len pointer arrays)
20/// plus the per-SSM-layer `h_state_ptrs` JIT slot.
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/// Exact scratch footprint for a ragged Q12 batch.
26///
27/// `total_tokens` is the packed `cu_seqlens` total while `max_chunk_len`
28/// sizes the per-stream metadata slots. Uniform callers should use
29/// [`q12_batched_scratch_bytes`]. Keeping this separate prevents VARLEN
30/// admission from pessimistically charging every stream at the longest
31/// request's length.
32pub fn q12_batched_scratch_bytes_varlen(
33    n: usize,
34    total_tokens: usize,
35    max_chunk_len: usize,
36    top_k: usize,
37    mrope: bool,
38) -> usize {
39    // MoE topk staging (indices+weights, both ×n streams), 64-byte aligned.
40    let moe = ((total_tokens * top_k * 4 * 2) + 63) & !63;
41    // Per-stream meta block — same formula as batch_kernel.rs.
42    let per_stream_meta = ((max_chunk_len * 16) + 64).max(4096);
43    // Stacked BatchedAttnMetadata (stage_batched.rs layout).
44    let pos = (total_tokens * 4 + 7) & !7;
45    let pos_streams = if mrope { 3 } else { 1 };
46    let slot = (total_tokens * 8 + 7) & !7;
47    let ptrs = ((n * std::mem::size_of::<u64>()) + 7) & !7;
48    // VARLEN cu_seqlens [n+1] i32 prefix-sum, staged after the pointer arrays
49    // (stage_batched.rs:122-126). This term scales with n, so omitting it made
50    // the SSOT under-count grow with batch size: at n>=4 the h_state_ptrs JIT
51    // slot overlapped a live per-stream pointer table → cross-stream KV/GDN
52    // bleed in decode (n<=3 clean, absorbed by over-provisioning slack).
53    let cu_seqlens = (((n + 1) * 4) + 7) & !7;
54    // VARLEN kv_lens [n] i32, staged immediately after cu_seqlens.
55    let kv_lens = ((n * 4) + 7) & !7;
56    let stage_meta = pos_streams * pos + slot + 2 * ptrs + cu_seqlens + kv_lens;
57    // h_state_ptrs JIT slot consumed per SSM layer (N device pointers).
58    let h_state_ptrs = n * std::mem::size_of::<u64>();
59    moe + n * per_stream_meta + stage_meta + h_state_ptrs
60}
spark_runtime/buffers/sizes_q12.rs

spark_runtime/buffers/
sizes_q12.rs
