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 qkv_output(&self) -> DevicePtr

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_ba(&self) -> DevicePtr

pub fn ssm_deinterleaved(&self) -> DevicePtr

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

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_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 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 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 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.

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.