spark_runtime/buffers/
accessors.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! `BufferArena` accessors. Split from `buffers.rs` (500-LoC cap).
4
5use super::{BufferArena, sizes::BufferSizes};
6use crate::gpu::{DevicePtr, GpuBackend};
7
8impl BufferArena {
9    pub fn hidden_states(&self) -> DevicePtr {
10        self.hidden_states
11    }
12    pub fn residual(&self) -> DevicePtr {
13        self.residual
14    }
15    pub fn norm_output(&self) -> DevicePtr {
16        self.norm_output
17    }
18    pub fn qkv_output(&self) -> DevicePtr {
19        self.qkv_output
20    }
21    pub fn attn_output(&self) -> DevicePtr {
22        self.attn_output
23    }
24    pub fn gate_logits(&self) -> DevicePtr {
25        self.gate_logits
26    }
27    pub fn gate_logits_f32(&self) -> DevicePtr {
28        self.gate_logits_f32
29    }
30    pub fn moe_router_in_f32(&self) -> DevicePtr {
31        self.moe_router_in_f32
32    }
33    pub fn moe_output(&self) -> DevicePtr {
34        self.moe_output
35    }
36    pub fn logits(&self) -> DevicePtr {
37        self.logits
38    }
39    pub fn ssm_qkvz(&self) -> DevicePtr {
40        self.ssm_qkvz
41    }
42    pub fn ssm_ba(&self) -> DevicePtr {
43        self.ssm_ba
44    }
45    /// Sequential [Q|K|V|Z] after deinterleaving.
46    pub fn ssm_deinterleaved(&self) -> DevicePtr {
47        self.ssm_deinterleaved
48    }
49    /// FP32 [gate, beta] for GDN (num_v_heads * 2 floats).
50    pub fn ssm_gates(&self) -> DevicePtr {
51        self.ssm_gates
52    }
53    /// FP32 conv1d output for SSM recurrent path (prevents BF16 precision drift).
54    pub fn ssm_conv_out_f32(&self) -> DevicePtr {
55        self.ssm_conv_out_f32
56    }
57    /// Scratch buffer for MoE routing + kernel metadata uploads.
58    pub fn scratch(&self) -> DevicePtr {
59        self.scratch
60    }
61    /// Mamba-2 SSD chunked-scan scratch (dt | dA_cumsum | CB). NULL if unused.
62    pub fn ssd_scratch(&self) -> DevicePtr {
63        self.ssd_scratch
64    }
65    /// Token IDs `[M]` u32 — stable across the layer loop (DeepSeek-V4 hash-MoE
66    /// reads `tid2eid[token_id]`). Upload the pass's token IDs here before the
67    /// layer loop; under CUDA-graph decode upload before each replay.
68    pub fn token_ids(&self) -> DevicePtr {
69        self.token_ids
70    }
71    /// Allocated byte size of the scratch buffer (#110: bounds-check
72    /// batched metadata-staging uploads against this).
73    pub fn scratch_bytes(&self) -> usize {
74        self.sizes.scratch
75    }
76    /// Batched expert gate projection output.
77    pub fn expert_gate_out(&self) -> DevicePtr {
78        self.expert_gate_out
79    }
80    /// Batched expert up projection output.
81    pub fn expert_up_out(&self) -> DevicePtr {
82        self.expert_up_out
83    }
84    /// Batched expert down projection output.
85    pub fn expert_down_out(&self) -> DevicePtr {
86        self.expert_down_out
87    }
88    /// Split-K decode attention workspace (F32 partials).
89    /// GDN FLA chunked-prefill scratch base (W|U|S|uc sub-divided by the caller).
90    /// `DevicePtr::NULL` unless this is a 128-dim-linear-head GDN model.
91    pub fn gdn_fla_scratch(&self) -> DevicePtr {
92        self.gdn_fla_scratch
93    }
94    /// Shared dense-FFN q8_1 activation scratch (Q4_K MMQ gate/up). NULL for MoE.
95    pub fn ffn_act_q8(&self) -> DevicePtr {
96        self.ffn_act_q8
97    }
98    /// Shared dense-FFN int8/NVFP4 activation scratch (a_i8 / packed). NULL for MoE.
99    pub fn ffn_act_a(&self) -> DevicePtr {
100        self.ffn_act_a
101    }
102    /// Shared dense-FFN int8/NVFP4 activation-scale scratch. NULL for MoE.
103    pub fn ffn_act_scale(&self) -> DevicePtr {
104        self.ffn_act_scale
105    }
106    /// Persistent FP8 block-scaled activation scratch for prefill projections.
107    /// Replaces a per-projection alloc/sync/free in the W8A8+FP32-epilogue path.
108    pub fn fp8_act(&self) -> DevicePtr {
109        self.fp8_act
110    }
111    /// Allocated byte size of `fp8_act` (debug bounds-check at call sites).
112    pub fn fp8_act_bytes(&self) -> usize {
113        self.sizes.fp8_act
114    }
115    /// Persistent per-128-block FP32 scales paired with `fp8_act`.
116    pub fn fp8_act_scale(&self) -> DevicePtr {
117        self.fp8_act_scale
118    }
119    /// Persistent BF16 transient-dequant scratch for native keep-packed Q2_0
120    /// prefill. Reused per projection: dequant into it, GEMM reads it (same
121    /// stream), no free. NULL unless `ATLAS_GGUF_NATIVE_Q2`.
122    pub fn q2_dequant_scratch(&self) -> DevicePtr {
123        self.q2_dequant_scratch
124    }
125    /// Allocated byte size of `q2_dequant_scratch` (debug bounds-check).
126    pub fn q2_dequant_scratch_bytes(&self) -> usize {
127        self.sizes.q2_dequant_scratch
128    }
129    /// Persistent q8_1 activation scratch for native Q2_0 MMQ prefill
130    /// (`ATLAS_GGUF_NATIVE_Q2_MMQ`). NULL unless the flag is set.
131    pub fn q2_act_q8(&self) -> DevicePtr {
132        self.q2_act_q8
133    }
134    /// Allocated byte size of `q2_act_q8` (debug bounds-check).
135    pub fn q2_act_q8_bytes(&self) -> usize {
136        self.sizes.q2_act_q8
137    }
138    pub fn splitk_workspace(&self) -> DevicePtr {
139        self.splitk_workspace
140    }
141    /// Grouped O-projection latent [M, o_groups*o_lora_rank] BF16 (V4-Flash).
142    pub fn o_latent(&self) -> DevicePtr {
143        self.o_latent
144    }
145    /// All-ones BF16 vector (max_dim) — weight for unweighted RMSNorm (q_b_norm).
146    pub fn norm_unit_w(&self) -> DevicePtr {
147        self.norm_unit_w
148    }
149    /// HC residual streams [M, hc_mult, hidden] BF16 (DeepSeek-V4 mHC).
150    pub fn hc_streams(&self) -> DevicePtr {
151        self.hc_streams
152    }
153
154    /// Low-rank mHC split-collapse scratch: `[T<=64, hc*H]` normed followed
155    /// by `[T<=64, rank]` low, both F32. See `sizes.rs`.
156    pub fn hc_lowrank_scratch(&self) -> DevicePtr {
157        self.hc_lowrank_scratch
158    }
159    /// QSA stage-2 prefill-selection scratch, shared by the indexer layers
160    /// (serial). Layout managed by `layers::qsa`; see `sizes.rs`.
161    pub fn qsa_select_scratch(&self) -> DevicePtr {
162        self.qsa_select_scratch
163    }
164    /// HC `post` mixing weights [M, hc_mult] F32.
165    pub fn hc_post(&self) -> DevicePtr {
166        self.hc_post
167    }
168    /// HC `comb` Sinkhorn matrix [M, hc_mult, hc_mult] F32.
169    pub fn hc_comb(&self) -> DevicePtr {
170        self.hc_comb
171    }
172    pub fn max_batch_tokens(&self) -> usize {
173        self.max_batch_tokens
174    }
175    /// Derived batched-decode metadata layout (rows/offsets). Byte-identical
176    /// to the legacy fixed 32-row layout for every serve `max_batch_size <= 32`.
177    pub fn decode_meta(&self) -> super::DecodeMetaLayout {
178        self.decode_meta
179    }
180    pub fn sizes(&self) -> &BufferSizes {
181        &self.sizes
182    }
183
184    /// Env-gated (`ATLAS_SSM_SAVE_DUMP`) per-buffer checksum probe.
185    ///
186    /// CBD: localize a stale/uninitialized decode-scratch buffer on the
187    /// prefix-cache skip path. Dumps sum/ssq/sabs over the FULL allocation
188    /// (so leftover-from-prior-occupant bytes in unwritten rows are visible)
189    /// for every reusable buffer. Treats raw bytes as f32 lanes — exact
190    /// numeric meaning is irrelevant; we only need a stable fingerprint that
191    /// differs iff the bytes differ. Synchronizes the stream first.
192    /// LoRA compressed activation scratch `xa = x@Aᵀ` [M, max_rank] BF16.
193    /// `DevicePtr::NULL` when no adapter is configured.
194    pub fn lora_xa(&self) -> DevicePtr {
195        self.lora_xa
196    }
197    /// Allocated byte size of `lora_xa` (0 when no adapter).
198    pub fn lora_xa_bytes(&self) -> usize {
199        self.sizes.lora_xa
200    }
201    /// LoRA expand scratch `delta = xa@Bᵀ` [M, max(hidden, intermediate)]
202    /// BF16. `DevicePtr::NULL` when no adapter is configured.
203    pub fn lora_delta(&self) -> DevicePtr {
204        self.lora_delta
205    }
206    /// Allocated byte size of `lora_delta` (0 when no adapter).
207    pub fn lora_delta_bytes(&self) -> usize {
208        self.sizes.lora_delta
209    }
210    /// LoRA hidden-activation scratch [M, intermediate_size] BF16 for the
211    /// runtime FFN delta path. `DevicePtr::NULL` when no adapter.
212    pub fn lora_hact(&self) -> DevicePtr {
213        self.lora_hact
214    }
215    /// Allocated byte size of `lora_hact` (0 when no adapter).
216    pub fn lora_hact_bytes(&self) -> usize {
217        self.sizes.lora_hact
218    }
219    /// LoRA per-request routing slots `[max_batch_tokens]` i32 for the prefill
220    /// path — one adapter SLOT index per prefilling token. `DevicePtr::NULL`
221    /// when no adapter is configured.
222    pub fn lora_seq_slot(&self) -> DevicePtr {
223        self.lora_seq_slot
224    }
225
226    pub fn debug_buffer_checksum(&self, gpu: &dyn GpuBackend, stream: u64, tag: &str) {
227        gpu.synchronize(stream).ok();
228        let probe = |name: &str, ptr: DevicePtr, bytes: usize| {
229            let mut hb = vec![0u8; bytes];
230            if gpu.copy_d2h(ptr, &mut hb).is_err() {
231                return;
232            }
233            let (mut sum, mut ssq, mut sabs) = (0f64, 0f64, 0f64);
234            for c in hb.chunks_exact(4) {
235                let v = f32::from_le_bytes([c[0], c[1], c[2], c[3]]) as f64;
236                if v.is_finite() {
237                    sum += v;
238                    ssq += v * v;
239                    sabs += v.abs();
240                }
241            }
242            tracing::warn!(
243                "ATLAS_BUF_CKSUM[{tag}] {name} bytes={bytes} sum={sum:.6} ssq={ssq:.6} sabs={sabs:.6}"
244            );
245        };
246        probe(
247            "hidden_states",
248            self.hidden_states,
249            self.sizes.hidden_states,
250        );
251        probe("residual", self.residual, self.sizes.residual);
252        probe("norm_output", self.norm_output, self.sizes.norm_output);
253        probe("qkv_output", self.qkv_output, self.sizes.qkv_output);
254        probe("attn_output", self.attn_output, self.sizes.attn_output);
255        probe("gate_logits", self.gate_logits, self.sizes.gate_logits);
256        probe("moe_output", self.moe_output, self.sizes.moe_output);
257        probe("ssm_qkvz", self.ssm_qkvz, self.sizes.ssm_qkvz);
258        probe("ssm_ba", self.ssm_ba, self.sizes.ssm_ba);
259        probe(
260            "ssm_deinterleaved",
261            self.ssm_deinterleaved,
262            self.sizes.ssm_deinterleaved,
263        );
264        probe("ssm_gates", self.ssm_gates, self.sizes.ssm_gates);
265        probe(
266            "ssm_conv_out_f32",
267            self.ssm_conv_out_f32,
268            self.sizes.ssm_conv_out_f32,
269        );
270        probe(
271            "expert_gate_out",
272            self.expert_gate_out,
273            self.sizes.expert_gate_out,
274        );
275        probe(
276            "expert_up_out",
277            self.expert_up_out,
278            self.sizes.expert_up_out,
279        );
280        probe(
281            "expert_down_out",
282            self.expert_down_out,
283            self.sizes.expert_down_out,
284        );
285        probe(
286            "splitk_workspace",
287            self.splitk_workspace,
288            self.sizes.splitk_workspace,
289        );
290    }
291
292    /// Zero only buffers that carry residual state between requests.
293    ///
294    /// During prefill, every buffer except hidden_states and residual is fully
295    /// overwritten before being read within the layer loop:
296    /// - norm_output, qkv_output, attn_output: written by each layer's projection
297    /// - gate_logits, moe_output: written by MoE gate/output
298    /// - ssm_*: written by SSM projection
299    /// - expert_*: written by expert compute
300    /// - logits: written by LM head on last token
301    /// - scratch: overwritten by metadata upload and MoE routing
302    /// - splitk_workspace: written by attention kernel
303    ///
304    /// This reduces per-chunk memset from 17 calls to 2, saving ~15 memset
305    /// launches × bandwidth on the LPDDR5X bus per prefill chunk.
306    pub fn zero_prefill_essentials(&self, gpu: &dyn GpuBackend, stream: u64) -> anyhow::Result<()> {
307        gpu.memset_async(self.hidden_states, 0, self.sizes.hidden_states, stream)?;
308        gpu.memset_async(self.residual, 0, self.sizes.residual, stream)?;
309        // MoE buffers: gate_logits may carry stale expert indices from a prior
310        // request with different token count, causing out-of-bounds expert access
311        // (CUDA error 700 at layer 38+ on 122B). Zero to prevent.
312        gpu.memset_async(self.gate_logits, 0, self.sizes.gate_logits, stream)?;
313        gpu.memset_async(self.expert_gate_out, 0, self.sizes.expert_gate_out, stream)?;
314        gpu.memset_async(self.expert_up_out, 0, self.sizes.expert_up_out, stream)?;
315        gpu.memset_async(self.expert_down_out, 0, self.sizes.expert_down_out, stream)?;
316        gpu.memset_async(self.moe_output, 0, self.sizes.moe_output, stream)?;
317        Ok(())
318    }
319
320    /// `zero_all`, but only the first `tokens` rows of every token-major arena.
321    ///
322    /// 🔴 Every buffer `zero_all` wipes is `[max_batch_tokens, row]`-major — verified against
323    /// the allocated sizes: `size / max_batch_tokens` is exactly one token's row for each of
324    /// them (`qkv_output` 3x8192 BF16, `attn_output` 64x256, `expert_gate_out` topk*2048,
325    /// `expert_down_out` topk*4096, ...). A decode step carrying `tokens` tokens can therefore
326    /// only ever read rows `0..tokens`, and zeroing the rest is dead bandwidth.
327    ///
328    /// Measured on GLM-5.3-Flash, 2 x GB10, `max_batch_tokens = 4096` (nsys, 2026-08-28):
329    /// `zero_all` issues 18 memsets totalling **1.59 GB and 8.01 ms on every single decode
330    /// token** — 9.4 % of an 85 ms step, all of it GPU-idle time before the first kernel.
331    ///
332    /// `logits`, `scratch` and `splitk_workspace` are NOT token-major (metadata arenas /
333    /// vocab-sized), so they keep the full wipe. They are 30 MB of the 1590.
334    pub fn zero_all_rows(
335        &self,
336        gpu: &dyn GpuBackend,
337        stream: u64,
338        tokens: usize,
339    ) -> anyhow::Result<()> {
340        let m = self.max_batch_tokens.max(1);
341        // A row-scaled length, falling back to the full wipe if the arena is not an exact
342        // multiple of `max_batch_tokens` (i.e. not token-major after all).
343        let head = |n: usize| {
344            if tokens >= m || m == 0 || !n.is_multiple_of(m) {
345                n
346            } else {
347                n / m * tokens
348            }
349        };
350        for (ptr, n) in [
351            (self.hidden_states, self.sizes.hidden_states),
352            (self.residual, self.sizes.residual),
353            (self.norm_output, self.sizes.norm_output),
354            (self.qkv_output, self.sizes.qkv_output),
355            (self.attn_output, self.sizes.attn_output),
356            (self.gate_logits, self.sizes.gate_logits),
357            (self.moe_output, self.sizes.moe_output),
358            (self.ssm_qkvz, self.sizes.ssm_qkvz),
359            (self.ssm_ba, self.sizes.ssm_ba),
360            (self.ssm_deinterleaved, self.sizes.ssm_deinterleaved),
361            (self.ssm_gates, self.sizes.ssm_gates),
362            (self.ssm_conv_out_f32, self.sizes.ssm_conv_out_f32),
363            (self.expert_gate_out, self.sizes.expert_gate_out),
364            (self.expert_up_out, self.sizes.expert_up_out),
365            (self.expert_down_out, self.sizes.expert_down_out),
366        ] {
367            gpu.memset_async(ptr, 0, head(n), stream)?;
368        }
369        // Not token-major — full wipe, 30 MB of the 1590.
370        gpu.memset_async(
371            self.splitk_workspace,
372            0,
373            self.sizes.splitk_workspace,
374            stream,
375        )?;
376        gpu.memset_async(self.logits, 0, self.sizes.logits, stream)?;
377        gpu.memset_async(self.scratch, 0, self.sizes.scratch, stream)?;
378        Ok(())
379    }
380
381    /// Zero all reusable buffers to eliminate stale data between requests.
382    /// Ensures deterministic computation regardless of request history.
383    pub fn zero_all(&self, gpu: &dyn GpuBackend, stream: u64) -> anyhow::Result<()> {
384        gpu.memset_async(self.hidden_states, 0, self.sizes.hidden_states, stream)?;
385        gpu.memset_async(self.residual, 0, self.sizes.residual, stream)?;
386        gpu.memset_async(self.norm_output, 0, self.sizes.norm_output, stream)?;
387        gpu.memset_async(self.qkv_output, 0, self.sizes.qkv_output, stream)?;
388        gpu.memset_async(self.attn_output, 0, self.sizes.attn_output, stream)?;
389        gpu.memset_async(self.gate_logits, 0, self.sizes.gate_logits, stream)?;
390        gpu.memset_async(self.moe_output, 0, self.sizes.moe_output, stream)?;
391        gpu.memset_async(self.ssm_qkvz, 0, self.sizes.ssm_qkvz, stream)?;
392        gpu.memset_async(self.ssm_ba, 0, self.sizes.ssm_ba, stream)?;
393        gpu.memset_async(
394            self.ssm_deinterleaved,
395            0,
396            self.sizes.ssm_deinterleaved,
397            stream,
398        )?;
399        gpu.memset_async(self.ssm_gates, 0, self.sizes.ssm_gates, stream)?;
400        gpu.memset_async(
401            self.ssm_conv_out_f32,
402            0,
403            self.sizes.ssm_conv_out_f32,
404            stream,
405        )?;
406        gpu.memset_async(
407            self.splitk_workspace,
408            0,
409            self.sizes.splitk_workspace,
410            stream,
411        )?;
412        gpu.memset_async(self.expert_gate_out, 0, self.sizes.expert_gate_out, stream)?;
413        gpu.memset_async(self.expert_up_out, 0, self.sizes.expert_up_out, stream)?;
414        gpu.memset_async(self.expert_down_out, 0, self.sizes.expert_down_out, stream)?;
415        gpu.memset_async(self.logits, 0, self.sizes.logits, stream)?;
416        gpu.memset_async(self.scratch, 0, self.sizes.scratch, stream)?;
417        Ok(())
418    }
419}
spark_runtime/buffers/accessors.rs

spark_runtime/buffers/
accessors.rs
