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