spark_runtime/buffers/
sizes_rowwise.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Row-wise FP8 GDN prefill BF16-weight slab sizing, split out of `sizes.rs`
4//! (≤500 LoC cap). Env-gated — 0 (→ NULL) unless `ATLAS_FP8_ROWWISE=1`, so
5//! every other recipe's ledger is byte-identical to before this entry existed.
6//!
7//! **WHY (#917 H100 receipt, 2026-09-11, `Qwen/Qwen3.8-27B-FP8`).** The
8//! `ATLAS_FP8_ROWWISE` GDN prefill arms dequantise their per-row FP8 weights
9//! to BF16 once and multiply with cuBLASLt, because
10//! `cublaslt::fp8_gemm_act_weight_t_rowwise` returns NOT_SUPPORTED on sm_121
11//! (measured 2026-08-15). That dequant used to be a lazy `gpu.alloc` memoised
12//! by weight pointer — `167772160` bytes for the fused `[QKV|Z]` weight PER
13//! LAYER, invisible to `--gpu-memory-utilization`, which is the same defect
14//! class that killed a 28-token prefill at layer 36 with `cuMemAlloc_v2
15//! failed: status 2`. Sizing it here makes it ONE arena allocation the
16//! preflight fitter can see (`preflight::headroom`'s `arena` term is
17//! `BufferSizes::total_bytes()`), and the arms bump-carve their slices from
18//! it instead of allocating.
19
20use atlas_core::config::ModelConfig;
21
22/// BF16 bytes ONE GDN layer's row-wise prefill arms dequantise and keep:
23/// the fused `in_proj_qkvz` `[ssm_qkvz_size, hidden]` and `out_proj`
24/// `[hidden, value_dim]`, 2 bytes per element.
25///
26/// EXACT, not an upper bound: these are the two weights
27/// `set_fp8_rowwise_prefill_weights` installs and the only two the row-wise
28/// arms dequantise. `value_dim = linear_num_value_heads * linear_value_head_dim`
29/// — the same extent `trait_prefill_block.rs` passes as the `out_proj` K.
30///
31/// Qwen3.8-27B (hidden 5120, 16x128 key heads, 48x128 value heads):
32/// `16384*5120*2 + 5120*6144*2 = 167772160 + 62914560 = 230686720` B.
33pub fn ssm_rowwise_w_bf16_layer_bytes(config: &ModelConfig) -> usize {
34    let bf16 = 2;
35    let value_dim = config.linear_num_value_heads * config.linear_value_head_dim;
36    let qkvz = config.ssm_qkvz_size() * config.hidden_size * bf16;
37    let out_proj = config.hidden_size * value_dim * bf16;
38    qkvz + out_proj
39}
40
41/// Arena bytes for the row-wise GDN prefill BF16-weight slab: the per-layer
42/// pair summed over every linear-attention layer, or 0 when the lever is off.
43///
44/// The env predicate is character for character
45/// `weight_loader::qwen35_dense::rowwise_fp8::rowwise_fp8_enabled` — `== Ok("1")`,
46/// NOT presence. The two must agree: the loader installs the per-row weights
47/// on that predicate and the arms then require this slab, so `ATLAS_FP8_ROWWISE=0`
48/// must leave BOTH off.
49pub fn ssm_rowwise_w_bf16_bytes(config: &ModelConfig) -> usize {
50    ssm_rowwise_w_bf16_bytes_for(
51        config,
52        std::env::var("ATLAS_FP8_ROWWISE").as_deref() == Ok("1"),
53    )
54}
55
56/// [`ssm_rowwise_w_bf16_bytes`] with the lever passed in, so the sizing
57/// arithmetic is pinnable at exact integers without touching the process
58/// environment.
59pub fn ssm_rowwise_w_bf16_bytes_for(config: &ModelConfig, rowwise_enabled: bool) -> usize {
60    if !rowwise_enabled {
61        return 0;
62    }
63    config.num_ssm_layers() * ssm_rowwise_w_bf16_layer_bytes(config)
64}
spark_runtime/buffers/sizes_rowwise.rs

spark_runtime/buffers/
sizes_rowwise.rs
