metrale_gpu_runtime/buffers/
rowwise_slab.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The row-wise FP8 GDN prefill BF16-weight slab: its size and the bump carve the prefill arms take their per-layer slices from.
4//!
5//! The slab is one arena allocation, so the preflight reserve (`total_bytes`)
6//! counts it.
7//!
8//! Owner: gpu-runtime.
9//! Invariants:
10//! - A carve never allocates: it returns a slice of the slab or an error (slab
11//!   absent, or exhausted).
12//! - Successful carves do not overlap and lie inside the slab.
13
14use super::BufferArena;
15use crate::gpu::DevicePtr;
16
17impl BufferArena {
18    /// 2026-09-25: Allocated byte size of the slab; 0 unless `METRALE_FP8_ROWWISE=1`
19    /// when the arena was sized.
20    pub fn ssm_rowwise_w_bf16_bytes(&self) -> usize {
21        self.sizes.ssm_rowwise_w_bf16
22    }
23
24    /// 2026-09-25: Carve the next `bytes` of the slab. Slices are never returned:
25    /// a GDN layer carves its `in_proj_qkvz` and `out_proj` copies on its first
26    /// prefill and keeps them for the arena's life. The slab holds exactly
27    /// `num_ssm_layers` such pairs (`sizes_rowwise.rs`), so exhaustion is a
28    /// sizing bug and is returned as an error, not served by a new allocation.
29    /// The cursor is not rolled back after that error.
30    ///
31    /// `Relaxed` ordering: the scheduler runs one forward at a time, the same
32    /// assumption `cublaslt::Ctx` states for its `Send`/`Sync`.
33    pub fn take_ssm_rowwise_w_bf16(&self, bytes: usize) -> anyhow::Result<DevicePtr> {
34        use std::sync::atomic::Ordering;
35        if self.ssm_rowwise_w_bf16 == DevicePtr::NULL {
36            anyhow::bail!(
37                "row-wise GDN prefill BF16-weight slab is absent: the arena was sized without                  METRALE_FP8_ROWWISE=1 but a row-wise prefill arm asked for {bytes} B. Both the                  loader's weight install and this ledger entry read the same lever, so they                  cannot legitimately disagree"
38            );
39        }
40        let base = self
41            .ssm_rowwise_w_bf16_used
42            .fetch_add(bytes, Ordering::Relaxed);
43        let end = base + bytes;
44        if end > self.sizes.ssm_rowwise_w_bf16 {
45            anyhow::bail!(
46                "row-wise GDN prefill BF16-weight slab exhausted: wanted {bytes} B at offset                  {base}, slab is {} B. `sizes_rowwise::ssm_rowwise_w_bf16_bytes` sizes it for                  num_ssm_layers x (in_proj_qkvz + out_proj)",
47                self.sizes.ssm_rowwise_w_bf16
48            );
49        }
50        Ok(self.ssm_rowwise_w_bf16.offset(base))
51    }
52}
metrale_gpu_runtime/buffers/rowwise_slab.rs

metrale_gpu_runtime/buffers/
rowwise_slab.rs
