metrale_gpu_runtime/
buffers.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The per-model GPU buffer arena: every intermediate tensor of a forward pass, allocated once and reused by every step.
4//!
5//! Owner: gpu-runtime.
6//! Invariants:
7//! - Every buffer is allocated at its [`BufferSizes`] size. The gated ones
8//!   (`ssd_scratch`, `gdn_fla_scratch`, `ffn_*`, `moe_fp8_scratch`, `q2_*`,
9//!   `lora_*`, `ssm_rowwise_w_bf16`) are `DevicePtr::NULL` when that size is 0.
10//! - No device memory is allocated after construction, and `release` frees
11//!   every pointer field.
12
13use crate::gpu::{DevicePtr, GpuBackend};
14use anyhow::Result;
15use metrale_config::ModelConfig;
16
17mod accessors;
18mod debug_checksum;
19pub mod decode_meta;
20mod moe_fp8_scratch;
21mod release;
22mod rowwise_slab;
23mod sizes;
24mod sizes_q12;
25mod sizes_q2;
26mod sizes_rowwise;
27pub use decode_meta::{DECODE_META_MAX_ROWS, DECODE_META_MIN_ROWS, DecodeMetaLayout};
28pub use moe_fp8_scratch::MoeFp8Scratch;
29pub use sizes::{BufferSizes, GATEUP_FUSED_MAX_M, VERIFY_ROW_CAP};
30pub use sizes_q2::q2_dequant_scratch_bytes;
31pub use sizes_q12::{
32    Q12_SIZING_STREAMS, q12_batched_scratch_bytes, q12_batched_scratch_bytes_varlen,
33};
34pub use sizes_rowwise::{
35    ssm_rowwise_w_bf16_bytes, ssm_rowwise_w_bf16_bytes_for, ssm_rowwise_w_bf16_layer_bytes,
36};
37
38/// 2026-09-25: The GPU buffers of one forward pass, `M = max_batch_tokens`.
39/// [`BufferSizes::from_config`] holds each buffer's size formula.
40pub struct BufferArena {
41    hidden_states: DevicePtr,
42    residual: DevicePtr,
43    norm_output: DevicePtr,
44    qkv_output: DevicePtr,
45    attn_output: DevicePtr,
46    gate_logits: DevicePtr,
47    gate_logits_f32: DevicePtr,
48    moe_router_in_f32: DevicePtr,
49    moe_output: DevicePtr,
50    logits: DevicePtr,
51    ssm_qkvz: DevicePtr,
52    ssm_ba: DevicePtr,
53    ssm_deinterleaved: DevicePtr,
54    ssm_gates: DevicePtr,
55    ssm_conv_out_f32: DevicePtr,
56    scratch: DevicePtr,
57    expert_gate_out: DevicePtr,
58    expert_up_out: DevicePtr,
59    expert_down_out: DevicePtr,
60    splitk_workspace: DevicePtr,
61    o_latent: DevicePtr,
62    norm_unit_w: DevicePtr,
63    hc_streams: DevicePtr,
64    hc_post: DevicePtr,
65    hc_comb: DevicePtr,
66    hc_lowrank_scratch: DevicePtr,
67    qsa_select_scratch: DevicePtr,
68    gdn_fla_scratch: DevicePtr,
69    ssd_scratch: DevicePtr,
70    token_ids: DevicePtr,
71    ffn_act_q8: DevicePtr,
72    ffn_act_a: DevicePtr,
73    ffn_act_scale: DevicePtr,
74    ffn_act_scale_kmajor: DevicePtr,
75    ffn_gate_up_fused: DevicePtr,
76    fp8_act: DevicePtr,
77    /// 2026-09-25: The grouped FP8 MoE activation, scale and tile-worklist slab that every
78    /// MoE layer reuses in turn; see [`BufferArena::moe_fp8_scratch`].
79    moe_fp8_scratch: DevicePtr,
80    fp8_act_scale: DevicePtr,
81    fp8_act_scale_kmajor: DevicePtr,
82    q2_dequant_scratch: DevicePtr,
83    lora_xa: DevicePtr,
84    lora_delta: DevicePtr,
85    lora_hact: DevicePtr,
86    lora_seq_slot: DevicePtr,
87    q2_act_q8: DevicePtr,
88    ssm_rowwise_w_bf16: DevicePtr,
89    /// 2026-09-25: Bytes carved from `ssm_rowwise_w_bf16` so far; it only grows.
90    ssm_rowwise_w_bf16_used: std::sync::atomic::AtomicUsize,
91    max_batch_tokens: usize,
92    /// 2026-09-25: Decode-metadata layout for the serve `max_batch_size`; the
93    /// batched-decode upload (`upload_batch_metadata_fixed` / `_at` in
94    /// metrale-model-engine) takes its offsets from it.
95    decode_meta: DecodeMetaLayout,
96    sizes: BufferSizes,
97}
98
99impl BufferArena {
100    /// 2026-09-25: Size the buffers with [`BufferSizes::from_config`] and allocate them.
101    pub fn new(
102        config: &ModelConfig,
103        max_batch_tokens: usize,
104        max_seq_len: usize,
105        kv_block_size: usize,
106        max_batch_size: usize,
107        gpu: &dyn GpuBackend,
108    ) -> Result<Self> {
109        let sizes = BufferSizes::from_config(
110            config,
111            max_batch_tokens,
112            max_seq_len,
113            kv_block_size,
114            max_batch_size,
115        );
116        Self::from_sizes(config, sizes, max_batch_tokens, max_batch_size, gpu)
117    }
118
119    /// 2026-09-25: [`BufferArena::new`] with the sizes passed in.
120    ///
121    /// `BufferSizes::from_config` reads the environment for the env-gated
122    /// entries (`q2_*`, `ssm_rowwise_w_bf16`), so a test that needs one of
123    /// them sized builds the sizes itself and calls this instead of setting a
124    /// process-global variable. On an allocation error the buffers already
125    /// allocated are not freed.
126    pub fn from_sizes(
127        config: &ModelConfig,
128        sizes: BufferSizes,
129        max_batch_tokens: usize,
130        max_batch_size: usize,
131        gpu: &dyn GpuBackend,
132    ) -> Result<Self> {
133        let decode_meta = DecodeMetaLayout::for_max_batch_size(max_batch_size);
134
135        let hidden_states = gpu.alloc(sizes.hidden_states)?;
136        let residual = gpu.alloc(sizes.residual)?;
137        let norm_output = gpu.alloc(sizes.norm_output)?;
138        let qkv_output = gpu.alloc(sizes.qkv_output)?;
139        let attn_output = gpu.alloc(sizes.attn_output)?;
140        let gate_logits = gpu.alloc(sizes.gate_logits)?;
141        let gate_logits_f32 = gpu.alloc(sizes.gate_logits_f32)?;
142        let moe_router_in_f32 = gpu.alloc(sizes.moe_router_in_f32)?;
143        let moe_output = gpu.alloc(sizes.moe_output)?;
144        let logits = gpu.alloc(sizes.logits)?;
145        let ssm_qkvz = gpu.alloc(sizes.ssm_qkvz)?;
146        let ssm_ba = gpu.alloc(sizes.ssm_ba)?;
147        let ssm_deinterleaved = gpu.alloc(sizes.ssm_deinterleaved)?;
148        let ssm_gates = gpu.alloc(sizes.ssm_gates)?;
149        let ssm_conv_out_f32 = gpu.alloc(sizes.ssm_conv_out_f32)?;
150        let scratch = gpu.alloc(sizes.scratch)?;
151        let expert_gate_out = gpu.alloc(sizes.expert_gate_out)?;
152        let expert_up_out = gpu.alloc(sizes.expert_up_out)?;
153        let expert_down_out = gpu.alloc(sizes.expert_down_out)?;
154        let splitk_workspace = gpu.alloc(sizes.splitk_workspace)?;
155        let o_latent = gpu.alloc(sizes.o_latent)?;
156        // 2026-09-25: All zeros: the `rms_norm` kernel scales by 1 + weight, so a
157        // zero weight is a plain normalize.
158        let norm_unit_w = gpu.alloc(sizes.norm_unit_w)?;
159        gpu.memset(norm_unit_w, 0, sizes.norm_unit_w)?;
160        let hc_streams = gpu.alloc(sizes.hc_streams)?;
161        let hc_post = gpu.alloc(sizes.hc_post)?;
162        let hc_comb = gpu.alloc(sizes.hc_comb)?;
163        let hc_lowrank_scratch = gpu.alloc(sizes.hc_lowrank_scratch)?;
164        let qsa_select_scratch = gpu.alloc(sizes.qsa_select_scratch)?;
165        // 2026-09-25: A gated entry of 0 bytes is left NULL, which its callers
166        // read as "not available".
167        let ssd_scratch = if sizes.ssd_scratch > 0 {
168            gpu.alloc(sizes.ssd_scratch)?
169        } else {
170            DevicePtr::NULL
171        };
172        let gdn_fla_scratch = if sizes.gdn_fla_scratch > 0 {
173            gpu.alloc(sizes.gdn_fla_scratch)?
174        } else {
175            DevicePtr::NULL
176        };
177        let token_ids = gpu.alloc(sizes.token_ids)?;
178        let ffn_act_q8 = if sizes.ffn_act_q8 > 0 {
179            gpu.alloc(sizes.ffn_act_q8)?
180        } else {
181            DevicePtr::NULL
182        };
183        let ffn_act_a = if sizes.ffn_act_a > 0 {
184            gpu.alloc(sizes.ffn_act_a)?
185        } else {
186            DevicePtr::NULL
187        };
188        let ffn_act_scale = if sizes.ffn_act_scale > 0 {
189            gpu.alloc(sizes.ffn_act_scale)?
190        } else {
191            DevicePtr::NULL
192        };
193        let ffn_act_scale_kmajor = if sizes.ffn_act_scale_kmajor > 0 {
194            gpu.alloc(sizes.ffn_act_scale_kmajor)?
195        } else {
196            DevicePtr::NULL
197        };
198        let ffn_gate_up_fused = if sizes.ffn_gate_up_fused > 0 {
199            gpu.alloc(sizes.ffn_gate_up_fused)?
200        } else {
201            DevicePtr::NULL
202        };
203        let fp8_act = gpu.alloc(sizes.fp8_act)?;
204        let moe_fp8_scratch = if sizes.moe_fp8_scratch > 0 {
205            gpu.alloc(sizes.moe_fp8_scratch)?
206        } else {
207            DevicePtr::NULL
208        };
209        let fp8_act_scale = gpu.alloc(sizes.fp8_act_scale)?;
210        let fp8_act_scale_kmajor = gpu.alloc(sizes.fp8_act_scale_kmajor)?;
211        let q2_dequant_scratch = if sizes.q2_dequant_scratch > 0 {
212            gpu.alloc(sizes.q2_dequant_scratch)?
213        } else {
214            DevicePtr::NULL
215        };
216        let lora_xa = if sizes.lora_xa > 0 {
217            gpu.alloc(sizes.lora_xa)?
218        } else {
219            DevicePtr::NULL
220        };
221        let lora_delta = if sizes.lora_delta > 0 {
222            gpu.alloc(sizes.lora_delta)?
223        } else {
224            DevicePtr::NULL
225        };
226        let lora_hact = if sizes.lora_hact > 0 {
227            gpu.alloc(sizes.lora_hact)?
228        } else {
229            DevicePtr::NULL
230        };
231        let lora_seq_slot = if sizes.lora_seq_slot > 0 {
232            gpu.alloc(sizes.lora_seq_slot)?
233        } else {
234            DevicePtr::NULL
235        };
236        let q2_act_q8 = if sizes.q2_act_q8 > 0 {
237            gpu.alloc(sizes.q2_act_q8)?
238        } else {
239            DevicePtr::NULL
240        };
241        let ssm_rowwise_w_bf16 = if sizes.ssm_rowwise_w_bf16 > 0 {
242            gpu.alloc(sizes.ssm_rowwise_w_bf16)?
243        } else {
244            DevicePtr::NULL
245        };
246
247        tracing::info!(
248            "Buffer arena: {} tokens × {:.1} MB total (attn_out={:.1}MB, ssm_deint={:.1}MB, kv_lora_rank={})",
249            max_batch_tokens,
250            sizes.total_bytes() as f64 / (1024.0 * 1024.0),
251            sizes.attn_output as f64 / (1024.0 * 1024.0),
252            sizes.ssm_deinterleaved as f64 / (1024.0 * 1024.0),
253            config.kv_lora_rank,
254        );
255
256        Ok(Self {
257            hidden_states,
258            residual,
259            norm_output,
260            qkv_output,
261            attn_output,
262            gate_logits,
263            gate_logits_f32,
264            moe_router_in_f32,
265            moe_output,
266            logits,
267            ssm_qkvz,
268            ssm_ba,
269            ssm_deinterleaved,
270            ssm_gates,
271            ssm_conv_out_f32,
272            scratch,
273            expert_gate_out,
274            expert_up_out,
275            expert_down_out,
276            splitk_workspace,
277            o_latent,
278            norm_unit_w,
279            hc_streams,
280            hc_post,
281            hc_comb,
282            hc_lowrank_scratch,
283            qsa_select_scratch,
284            gdn_fla_scratch,
285            ssd_scratch,
286            token_ids,
287            ffn_act_q8,
288            ffn_act_a,
289            ffn_act_scale,
290            ffn_act_scale_kmajor,
291            ffn_gate_up_fused,
292            fp8_act,
293            moe_fp8_scratch,
294            fp8_act_scale,
295            fp8_act_scale_kmajor,
296            q2_dequant_scratch,
297            lora_xa,
298            lora_delta,
299            lora_hact,
300            lora_seq_slot,
301            q2_act_q8,
302            ssm_rowwise_w_bf16,
303            ssm_rowwise_w_bf16_used: std::sync::atomic::AtomicUsize::new(0),
304            max_batch_tokens,
305            decode_meta,
306            sizes,
307        })
308    }
309}
310
311#[cfg(test)]
312mod tests;
metrale_gpu_runtime/buffers.rs

metrale_gpu_runtime/
buffers.rs
