metrale_gpu_runtime/buffers/
sizes_rowwise.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Size of the row-wise FP8 GDN prefill BF16-weight slab: 0 unless `METRALE_FP8_ROWWISE=1`.
4//!
5//! The row-wise GDN prefill arms multiply BF16 copies of their per-row FP8
6//! weights. Sizing the copies here makes them one arena allocation that the
7//! preflight reserve (`BufferSizes::total_bytes`) counts; the arms carve their
8//! slices from it (`rowwise_slab.rs`).
9//!
10//! Owner: gpu-runtime.
11//! Invariants:
12//! - [`ssm_rowwise_w_bf16_bytes`] reads the lever with the same predicate as
13//!   the loader's `rowwise_fp8_enabled` (`== "1"`).
14
15use metrale_config::ModelConfig;
16
17/// 2026-09-25: BF16 bytes of one GDN layer's row-wise copies: `in_proj_qkvz`
18/// `[ssm_qkvz_size, hidden]` and `out_proj` `[hidden, value_dim]`, the two
19/// weights `set_fp8_rowwise_prefill_weights` installs, with
20/// `value_dim = linear_num_value_heads * linear_value_head_dim`. At the
21/// Qwen3.8-27B shapes of `buffers/tests.rs` (hidden 5120, 16x128 key heads,
22/// 48x128 value heads) that is `167772160 + 62914560` B.
23pub fn ssm_rowwise_w_bf16_layer_bytes(config: &ModelConfig) -> usize {
24    let bf16 = 2;
25    let value_dim = config.linear_num_value_heads * config.linear_value_head_dim;
26    let qkvz = config.ssm_qkvz_size() * config.hidden_size * bf16;
27    let out_proj = config.hidden_size * value_dim * bf16;
28    qkvz + out_proj
29}
30
31/// 2026-09-25: Slab bytes: the per-layer pair times `num_ssm_layers`, or 0
32/// unless `METRALE_FP8_ROWWISE=1`. The predicate is the one in the loader's
33/// `rowwise_fp8_enabled` (`weight_loader/qwen35_dense/rowwise_fp8.rs`): the
34/// loader installs the row-wise weights on it and the arms then need this
35/// slab, so the two must agree.
36pub fn ssm_rowwise_w_bf16_bytes(config: &ModelConfig) -> usize {
37    ssm_rowwise_w_bf16_bytes_for(
38        config,
39        std::env::var("METRALE_FP8_ROWWISE").as_deref() == Ok("1"),
40    )
41}
42
43/// 2026-09-25: [`ssm_rowwise_w_bf16_bytes`] with the lever passed in, for tests
44/// that must not set the process environment.
45pub fn ssm_rowwise_w_bf16_bytes_for(config: &ModelConfig, rowwise_enabled: bool) -> usize {
46    if !rowwise_enabled {
47        return 0;
48    }
49    config.num_ssm_layers() * ssm_rowwise_w_bf16_layer_bytes(config)
50}
metrale_gpu_runtime/buffers/sizes_rowwise.rs

metrale_gpu_runtime/buffers/
sizes_rowwise.rs
