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;