spark_runtime/
buffers.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Pre-allocated GPU buffer arena for intermediate tensors.
4//!
5//! All buffer sizes derive from [`ModelConfig`] (SSOT). The arena is
6//! allocated once during initialization and reused across decode steps.
7
8use crate::gpu::{DevicePtr, GpuBackend};
9use anyhow::Result;
10use atlas_core::config::ModelConfig;
11
12mod accessors;
13pub mod decode_meta;
14mod rowwise_slab;
15mod sizes;
16mod sizes_q12;
17mod sizes_q2;
18mod sizes_rowwise;
19pub use decode_meta::{DECODE_META_MAX_ROWS, DECODE_META_MIN_ROWS, DecodeMetaLayout};
20pub use sizes::{BufferSizes, GATEUP_FUSED_MAX_M};
21pub use sizes_q2::q2_dequant_scratch_bytes;
22pub use sizes_q12::{
23    Q12_SIZING_STREAMS, q12_batched_scratch_bytes, q12_batched_scratch_bytes_varlen,
24};
25pub use sizes_rowwise::{
26    ssm_rowwise_w_bf16_bytes, ssm_rowwise_w_bf16_bytes_for, ssm_rowwise_w_bf16_layer_bytes,
27};
28
29/// Pre-allocated GPU buffers for a single forward pass.
30///
31/// Each buffer is sized for `max_batch_tokens` tokens through the model.
32/// Buffers are reused across steps — no per-step allocation.
33///
34/// Expert output buffers are sized for max(k_max, max_batch_tokens) to
35/// support both speculative decode (K=3) and batched MoE prefill. At N=512,
36/// this adds ~31 MB (vs the old grouped-GEMM approach that needed 260 MB
37/// and caused a 15% decode regression). The GEMV-based prefill kernels
38/// only touch k_max slots during decode, so the extra pages don't affect
39/// decode bandwidth on unified memory.
40pub struct BufferArena {
41    /// Hidden states: [M, hidden_size] in BF16.
42    hidden_states: DevicePtr,
43    /// Residual stream: [M, hidden_size] in BF16.
44    residual: DevicePtr,
45    /// Post-norm output: [M, hidden_size] in BF16.
46    norm_output: DevicePtr,
47    /// QKV projection output for full attention: [M, (Hq + 2*Hkv) * D] in BF16.
48    qkv_output: DevicePtr,
49    /// Attention output: [M, Hq * D] in BF16.
50    attn_output: DevicePtr,
51    /// MoE gate logits: [M, num_experts] in BF16.
52    gate_logits: DevicePtr,
53    /// MoE gate logits: [M, num_experts] in FP32 (ATLAS_FP32_GATE path).
54    gate_logits_f32: DevicePtr,
55    /// MoE-input norm output: [M, hidden_size] in FP32 (ATLAS_FP32_ROUTING).
56    moe_router_in_f32: DevicePtr,
57    /// MoE output: [M, hidden_size] in BF16.
58    moe_output: DevicePtr,
59    /// Logits: [M, vocab_size] in BF16.
60    logits: DevicePtr,
61    /// SSM QKVZ projection: [M, ssm_qkvz_size] in BF16.
62    ssm_qkvz: DevicePtr,
63    /// SSM beta-alpha projection: [M, ssm_ba_size] in BF16.
64    ssm_ba: DevicePtr,
65    /// SSM deinterleaved QKVZ: [M, ssm_qkvz_size] in BF16 (sequential [Q|K|V|Z]).
66    ssm_deinterleaved: DevicePtr,
67    /// SSM FP32 gates: [num_v_heads * 2] as FP32 (gate + beta for GDN).
68    ssm_gates: DevicePtr,
69    /// SSM conv1d output in FP32: [M, conv_dim] as FP32.
70    /// Prevents BF16 truncation in the SSM recurrent path (conv → GDN).
71    /// Without this, ~7 bits of precision are lost every token, causing
72    /// coherence degradation after 8k+ tokens.
73    ssm_conv_out_f32: DevicePtr,
74    /// Scratch space for kernel metadata (positions, slot_mapping, block_tables).
75    scratch: DevicePtr,
76    /// Expert gate projection output: [k2 * top_k, moe_intermediate_size] BF16.
77    expert_gate_out: DevicePtr,
78    /// Expert up projection output: [k2 * top_k, moe_intermediate_size] BF16.
79    expert_up_out: DevicePtr,
80    /// Expert down projection output: [k2 * top_k, hidden_size] BF16.
81    expert_down_out: DevicePtr,
82    /// Split-K decode attention workspace: partials from split CTAs (F32).
83    splitk_workspace: DevicePtr,
84    /// Grouped O-projection latent: [M, o_groups*o_lora_rank] BF16 (V4-Flash).
85    o_latent: DevicePtr,
86    /// Zero-filled BF16 weight (max_dim) for unweighted RMSNorm under the
87    /// offset-from-1 kernel convention (scale = 1+weight → 1.0). Used by q_b_norm.
88    norm_unit_w: DevicePtr,
89    /// HC residual streams: [M, hc_mult, hidden] BF16 (DeepSeek-V4 mHC).
90    hc_streams: DevicePtr,
91    /// HC `post` mixing weights: [M, hc_mult] F32.
92    hc_post: DevicePtr,
93    /// HC `comb` Sinkhorn matrix: [M, hc_mult, hc_mult] F32.
94    hc_comb: DevicePtr,
95    hc_lowrank_scratch: DevicePtr,
96    qsa_select_scratch: DevicePtr,
97    /// GDN FLA chunked-prefill scratch (W|U|S|uc sub-divided). NULL unless the
98    /// model is a 128-dim-linear-head GDN model (ATLAS_GDN_FLA path).
99    gdn_fla_scratch: DevicePtr,
100    /// Mamba-2 SSD chunked-scan scratch (dt | dA_cumsum | CB). NULL unless the model
101    /// has Mamba-2 SSM layers.
102    ssd_scratch: DevicePtr,
103    /// Token IDs `[M]` u32 — stable across the layer loop so DeepSeek-V4
104    /// hash-MoE layers can read `tid2eid[token_id]`.
105    token_ids: DevicePtr,
106    /// Shared FFN activation-quant scratch (dense-FFN MMQ/int8 prefill path).
107    /// Allocated once here instead of per-DenseFfnLayer (64× would leak ~18GB).
108    /// NULL unless the model is dense (`num_experts == 0`).
109    /// `ffn_act_q8`: q8_1 activations for the Q4_K MMQ gate/up GEMM.
110    /// `ffn_act_a` / `ffn_act_scale`: int8 (a_i8 / a_scale) — reused for NVFP4 packed/scale.
111    ffn_act_q8: DevicePtr,
112    ffn_act_a: DevicePtr,
113    ffn_act_scale: DevicePtr,
114    /// `[K/128, ceil16(M)]` transposed copy of `ffn_act_scale` — the VEC128
115    /// B-scale layout cuBLASLt documents (token index contiguous). NULL for MoE.
116    ffn_act_scale_kmajor: DevicePtr,
117    /// `[ceil16(GATEUP_FUSED_MAX_M), 2 * intermediate]` BF16 output of the
118    /// FUSED dense-FFN gate+up decode GEMM (#927). NULL for MoE.
119    ffn_gate_up_fused: DevicePtr,
120    /// Persistent FP8 block-scaled activation scratch for prefill projections.
121    fp8_act: DevicePtr,
122    /// Persistent per-128-block FP32 scales paired with `fp8_act`.
123    fp8_act_scale: DevicePtr,
124    /// `[K/128, ceil16(M)]` transpose of `fp8_act_scale` (cuBLASLt VEC128).
125    fp8_act_scale_kmajor: DevicePtr,
126    /// Persistent BF16 transient-dequant scratch for native keep-packed Q2_0
127    /// prefill. Reused per projection — replaces a per-matmul alloc/sync/free.
128    q2_dequant_scratch: DevicePtr,
129    /// LoRA shrink scratch `xa = x@Aᵀ`: [M, adapter_max_rank] BF16.
130    /// NULL when no adapter is configured.
131    lora_xa: DevicePtr,
132    /// LoRA expand scratch `delta = xa@Bᵀ`: [M, max(hidden, intermediate)]
133    /// BF16. NULL when no adapter is configured.
134    lora_delta: DevicePtr,
135    /// LoRA hidden-activation scratch: [M, intermediate_size] BF16 for the
136    /// runtime FFN delta path. NULL when no adapter is configured.
137    lora_hact: DevicePtr,
138    /// LoRA per-request routing slots `[M]` i32 for the prefill path (one
139    /// adapter SLOT index per prefilling token). NULL when no adapter.
140    lora_seq_slot: DevicePtr,
141    /// Persistent q8_1_mmq activation scratch for native Q2_0 MMQ prefill
142    /// (`ATLAS_GGUF_NATIVE_Q2_MMQ`). Shared by every kept-packed projection;
143    /// each seam quantizes its activation here then runs the packed MMQ GEMM.
144    q2_act_q8: DevicePtr,
145    /// Row-wise FP8 GDN prefill BF16-weight slab (`ATLAS_FP8_ROWWISE`). One
146    /// allocation for EVERY GDN layer's dequanted `in_proj_qkvz` + `out_proj`;
147    /// `take_ssm_rowwise_w_bf16` bump-carves a layer's slice on its first
148    /// prefill. NULL unless the lever is armed. See `sizes_rowwise.rs` for the
149    /// #917 receipt this replaces.
150    ssm_rowwise_w_bf16: DevicePtr,
151    /// Bytes already carved out of `ssm_rowwise_w_bf16`. Bump-only: a slice
152    /// lives as long as the arena does, exactly like the weight it holds.
153    ssm_rowwise_w_bf16_used: std::sync::atomic::AtomicUsize,
154    /// Maximum batch tokens this arena was sized for.
155    max_batch_tokens: usize,
156    /// Derived batched-decode metadata layout (rows = max(32, serve
157    /// max_batch_size)); byte-identical to the legacy fixed 32-row gaps for
158    /// every bs <= 32. SSOT consumed by `upload_batch_metadata_fixed`/`_at`.
159    decode_meta: DecodeMetaLayout,
160    /// Sizes in bytes for each buffer (for debug/logging).
161    sizes: BufferSizes,
162}
163
164impl BufferArena {
165    /// Allocate all intermediate buffers on the GPU.
166    pub fn new(
167        config: &ModelConfig,
168        max_batch_tokens: usize,
169        max_seq_len: usize,
170        kv_block_size: usize,
171        max_batch_size: usize,
172        gpu: &dyn GpuBackend,
173    ) -> Result<Self> {
174        let sizes = BufferSizes::from_config(
175            config,
176            max_batch_tokens,
177            max_seq_len,
178            kv_block_size,
179            max_batch_size,
180        );
181        Self::from_sizes(config, sizes, max_batch_tokens, max_batch_size, gpu)
182    }
183
184    /// [`BufferArena::new`] with the ledger handed in instead of derived.
185    ///
186    /// `BufferSizes::from_config` reads the process environment for the
187    /// env-gated entries (`q2_*`, `ssm_rowwise_w_bf16`), and `set_var` is
188    /// process-global and unsafe — so a test that wants one of those arms
189    /// ARMED builds the sizes by hand and comes in here rather than racing
190    /// every other test in the binary. Production still goes through `new`.
191    pub fn from_sizes(
192        config: &ModelConfig,
193        sizes: BufferSizes,
194        max_batch_tokens: usize,
195        max_batch_size: usize,
196        gpu: &dyn GpuBackend,
197    ) -> Result<Self> {
198        let decode_meta = DecodeMetaLayout::for_max_batch_size(max_batch_size);
199
200        let hidden_states = gpu.alloc(sizes.hidden_states)?;
201        let residual = gpu.alloc(sizes.residual)?;
202        let norm_output = gpu.alloc(sizes.norm_output)?;
203        let qkv_output = gpu.alloc(sizes.qkv_output)?;
204        let attn_output = gpu.alloc(sizes.attn_output)?;
205        let gate_logits = gpu.alloc(sizes.gate_logits)?;
206        let gate_logits_f32 = gpu.alloc(sizes.gate_logits_f32)?;
207        let moe_router_in_f32 = gpu.alloc(sizes.moe_router_in_f32)?;
208        let moe_output = gpu.alloc(sizes.moe_output)?;
209        let logits = gpu.alloc(sizes.logits)?;
210        let ssm_qkvz = gpu.alloc(sizes.ssm_qkvz)?;
211        let ssm_ba = gpu.alloc(sizes.ssm_ba)?;
212        let ssm_deinterleaved = gpu.alloc(sizes.ssm_deinterleaved)?;
213        let ssm_gates = gpu.alloc(sizes.ssm_gates)?;
214        let ssm_conv_out_f32 = gpu.alloc(sizes.ssm_conv_out_f32)?;
215        let scratch = gpu.alloc(sizes.scratch)?;
216        let expert_gate_out = gpu.alloc(sizes.expert_gate_out)?;
217        let expert_up_out = gpu.alloc(sizes.expert_up_out)?;
218        let expert_down_out = gpu.alloc(sizes.expert_down_out)?;
219        let splitk_workspace = gpu.alloc(sizes.splitk_workspace)?;
220        let o_latent = gpu.alloc(sizes.o_latent)?;
221        // Zero-filled "weight" for unweighted RMSNorm under the offset-from-1
222        // convention used by the rms_norm kernel (scale = 1 + weight). Weight = 0
223        // → scale = 1.0, i.e. a pure normalize (DeepSeek-V4 q_b_norm).
224        let norm_unit_w = gpu.alloc(sizes.norm_unit_w)?;
225        gpu.memset(norm_unit_w, 0, sizes.norm_unit_w)?;
226        let hc_streams = gpu.alloc(sizes.hc_streams)?;
227        let hc_post = gpu.alloc(sizes.hc_post)?;
228        let hc_comb = gpu.alloc(sizes.hc_comb)?;
229        let hc_lowrank_scratch = gpu.alloc(sizes.hc_lowrank_scratch)?;
230        let qsa_select_scratch = gpu.alloc(sizes.qsa_select_scratch)?;
231        // GDN FLA scratch: only allocate for the 128-dim-linear-head GDN path
232        // (size 0 → NULL → ATLAS_GDN_FLA dispatch stays disabled).
233        let ssd_scratch = if sizes.ssd_scratch > 0 {
234            gpu.alloc(sizes.ssd_scratch)?
235        } else {
236            DevicePtr::NULL
237        };
238        let gdn_fla_scratch = if sizes.gdn_fla_scratch > 0 {
239            gpu.alloc(sizes.gdn_fla_scratch)?
240        } else {
241            DevicePtr::NULL
242        };
243        let token_ids = gpu.alloc(sizes.token_ids)?;
244        // Shared dense-FFN activation-quant scratch (MMQ/int8 prefill). Sized 0
245        // for MoE models → NULL → per-layer ensure_* path stays inert.
246        let ffn_act_q8 = if sizes.ffn_act_q8 > 0 {
247            gpu.alloc(sizes.ffn_act_q8)?
248        } else {
249            DevicePtr::NULL
250        };
251        let ffn_act_a = if sizes.ffn_act_a > 0 {
252            gpu.alloc(sizes.ffn_act_a)?
253        } else {
254            DevicePtr::NULL
255        };
256        let ffn_act_scale = if sizes.ffn_act_scale > 0 {
257            gpu.alloc(sizes.ffn_act_scale)?
258        } else {
259            DevicePtr::NULL
260        };
261        let ffn_act_scale_kmajor = if sizes.ffn_act_scale_kmajor > 0 {
262            gpu.alloc(sizes.ffn_act_scale_kmajor)?
263        } else {
264            DevicePtr::NULL
265        };
266        let ffn_gate_up_fused = if sizes.ffn_gate_up_fused > 0 {
267            gpu.alloc(sizes.ffn_gate_up_fused)?
268        } else {
269            DevicePtr::NULL
270        };
271        let fp8_act = gpu.alloc(sizes.fp8_act)?;
272        let fp8_act_scale = gpu.alloc(sizes.fp8_act_scale)?;
273        let fp8_act_scale_kmajor = gpu.alloc(sizes.fp8_act_scale_kmajor)?;
274        // Q2_0 prefill dequant scratch. 0 → NULL unless ATLAS_GGUF_NATIVE_Q2.
275        let q2_dequant_scratch = if sizes.q2_dequant_scratch > 0 {
276            gpu.alloc(sizes.q2_dequant_scratch)?
277        } else {
278            DevicePtr::NULL
279        };
280        // LoRA scratch: only allocate when an adapter is configured
281        // (size 0 → NULL; cuMemAlloc rejects 0-byte allocations).
282        let lora_xa = if sizes.lora_xa > 0 {
283            gpu.alloc(sizes.lora_xa)?
284        } else {
285            DevicePtr::NULL
286        };
287        let lora_delta = if sizes.lora_delta > 0 {
288            gpu.alloc(sizes.lora_delta)?
289        } else {
290            DevicePtr::NULL
291        };
292        let lora_hact = if sizes.lora_hact > 0 {
293            gpu.alloc(sizes.lora_hact)?
294        } else {
295            DevicePtr::NULL
296        };
297        let lora_seq_slot = if sizes.lora_seq_slot > 0 {
298            gpu.alloc(sizes.lora_seq_slot)?
299        } else {
300            DevicePtr::NULL
301        };
302        // Q2_0 MMQ prefill q8_1 activation scratch. 0 → NULL unless ATLAS_GGUF_NATIVE_Q2_MMQ.
303        let q2_act_q8 = if sizes.q2_act_q8 > 0 {
304            gpu.alloc(sizes.q2_act_q8)?
305        } else {
306            DevicePtr::NULL
307        };
308        // Row-wise GDN prefill BF16 weights. 0 → NULL unless ATLAS_FP8_ROWWISE.
309        let ssm_rowwise_w_bf16 = if sizes.ssm_rowwise_w_bf16 > 0 {
310            gpu.alloc(sizes.ssm_rowwise_w_bf16)?
311        } else {
312            DevicePtr::NULL
313        };
314
315        tracing::info!(
316            "Buffer arena: {} tokens × {:.1} MB total (attn_out={:.1}MB, ssm_deint={:.1}MB, kv_lora_rank={})",
317            max_batch_tokens,
318            sizes.total_bytes() as f64 / (1024.0 * 1024.0),
319            sizes.attn_output as f64 / (1024.0 * 1024.0),
320            sizes.ssm_deinterleaved as f64 / (1024.0 * 1024.0),
321            config.kv_lora_rank,
322        );
323
324        Ok(Self {
325            hidden_states,
326            residual,
327            norm_output,
328            qkv_output,
329            attn_output,
330            gate_logits,
331            gate_logits_f32,
332            moe_router_in_f32,
333            moe_output,
334            logits,
335            ssm_qkvz,
336            ssm_ba,
337            ssm_deinterleaved,
338            ssm_gates,
339            ssm_conv_out_f32,
340            scratch,
341            expert_gate_out,
342            expert_up_out,
343            expert_down_out,
344            splitk_workspace,
345            o_latent,
346            norm_unit_w,
347            hc_streams,
348            hc_post,
349            hc_comb,
350            hc_lowrank_scratch,
351            qsa_select_scratch,
352            gdn_fla_scratch,
353            ssd_scratch,
354            token_ids,
355            ffn_act_q8,
356            ffn_act_a,
357            ffn_act_scale,
358            ffn_act_scale_kmajor,
359            ffn_gate_up_fused,
360            fp8_act,
361            fp8_act_scale,
362            fp8_act_scale_kmajor,
363            q2_dequant_scratch,
364            lora_xa,
365            lora_delta,
366            lora_hact,
367            lora_seq_slot,
368            q2_act_q8,
369            ssm_rowwise_w_bf16,
370            ssm_rowwise_w_bf16_used: std::sync::atomic::AtomicUsize::new(0),
371            max_batch_tokens,
372            decode_meta,
373            sizes,
374        })
375    }
376}
377
378/// Release every buffer this arena owns.
379///
380/// The destructure below is **exhaustive on purpose — no `..`**. A buffer added
381/// to `BufferArena` without a matching free is a leak that only shows up as the
382/// next model failing to fit, so the compiler is made to refuse the addition
383/// instead. If this line stops compiling, the fix is to free the new field, not
384/// to add a wildcard.
385impl atlas_core::scope::ModelResource<dyn GpuBackend> for BufferArena {
386    fn label(&self) -> &'static str {
387        "buffer arena"
388    }
389
390    fn release(&mut self, gpu: &dyn GpuBackend) -> anyhow::Result<()> {
391        let Self {
392            // Not allocations — named rather than wildcarded so the
393            // exhaustiveness check above keeps its teeth.
394            sizes: _,
395            max_batch_tokens: _,
396            // Layout, not an allocation — derived from `--max-batch-size`.
397            decode_meta: _,
398            hidden_states,
399            residual,
400            norm_output,
401            qkv_output,
402            attn_output,
403            gate_logits,
404            gate_logits_f32,
405            moe_router_in_f32,
406            moe_output,
407            logits,
408            ssm_qkvz,
409            ssm_ba,
410            ssm_deinterleaved,
411            ssm_gates,
412            ssm_conv_out_f32,
413            scratch,
414            expert_gate_out,
415            expert_up_out,
416            expert_down_out,
417            splitk_workspace,
418            o_latent,
419            norm_unit_w,
420            hc_streams,
421            hc_post,
422            hc_comb,
423            hc_lowrank_scratch,
424            qsa_select_scratch,
425            gdn_fla_scratch,
426            ssd_scratch,
427            token_ids,
428            ffn_act_q8,
429            ffn_act_a,
430            ffn_act_scale,
431            ffn_act_scale_kmajor,
432            ffn_gate_up_fused,
433            fp8_act,
434            fp8_act_scale,
435            fp8_act_scale_kmajor,
436            lora_xa,
437            lora_delta,
438            lora_hact,
439            lora_seq_slot,
440            q2_dequant_scratch,
441            q2_act_q8,
442            ssm_rowwise_w_bf16,
443            // A cursor into the slab above, not an allocation.
444            ssm_rowwise_w_bf16_used: _,
445        } = self;
446        // Every pointer, then NULL it: `release` must be idempotent because a
447        // `Drop` backstop may call it again, and `free` already no-ops on NULL.
448        let owned = [
449            *hidden_states,
450            *residual,
451            *norm_output,
452            *qkv_output,
453            *attn_output,
454            *gate_logits,
455            *gate_logits_f32,
456            *moe_router_in_f32,
457            *moe_output,
458            *logits,
459            *ssm_qkvz,
460            *ssm_ba,
461            *ssm_deinterleaved,
462            *ssm_gates,
463            *ssm_conv_out_f32,
464            *scratch,
465            *expert_gate_out,
466            *expert_up_out,
467            *expert_down_out,
468            *splitk_workspace,
469            *o_latent,
470            *norm_unit_w,
471            *hc_streams,
472            *hc_lowrank_scratch,
473            *qsa_select_scratch,
474            *hc_post,
475            *hc_comb,
476            *gdn_fla_scratch,
477            *ssd_scratch,
478            *token_ids,
479            *ffn_act_q8,
480            *ffn_act_a,
481            *ffn_act_scale,
482            *ffn_act_scale_kmajor,
483            *ffn_gate_up_fused,
484            *fp8_act,
485            *fp8_act_scale,
486            *fp8_act_scale_kmajor,
487            *lora_xa,
488            *lora_delta,
489            *lora_hact,
490            *lora_seq_slot,
491            *q2_dequant_scratch,
492            *q2_act_q8,
493            *ssm_rowwise_w_bf16,
494        ];
495        let mut first_error = None;
496        for ptr in owned {
497            if let Err(e) = gpu.free(ptr)
498                && first_error.is_none()
499            {
500                first_error = Some(e);
501            }
502        }
503        *hidden_states = DevicePtr::NULL;
504        *residual = DevicePtr::NULL;
505        *norm_output = DevicePtr::NULL;
506        *qkv_output = DevicePtr::NULL;
507        *attn_output = DevicePtr::NULL;
508        *gate_logits = DevicePtr::NULL;
509        *gate_logits_f32 = DevicePtr::NULL;
510        *moe_router_in_f32 = DevicePtr::NULL;
511        *moe_output = DevicePtr::NULL;
512        *logits = DevicePtr::NULL;
513        *ssm_qkvz = DevicePtr::NULL;
514        *ssm_ba = DevicePtr::NULL;
515        *ssm_deinterleaved = DevicePtr::NULL;
516        *ssm_gates = DevicePtr::NULL;
517        *ssm_conv_out_f32 = DevicePtr::NULL;
518        *scratch = DevicePtr::NULL;
519        *expert_gate_out = DevicePtr::NULL;
520        *expert_up_out = DevicePtr::NULL;
521        *expert_down_out = DevicePtr::NULL;
522        *splitk_workspace = DevicePtr::NULL;
523        *o_latent = DevicePtr::NULL;
524        *norm_unit_w = DevicePtr::NULL;
525        *hc_streams = DevicePtr::NULL;
526        *hc_lowrank_scratch = DevicePtr::NULL;
527        *qsa_select_scratch = DevicePtr::NULL;
528        *hc_post = DevicePtr::NULL;
529        *hc_comb = DevicePtr::NULL;
530        *gdn_fla_scratch = DevicePtr::NULL;
531        *ssd_scratch = DevicePtr::NULL;
532        *token_ids = DevicePtr::NULL;
533        *ffn_act_q8 = DevicePtr::NULL;
534        *ffn_act_a = DevicePtr::NULL;
535        *ffn_act_scale = DevicePtr::NULL;
536        *ffn_act_scale_kmajor = DevicePtr::NULL;
537        *ffn_gate_up_fused = DevicePtr::NULL;
538        *fp8_act = DevicePtr::NULL;
539        *fp8_act_scale = DevicePtr::NULL;
540        *fp8_act_scale_kmajor = DevicePtr::NULL;
541        *lora_xa = DevicePtr::NULL;
542        *lora_delta = DevicePtr::NULL;
543        *lora_hact = DevicePtr::NULL;
544        *lora_seq_slot = DevicePtr::NULL;
545        *q2_dequant_scratch = DevicePtr::NULL;
546        *q2_act_q8 = DevicePtr::NULL;
547        *ssm_rowwise_w_bf16 = DevicePtr::NULL;
548        match first_error {
549            Some(e) => Err(e),
550            None => Ok(()),
551        }
552    }
553}
554
555#[cfg(test)]
556mod tests;