spark_runtime/buffers/
rowwise_slab.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The row-wise FP8 GDN prefill BF16-weight slab: its size accessor and the
4//! bump carve the prefill arms take their per-layer slices from. Split from
5//! `accessors.rs` (≤500 LoC cap) rather than wedged into it, because the carve
6//! is a tiny allocator and not a getter.
7//!
8//! **WHY IT EXISTS (#917 H100 receipt, 2026-09-11).** The `ATLAS_FP8_ROWWISE`
9//! GDN arms used to get their BF16 weight from a lazy `gpu.alloc` memoised by
10//! weight pointer — `167772160` B for the fused `[QKV|Z]` weight per layer,
11//! with no `BufferSizes` entry, so `--gpu-memory-utilization` could not see it
12//! and a 28-token prefill died at layer 36 with `cuMemAlloc_v2 failed:
13//! status 2`. The bytes are the same; the difference is that they are now ONE
14//! ledgered allocation the preflight fitter prices.
15
16use super::BufferArena;
17use crate::gpu::DevicePtr;
18
19impl BufferArena {
20    /// Allocated byte size of the row-wise GDN prefill BF16-weight slab.
21    /// 0 when `ATLAS_FP8_ROWWISE` was not armed at boot.
22    pub fn ssm_rowwise_w_bf16_bytes(&self) -> usize {
23        self.sizes.ssm_rowwise_w_bf16
24    }
25
26    /// Carve the next `bytes` of the row-wise GDN prefill BF16-weight slab.
27    ///
28    /// Bump-only and never returned: each GDN layer takes its `in_proj_qkvz`
29    /// and `out_proj` slices on its FIRST prefill and holds them for the life
30    /// of the arena, because a dequanted weight is as immutable as the weight
31    /// it came from. Sized in `sizes_rowwise::ssm_rowwise_w_bf16_bytes` for
32    /// exactly `num_ssm_layers` of those pairs, so exhaustion means the sizing
33    /// and the callers disagree — a bug, reported as one rather than papered
34    /// over with a fresh allocation (that is the #917 defect this replaces).
35    ///
36    /// `Relaxed` is enough: the scheduler drives one forward at a time (the
37    /// same single-threaded invariant `cublaslt::Ctx` documents), so this is a
38    /// counter that happens to be atomic rather than a contended one.
39    pub fn take_ssm_rowwise_w_bf16(&self, bytes: usize) -> anyhow::Result<DevicePtr> {
40        use std::sync::atomic::Ordering;
41        if self.ssm_rowwise_w_bf16 == DevicePtr::NULL {
42            anyhow::bail!(
43                "row-wise GDN prefill BF16-weight slab is absent: the arena was sized without                  ATLAS_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"
44            );
45        }
46        let base = self
47            .ssm_rowwise_w_bf16_used
48            .fetch_add(bytes, Ordering::Relaxed);
49        let end = base + bytes;
50        if end > self.sizes.ssm_rowwise_w_bf16 {
51            anyhow::bail!(
52                "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)",
53                self.sizes.ssm_rowwise_w_bf16
54            );
55        }
56        Ok(self.ssm_rowwise_w_bf16.offset(base))
57    }
58}
spark_runtime/buffers/rowwise_slab.rs

spark_runtime/buffers/
rowwise_slab.rs
