pub struct BufferArena { /* private fields */ }Expand description
2026-09-25: The GPU buffers of one forward pass, M = max_batch_tokens.
BufferSizes::from_config holds each buffer’s size formula.
Implementations§
Source§impl BufferArena
impl BufferArena
pub fn residual(&self) -> DevicePtr
pub fn norm_output(&self) -> DevicePtr
Sourcepub fn norm_output_bytes(&self) -> usize
pub fn norm_output_bytes(&self) -> usize
2026-09-25: Allocated byte size of norm_output, the bound the attention
prefill o_proj’s cuBLASLt arm (paged_oproj.rs) checks its ceil16(M)
rows against.
pub fn qkv_output(&self) -> DevicePtr
Sourcepub fn qkv_output_bytes(&self) -> usize
pub fn qkv_output_bytes(&self) -> usize
2026-09-25: Allocated byte size of qkv_output, the bound the multi-seq
decode W8A8 arm (multi_seq/w8a8_decode.rs) and the cuBLASLt Q/K/V
prefill arm check their ceil16(M) rows against.
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
Sourcepub fn ssm_qkvz_bytes(&self) -> usize
pub fn ssm_qkvz_bytes(&self) -> usize
2026-09-25: Allocated byte size of ssm_qkvz, the QKVZ projection’s
destination on an interleaved model, checked by the cuBLASLt arm against
its ceil16(M) rows.
pub fn ssm_ba(&self) -> DevicePtr
Sourcepub fn ssm_deinterleaved(&self) -> DevicePtr
pub fn ssm_deinterleaved(&self) -> DevicePtr
2026-09-25: Sequential [Q|K|V|Z] rows after deinterleaving.
Sourcepub fn ssm_deinterleaved_bytes(&self) -> usize
pub fn ssm_deinterleaved_bytes(&self) -> usize
2026-09-25: Allocated byte size of ssm_deinterleaved, the QKVZ
projection’s destination on a sequential model; same padded-M check.
Sourcepub fn ssm_gates(&self) -> DevicePtr
pub fn ssm_gates(&self) -> DevicePtr
2026-09-25: GDN gates in FP32: per token, linear_num_value_heads gate
values, then as many beta values.
Sourcepub fn ssm_conv_out_f32(&self) -> DevicePtr
pub fn ssm_conv_out_f32(&self) -> DevicePtr
2026-09-25: FP32 conv1d output of the SSM layers.
Sourcepub fn scratch(&self) -> DevicePtr
pub fn scratch(&self) -> DevicePtr
2026-09-25: Scratch for MoE routing and the metadata uploads (layouts in sizes.rs).
Sourcepub fn ssd_scratch(&self) -> DevicePtr
pub fn ssd_scratch(&self) -> DevicePtr
2026-09-25: Mamba-2 SSD chunked-scan scratch (dt | dA_cumsum | CB); NULL without Mamba-2 layers.
Sourcepub fn token_ids(&self) -> DevicePtr
pub fn token_ids(&self) -> DevicePtr
2026-09-25: Token ids [M] u32 of the pass, read by the DeepSeek-V4
hash-routed MoE layers. The caller uploads them before the layer loop,
and under CUDA-graph decode before each replay.
Sourcepub fn scratch_bytes(&self) -> usize
pub fn scratch_bytes(&self) -> usize
2026-09-25: Allocated byte size of scratch, the bound for metadata staging.
pub fn expert_gate_out(&self) -> DevicePtr
pub fn expert_up_out(&self) -> DevicePtr
Sourcepub fn expert_gate_out_bytes(&self) -> usize
pub fn expert_gate_out_bytes(&self) -> usize
2026-09-25: Allocated byte size of expert_gate_out, equal to that of
expert_up_out (debug-asserted).
pub fn moe_output_bytes(&self) -> usize
pub fn expert_down_out(&self) -> DevicePtr
pub fn expert_down_out_bytes(&self) -> usize
pub fn gate_logits_bytes(&self) -> usize
pub fn logits_bytes(&self) -> usize
pub fn attn_output_bytes(&self) -> usize
Sourcepub fn gdn_fla_scratch(&self) -> DevicePtr
pub fn gdn_fla_scratch(&self) -> DevicePtr
2026-09-25: GDN FLA prefill scratch (W | U | S | uc | gc, carved by the
caller); DevicePtr::NULL unless both linear-attention head dims are 128.
Sourcepub fn ffn_act_q8(&self) -> DevicePtr
pub fn ffn_act_q8(&self) -> DevicePtr
2026-09-25: Shared dense-FFN q8_1 activation scratch (Q4_K MMQ); NULL for MoE.
Sourcepub fn ffn_act_a(&self) -> DevicePtr
pub fn ffn_act_a(&self) -> DevicePtr
2026-09-25: Shared dense-FFN int8 / NVFP4 / FP8 activation scratch; NULL for MoE.
Sourcepub fn ffn_act_scale(&self) -> DevicePtr
pub fn ffn_act_scale(&self) -> DevicePtr
2026-09-25: The scales paired with ffn_act_a; NULL for MoE.
Sourcepub fn ffn_gate_up_fused(&self) -> DevicePtr
pub fn ffn_gate_up_fused(&self) -> DevicePtr
2026-09-25: [ceil16(GATEUP_FUSED_MAX_M), 2 * intermediate] BF16 output
of the fused dense-FFN gate+up decode GEMM: gate at column 0, up at
column intermediate. NULL for MoE.
Sourcepub fn ffn_gate_up_fused_bytes(&self) -> usize
pub fn ffn_gate_up_fused_bytes(&self) -> usize
2026-09-25: Allocated byte size of ffn_gate_up_fused; the fused arm
selects itself only when its [ceil16(m), 2 * intermediate] output fits.
pub fn ffn_act_a_bytes(&self) -> usize
pub fn ffn_act_scale_bytes(&self) -> usize
Sourcepub fn ffn_act_scale_kmajor(&self) -> DevicePtr
pub fn ffn_act_scale_kmajor(&self) -> DevicePtr
2026-09-25: [K/128, ceil16(M)] transpose of the dense-FFN FP8 activation
scales for the cuBLASLt block-scaled GEMM; NULL for MoE.
pub fn ffn_act_scale_kmajor_bytes(&self) -> usize
Sourcepub fn fp8_act(&self) -> DevicePtr
pub fn fp8_act(&self) -> DevicePtr
2026-09-25: FP8 activation scratch for the prefill projections.
pub fn fp8_act_bytes(&self) -> usize
Sourcepub fn fp8_act_scale(&self) -> DevicePtr
pub fn fp8_act_scale(&self) -> DevicePtr
2026-09-25: One FP32 scale per 128 elements of fp8_act.
pub fn fp8_act_scale_bytes(&self) -> usize
Sourcepub fn fp8_act_scale_kmajor(&self) -> DevicePtr
pub fn fp8_act_scale_kmajor(&self) -> DevicePtr
2026-09-25: [K/128, ceil16(M)] transpose of fp8_act_scale for the
cuBLASLt block-scaled projections.
pub fn fp8_act_scale_kmajor_bytes(&self) -> usize
Sourcepub fn q2_dequant_scratch(&self) -> DevicePtr
pub fn q2_dequant_scratch(&self) -> DevicePtr
2026-09-25: BF16 dequant scratch for keep-packed Q2_0 prefill, reused by
every projection; NULL unless METRALE_GGUF_NATIVE_Q2=1.
pub fn q2_dequant_scratch_bytes(&self) -> usize
Sourcepub fn q2_act_q8(&self) -> DevicePtr
pub fn q2_act_q8(&self) -> DevicePtr
2026-09-25: q8_1 activation scratch for keep-packed Q2_0 MMQ prefill; NULL
unless METRALE_GGUF_NATIVE_Q2_MMQ=1.
pub fn q2_act_q8_bytes(&self) -> usize
pub fn splitk_workspace(&self) -> DevicePtr
Sourcepub fn o_latent(&self) -> DevicePtr
pub fn o_latent(&self) -> DevicePtr
2026-09-25: Grouped O-projection latent [M, o_groups * o_lora_rank] BF16.
Sourcepub fn norm_unit_w(&self) -> DevicePtr
pub fn norm_unit_w(&self) -> DevicePtr
2026-09-25: All-zero BF16 weight: rms_norm with it is a plain normalize.
Sourcepub fn hc_streams(&self) -> DevicePtr
pub fn hc_streams(&self) -> DevicePtr
2026-09-25: Hyper-connection residual streams [M, hc_mult, hidden] FP32.
Sourcepub fn hc_lowrank_scratch(&self) -> DevicePtr
pub fn hc_lowrank_scratch(&self) -> DevicePtr
2026-09-25: Low-rank hyper-connection scratch; its two layouts are
described in sizes.rs.
Sourcepub fn qsa_select_scratch(&self) -> DevicePtr
pub fn qsa_select_scratch(&self) -> DevicePtr
2026-09-25: QSA prefill-selection scratch shared by the indexer layers;
layers/qsa_select.rs carves it as sizes.rs sizes it.
Sourcepub fn hc_comb(&self) -> DevicePtr
pub fn hc_comb(&self) -> DevicePtr
2026-09-25: Hyper-connection comb matrix [M, hc_mult, hc_mult] F32.
pub fn max_batch_tokens(&self) -> usize
Sourcepub fn decode_meta(&self) -> DecodeMetaLayout
pub fn decode_meta(&self) -> DecodeMetaLayout
2026-09-25: The batched-decode metadata layout.
pub fn sizes(&self) -> &BufferSizes
Sourcepub fn lora_xa(&self) -> DevicePtr
pub fn lora_xa(&self) -> DevicePtr
2026-09-25: LoRA shrink scratch xa = x @ Aᵀ, [M, adapter_max_rank]
BF16; DevicePtr::NULL without an adapter.
pub fn lora_xa_bytes(&self) -> usize
Sourcepub fn lora_delta(&self) -> DevicePtr
pub fn lora_delta(&self) -> DevicePtr
2026-09-25: LoRA expand scratch delta = xa @ Bᵀ BF16; DevicePtr::NULL
without an adapter.
pub fn lora_delta_bytes(&self) -> usize
Sourcepub fn lora_hact(&self) -> DevicePtr
pub fn lora_hact(&self) -> DevicePtr
2026-09-25: LoRA hidden-activation scratch [M, intermediate_size] BF16;
DevicePtr::NULL without an adapter.
pub fn lora_hact_bytes(&self) -> usize
Sourcepub fn lora_seq_slot(&self) -> DevicePtr
pub fn lora_seq_slot(&self) -> DevicePtr
2026-09-25: LoRA adapter slot per prefill token, [M] i32;
DevicePtr::NULL without an adapter.
Sourcepub fn zero_prefill_essentials(
&self,
gpu: &dyn GpuBackend,
stream: u64,
) -> Result<()>
pub fn zero_prefill_essentials( &self, gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>
2026-09-25: Zero hidden_states, residual, gate_logits, the three
expert outputs and moe_output: seven memsets, where zero_all issues
eighteen. The prefill paths call this on a sequence’s first chunk when
they do not call zero_all.
Sourcepub fn zero_all_rows(
&self,
gpu: &dyn GpuBackend,
stream: u64,
tokens: usize,
) -> Result<()>
pub fn zero_all_rows( &self, gpu: &dyn GpuBackend, stream: u64, tokens: usize, ) -> Result<()>
2026-09-25: zero_all limited to the first tokens rows of each
token-major buffer: for a buffer whose size is a multiple of
max_batch_tokens, size / max_batch_tokens * tokens bytes; any other
buffer, and every buffer when tokens >= max_batch_tokens, is zeroed
whole. splitk_workspace, logits and scratch are zeroed whole. The
caller must read only rows 0..tokens. Measured 2026-08-28 on
GLM-5.3-Flash (nsys, max_batch_tokens = 4096): zero_all took 8.01 ms
of an 85 ms decode step.
Sourcepub fn zero_all(&self, gpu: &dyn GpuBackend, stream: u64) -> Result<()>
pub fn zero_all(&self, gpu: &dyn GpuBackend, stream: u64) -> Result<()>
2026-09-25: Zero, in full, the eighteen buffers zero_all_rows covers.
The others (gate_logits_f32, moe_router_in_f32, o_latent,
norm_unit_w, the hyper-connection, QSA, FLA and SSD scratch,
token_ids, ffn_*, fp8_*, q2_*, lora_* and the row-wise slab)
keep their contents.
Source§impl BufferArena
impl BufferArena
Sourcepub fn debug_buffer_checksum(
&self,
gpu: &dyn GpuBackend,
stream: u64,
tag: &str,
)
pub fn debug_buffer_checksum( &self, gpu: &dyn GpuBackend, stream: u64, tag: &str, )
2026-09-25: Synchronize stream, then log at warn level the sum, sum of
squares and sum of absolute values of each listed buffer read as f32
lanes over its whole allocation (non-finite lanes skipped), tagged with
tag. A failed synchronize or copy is ignored; a failed copy skips that
buffer. The model engine calls it at the first decode step when
ssm_save_dump is set.
Source§impl BufferArena
impl BufferArena
Sourcepub fn moe_fp8_scratch(
&self,
c: &ModelConfig,
rows: usize,
) -> Result<MoeFp8Scratch>
pub fn moe_fp8_scratch( &self, c: &ModelConfig, rows: usize, ) -> Result<MoeFp8Scratch>
2026-09-25: Fixed offsets for a given shape, without allocating or mutating ownership.
§Errors
When rows exceeds the arena’s batch, the model has no experts, the shape
needs more than the slab holds, or the slab was released.
Source§impl BufferArena
impl BufferArena
Sourcepub fn ssm_rowwise_w_bf16_bytes(&self) -> usize
pub fn ssm_rowwise_w_bf16_bytes(&self) -> usize
2026-09-25: Allocated byte size of the slab; 0 unless METRALE_FP8_ROWWISE=1
when the arena was sized.
Sourcepub fn take_ssm_rowwise_w_bf16(&self, bytes: usize) -> Result<DevicePtr>
pub fn take_ssm_rowwise_w_bf16(&self, bytes: usize) -> Result<DevicePtr>
2026-09-25: Carve the next bytes of the slab. Slices are never returned:
a GDN layer carves its in_proj_qkvz and out_proj copies on its first
prefill and keeps them for the arena’s life. The slab holds exactly
num_ssm_layers such pairs (sizes_rowwise.rs), so exhaustion is a
sizing bug and is returned as an error, not served by a new allocation.
The cursor is not rolled back after that error.
Relaxed ordering: the scheduler runs one forward at a time, the same
assumption cublaslt::Ctx states for its Send/Sync.
Source§impl BufferArena
impl BufferArena
Sourcepub 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>
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>
2026-09-25: Size the buffers with BufferSizes::from_config and allocate them.
Sourcepub fn from_sizes(
config: &ModelConfig,
sizes: BufferSizes,
max_batch_tokens: usize,
max_batch_size: usize,
gpu: &dyn GpuBackend,
) -> Result<Self>
pub fn from_sizes( config: &ModelConfig, sizes: BufferSizes, max_batch_tokens: usize, max_batch_size: usize, gpu: &dyn GpuBackend, ) -> Result<Self>
2026-09-25: BufferArena::new with the sizes passed in.
BufferSizes::from_config reads the environment for the env-gated
entries (q2_*, ssm_rowwise_w_bf16), so a test that needs one of
them sized builds the sizes itself and calls this instead of setting a
process-global variable. On an allocation error the buffers already
allocated are not freed.
Trait Implementations§
Source§impl ModelResource<dyn GpuBackend> for BufferArena
2026-09-25: The destructure in release names every field, with no .., so
a field added to BufferArena does not compile until it is handled here.
Free the new buffer rather than adding a wildcard: a missed free is a leak
that shows only as the next model failing to fit.
impl ModelResource<dyn GpuBackend> for BufferArena
2026-09-25: The destructure in release names every field, with no .., so
a field added to BufferArena does not compile until it is handled here.
Free the new buffer rather than adding a wildcard: a missed free is a leak
that shows only as the next model failing to fit.