metrale_gpu_runtime/buffers/
accessors.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `BufferArena` accessors (device pointers and allocated byte sizes) and the zeroing helpers.
4//!
5//! Owner: gpu-runtime.
6//! Invariants:
7//! - Every `*_bytes` accessor returns the allocated size of its buffer, from
8//!   `BufferSizes`.
9//! - The zeroing helpers only enqueue `memset_async` on the given stream; the
10//!   first failed enqueue is returned and later buffers are not zeroed.
11
12use super::{BufferArena, sizes::BufferSizes};
13use crate::gpu::{DevicePtr, GpuBackend};
14
15impl BufferArena {
16    pub fn hidden_states(&self) -> DevicePtr {
17        self.hidden_states
18    }
19    pub fn residual(&self) -> DevicePtr {
20        self.residual
21    }
22    pub fn norm_output(&self) -> DevicePtr {
23        self.norm_output
24    }
25    /// 2026-09-25: Allocated byte size of `norm_output`, the bound the attention
26    /// prefill o_proj's cuBLASLt arm (`paged_oproj.rs`) checks its `ceil16(M)`
27    /// rows against.
28    pub fn norm_output_bytes(&self) -> usize {
29        self.sizes.norm_output
30    }
31    pub fn qkv_output(&self) -> DevicePtr {
32        self.qkv_output
33    }
34    /// 2026-09-25: Allocated byte size of `qkv_output`, the bound the multi-seq
35    /// decode W8A8 arm (`multi_seq/w8a8_decode.rs`) and the cuBLASLt Q/K/V
36    /// prefill arm check their `ceil16(M)` rows against.
37    pub fn qkv_output_bytes(&self) -> usize {
38        self.sizes.qkv_output
39    }
40    pub fn attn_output(&self) -> DevicePtr {
41        self.attn_output
42    }
43    pub fn gate_logits(&self) -> DevicePtr {
44        self.gate_logits
45    }
46    pub fn gate_logits_f32(&self) -> DevicePtr {
47        self.gate_logits_f32
48    }
49    pub fn moe_router_in_f32(&self) -> DevicePtr {
50        self.moe_router_in_f32
51    }
52    pub fn moe_output(&self) -> DevicePtr {
53        self.moe_output
54    }
55    pub fn logits(&self) -> DevicePtr {
56        self.logits
57    }
58    pub fn ssm_qkvz(&self) -> DevicePtr {
59        self.ssm_qkvz
60    }
61    /// 2026-09-25: Allocated byte size of `ssm_qkvz`, the QKVZ projection's
62    /// destination on an interleaved model, checked by the cuBLASLt arm against
63    /// its `ceil16(M)` rows.
64    pub fn ssm_qkvz_bytes(&self) -> usize {
65        self.sizes.ssm_qkvz
66    }
67    pub fn ssm_ba(&self) -> DevicePtr {
68        self.ssm_ba
69    }
70    /// 2026-09-25: Sequential `[Q|K|V|Z]` rows after deinterleaving.
71    pub fn ssm_deinterleaved(&self) -> DevicePtr {
72        self.ssm_deinterleaved
73    }
74    /// 2026-09-25: Allocated byte size of `ssm_deinterleaved`, the QKVZ
75    /// projection's destination on a sequential model; same padded-M check.
76    pub fn ssm_deinterleaved_bytes(&self) -> usize {
77        self.sizes.ssm_deinterleaved
78    }
79    /// 2026-09-25: GDN gates in FP32: per token, `linear_num_value_heads` gate
80    /// values, then as many beta values.
81    pub fn ssm_gates(&self) -> DevicePtr {
82        self.ssm_gates
83    }
84    /// 2026-09-25: FP32 conv1d output of the SSM layers.
85    pub fn ssm_conv_out_f32(&self) -> DevicePtr {
86        self.ssm_conv_out_f32
87    }
88    /// 2026-09-25: Scratch for MoE routing and the metadata uploads (layouts in `sizes.rs`).
89    pub fn scratch(&self) -> DevicePtr {
90        self.scratch
91    }
92    /// 2026-09-25: Mamba-2 SSD chunked-scan scratch (dt | dA_cumsum | CB); NULL without Mamba-2 layers.
93    pub fn ssd_scratch(&self) -> DevicePtr {
94        self.ssd_scratch
95    }
96    /// 2026-09-25: Token ids `[M]` u32 of the pass, read by the DeepSeek-V4
97    /// hash-routed MoE layers. The caller uploads them before the layer loop,
98    /// and under CUDA-graph decode before each replay.
99    pub fn token_ids(&self) -> DevicePtr {
100        self.token_ids
101    }
102    /// 2026-09-25: Allocated byte size of `scratch`, the bound for metadata staging.
103    pub fn scratch_bytes(&self) -> usize {
104        self.sizes.scratch
105    }
106    pub fn expert_gate_out(&self) -> DevicePtr {
107        self.expert_gate_out
108    }
109    pub fn expert_up_out(&self) -> DevicePtr {
110        self.expert_up_out
111    }
112    /// 2026-09-25: Allocated byte size of `expert_gate_out`, equal to that of
113    /// `expert_up_out` (debug-asserted).
114    pub fn expert_gate_out_bytes(&self) -> usize {
115        debug_assert_eq!(self.sizes.expert_gate_out, self.sizes.expert_up_out);
116        self.sizes.expert_gate_out
117    }
118    pub fn moe_output_bytes(&self) -> usize {
119        self.sizes.moe_output
120    }
121    pub fn expert_down_out(&self) -> DevicePtr {
122        self.expert_down_out
123    }
124    pub fn expert_down_out_bytes(&self) -> usize {
125        self.sizes.expert_down_out
126    }
127    pub fn gate_logits_bytes(&self) -> usize {
128        self.sizes.gate_logits
129    }
130    pub fn logits_bytes(&self) -> usize {
131        self.sizes.logits
132    }
133    pub fn attn_output_bytes(&self) -> usize {
134        self.sizes.attn_output
135    }
136    /// 2026-09-25: GDN FLA prefill scratch (W | U | S | uc | gc, carved by the
137    /// caller); `DevicePtr::NULL` unless both linear-attention head dims are 128.
138    pub fn gdn_fla_scratch(&self) -> DevicePtr {
139        self.gdn_fla_scratch
140    }
141    /// 2026-09-25: Shared dense-FFN q8_1 activation scratch (Q4_K MMQ); NULL for MoE.
142    pub fn ffn_act_q8(&self) -> DevicePtr {
143        self.ffn_act_q8
144    }
145    /// 2026-09-25: Shared dense-FFN int8 / NVFP4 / FP8 activation scratch; NULL for MoE.
146    pub fn ffn_act_a(&self) -> DevicePtr {
147        self.ffn_act_a
148    }
149    /// 2026-09-25: The scales paired with `ffn_act_a`; NULL for MoE.
150    pub fn ffn_act_scale(&self) -> DevicePtr {
151        self.ffn_act_scale
152    }
153    /// 2026-09-25: `[ceil16(GATEUP_FUSED_MAX_M), 2 * intermediate]` BF16 output
154    /// of the fused dense-FFN gate+up decode GEMM: gate at column 0, up at
155    /// column `intermediate`. NULL for MoE.
156    pub fn ffn_gate_up_fused(&self) -> DevicePtr {
157        self.ffn_gate_up_fused
158    }
159    /// 2026-09-25: Allocated byte size of `ffn_gate_up_fused`; the fused arm
160    /// selects itself only when its `[ceil16(m), 2 * intermediate]` output fits.
161    pub fn ffn_gate_up_fused_bytes(&self) -> usize {
162        self.sizes.ffn_gate_up_fused
163    }
164    pub fn ffn_act_a_bytes(&self) -> usize {
165        self.sizes.ffn_act_a
166    }
167    pub fn ffn_act_scale_bytes(&self) -> usize {
168        self.sizes.ffn_act_scale
169    }
170    /// 2026-09-25: `[K/128, ceil16(M)]` transpose of the dense-FFN FP8 activation
171    /// scales for the cuBLASLt block-scaled GEMM; NULL for MoE.
172    pub fn ffn_act_scale_kmajor(&self) -> DevicePtr {
173        self.ffn_act_scale_kmajor
174    }
175    pub fn ffn_act_scale_kmajor_bytes(&self) -> usize {
176        self.sizes.ffn_act_scale_kmajor
177    }
178    /// 2026-09-25: FP8 activation scratch for the prefill projections.
179    pub fn fp8_act(&self) -> DevicePtr {
180        self.fp8_act
181    }
182    pub fn fp8_act_bytes(&self) -> usize {
183        self.sizes.fp8_act
184    }
185    /// 2026-09-25: One FP32 scale per 128 elements of `fp8_act`.
186    pub fn fp8_act_scale(&self) -> DevicePtr {
187        self.fp8_act_scale
188    }
189    pub fn fp8_act_scale_bytes(&self) -> usize {
190        self.sizes.fp8_act_scale
191    }
192    /// 2026-09-25: `[K/128, ceil16(M)]` transpose of `fp8_act_scale` for the
193    /// cuBLASLt block-scaled projections.
194    pub fn fp8_act_scale_kmajor(&self) -> DevicePtr {
195        self.fp8_act_scale_kmajor
196    }
197    pub fn fp8_act_scale_kmajor_bytes(&self) -> usize {
198        self.sizes.fp8_act_scale_kmajor
199    }
200    /// 2026-09-25: BF16 dequant scratch for keep-packed Q2_0 prefill, reused by
201    /// every projection; NULL unless `METRALE_GGUF_NATIVE_Q2=1`.
202    pub fn q2_dequant_scratch(&self) -> DevicePtr {
203        self.q2_dequant_scratch
204    }
205    pub fn q2_dequant_scratch_bytes(&self) -> usize {
206        self.sizes.q2_dequant_scratch
207    }
208    /// 2026-09-25: q8_1 activation scratch for keep-packed Q2_0 MMQ prefill; NULL
209    /// unless `METRALE_GGUF_NATIVE_Q2_MMQ=1`.
210    pub fn q2_act_q8(&self) -> DevicePtr {
211        self.q2_act_q8
212    }
213    pub fn q2_act_q8_bytes(&self) -> usize {
214        self.sizes.q2_act_q8
215    }
216    pub fn splitk_workspace(&self) -> DevicePtr {
217        self.splitk_workspace
218    }
219    /// 2026-09-25: Grouped O-projection latent `[M, o_groups * o_lora_rank]` BF16.
220    pub fn o_latent(&self) -> DevicePtr {
221        self.o_latent
222    }
223    /// 2026-09-25: All-zero BF16 weight: `rms_norm` with it is a plain normalize.
224    pub fn norm_unit_w(&self) -> DevicePtr {
225        self.norm_unit_w
226    }
227    /// 2026-09-25: Hyper-connection residual streams `[M, hc_mult, hidden]` FP32.
228    pub fn hc_streams(&self) -> DevicePtr {
229        self.hc_streams
230    }
231
232    /// 2026-09-25: Low-rank hyper-connection scratch; its two layouts are
233    /// described in `sizes.rs`.
234    pub fn hc_lowrank_scratch(&self) -> DevicePtr {
235        self.hc_lowrank_scratch
236    }
237    /// 2026-09-25: QSA prefill-selection scratch shared by the indexer layers;
238    /// `layers/qsa_select.rs` carves it as `sizes.rs` sizes it.
239    pub fn qsa_select_scratch(&self) -> DevicePtr {
240        self.qsa_select_scratch
241    }
242    /// 2026-09-25: Hyper-connection `post` weights `[M, hc_mult]` F32.
243    pub fn hc_post(&self) -> DevicePtr {
244        self.hc_post
245    }
246    /// 2026-09-25: Hyper-connection `comb` matrix `[M, hc_mult, hc_mult]` F32.
247    pub fn hc_comb(&self) -> DevicePtr {
248        self.hc_comb
249    }
250    pub fn max_batch_tokens(&self) -> usize {
251        self.max_batch_tokens
252    }
253    /// 2026-09-25: The batched-decode metadata layout.
254    pub fn decode_meta(&self) -> super::DecodeMetaLayout {
255        self.decode_meta
256    }
257    pub fn sizes(&self) -> &BufferSizes {
258        &self.sizes
259    }
260
261    /// 2026-09-25: LoRA shrink scratch `xa = x @ Aᵀ`, `[M, adapter_max_rank]`
262    /// BF16; `DevicePtr::NULL` without an adapter.
263    pub fn lora_xa(&self) -> DevicePtr {
264        self.lora_xa
265    }
266    pub fn lora_xa_bytes(&self) -> usize {
267        self.sizes.lora_xa
268    }
269    /// 2026-09-25: LoRA expand scratch `delta = xa @ Bᵀ` BF16; `DevicePtr::NULL`
270    /// without an adapter.
271    pub fn lora_delta(&self) -> DevicePtr {
272        self.lora_delta
273    }
274    pub fn lora_delta_bytes(&self) -> usize {
275        self.sizes.lora_delta
276    }
277    /// 2026-09-25: LoRA hidden-activation scratch `[M, intermediate_size]` BF16;
278    /// `DevicePtr::NULL` without an adapter.
279    pub fn lora_hact(&self) -> DevicePtr {
280        self.lora_hact
281    }
282    pub fn lora_hact_bytes(&self) -> usize {
283        self.sizes.lora_hact
284    }
285    /// 2026-09-25: LoRA adapter slot per prefill token, `[M]` i32;
286    /// `DevicePtr::NULL` without an adapter.
287    pub fn lora_seq_slot(&self) -> DevicePtr {
288        self.lora_seq_slot
289    }
290
291    /// 2026-09-25: Zero `hidden_states`, `residual`, `gate_logits`, the three
292    /// expert outputs and `moe_output`: seven memsets, where `zero_all` issues
293    /// eighteen. The prefill paths call this on a sequence's first chunk when
294    /// they do not call `zero_all`.
295    pub fn zero_prefill_essentials(&self, gpu: &dyn GpuBackend, stream: u64) -> anyhow::Result<()> {
296        gpu.memset_async(self.hidden_states, 0, self.sizes.hidden_states, stream)?;
297        gpu.memset_async(self.residual, 0, self.sizes.residual, stream)?;
298        gpu.memset_async(self.gate_logits, 0, self.sizes.gate_logits, stream)?;
299        gpu.memset_async(self.expert_gate_out, 0, self.sizes.expert_gate_out, stream)?;
300        gpu.memset_async(self.expert_up_out, 0, self.sizes.expert_up_out, stream)?;
301        gpu.memset_async(self.expert_down_out, 0, self.sizes.expert_down_out, stream)?;
302        gpu.memset_async(self.moe_output, 0, self.sizes.moe_output, stream)?;
303        Ok(())
304    }
305
306    /// 2026-09-25: `zero_all` limited to the first `tokens` rows of each
307    /// token-major buffer: for a buffer whose size is a multiple of
308    /// `max_batch_tokens`, `size / max_batch_tokens * tokens` bytes; any other
309    /// buffer, and every buffer when `tokens >= max_batch_tokens`, is zeroed
310    /// whole. `splitk_workspace`, `logits` and `scratch` are zeroed whole. The
311    /// caller must read only rows `0..tokens`. Measured 2026-08-28 on
312    /// GLM-5.3-Flash (nsys, `max_batch_tokens = 4096`): `zero_all` took 8.01 ms
313    /// of an 85 ms decode step.
314    pub fn zero_all_rows(
315        &self,
316        gpu: &dyn GpuBackend,
317        stream: u64,
318        tokens: usize,
319    ) -> anyhow::Result<()> {
320        let m = self.max_batch_tokens.max(1);
321        let head = |n: usize| {
322            if tokens >= m || m == 0 || !n.is_multiple_of(m) {
323                n
324            } else {
325                n / m * tokens
326            }
327        };
328        for (ptr, n) in [
329            (self.hidden_states, self.sizes.hidden_states),
330            (self.residual, self.sizes.residual),
331            (self.norm_output, self.sizes.norm_output),
332            (self.qkv_output, self.sizes.qkv_output),
333            (self.attn_output, self.sizes.attn_output),
334            (self.gate_logits, self.sizes.gate_logits),
335            (self.moe_output, self.sizes.moe_output),
336            (self.ssm_qkvz, self.sizes.ssm_qkvz),
337            (self.ssm_ba, self.sizes.ssm_ba),
338            (self.ssm_deinterleaved, self.sizes.ssm_deinterleaved),
339            (self.ssm_gates, self.sizes.ssm_gates),
340            (self.ssm_conv_out_f32, self.sizes.ssm_conv_out_f32),
341            (self.expert_gate_out, self.sizes.expert_gate_out),
342            (self.expert_up_out, self.sizes.expert_up_out),
343            (self.expert_down_out, self.sizes.expert_down_out),
344        ] {
345            gpu.memset_async(ptr, 0, head(n), stream)?;
346        }
347        gpu.memset_async(
348            self.splitk_workspace,
349            0,
350            self.sizes.splitk_workspace,
351            stream,
352        )?;
353        gpu.memset_async(self.logits, 0, self.sizes.logits, stream)?;
354        gpu.memset_async(self.scratch, 0, self.sizes.scratch, stream)?;
355        Ok(())
356    }
357
358    /// 2026-09-25: Zero, in full, the eighteen buffers `zero_all_rows` covers.
359    /// The others (`gate_logits_f32`, `moe_router_in_f32`, `o_latent`,
360    /// `norm_unit_w`, the hyper-connection, QSA, FLA and SSD scratch,
361    /// `token_ids`, `ffn_*`, `fp8_*`, `q2_*`, `lora_*` and the row-wise slab)
362    /// keep their contents.
363    pub fn zero_all(&self, gpu: &dyn GpuBackend, stream: u64) -> anyhow::Result<()> {
364        gpu.memset_async(self.hidden_states, 0, self.sizes.hidden_states, stream)?;
365        gpu.memset_async(self.residual, 0, self.sizes.residual, stream)?;
366        gpu.memset_async(self.norm_output, 0, self.sizes.norm_output, stream)?;
367        gpu.memset_async(self.qkv_output, 0, self.sizes.qkv_output, stream)?;
368        gpu.memset_async(self.attn_output, 0, self.sizes.attn_output, stream)?;
369        gpu.memset_async(self.gate_logits, 0, self.sizes.gate_logits, stream)?;
370        gpu.memset_async(self.moe_output, 0, self.sizes.moe_output, stream)?;
371        gpu.memset_async(self.ssm_qkvz, 0, self.sizes.ssm_qkvz, stream)?;
372        gpu.memset_async(self.ssm_ba, 0, self.sizes.ssm_ba, stream)?;
373        gpu.memset_async(
374            self.ssm_deinterleaved,
375            0,
376            self.sizes.ssm_deinterleaved,
377            stream,
378        )?;
379        gpu.memset_async(self.ssm_gates, 0, self.sizes.ssm_gates, stream)?;
380        gpu.memset_async(
381            self.ssm_conv_out_f32,
382            0,
383            self.sizes.ssm_conv_out_f32,
384            stream,
385        )?;
386        gpu.memset_async(
387            self.splitk_workspace,
388            0,
389            self.sizes.splitk_workspace,
390            stream,
391        )?;
392        gpu.memset_async(self.expert_gate_out, 0, self.sizes.expert_gate_out, stream)?;
393        gpu.memset_async(self.expert_up_out, 0, self.sizes.expert_up_out, stream)?;
394        gpu.memset_async(self.expert_down_out, 0, self.sizes.expert_down_out, stream)?;
395        gpu.memset_async(self.logits, 0, self.sizes.logits, stream)?;
396        gpu.memset_async(self.scratch, 0, self.sizes.scratch, stream)?;
397        Ok(())
398    }
399}
metrale_gpu_runtime/buffers/accessors.rs

metrale_gpu_runtime/buffers/
accessors.rs
