Struct BufferSizes 

pub struct BufferSizes {

    pub hidden_states: usize,
    pub residual: usize,
    pub norm_output: usize,
    pub qkv_output: usize,
    pub attn_output: usize,
    pub gate_logits: usize,
    pub gate_logits_f32: usize,
    pub moe_router_in_f32: usize,
    pub moe_output: usize,
    pub logits: usize,
    pub ssm_qkvz: usize,
    pub ssm_ba: usize,
    pub ssm_deinterleaved: usize,
    pub ssm_gates: usize,
    pub ssm_conv_out_f32: usize,
    pub scratch: usize,
    pub expert_gate_out: usize,
    pub expert_up_out: usize,
    pub expert_down_out: usize,
    pub splitk_workspace: usize,
    pub gdn_fla_scratch: usize,
    pub ssd_scratch: usize,
    pub o_latent: usize,
    pub norm_unit_w: usize,
    pub hc_streams: usize,
    pub hc_post: usize,
    pub hc_comb: usize,
    pub hc_lowrank_scratch: usize,
    pub qsa_select_scratch: usize,
    pub token_ids: usize,
    pub ffn_act_q8: usize,
    pub ffn_act_a: usize,
    pub ffn_act_scale: usize,
    pub fp8_act: usize,
    pub fp8_act_scale: usize,
    pub lora_xa: usize,
    pub lora_delta: usize,
    pub lora_hact: usize,
    pub lora_seq_slot: usize,
    pub q2_dequant_scratch: usize,
    pub q2_act_q8: usize,
}

Byte sizes of each buffer, derived from ModelConfig.

Fields

hidden_states: usize

residual: usize

norm_output: usize

qkv_output: usize

attn_output: usize

gate_logits: usize

gate_logits_f32: usize

FP32 gate logits [m, num_experts] for the ATLAS_FP32_GATE routing path. Keeps the router GEMM accumulator unrounded into top-K so near-tied experts don’t flip on a BF16 store. Allocated whenever num_experts > 0.

moe_router_in_f32: usize

FP32 MoE-input norm output [m, hidden] for ATLAS_FP32_ROUTING — the full-precision router_in the gate GEMM consumes. Allocated when experts > 0.

moe_output: usize

logits: usize

ssm_qkvz: usize

ssm_ba: usize

ssm_deinterleaved: usize

ssm_gates: usize

ssm_conv_out_f32: usize

scratch: usize

expert_gate_out: usize

expert_up_out: usize

expert_down_out: usize

splitk_workspace: usize

gdn_fla_scratch: usize

GDN FLA chunked-prefill scratch (single buffer, sub-divided W|U|S|uc). 0 unless the model is a 128-dim-linear-head GDN model (ATLAS_GDN_FLA path).

ssd_scratch: usize

Mamba-2 SSD chunked-scan scratch (single buffer, sub-divided dt | dA_cumsum | CB). 0 unless the model has Mamba-2 SSM layers. Shared across layers: they run sequentially on one stream, so one allocation serves all 40.

o_latent: usize

Grouped O-projection latent: [M, o_groups*o_lora_rank] BF16 (V4-Flash). 256 (placeholder) when o_groups == 0.

norm_unit_w: usize

Zero-filled BF16 weight (length max_dim) for unweighted RMSNorm under the offset-from-1 kernel convention (scale = 1+weight → 1.0). DeepSeek-V4 q_b_norm.

hc_streams: usize

HC residual streams: [M, hc_mult, hidden] BF16 (DeepSeek-V4 mHC). 256 (placeholder) when hc_mult == 0.

hc_post: usize

HC post mixing weights: [M, hc_mult] F32.

hc_comb: usize

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

hc_lowrank_scratch: usize

Low-rank mHC split-collapse scratch (Qwen3.8-Flash-Next): the staged normed vector [T, hc_mult*hidden] F32 plus the rank vector [T, hc_lowrank] F32, for SMALL T only — decode runs the collapse as three multi-block launches because grid=[1] starves the fused kernel (measured 2.0 ms/call, one SM’s bandwidth). Sized for 64 tokens; the dispatcher falls back to the fused kernel above that.

qsa_select_scratch: usize

QSA stage-2 prefill-selection scratch (Qwen3.8-Flash-Next), SHARED across the 12 indexer layers (they run serially). Layout, slabbed at 2048 selective rows: qk [2048, (n_heads+1)*hd] BF16, q_post [2048, n_heads, hd] F32, scores [2048, max_seq/ratio] F32, lists [2048, topk] i32. 256 (placeholder) when no indexer.

token_ids: usize

Token IDs [M] u32 for the current pass — stable across the layer loop so DeepSeek-V4 hash-MoE layers can read tid2eid[token_id]. Always allocated (small); unused by models without hash routing.

ffn_act_q8: usize

Dense-FFN activation-quant scratch, SHARED across all layers by the MMQ (Q4_K), int8 (W4A8), and NVFP4 (W4A4) prefill paths. Was previously a per-DenseFfnLayer field → 64× duplication (18 GB on Qwen3.6-27B) that OOM’d chunked prefill layer-by-layer. Sized for the largest projection K. ffn_act_q8: q8_1_mmq activations m*kpad*4 + 1MB (Q4_K path). ffn_act_a: int8 [m,K] / NVFP4 packed [m,K/2] activations. ffn_act_scale: int8 [m,K/32]*4 / NVFP4 [m,K/16] group scales. 0 for MoE models (dense FFN prefill path is Dense-only).

ffn_act_a: usize

ffn_act_scale: usize

fp8_act: usize

FP8 block-scaled activation scratch for prefill projections (qkv / o / ssm-qkvz). Persistent so the W8A8+FP32-epilogue path stops doing a per-projection cuMemAlloc + cuStreamSynchronize + cuMemFree. 1 byte/elem.

fp8_act_scale: usize

Per-128-block FP32 scales paired with fp8_act (one f32 per 128 elems).

lora_xa: usize

LoRA shrink output xa = x@Aᵀ: [m, adapter_max_rank] BF16. 0 (→ NULL alloc) when no adapter is configured (adapter_max_rank == 0).

lora_delta: usize

LoRA expand output delta = xa@Bᵀ: [m, max target n_out] BF16, where max n_out = max(hidden, intermediate) — covers k/v/o/gate/up/down in v0 (q_proj is excluded). 0 (→ NULL) when no adapter.

lora_hact: usize

LoRA hidden-activation scratch [m, intermediate_size] BF16 for the runtime delta path on FFN projections. 0 (→ NULL) when no adapter.

lora_seq_slot: usize

LoRA per-request routing slots [m] i32 — one adapter SLOT index per prefilling token (all equal for a single-request prefill; resolves -1→active before upload). Dedicated buffer (not a packed meta offset) so the m-element prefill slot array never collides with the per-path positions/slots/block_table region. 0 (→ NULL) when no adapter (adapter_max_rank == 0).

q2_dequant_scratch: usize

Native keep-packed Q2_0 prefill transient-dequant scratch (ATLAS_GGUF_NATIVE_Q2=1). ONE persistent BF16 [N,K] buffer sized to the LARGEST keep-packed projection, REUSED for every per-projection dequant so prefill stops doing a per-matmul cuMemAlloc + cuStreamSynchronize + cuMemFree (the multi-second fixed cost behind the 3.7 s / 28-token TTFT regression). 0 (→ NULL) unless the flag is set.

q2_act_q8: usize

Native Q2_0 MMQ prefill q8_1 activation scratch (ATLAS_GGUF_NATIVE_Q2_MMQ=1). ONE persistent q8_1_mmq buffer (m*kpad*4 + 1MB) shared by every kept-packed projection (FFN gate/up/down, attn q/k/v/o, GDN qkvz): each seam quantizes its BF16 activation into this buffer then runs the packed MMQ GEMM — so the 2-bit weight is never dequantized to a BF16 scratch (kills the ~2s dequant tax AND the shared-q2_dequant_scratch co-dispatch race). Sized to the widest projection K = max(hidden, intermediate, q_heads*head_dim). 0 (→ NULL) unless the MMQ sub-flag is set.

Implementations

impl BufferSizes

pub fn from_config( config: &ModelConfig, max_batch_tokens: usize, max_seq_len: usize, kv_block_size: usize, max_batch_size: usize, ) -> Self

Compute all buffer sizes from model config and max batch tokens.

All sizes in bytes. BF16 = 2 bytes per element. Logits buffer is capped: only needed for decode (1 token) or speculative verification (K tokens), never for full prefill.

max_seq_len and kv_block_size are needed to size the scratch buffer for block table metadata during batched decode / verify.

pub fn total_bytes(&self) -> usize

Total bytes across all buffers.

Trait Implementations

impl Clone for BufferSizes

fn clone(&self) -> BufferSizes

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for BufferSizes

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.