Module buffers

Module buffers 

Source
Expand description

2026-09-25: The per-model GPU buffer arena: every intermediate tensor of a forward pass, allocated once and reused by every step.

Owner: gpu-runtime. Invariants:

  • Every buffer is allocated at its BufferSizes size. The gated ones (ssd_scratch, gdn_fla_scratch, ffn_*, moe_fp8_scratch, q2_*, lora_*, ssm_rowwise_w_bf16) are DevicePtr::NULL when that size is 0.
  • No device memory is allocated after construction, and release frees every pointer field.

Re-exports§

pub use decode_meta::DECODE_META_MAX_ROWS;
pub use decode_meta::DECODE_META_MIN_ROWS;
pub use decode_meta::DecodeMetaLayout;

Modules§

decode_meta
2026-09-25: Fixed-stride batched-decode metadata layout, derived from the serve max_batch_size.

Structs§

BufferArena
2026-09-25: The GPU buffers of one forward pass, M = max_batch_tokens. BufferSizes::from_config holds each buffer’s size formula.
BufferSizes
2026-09-25: Byte size of each arena buffer. M is max_batch_tokens and ceil16(M) is M rounded up to a multiple of 16.
MoeFp8Scratch
2026-09-25: Disjoint regions; shared/gate/up/down phases reuse these on one stream.

Constants§

GATEUP_FUSED_MAX_M
2026-09-25: The widest M the fused dense-FFN gate+up decode GEMM serves, and so the row extent of ffn_gate_up_fused. Declared here because the arena is sized in this crate; the dispatch (layers/dense_ffn_gateup_fused.rs in metrale-model-layers) reads it for its 5..=GATEUP_FUSED_MAX_M band, so the band and the buffer cannot disagree.
Q12_SIZING_STREAMS
2026-09-25: Streams scratch is provisioned for in the kernel-batched prefill staging. A batch whose footprint does not fit is refused by check_kernel_batched_eligible (metrale-model-engine) and runs per stream, so this sets how often the batched path is available, not safety.
VERIFY_ROW_CAP
2026-09-29: Most rows (R = Σ ks) one batched MTP verify scores: 128 sequences at one draft (the serve’s widest batch, DECODE_META_MAX_ROWS). It sizes the verify block-table rows (sizes/regions.rs) and the logits floor here; model-engine’s VERIFY_ROW_CAP and the scheduler’s VERIFY_ROW_BUDGET are this value. Was 160 (32 sequences at up to five rows).

Functions§

q2_dequant_scratch_bytes
2026-09-25: Bytes of the keep-packed Q2_0 dequant scratch: hidden_size times the widest of intermediate, QKVZ, Mamba-2 in_proj, the q projection and 2 · kv_heads · head_dim, in BF16. This covers a weight [N, K] with hidden_size as one dimension and the other at most that width; the dequant callers debug-assert N * K * 2 against the allocated size. It does not depend on the batch.
q12_batched_scratch_bytes
2026-09-25: Scratch bytes of the kernel-batched prefill staging for n streams of chunk_len tokens each. Both BufferSizes::from_config and check_kernel_batched_eligible call it (or the varlen form), so they agree on whether a batch fits. The terms follow prefill_b/batch_kernel.rs (MoE top-K area, n per-stream metadata blocks), prefill_b/stage_batched.rs (positions, three streams under MRoPE, slots, the block-table and seq_len pointer arrays, cu_seqlens, kv_lens) and the h_state_ptrs table.
q12_batched_scratch_bytes_varlen
2026-09-25: q12_batched_scratch_bytes for a ragged batch: total_tokens is the packed total and max_chunk_len sizes the per-stream metadata slots, so a varlen batch is not charged the longest stream’s length for every stream.
ssm_rowwise_w_bf16_bytes
2026-09-25: Slab bytes: the per-layer pair times num_ssm_layers, or 0 unless METRALE_FP8_ROWWISE=1. The predicate is the one in the loader’s rowwise_fp8_enabled (weight_loader/qwen35_dense/rowwise_fp8.rs): the loader installs the row-wise weights on it and the arms then need this slab, so the two must agree.
ssm_rowwise_w_bf16_bytes_for
2026-09-25: ssm_rowwise_w_bf16_bytes with the lever passed in, for tests that must not set the process environment.
ssm_rowwise_w_bf16_layer_bytes
2026-09-25: BF16 bytes of one GDN layer’s row-wise copies: in_proj_qkvz [ssm_qkvz_size, hidden] and out_proj [hidden, value_dim], the two weights set_fp8_rowwise_prefill_weights installs, with value_dim = linear_num_value_heads * linear_value_head_dim. At the Qwen3.8-27B shapes of buffers/tests.rs (hidden 5120, 16x128 key heads, 48x128 value heads) that is 167772160 + 62914560 B.