Struct BufferArena 

pub struct BufferArena { /* private fields */ }

Pre-allocated GPU buffers for a single forward pass.

Each buffer is sized for max_batch_tokens tokens through the model. Buffers are reused across steps — no per-step allocation.

Expert output buffers are sized for max(k_max, max_batch_tokens) to support both speculative decode (K=3) and batched MoE prefill. At N=512, this adds ~31 MB (vs the old grouped-GEMM approach that needed 260 MB and caused a 15% decode regression). The GEMV-based prefill kernels only touch k_max slots during decode, so the extra pages don’t affect decode bandwidth on unified memory.

Implementations

impl BufferArena

pub fn hidden_states(&self) -> DevicePtr

pub fn residual(&self) -> DevicePtr

pub fn norm_output(&self) -> DevicePtr

pub fn norm_output_bytes(&self) -> usize

Allocated byte size of norm_output. Bounds-check for the attention prefill o_proj’s cuBLASLt arm, which writes ceil16(M) rows — and a prefill token count is not a multiple of 16 (#927).

pub fn qkv_output(&self) -> DevicePtr

pub fn qkv_output_bytes(&self) -> usize

Allocated byte size of qkv_output. Bounds-check for the multi-seq decode W8A8 arm, which writes ceil16(M) rows at a per_seq_qkv row pitch — the padded rows land in slots the step does not use, which is in-bounds only while the buffer holds them (#927).

pub fn attn_output(&self) -> DevicePtr

pub fn gate_logits(&self) -> DevicePtr

pub fn gate_logits_f32(&self) -> DevicePtr

pub fn moe_router_in_f32(&self) -> DevicePtr

pub fn moe_output(&self) -> DevicePtr

pub fn logits(&self) -> DevicePtr

pub fn ssm_qkvz(&self) -> DevicePtr

pub fn ssm_qkvz_bytes(&self) -> usize

Allocated byte size of ssm_qkvz — the QKVZ projection’s destination on an INTERLEAVED model. Bounds-check for the cuBLASLt arm, which writes ceil16(M) rows (see sizes.rs).

pub fn ssm_ba(&self) -> DevicePtr

pub fn ssm_deinterleaved(&self) -> DevicePtr

Sequential [Q|K|V|Z] after deinterleaving.

pub fn ssm_deinterleaved_bytes(&self) -> usize

Allocated byte size of ssm_deinterleaved — the QKVZ projection’s destination on a SEQUENTIAL model. Same padded-M bounds check.

pub fn ssm_gates(&self) -> DevicePtr

FP32 [gate, beta] for GDN (num_v_heads * 2 floats).

pub fn ssm_conv_out_f32(&self) -> DevicePtr

FP32 conv1d output for SSM recurrent path (prevents BF16 precision drift).

pub fn scratch(&self) -> DevicePtr

Scratch buffer for MoE routing + kernel metadata uploads.

pub fn ssd_scratch(&self) -> DevicePtr

Mamba-2 SSD chunked-scan scratch (dt | dA_cumsum | CB). NULL if unused.

pub fn token_ids(&self) -> DevicePtr

Token IDs [M] u32 — stable across the layer loop (DeepSeek-V4 hash-MoE reads tid2eid[token_id]). Upload the pass’s token IDs here before the layer loop; under CUDA-graph decode upload before each replay.

pub fn scratch_bytes(&self) -> usize

Allocated byte size of the scratch buffer (#110: bounds-check batched metadata-staging uploads against this).

pub fn expert_gate_out(&self) -> DevicePtr

Batched expert gate projection output.

pub fn expert_up_out(&self) -> DevicePtr

Batched expert up projection output.

pub fn expert_gate_out_bytes(&self) -> usize

Allocated byte size of expert_gate_out / expert_up_out (identical by construction). Debug bounds-check for GEMM paths that write PADDED M rows — the FP8 block-scaled cuBLASLt matmul rounds M up to 16.

pub fn moe_output_bytes(&self) -> usize

Allocated byte size of moe_output (same padded-M debug check).

pub fn expert_down_out(&self) -> DevicePtr

Batched expert down projection output.

pub fn gdn_fla_scratch(&self) -> DevicePtr

Split-K decode attention workspace (F32 partials). GDN FLA chunked-prefill scratch base (W|U|S|uc sub-divided by the caller). DevicePtr::NULL unless this is a 128-dim-linear-head GDN model.

pub fn ffn_act_q8(&self) -> DevicePtr

Shared dense-FFN q8_1 activation scratch (Q4_K MMQ gate/up). NULL for MoE.

pub fn ffn_act_a(&self) -> DevicePtr

Shared dense-FFN int8/NVFP4 activation scratch (a_i8 / packed). NULL for MoE.

pub fn ffn_act_scale(&self) -> DevicePtr

Shared dense-FFN int8/NVFP4 activation-scale scratch. NULL for MoE.

pub fn ffn_gate_up_fused(&self) -> DevicePtr

[ceil16(GATEUP_FUSED_MAX_M), 2 * intermediate] BF16 output of the FUSED dense-FFN gate+up decode GEMM (#927): a row is [gate | up], gate at column 0 and up at column intermediate. NULL for MoE.

pub fn ffn_gate_up_fused_bytes(&self) -> usize

Allocated byte size of ffn_gate_up_fused — the bound the fused arm checks the padded [ceil16(m), 2*inter] extent against before it can select itself.

pub fn ffn_act_a_bytes(&self) -> usize

Allocated byte size of ffn_act_a (debug bounds-check at call sites).

pub fn ffn_act_scale_bytes(&self) -> usize

Allocated byte size of ffn_act_scale (debug bounds-check at call sites).

pub fn ffn_act_scale_kmajor(&self) -> DevicePtr

Transposed ([K/128, ceil16(M)]) dense-FFN activation scales — the VEC128 B-scale layout the cuBLASLt block-scaled FP8 GEMM documents. NULL for MoE.

pub fn ffn_act_scale_kmajor_bytes(&self) -> usize

Allocated byte size of ffn_act_scale_kmajor (bounds-check at call sites).

pub fn fp8_act(&self) -> DevicePtr

Persistent FP8 block-scaled activation scratch for prefill projections. Replaces a per-projection alloc/sync/free in the W8A8+FP32-epilogue path.

pub fn fp8_act_bytes(&self) -> usize

Allocated byte size of fp8_act (debug bounds-check at call sites).

pub fn fp8_act_scale(&self) -> DevicePtr

Persistent per-128-block FP32 scales paired with fp8_act.

pub fn fp8_act_scale_bytes(&self) -> usize

Allocated byte size of fp8_act_scale (debug bounds-check at call sites).

pub fn fp8_act_scale_kmajor(&self) -> DevicePtr

Transposed ([K/128, ceil16(M)]) copy of fp8_act_scale — the VEC128 B-scale layout the cuBLASLt block-scaled FP8 GEMM documents. The prefill-projection sibling of ffn_act_scale_kmajor.

pub fn fp8_act_scale_kmajor_bytes(&self) -> usize

Allocated byte size of fp8_act_scale_kmajor (bounds-check at call sites).

pub fn q2_dequant_scratch(&self) -> DevicePtr

Persistent BF16 transient-dequant scratch for native keep-packed Q2_0 prefill. Reused per projection: dequant into it, GEMM reads it (same stream), no free. NULL unless ATLAS_GGUF_NATIVE_Q2.

pub fn q2_dequant_scratch_bytes(&self) -> usize

Allocated byte size of q2_dequant_scratch (debug bounds-check).

pub fn q2_act_q8(&self) -> DevicePtr

Persistent q8_1 activation scratch for native Q2_0 MMQ prefill (ATLAS_GGUF_NATIVE_Q2_MMQ). NULL unless the flag is set.

pub fn q2_act_q8_bytes(&self) -> usize

Allocated byte size of q2_act_q8 (debug bounds-check).

pub fn splitk_workspace(&self) -> DevicePtr

pub fn o_latent(&self) -> DevicePtr

Grouped O-projection latent [M, o_groups*o_lora_rank] BF16 (V4-Flash).

pub fn norm_unit_w(&self) -> DevicePtr

All-ones BF16 vector (max_dim) — weight for unweighted RMSNorm (q_b_norm).

pub fn hc_streams(&self) -> DevicePtr

HC residual streams [M, hc_mult, hidden] BF16 (DeepSeek-V4 mHC).

pub fn hc_lowrank_scratch(&self) -> DevicePtr

Low-rank mHC split-collapse scratch: [T<=64, hc*H] normed followed by [T<=64, rank] low, both F32. See sizes.rs.

pub fn qsa_select_scratch(&self) -> DevicePtr

QSA stage-2 prefill-selection scratch, shared by the indexer layers (serial). Layout managed by layers::qsa; see sizes.rs.

pub fn hc_post(&self) -> DevicePtr

HC post mixing weights [M, hc_mult] F32.

pub fn hc_comb(&self) -> DevicePtr

HC comb Sinkhorn matrix [M, hc_mult, hc_mult] F32.

pub fn max_batch_tokens(&self) -> usize

pub fn decode_meta(&self) -> DecodeMetaLayout

Derived batched-decode metadata layout (rows/offsets). Byte-identical to the legacy fixed 32-row layout for every serve max_batch_size <= 32.

pub fn sizes(&self) -> &BufferSizes

pub fn lora_xa(&self) -> DevicePtr

Env-gated (ATLAS_SSM_SAVE_DUMP) per-buffer checksum probe.

CBD: localize a stale/uninitialized decode-scratch buffer on the prefix-cache skip path. Dumps sum/ssq/sabs over the FULL allocation (so leftover-from-prior-occupant bytes in unwritten rows are visible) for every reusable buffer. Treats raw bytes as f32 lanes — exact numeric meaning is irrelevant; we only need a stable fingerprint that differs iff the bytes differ. Synchronizes the stream first. LoRA compressed activation scratch xa = x@Aᵀ [M, max_rank] BF16. DevicePtr::NULL when no adapter is configured.

pub fn lora_xa_bytes(&self) -> usize

Allocated byte size of lora_xa (0 when no adapter).

pub fn lora_delta(&self) -> DevicePtr

LoRA expand scratch delta = xa@Bᵀ [M, max(hidden, intermediate)] BF16. DevicePtr::NULL when no adapter is configured.

pub fn lora_delta_bytes(&self) -> usize

Allocated byte size of lora_delta (0 when no adapter).

pub fn lora_hact(&self) -> DevicePtr

LoRA hidden-activation scratch [M, intermediate_size] BF16 for the runtime FFN delta path. DevicePtr::NULL when no adapter.

pub fn lora_hact_bytes(&self) -> usize

Allocated byte size of lora_hact (0 when no adapter).

pub fn lora_seq_slot(&self) -> DevicePtr

LoRA per-request routing slots [max_batch_tokens] i32 for the prefill path — one adapter SLOT index per prefilling token. DevicePtr::NULL when no adapter is configured.

pub fn debug_buffer_checksum( &self, gpu: &dyn GpuBackend, stream: u64, tag: &str, )

pub fn zero_prefill_essentials( &self, gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>

Zero only buffers that carry residual state between requests.

During prefill, every buffer except hidden_states and residual is fully overwritten before being read within the layer loop:

• norm_output, qkv_output, attn_output: written by each layer’s projection
• gate_logits, moe_output: written by MoE gate/output
• ssm_*: written by SSM projection
• expert_*: written by expert compute
• logits: written by LM head on last token
• scratch: overwritten by metadata upload and MoE routing
• splitk_workspace: written by attention kernel

This reduces per-chunk memset from 17 calls to 2, saving ~15 memset launches × bandwidth on the LPDDR5X bus per prefill chunk.

pub fn zero_all_rows( &self, gpu: &dyn GpuBackend, stream: u64, tokens: usize, ) -> Result<()>

zero_all, but only the first tokens rows of every token-major arena.

🔴 Every buffer zero_all wipes is [max_batch_tokens, row]-major — verified against the allocated sizes: size / max_batch_tokens is exactly one token’s row for each of them (qkv_output 3x8192 BF16, attn_output 64x256, expert_gate_out topk2048, expert_down_out topk4096, …). A decode step carrying tokens tokens can therefore only ever read rows 0..tokens, and zeroing the rest is dead bandwidth.

Measured on GLM-5.3-Flash, 2 x GB10, max_batch_tokens = 4096 (nsys, 2026-08-28): zero_all issues 18 memsets totalling 1.59 GB and 8.01 ms on every single decode token — 9.4 % of an 85 ms step, all of it GPU-idle time before the first kernel.

logits, scratch and splitk_workspace are NOT token-major (metadata arenas / vocab-sized), so they keep the full wipe. They are 30 MB of the 1590.

pub fn zero_all(&self, gpu: &dyn GpuBackend, stream: u64) -> Result<()>

Zero all reusable buffers to eliminate stale data between requests. Ensures deterministic computation regardless of request history.

impl BufferArena

pub fn ssm_rowwise_w_bf16_bytes(&self) -> usize

Allocated byte size of the row-wise GDN prefill BF16-weight slab. 0 when ATLAS_FP8_ROWWISE was not armed at boot.

pub fn take_ssm_rowwise_w_bf16(&self, bytes: usize) -> Result<DevicePtr>

Carve the next bytes of the row-wise GDN prefill BF16-weight slab.

Bump-only and never returned: each GDN layer takes its in_proj_qkvz and out_proj slices on its FIRST prefill and holds them for the life of the arena, because a dequanted weight is as immutable as the weight it came from. Sized in sizes_rowwise::ssm_rowwise_w_bf16_bytes for exactly num_ssm_layers of those pairs, so exhaustion means the sizing and the callers disagree — a bug, reported as one rather than papered over with a fresh allocation (that is the #917 defect this replaces).

Relaxed is enough: the scheduler drives one forward at a time (the same single-threaded invariant cublaslt::Ctx documents), so this is a counter that happens to be atomic rather than a contended one.

impl BufferArena

pub fn new( config: &ModelConfig, max_batch_tokens: usize, max_seq_len: usize, kv_block_size: usize, max_batch_size: usize, gpu: &dyn GpuBackend, ) -> Result<Self>

Allocate all intermediate buffers on the GPU.

pub fn from_sizes( config: &ModelConfig, sizes: BufferSizes, max_batch_tokens: usize, max_batch_size: usize, gpu: &dyn GpuBackend, ) -> Result<Self>

BufferArena::new with the ledger handed in instead of derived.

BufferSizes::from_config reads the process environment for the env-gated entries (q2_*, ssm_rowwise_w_bf16), and set_var is process-global and unsafe — so a test that wants one of those arms ARMED builds the sizes by hand and comes in here rather than racing every other test in the binary. Production still goes through new.

Trait Implementations

impl ModelResource<dyn GpuBackend> for BufferArena

Release every buffer this arena owns.

The destructure below is exhaustive on purpose — no ... A buffer added to BufferArena without a matching free is a leak that only shows up as the next model failing to fit, so the compiler is made to refuse the addition instead. If this line stops compiling, the fix is to free the new field, not to add a wildcard.

fn label(&self) -> &'static str

Human name, for the teardown report and for attributing a failure.

fn release(&mut self, gpu: &dyn GpuBackend) -> Result<()>

Release everything this owns. Must be idempotent: the host calls it, and a Drop backstop may call it again.