Beyond Masked Sparsity: SNACK Enables Truly Sparse Neural Networks on GPU

Beyond Masked Sparsity: SNACK Enables Truly Sparse Neural Networks on GPU

arXiv:2610.04093v1 [cs.LG] 2 Oct 2026

Jafar Badour University of Twente jafar.badour@utwente.nl

Maurice van Keulen University of Twente m.vankeulen@utwente.nl

Elena Mocanu University of Twente e.mocanu@utwente.nl

Abstract Deep neural networks continue to grow in parameter count, driving up training and inference cost on GPUs. Sparse neural networks and Dynamic Sparse Training (DST) promise to reduce these costs, but most implementations rely on binary masks over dense tensors and recover little of the theoretical compute, memory, or energy savings. We propose SNACK, a truly sparse GPU layer that stores and computes only non-zero connections. SNACK exposes a simple PyTorch API for restructuring connections and backpropagating gradients entirely in the sparse paradigm, and ships SNACK-COO, a custom COO-format SpMM CUDA kernel with a batch-to-Streaming-Multiprocessor mapping tuned for the small-batch, high-sparsity regime typical of large-model training and single-stream inference. At the kernel level, SNACK is up to 7× faster than the masked dense baseline (Dense+Mask) and competitive with cuSPARSE, Sputnik, and FlashSparse at 95% sparsity. At 90% sparsity, a single SNACK layer accelerates training by 8× and 3.7×, and inference by 4× and 2×, over Dense+Mask and fully dense layers, respectively, while using 72% less memory than dense and substantially less energy. End-to-end, SNACK reduces GPT-2 peak training memory by up to 40% and graph-style inference latency by 4.8× over Dense+Mask at 99% sparsity.

1

Introduction

Training and serving large neural networks is dominated by dense matrix multiplications whose cost scales with model size, raising real concerns about energy, hardware, and carbon emissions [46, 44]. Large language models in particular routinely involve weight matrices of size around 10000 × 10000, with both training and inference dominated by these multiplications. Sparse neural networks (SNNs) and Dynamic Sparse Training (DST) [33, 8] aim to cut this cost by keeping most weights at zero throughout training, in principle reducing parameters and FLOPs quadratically. In practice, however, most SNN implementations simulate sparsity by applying a binary mask to a dense tensor [13, 16]: the dense matmul still runs underneath, and almost none of the theoretical compute, memory, or energy savings materialize on the GPU. This leaves a gap between the promise of truly sparse neural networks (where weights live in a sparse format end-to-end [29]) and what is actually deployed. Mature high-performance SpMM(Sparse Matrix Multiplication) kernels such as cuSPARSE [37], Sputnik [14], and FlashSparse [45] cover several SpMM regimes well, but the very small batch sizes (down to batch = 1) typical of memory-constrained DST training and latency-bound single-stream inference on commodity GPUs remain underserved, so many fall back to a dense matmul times a binary mask, recovering little of the theoretical savings. We argue that this fallback should and can be avoided: with a small-batch, high-sparsity kernel, a truly sparse layer saves real time, memory, and energy over the masked-dense baseline. We further argue that for other regimes, current and future SpMM kernels can be plugged into our layer as backends, while still benefiting from its sparse gradient handling and sparse initialization. In this paper, we close that gap with SNACK (Sparse Network Active CUDA Kernel), which contributes two pieces. First, SNACK-COO: a custom COO-format SpMM CUDA kernel with a batch-to-Streaming-Multiprocessor mapping tuned for 40th Conference on Neural Information Processing Systems (NeurIPS 2026).

the small-batch, high-sparsity regime. Second, the SNACK layer: a truly sparse, backend-agnostic drop-in for nn.Linear with a sparse outer-product gradient and a memory-cheap sparse initializer, into which third-party SpMM kernels (cuSPARSE, Sputnik, FlashSparse) or our own SNACK-COO plug in as interchangeable backends. We demonstrate this second contribution end-to-end on a consumer-class RTX 3070 (6 GB): dropping Sputnik into the SNACK layer (SNACK(Sputnik)) trains GPT-2 to validation perplexity statistically equivalent to Dense+Mask (the two backends cannot be seed-aligned because Sputnik’s path through Algorithm 2 consumes a different RNG (Random Number Generator) stream from the dense Gaussian-then-mask path; §5) while using up to 40% less peak memory at η = 0.99, so the same model trains and serves on a card that would otherwise be excluded by the dense weight-tensor memory bill. On a single layer at 90% sparsity (Section 4), SNACK delivers: • vs. Dense+Mask: 4× / 8× faster inference/training, 72% less energy, 73% less memory. • vs. Dense: 2× / 3.7× faster inference/training, 35% less energy, 72% less memory. End-to-end (Section 5), SNACK further cuts GPT-2 peak training memory by up to 40% and graphstyle inference latency by 4.8× over Dense+Mask at 99% sparsity. We anticipate that SNACK will serve as a starting point for further work on truly-sparse layers and per-regime SpMM kernels, both on the consumer-class accelerators where the memory bill is the binding constraint and on the datacenter-class hardware where the kernel-choice frontier is still evolving.

2

Related Works

We organize prior work along three axes that bound our contribution: (i) sparsity-inducing training methods, (ii) GPU SpMM systems, and (iii) end-to-end truly sparse implementations. An extended discussion (covering early pruning, additional sparse-format details, and CPU/IPU truly-sparse work) is provided in Appendix E. Table 1: DST-oriented methods that maintain a sparse representation throughout training. “Truly Sparse” means no dense weight tensor is materialized at any point; “Unstructured” means the method does not require a block, N :M , or other hardware-friendly sparsity pattern; “Hardware-Aware” means the method requires structured patterns or specialized accelerator support. SNACK is, to our knowledge, the first truly sparse PyTorch building block on commodity GPUs that supports unstructured sparsity end-to-end and exposes both forward and backward passes in the sparse paradigm. Method Truly SNNs at Scale [6] SET-MLP [29] QuickSelection [2] Sparse-IFT [49] PopSparse [26] BLaST [39] SNACK (ours)

Truly Sparse

Unstructured

✓ ✓ ✓

✓ ✓ ✓

Hardware-Aware

✓ ✓ ✓

✓ ✓

✓

Sparsity Phase

CPU/GPU/IPU

Training Training Training (Autoencoder) Training & Inference Training & Inference Training & Inference Training & Inference

CPU CPU CPU GPU IPU GPU GPU

Sparsity-inducing training. Sparsity in deep networks has been induced via post-hoc pruning [18, 9, 10, 60, 27], pruning-at-init [25, 47], and Dynamic Sparse Training (DST) [4, 33, 7, 8, 35], which adapts connectivity throughout training. With rare exceptions, all of these maintain a dense weight tensor with a binary mask [28, 13, 16], recovering little of the theoretical compute or memory savings; [55] attacks this from the optimizer side by storing optimizer states only for active parameters, which is complementary to the kernel-level fix we pursue. Hardware-friendly transformations such as Sparse-IFT [49] and SPDF [48] preserve dense FLOP counts to remain compatible with GEMM, which is orthogonal to the truly sparse direction we pursue. At LLM scale, [1] pretrains highsparsity Llama variants for efficient deployment, and orthogonally, [59] reduces optimizer-state memory via low-rank gradient projections; both target the same memory wall as SNACK but along different axes (sparse pretraining recipe and low-rank gradients, respectively, vs. a truly sparse weight tensor). End-to-end, these methods are typically validated on language modeling with GPT-2 [41] on WikiText-103 [31] and similar benchmarks, which we adopt in Section 5. We discuss the broader literature in Appendix E. GPU SpMM systems. cuSPARSE [37] and Sputnik [14] optimize SpMM for the moderate-sparsity, large-batch regime that dominates dense neural network training; Sputnik exceeds dense throughput starting at ≈ 71% sparsity. FlashSparse [45] maps unstructured SpMM onto NVIDIA Tensor Cores, 2

targeting graph workloads at moderate sparsity, e.g. GAMLP-style models [57] on the Open Graph Benchmark [20]. BLaST [39] pairs a fused Triton block-SpMM kernel with a block-prune-and-grow recipe to sparsify Transformer MLPs up to 95% on multi-GPU LLM pretraining, but requires structured (block) sparsity and keeps the dense weight tensor live during training. MOSparse [51] addresses sparse-sparse multiplication (SpGEMM, a different operation than SpMM), and SparseRL [30] synthesizes SpMV kernels via reinforcement learning. Background on CSR/BSR/COO formats and earlier kernel-engineering work [3, 32, 54] is given in Appendix E. Truly sparse implementations and positioning. Truly sparse training (no dense weight tensor at any point) has been demonstrated on CPU using SciPy/Cython data structures [29, 6, 2] and on IPU hardware in PopSparse [26], which relies on block sparsity (b ∈ {4, 8, 16}) to align with vectorized IPU instructions. On commodity GPUs, hardware-supported 2:4 structured sparsity (Ampere/Hopper Tensor Cores) caps usable sparsity at 50% and requires re-pruning, so DST methods such as RigL [8] and SET [33], which produce unstructured masks, instead fall back to dense+mask. As summarized in Table 1, the combination of truly sparse, unstructured, commodity-GPU with no hardware-specific sparsity support is, to our knowledge, unique to SNACK. Throughout the paper, “sparsity” refers to unstructured sparsity unless stated otherwise.

3

Sparse Network Active CUDA Kernel (SNACK)

SNACK is a sparse-by-construction layer block (Figure 1) that drops in for nn.Linear and keeps weights in COO format end-to-end. The four subsections below cover the SpMM kernel (Sections 3.1 and 3.2), the uniform initializer (Algorithm 2), sparse gradient computation (Section 3.4), and the PyTorch wrapping (Section 3.5). Appendix F provides the full technical companion: a parallelefficiency proof for SNACK-COO, a geometric-distribution analysis of the initializer’s expected cost, the complete chain-rule and mini-batch gradient derivations, the CUDA/autograd implementation plumbing, and a summary of how end-to-end true sparsity is preserved across forward, backward, and topology updates.

Figure 1: SNACK layer overview. Two stacked SNACK blocks (A and B) act as drop-in replacements for standard nn.Linear layers. Each block stores its weight matrix SpM(·) on the GPU as the COO triple (Av , Ar , Ac ) over its non-zeros only: the highlighted (blue) edges are the active sparse subnetwork that consumes memory and FLOPs, while the faded (grey) edges depict the underlying dense linear layer that SNACK never materializes during training or inference. The forward pass is computed by SNACK-COO directly on the triple, the backward pass uses a sparse outer product over supp(W ), and DST methods may add or drop active connections in place, so SNACK and dense blocks can be freely interleaved in a network without changing any caller code. 3.1

Sparse Matrix Multiplication

We adopt a unified notation: A ∈ Rm×k is the Algorithm 1: Sparse Matrix Multiplication sparse weight matrix with sparsity η (fraction of k×n zero entries), B ∈ R is the dense activa- Input: Sparse matrix A ∈ Rm×k , activation matrix tion/input matrix where n is the batch dimension, B ∈ Rk×n and C = AB ∈ Rm×n is the output. Skipping ze- Output: Matrix C = A × B ros reduces the dense O(mkn) cost to (1 − η) mkn, 1. Initialize result matrix C ∈ Rm×n with all zeros so runtime is proportional to the number of non- 2. for each nonzero element (i, j, Az ) in A do a. for each nonzero element (j, l, Bz ) in B with zeros nnz = (1 − η) mk. Algorithm 1 expresses this in the Coordinate (COO) matching j Ci,l ← Ci,l + Az · Bz format [3], which stores each non-zero Ai,j as a 3. returni. C triple, i.e. as three arrays (Av , Ar , Ac ) over the nnz non-zeros (full definition in Appendix B); CSR or BSR are interchangeable. Standard COO 3

implementations distribute rows of A across CUDA cores and treat n as an outer loop, which underutilizes SMs at small n and motivates Section 3.2. 3.2

SNACK-COO: Batch-to-SM Mapping

The standard COO scheme assigns one row of A per CUDA thread and treats the batch dimension as an outer loop, yielding O(nm) thread tasks and leaving SMs idle whenever n ≪ m. Our proposed SNACK-COO instead launches a 2D block grid (Nb , n) with Nb = ⌈nnz/T ⌉ for a thread-block size T (typically 256): block (i, b) owns the i-th chunk of T non-zeros and computes its contribution to output column b. Thread τ in block (i, b) reads the COO triple at index k = iT + τ and atomically accumulates CAr [k], b += Av [k]·BAc [k], b . The task pool becomes Nb ×n blocks (i.e. nnz ×n threadtasks), so SM utilization tracks the non-zero workload rather than the matrix height, which is decisive at small n. Appendix F formalizes this with the parallel-efficiency inequality ϵSNACK-COO > ϵstandard p p for n ≪ m, and shows why batch-wise blocks reduce atomic contention compared to row-wise distribution. 3.3 Uniform Connection Initializer Most DST methods initialize the mask by drawing an m × k tensor of random numbers and thresholding, which materializes a dense intermediate and can OOM for large layers. Algorithm 2 avoids this: it samples nnz = (1 − η)mk unique index pairs uniformly at random and never builds a dense matrix. The expected cost is O(ζ · nnz), where ζ is the expected number of trials before drawing a fresh pair. At step i, the probability of drawing a fresh pair is pi = 1 − i/(mk), so by the geometric distribution the expected number of trials at that step is 1/pi = mk/(mk − i). Summing over the nnz unique pairs and applyPb ing the harmonic approximation i=a 1/i ≈ Pnnz mk ln(b/a) yields i=0 mk−i ≈ −mk ln η for η ∈ (0, 1), i.e. essentially linear in nnz. Collisions dominate as η → 0 (dense regime) but vanish as η → 1, so the rejection penalty ζ = − ln η/(1 − η) is bounded and small in the high-sparsity regime DST targets. The full derivation is in Appendix F.

Algorithm 2: Uniform Initializer Input: Sparse matrix of size m × k, target sparsity η Output: Set C = {ct = ⟨i, j⟩} of (1 − η) · m · k unique indices 1. Initialize empty index set C, hash set S 2. for t = 1 to (1 − η) · m · k do a. repeat i. i ← getrand(m) + 1 ii. j ← getrand(k) + 1 iii. if (i, j) ∈ S then continue iv. else add (i, j) to S ct ← ⟨i, j⟩ break b. until a new pair is added 3. return C

3.4 Sparse Gradient Computation For a SNACK layer with sparse weight W ∈ Rm×n , input X ∈ Rn×b , and pre-activation A = W X, ∂L ∂L ∂L ∂L = ( ∂A ) X ⊤ and ∂X = W ⊤ ( ∂A ). The input-gradient term is sparse-timesthe chain rule yields ∂W ⊤ dense and reuses our SpMM directly since W inherits W ’s sparsity. The weight-gradient term is the cost driver: a naive outer product produces a dense m × n matrix the size of the unpruned W . Pb ∂L SNACK instead computes ( ∂W )i,j = t=1 ∇Yi,t Xj,t only at indices (i, j) ∈ supp(W ), dropping the cost from O(mnb) to O(nnz · b) while preserving SGD correctness (frozen entries see zero gradient anyway). Activation functions operate element-wise and are agnostic to W ’s sparsity [17]. Appendix F gives the complete derivation: the chain rule for ∂L/∂W and ∂L/∂X, the entry-wise expansion that motivates restricting work to supp(W ), and the generalization to mini-batches as a sum of b rank-1 sparse outer products. The matrix on the right is ∇Y ⊗ X = (∇Y ) X ⊤ ,   the entry-wise expansion of the outer product. Dense X1 ∇Y1 . . . X1 ∇Yn backprop fills all mn entries; SNACK computes only  X2 ∇Y1 . . . X2 ∇Yn   those at (i, j) ∈ supp(W ). For a single non-zero (∇Y ) X ⊤ =    . . .   . . . ⟨2, 1⟩, the boxed cell X2 ∇Y1 is the only entry ma. . . terialized. Across forward, gradient, and topology Xm ∇Y1 . . . Xm ∇Yn updates, we maintain end-to-end true sparsity: no Boxed entry computed iff ⟨2, 1⟩ ∈ supp(W ). step ever instantiates a dense m × n tensor. 4

SpMM/DenseMM (batch = 1)

SpMM/DenseMM (batch = 2)

5,150.9 s

CUDA time

801.2 s 124.6 s 19.4 s 3.0 s 80 DenseMM DenseMM+Mask

90 Sparsity SNACK-COO SpMM_COO(Sputnik)

100 SpMM_CSR(Sputnik DL-opt) SpMM_COO(cuSPARSE)

80

90 Sparsity

SpMM_CSR(cuSPARSE) SpMM_FlashSparse

SpMM_COO(Torch) SpMM_CSR(Torch)

100 SpMM_COO(Cupy) SpMM_CSR(Cupy)

Figure 2: SpMM CUDA time vs. sparsity at fixed batch sizes {1, 2} on a 5000×5000 matrix. Each panel sweeps sparsity for one batch size. SNACK-COO (ours, orange) stays below the Dense+Mask baseline across these small-batch settings and remains competitive with the dense reference. 3.5

Implementation

SNACK is exposed as a PyTorch nn.Module that drops in for nn.Linear (Figure 1): forward and backward are wired through a custom torch.autograd.Function backed by the SNACK-COO ∂L CUDA kernel and a sparse outer-product CUDA kernel for ∂W . Weights live entirely on the GPU as the COO triple (Av , Ar , Ac ) and are never materialized as a dense m × n tensor at any point: initialization, forward, backward, or topology updates. The SpMM backend is swappable: SNACKCOO is the default for the small-batch DST regime, but any third-party SpMM (Sputnik, FlashSparse, cuSPARSE) can be plugged in for regimes where it is faster, while still inheriting SNACK’s sparse outer-product gradient and uniform initializer. Wrapper, kernel, and autograd plumbing details are in Appendix F; code is available at https://github.com/JafarBadour/Snack-SNN.

4

Evaluation and Results

We evaluate SNACK at three increasingly demanding levels: (i) the SpMM kernel in isolation (Section 4.1), (ii) a single sparse layer in a 100-iteration MLP-style training cycle (Section 4.2), and (iii) end-to-end on GPT-2 and a graph workload (Section 5). Unless stated otherwise, all micro-benchmarks run on an RTX 3070 (6 GB) at sparsities η ∈ {0.9, 0.95, 0.99} and small leading dimensions, the regime SNACK-COO targets. We report CUDA time, peak GPU memory, energy, and parameter count; per-input-size sweeps and additional figures are deferred to Appendix D. 4.1

SpMM Kernel: SNACK-COO vs. State-of-the-Art

We benchmark SNACK-COO against eleven SpMM implementations [12, 5, 38, 40] from PyTorch (CSR/COO), CuPy (CSR/COO), JAX (BSR), cuSPARSE (CSR/COO), Sputnik (CSR/COO), and FlashSparse on a 5000 × 5000 sparse matrix at 95% sparsity, against a (batch × 5000) dense matrix, averaged over 1000 multiplications. Table 2 reports CUDA time at batch 1 across four matrix sizes: Our SNACK-COO is the fastest sparse method on every leading-dimension size we tested and the only sparse implementation that beats Dense+Mask, achieving a 7× speedup over Dense+Mask and a 3.5× speedup over the dense baseline at 5000 × 5000. Figure 2 reports sparsity sweeps at fixed batch sizes {1, 2} and shows the same advantage across the small-batch regime. Appendix D.2 replicates this sweep on a datacenter-class NVIDIA A100 across six leading-dimension sizes (500– 10000) at batch {1, 4}: at batch 1 SNACK-COO retains the per-column lead at every size (7.5× over Dense+Mask, 3.7× over dense at 5000×5000), with SpMM_CSR(Cupy) closing to within 13% only at {8500, 10000}; at batch 4 SNACK-COO leads at {500, 1000} and is overtaken from 5000×5000 upward by the two Sputnik variants (within ±10% of each other at every size), the regime crossover predicted by the parallel-efficiency analysis (Appendix F) and motivating the swappable backend of Section 3.5. 4.2

Single-Layer Training Cycle

We replace one 10000 × 10000 linear projection with SNACK and run 100 SGD iterations against Dense+Mask (the simulated-sparsity baseline) and Dense (Only). Figure 3 sweeps sparsity (and, for the two latency / per-batch energy panels, mini-batch size) across seven GPU-level metrics; Table 3 5

Table 2: Performance comparison of sparse matrix multiplication implementations. SpM sparsity is 95%, batch size of 1. Our SNACK-COO method outperforms other methods (7x reduction in time vs DenseMM+Mask). † Simulates sparsity. Method

is sparse

SpMM_CSR(Torch) SpMM_CSR(Cupy) SpMM_COO(Torch) SNACK-COO (ours) SpMM_COO(Cupy) SpMM_BSR(Jax) SpMM_CSR(cuSPARSE) SpMM_COO(cuSPARSE) SpMM_CSR(Sputnik DL-opt) SpMM_COO(Sputnik) FlashSparse DenseMM+Mask

Yes

DenseMM(baseline)

no

†

CUDA (ms) 500x500

1000x1000

5000x5000

10000x10000

27.05 11.95 27.41 2.43 207.86 88.68 27.14 26.56 3.21 3.22 29.85 4.32

26.84 11.66 29.10 2.38 215.45 75.78 27.47 29.15 3.31 3.22 39.11 4.24

231.83 12.27 341.60 7.28 1231.38 126.39 225.09 333.83 8.93 8.89 360.35 50.46

885.33 17.90 1396.26 15.17 4946.64 320.32 858.96 1362.39 25.37 25.36 1180.53 193.91

2.17

2.13

25.36

97.44

Figure 3: Layer-level GPU profile, 10000×10000 projection (100 SGD iterations). Top row: per-iteration CUDA time (log µs; the multiple lines per method are mini-batches b ∈ {1, 2, 4, 8, 16, 32, 64}, annotated bz=b); average total energy (mWh); total energy per batch (mWh, same b-sweep). Bottom row: total parameters (sparse values + COO indices for SNACK), trainable parameters, peak GPU memory (MB), and steady-state power (W); these four panels report a single curve per method. Only SNACK’s parameter, memory, and energy curves fall linearly with 1 − η; Dense+Mask stays flat because it still allocates and multiplies through a dense W . CUDA-time and energy axes are log-scaled because of the wide dynamic range across η.

pins the same numbers at batch 4. At η = 0.9, SNACK uses only 30% of Dense’s parameters (10% trainable plus their COO indices) and 28% of its memory; at η = 0.99 it is 17× faster and ≈12× more energy-efficient than Dense+Mask. Crucially, SNACK’s parameter, memory, and energy costs scale linearly with 1 − η, while Dense+Mask remains essentially flat: direct evidence that masking forfeits the savings sparsity should provide. 4.3

Numerical Equivalence and DST Accuracy

SNACK is a drop-in replacement: forward outputs and gradients match Dense+Mask up to relative error ϵ ≈ 10−4 (attributable to floating-point reordering), so any SNN trained with Dense+Mask can be swapped to SNACK without changes to the optimizer or the loss. Plugged into the DST pipeline of [22], RigL(SNACK) matches RigL on ResNet-50/ImageNet at η ∈ {0.8, 0.9, 0.95, 0.99} within 0.2% top-1 accuracy (see Appendix C). 6

Table 3: Performance and environmental impact comparison of different implementations on a 10000×10000 matrix at batch size 4. Energy (mWh) and CO2 (mg) use local emission factors. Dense (Only) is sparsity-independent and is reported once at the bottom as a reference (mean over the three sparsity-block runs; per-run variation is < 2% on every column except memory). Sparsity %

Params (k)

Memory (MB)

CUDA τ (µs)

Power (W)

Energy (mWh)

CO2 (mg)

Dense+Mask SNACK (ours)

90 90

200,000 30,000

4371 1173

1085.6 484.4

94.6 81.7

3.08 0.85

1.38 0.38

Dense+Mask SNACK (ours)

95 95

200,000 15,000

4140 947

1085.6 254.6

94.7 71.6

3.15 0.45

1.42 0.20

Dense+Mask SNACK (ours)

99 99

200,000 3,000

4447 693

1085.3 63.6

94.4 66.3

3.07 0.26

1.38 0.12

Dense (Only)

–

100,000

4192

497.3

105.4

1.30

0.59

Type

MEMORY USAGE COMPARISON SNACK (Sparsity 0.9) SNACK (Sparsity 0.95) SNACK (Sparsity 0.99) Dense (Only) Dense+Mask (Sparsity 0.9) Dense+Mask (Sparsity 0.95) Dense+Mask (Sparsity 0.99) GPU Memory Limit (~8GB)

7000

Memory Usage (MB)

6000 5000 4000 3000 2000 1000 5000

10000

15000

20000

Input Size

25000

Figure 4: Memory vs. input size at hidden dimension 10000. Dense and Dense+Mask OOM at input ≥ 15000; Dense+Mask at η = 0.9 fails earlier (≥ 12500) due to lazy allocation. SNACK scales gracefully because no step ever materializes a dense W .

4.4

Memory Headroom at Larger Layers

We sweep input size at hidden dimension 10000, recording peak memory at η ∈ {0.9, 0.95, 0.99} and the first input size that triggers CUDA OOM (Out Of Memory) (Figure 4). Dense and Dense+Mask OOM at input ≥ 15000 (inside the hidden-size range of modern LLMs [50]) because the binary mask still requires a dense W . SNACK avoids OOM at every sparsity tested by storing only non-zeros, fitting the same workload in roughly an order of magnitude less memory at η = 0.99. Granular perinput-size profiles (2500 –12500) and a numerical breakdown at 12500 × 10000 are in Appendix D.

5

End-to-end Experiments

We test SNACK end-to-end on (i) GPT-2 [41] dynamic sparse training and inference on WikiText103 [31] via HuggingFace transformers [53], and (ii) batch-1 GAMLP-style [57] inference on ogbn-products [20]. Backends: dense (cuBLAS), Dense+Mask, SNACK(Sputnik); GAMLP also tests SNACK-COO. Sparsities are layer-local (only c_fc/c_proj per GPT-2 block and the 5000×5000 GAMLP feature head are replaced; attention, embeddings, norms, LM (Language Model) head, and graph propagation stay dense), so for GPT-2 (small) η=0.9 removes only ≈41% of total parameters — an architectural ceiling, not a kernel limit. All GPT-2 runs use a consumer-class RTX 3070 (6 GB) to demonstrate viability on commodity hardware where the weight-tensor memory bill is the binding constraint, so we lead on memory (and, for inference, latency). The two backends cannot be seedaligned — SNACK(Sputnik) draws only (1−η)mk non-zeros via Algorithm 2, Dense+Mask draws Gaussian-then-mask, distinct RNG streams with no alignment hook — so we report distributional equivalence over independent seeds, not bit-identity. GPT-2 training (Table 4): an explicit time/energy vs. memory trade. SNACK(Sputnik) matches (and slightly better) Dense+Mask perplexity at every η (∆ppl ∈ [−1.62, +0.39] (ppl here is perplexity), due to unseedable initializer divergence) and always cuts peak memory (1366 vs. 2261 MB at 7

Table 4: GPT-2 (small) dynamic sparse training across target sparsity levels. We compare Dense+Mask and SNACK(Sputnik) (ours), reporting validation perplexity (WikiText-103 validation split), peak GPU memory (MB), per-step energy (mJ), and per-token training latency (ms); dense is the no-sparsity reference. Best per-metric per-sparsity value is in bold. Sparsity

Method

Perplexity

Memory

Energy_spent

Latency_token_ms

0.25 0.25

Dense+Mask SNACK(Sputnik) (ours)

29.92 29.73

2260.84 2506.92

6916.40 345236.72

41.64 582.63

0.50 0.50

Dense+Mask SNACK(Sputnik) (ours)

36.50 36.89

2260.84 2126.79

7021.94 190115.90

41.59 453.23

0.75 0.75

Dense+Mask SNACK(Sputnik) (ours)

43.08 43.05

2260.84 1746.65

7127.47 34995.07

41.53 332.92

0.85 0.85

Dense+Mask SNACK(Sputnik) (ours)

62.64 62.72

2260.84 1601.57

7078.19 18627.35

41.69 176.56

0.90 0.90

Dense+Mask SNACK(Sputnik) (ours)

77.63 77.25

2260.84 1498.08

7054.77 14217.54

41.72 127.16

0.95 0.95

Dense+Mask SNACK(Sputnik) (ours)

103.97 103.16

2260.84 1420.60

7011.71 10171.30

41.71 86.13

0.99 0.99

Dense+Mask SNACK(Sputnik) (ours)

131.90 130.28

2260.84 1366.27

6978.18 8166.11

41.76 70.67

-

dense

29.73

2031.74

6683.58

39.93

η=0.99, ≈40% monotone in η); the cost is per-step latency and energy, higher at low η (582 ms / 345k mJ) and at parity by high η (71 ms / 8k mJ between η=0.25 and 0.99) as the fixed launch / format-conversion overhead amortizes against (1−η)mnb arithmetic — the trade flips toward SNACK exactly where DST operates (per-step energy ratio ≈50 ×→≈1.17× across the sweep). The point of the trade is not faster steps, it is feasibility: the same recipe places a 124 M-parameter GPT-2 onto cards that cannot host Dense+Mask at all — on the RTX 3070 (6 GB) used here SNACK(Sputnik) at η=0.99 frees ≈900 MB beyond Dense+Mask’s ≈3.7 GB headroom, and the same recipe in principle drops onto a 2 GB-class accelerator on which Dense+Mask OOMs immediately. Table 5: GPT-2 (small) inference across target sparsity levels at batch size 1 on a consumer-class RTX 3070 (6 GB). We compare Dense+Mask and SNACK(Sputnik) (ours), reporting peak GPU memory (MB), per-token energy (mJ), and per-token latency (ms); dense is the no-sparsity reference. Best per-metric per-sparsity value is in bold. Sparsity

Method

Memory

Energy_spent

Latency_token_ms

0.25 0.25

Dense+Mask SNACK(Sputnik) (ours)

1255.59 1609.92

761.16 2237.78

5.42 11.91

0.5 0.5

Dense+Mask SNACK(Sputnik) (ours)

1255.59 1284.49

749.48 1507.93

5.33 9.31

0.75 0.75

Dense+Mask SNACK(Sputnik) (ours)

1255.59 959.05

737.79 778.07

5.25 6.70

0.85 0.85

Dense+Mask SNACK(Sputnik) (ours)

1255.59 807.59

740.36 685.85

5.44 6.89

0.9 0.9

Dense+Mask SNACK(Sputnik) (ours)

1250.87 707.97

730.59 627.04

5.33 6.95

0.95 0.95

Dense+Mask SNACK(Sputnik) (ours)

1255.59 632.92

725.91 593.18

5.33 6.82

0.99 0.99

Dense+Mask SNACK(Sputnik) (ours)

1255.59 578.59

726.93 565.64

5.41 6.87

-

dense

1017.14

605.89

5.09

GPT-2 inference (Table 5): SNACK(Sputnik) drops peak memory monotonically from 1610 to 579 MB (2.2× over Dense+Mask’s flat ≈1255 MB), keeps latency within ≈1.3× of dense 8

Table 6: Graph-style sparse inference on a 5000×5000 GAMLP-shaped workload at batch size 1, averaged over 2,000 samples after warm-up. We compare Dense+Mask and SNACK(SNACK-COO) (ours), reporting peak memory (MB), per-sample energy (mJ), and inference latency (ms); dense is the no-sparsity reference. Best per-metric per-sparsity value is in bold. Method

Sparsity

Memory

Energy_spent

Inference_latency_ms

Dense+Mask SNACK(SNACK-COO) (ours) Dense+Mask SNACK(SNACK-COO) (ours) Dense+Mask SNACK(SNACK-COO) (ours)

0.9 0.9 0.95 0.95 0.99 0.99

408.56 128.60 408.56 118.60 309.89 110.60

52.79 11.82 53.63 10.14 50.86 9.82

1.02 0.25 1.01 0.21 1.01 0.21

dense

baseline

367.71

47.51

0.91

at η ≥ 0.75, and crosses into a per-token energy win at η ≥ 0.85 (−22% at η=0.99): the time/energy/memory trade-off no longer trades — SNACK is at parity (latency) or ahead (memory; energy at high η) on every axis. Deployment consequence: the same model fits in 579 MB at η=0.99 — under half of Dense+Mask’s 1255 MB — freeing room for two co-resident copies, doubled context, or sub-1 GB accelerators where Dense+Mask cannot fit. GAMLP graph inference (Table 6). Training runs at batch 4096 via SNACK(Sputnik); the trained sparse weights are then served at batch 1 through SNACK-COO (benchmarks are in Appendix G). SNACK-COO cuts peak memory 2.8× and latency 4.8× over Dense+Mask at η=0.99 (110.6 vs. 309.9 MB; 0.21 vs. 1.01 ms), and is 3.6–4.3× faster and 2.8–3.3× smaller than the dense reference at every η tested — confirming the per-regime story: SNACK-COO wins at small leading dimension and high η, SNACK(Sputnik) covers larger-batch training.

6

Conclusion, Limitations, and Future Work

We presented SNACK, a backend-agnostic, truly sparse PyTorch layer for commodity GPUs, and SNACK-COO, a COO-format SpMM kernel whose batch-to-Streaming-Multiprocessor mapping targets the small-batch, high-sparsity regime current libraries underserve. The layer preserves true sparsity through forward, the sparse outer-product backward DST topology updates via the memory-cheap uniform initializer, and admits cuSPARSE / Sputnik / FlashSparse as drop-in backends. Layer-wise, this yields up to 7× kernel-level speedup over Dense+Mask, up to 40% peak-memory reduction in GPT-2 dynamic sparse training, and 4.8× lower graph-inference latency at 99% sparsity, all measured on a single consumer-class RTX 3070 (6 GB) chosen specifically to demonstrate viability on commodity hardware, at validation-perplexity / accuracy parity with Dense+Mask on the workloads we test. By decoupling the sparse-PyTorch surface from any single SpMM, SNACK turns each new kernel into a transparent layer-level speedup, redirecting research effort toward the kernels themselves. Limitations. SNACK-COO is tuned for n ≪ m and η ⩾ 0.9 in fp32; outside this regime, kernels for larger batches or moderate sparsity (Sputnik, FlashSparse, BLaST) are faster, and the layer relies on backend swapping rather than a single dominant kernel. End-to-end gains are bounded architecturally: only sparsified projections recover compute (e.g. ≈41% of GPT-2-small at η = 0.9). Although SNACK, as a framework, due to its flexibility and modularity, can easily integrate Sputnik, FlashSparse, or CuSparse, its advantage on large models is investigated just for one architecture, leaving the other models largely under-explored. Yet, the potential of using it in conjunction with a large batch size is a promising research direction. Future Work. Immediate extensions: Tensor-Core SNACK-COO for fp16/bf16, multi-GPU SNACK with sparsity-aware collectives, and a runtime backend dispatcher. The substrate also generalises to sparse attention, MoE routers, quantization-×-sparsity, and a Triton/torch.compile port for AMD/Intel.

References [1] Abhinav Agarwalla, Abhay Gupta, Alexandre Marques, Shubhra Pandit, Michael Goin, Eldar Kurtic, Kevin Leong, Tuan Nguyen, Mahmoud Salem, Dan Alistarh, et al. Enabling high9

sparsity foundational Llama models with efficient pretraining and deployment. arXiv preprint arXiv:2405.03594, 2024. [2] Zahra Atashgahi, Ghada Sokar, Tim van der Lee, Elena Mocanu, Decebal Constantin Mocanu, Raymond Veldhuis, and Mykola Pechenizkiy. Quick and robust feature selection: the strength of energy-efficient sparse training for autoencoders. Mach. Learn., 111(1):377–414, January 2022. ISSN 0885-6125. [3] Nathan Bell and Michael Garland. Implementing sparse matrix-vector multiplication on throughput-oriented processors. In Proceedings of the conference on high performance computing networking, storage and analysis, pages 1–11, 2009. [4] Guillaume Bellec, David Kappel, Wolfgang Maass, and Robert Legenstein. Deep rewiring: Training very sparse deep networks. In International Conference on Learning Representations, 2018. [5] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/jax-ml/jax. [6] Selima Curci, Decebal Constantin Mocanu, and Mykola Pechenizkiyi. Truly sparse neural networks at scale. arXiv preprint arXiv:2102.01732, 2021. [7] Tim Dettmers and Luke Zettlemoyer. Sparse networks from scratch: Faster training without losing performance. arXiv preprint arXiv:1907.04840, 2019. [8] Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In International conference on machine learning, pages 2943–2952. PMLR, 2020. [9] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=rJl-b3RcF7. [10] Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International conference on machine learning, pages 10323–10337. PMLR, 2023. [11] Elias Frantar, Carlos Riquelme Ruiz, Neil Houlsby, Dan Alistarh, and Utku Evci. Scaling laws for sparsely-connected foundation models. In B. Kim, Y. Yue, S. Chaudhuri, K. Fragkiadaki, M. Khan, and Y. Sun, editors, International Conference on Representation Learning, volume 2024, pages 42175–42194, 2024. URL https://proceedings.iclr.cc/paper_files/ paper/2024/file/b83abe0c353d5dddf939e3533303220f-Paper-Conference.pdf. [12] Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. Systems for Machine Learning, 4(9):1–3, 2018. [13] Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint arXiv:1902.09574, 2019. [14] Trevor Gale, Matei Zaharia, Cliff Young, and Erich Elsen. Sparse GPU kernels for deep learning. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–14. IEEE, 2020. [15] John R Gilbert, Cleve Moler, and Robert Schreiber. Sparse matrices in matlab: Design and implementation. SIAM journal on matrix analysis and applications, 13(1):333–356, 1992. [16] Laura Graesser, Utku Evci, Erich Elsen, and Pablo Samuel Castro. The state of sparse training in deep reinforcement learning. In International Conference on Machine Learning, pages 7766–7792. PMLR, 2022. [17] Martin T Hagan, Howard B Demuth, and Mark Beale. Neural network design. PWS Publishing Co., 1997. 10

[18] Babak Hassibi, David G Stork, and Gregory J Wolff. Optimal brain surgeon and general network pruning. In IEEE international conference on neural networks, pages 293–299. IEEE, 1993. [19] Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: pruning and growth for efficient inference and training in neural networks. J. Mach. Learn. Res., 22(1), January 2021. ISSN 1532-4435. [20] Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:22118–22133, 2020. [21] Eun-Jin Im. Optimizing the performance of sparse matrix-vector multiplication. University of California, Berkeley, 2000. [22] Mike Lasby, Anna Golubeva, Utku Evci, Mihai Nica, and Yani Ioannou. Dynamic sparse training with structured sparsity. arXiv preprint arXiv:2305.02299, 2023. [23] Mike Lasby, Anna Golubeva, Utku Evci, Mihai Nica, and Yani Ioannou. Dynamic sparse training with structured sparsity, 2024. URL https://arxiv.org/abs/2305.02299. [24] Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. In D. Touretzky, editor, Advances in Neural Information Processing Systems, volume 2. MorganKaufmann, 1989. URL https://proceedings.neurips.cc/paper_files/paper/1989/ file/6c9882bbac1c7093bd25041881277658-Paper.pdf. [25] Namhoon Lee, Thalaiyasingam Ajanthan, and Philip Torr. SNIP: Single-shot network pruning based on connection sensitivity. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=B1VZqjAcYX. [26] Zhiyi Li, Douglas Orr, Valeriu Ohan, Tom Murray, Adam Sanders, Deniz Beker, Dominic Masters, et al. Popsparse: Accelerated block sparse matrix multiplication on ipu. arXiv preprint arXiv:2303.16999, 2023. [27] Tailin Liang, John Glossner, Lei Wang, Shaobo Shi, and Xiaotong Zhang. Pruning and quantization for deep neural network acceleration: A survey. Neurocomputing, 461:370–403, 2021. [28] Junjie Liu, Zhe Xu, Runbin Shi, Ray CC Cheung, and Hayden KH So. Dynamic sparse training: Find efficient sparse network from scratch with trainable masked layers. arXiv preprint arXiv:2005.06870, 2020. [29] Shiwei Liu, Lu Yin, Decebal Constantin Mocanu, and Mykola Pechenizkiy. Do we actually need dense over-parameterization? in-time over-parameterization in sparse training. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 6989–7000. PMLR, 18–24 Jul 2021. URL https://proceedings.mlr.press/v139/liu21y.html. [30] Sijia Luo, Xiaokang Zhang, Yuxuan Hu, Bohan Zhang, Ke Wang, Jinbo Su, Mengshu Sun, Lei Liang, and Jing Zhang. Sparse-rl: Breaking the memory wall in llm reinforcement learning via stable sparse rollouts. arXiv preprint arXiv:2601.10079, 2026. [31] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations (ICLR), 2017. [32] Duane Merrill and Michael Garland. Merge-based parallel sparse matrix-vector multiplication. In SC’16: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 678–689. IEEE, 2016. [33] Decebal Constantin Mocanu, Elena Mocanu, Peter Stone, Phuong H Nguyen, Madeleine Gibescu, and Antonio Liotta. Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science. Nature communications, 9(1):2383, 2018. 11

[34] Decebal Constantin Mocanu, Elena Mocanu, Tiago Pinto, Selima Curci, Phuong H. Nguyen, Madeleine Gibescu, Damien Ernst, and Zita A. Vale. Sparse training theory for scalable and efficient agents. In Proceedings of the 20th International Conference on Autonomous Agents and MultiAgent Systems, AAMAS ’21, page 34–38, Richland, SC, 2021. International Foundation for Autonomous Agents and Multiagent Systems. ISBN 9781450383073. [35] Hesham Mostafa and Xin Wang. Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization. In International Conference on Machine Learning, pages 4646–4655. PMLR, 2019. [36] Francisco Muñoz-Martínez, Raveesh Garg, Michael Pellauer, José L Abellán, Manuel E Acacio, and Tushar Krishna. Flexagon: A multi-dataflow sparse-sparse matrix multiplication accelerator for efficient dnn processing. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, pages 252–265, 2023. [37] Maxim Naumov, L Chien, Philippe Vandermersch, and Ujval Kapasi. Cusparse library: A set of basic linear algebra subroutines for sparse matrices. In GPU Technology Conference, number 5, page 17, 2010. [38] ROYUD Nishino and Shohei Hido Crissman Loomis. Cupy: A numpy-compatible library for nvidia gpu calculations. 31st confernce on neural information processing systems, 151(7), 2017. [39] Patrik Okanovic, Sameer Deshmukh, Grzegorz Kwasniewski, Yi Zhu, Haruto Fujii, Sakina Fatima, Maciej Besta, Kentaro Katayama, Takumi Honda, Yusuke Nagasaka, et al. Blast: High performance inference and pretraining using block sparse transformers. arXiv preprint arXiv:2507.03117, 2025. [40] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: an imperative style, high-performance deep learning library. Curran Associates Inc., Red Hook, NY, USA, 2019. [41] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. Technical report, OpenAI, 2019. [42] Yousef Saad. Iterative methods for sparse linear systems. SIAM, 2003. [43] Erik Schultheis and Rohit Babbar. Towards memory-efficient training for extremely large output spaces–learning with 670k labels on a single commodity gpu. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 689–704. Springer, 2023. [44] Roy Schwartz, Jesse Dodge, Noah A Smith, and Oren Etzioni. Green ai. Communications of the ACM, 63(12):54–63, 2020. [45] Jinliang Shi, Shigang Li, Youxuan Xu, Rongtian Fu, Xueying Wang, and Tong Wu. FlashSparse: Minimizing computation redundancy for fast sparse matrix multiplications on tensor cores. In Proceedings of the 30th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming (PPoPP), 2025. [46] Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in nlp. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 3645–3650, 2019. [47] Hidenori Tanaka, Daniel Kunin, Daniel L. Yamins, and Surya Ganguli. Pruning neural networks without any data by iteratively conserving synaptic flow. In Advances in Neural Information Processing Systems, volume 33, pages 6377–6389, 2020. [48] Vithursan Thangarasa, Abhay Gupta, William Marshall, Tianda Li, Kevin Leong, Dennis DeCoste, Sean Lie, and Shreyas Saxena. Spdf: Sparse pre-training and dense fine-tuning for large language models. In Uncertainty in Artificial Intelligence, pages 2134–2146. PMLR, 2023. 12

[49] Vithursan Thangarasa, Shreyas Saxena, Abhay Gupta, and Sean Lie. Sparse-ift: Sparse iso-flop transformations for maximizing training efficiency. arXiv preprint arXiv:2303.11525, 2023. [50] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023. [51] Yizhuo Wang, Hongpeng Lin, Bingxin Wei, Jianhua Gao, and Weixing Ji. Optimizing general sparse matrix-matrix multiplication on the gpu. ACM Transactions on Architecture and Code Optimization, 22(4):1–25, 2025. [52] Wei Wen, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Learning structured sparsity in deep neural networks. Advances in neural information processing systems, 29, 2016. [53] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-theart natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38–45, 2020. [54] Carl Yang, Aydın Buluç, and John D Owens. Design principles for sparse matrix multiplication on the gpu. In European Conference on Parallel Processing, pages 672–687. Springer, 2018. [55] Geng Yuan, Xiaolong Ma, Wei Niu, Zhengang Li, Zhenglun Kong, Ning Liu, Yifan Gong, Zheng Zhan, Chaoyang He, Qing Jin, et al. MEST: Accurate and fast memory-economic sparse training framework on the edge. In Advances in Neural Information Processing Systems, volume 34, pages 20838–20850, 2021. [56] Guowei Zhang, Nithya Attaluri, Joel S Emer, and Daniel Sanchez. Gamma: Leveraging gustavson’s algorithm to accelerate sparse matrix multiplication. In Proceedings of the 26th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, pages 687–701, 2021. [57] Wentao Zhang, Ziqi Yin, Zeang Sheng, Yang Li, Wen Ouyang, Xiaosen Li, Yangyu Tao, Zhi Yang, and Bin Cui. Graph attention multi-layer perceptron. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 4560–4570, 2022. [58] Zhekai Zhang, Hanrui Wang, Song Han, and William J. Dally. Sparch: Efficient architecture for sparse matrix multiplication. In 2020 IEEE International Symposium on High Performance Computer Architecture (HPCA), pages 261–274, 2020. doi: 10.1109/HPCA47549.2020.00030. [59] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. GaLore: Memory-efficient LLM training by gradient low-rank projection. In Forty-first International Conference on Machine Learning, 2024. [60] Michael Zhu and Suyog Gupta. To prune, or not to prune: exploring the efficacy of pruning for model compression. arXiv preprint arXiv:1710.01878, 2017.

13

Appendix A

DST Background

Sparse training of deep neural networks, also known as dynamic sparse training (DST), encompasses methods that maintain sparsity throughout training by starting from a sparse initialization and iteratively updating the network’s connectivity. During training, a fraction of weights is periodically removed and regrown, allowing the sparse structure to adapt to the learning dynamics. Recent work [13, 19, 34, 16, 11] provides a comprehensive overview of DST algorithms, theoretical foundations, and practical trade-offs, demonstrating their potential to scale deep learning efficiently without significant loss in performance.

B

Sparse Matrices & Sparse Operations

Sparse Matrices are the cornerstone of many DST and Pruning methods. Sparse Matrices are data structures containing compact data representations of a dense matrix with missing elements (generally represented by zeroes). Sparse matrix multiplication (SpMM) is a popular operation in many engineering applications. It is especially crucial when having large data structures. Experiments added early on by [15] in MATLAB explored having sparse vectors in each row. Assuming we have the following matrix: " # 0 0 3 A= 4 0 0 0 5 0 Sparse matrix storage data structure is essential in optimizing memory and computational resources. There are many structures for SpM representation. A prevalent structure is the Compressed Sparse Row (CSR) format. It is used for optimized row-wise storage, reducing memory overhead since memory is optimized to be accessed in contiguous segments [42]. A similar format is Compressed Sparse Column (CSC), which is optimized for column-wise storage. The Coordinate List format (COO) is a simple representation of the storage data structure of (SpM), where non-zero values and their coordinates are explicitly stored [3]. Examples on how each representation is shown in Table 7. Table 7: Sparse matrix representations in COO, CSR, and CSC formats. Format Data Indices Index Pointer COO CSR CSC

[3, 4, 5] [3, 4, 5] [4, 5, 3]

row = [0, 1, 2] col = [2, 0, 1] row = [1, 2, 0]

col = [2, 0, 1] indptr = [0, 1, 2, 3] indptr = [0, 1, 2, 3]

Another format is the Block Sparse Row (BSR). It extends CSR by grouping non-zero elements into blocks. Blocks are usually n × n in size, and if a block does not contain a non-zero element, it gets ignored. It is less efficient, memory-wise, than other formats but more computationally fast as it improves cache efficiency and parallelism [21]. There are many theoretical architectures for sparse accelerators, such as SpArch, which features a fixed outer product dataflow. Flexagon advances this by supporting multiple dataflows (inner product, outer product) within a reconfigurable architecture [58, 36]. Gamma [56] leverages Gustavson’s algorithm with specialized fiber cache and merger hardware to accelerate sparse matrix multiplication. Recent work by [43] demonstrates custom CUDA kernels for uniform sparse matrix multiplication in extreme classification scenarios, where each output neuron receives exactly the same number of inputs, achieving memory-efficient training with 500K labels on commodity GPUs.

C

Experiment within Condensed Sparsity Training Phase

The flexibility of the SNACK layer allows for fast integration into any dynamic sparse training pipeline. In this section, we provide results integrating the SNACK layer into a CNN-based architecture. 14

Table 8: Top-1 ImageNet test accuracy of ResNet-50 trained with RigL at high sparsities. RigL* is the method used by [23]; RigL (SNACK) is the same method but using SNACK instead of a dense layer with masked weights. Sparsity RigL* RigL (SNACK) (ours) Dense (%) ResNet-50 80 90 95 99

74.9 ± 0.1 72.8 ± 0.1 69.6 ± 0.1 51.4 ± 0.2

74.7 ± 0.2 72.8 ± 0.1 69.6 ± 0.2 51.4 ± 0.3

76.7 ±0.1

Table 9: Energy consumption (in kWh) of ResNet-50 trained with RigL at various sparsities. RigL* are from [23]; RigL (SNACK) replaces dense masked layers with truly sparse CUDA kernels. Sparsity RigL* RigL Dense % (SNACK) (ours) ResNet-50 80 90 95 99

5.170 ± 0.007 5.172 ± 0.005 5.156 ± 0.006 5.112 ± 0.007

5.217 ± 0.010 5.045 ± 0.010 5.033 ± 0.015 5.017 ± 0.020

4.26 ± 0.003 4.25 ± 0.004 4.24 ± 0.003 4.25 ± 0.003

Recently, [23] showed that DST can be effectively applied to computer vision tasks under a structured sparsity setup, achieving increased accuracy compared to its dense counterpart. The main motivation behind choosing to integrate SNACK with this work is to show how SNACK is integrable with any DST-based work and it is worthwhile to use it for energy cost reduction and restricted memory spaces environments. Starting from their reported results, we introduce RigL(SNACK) to train a ResNet-50 on ImageNet. In order to do so, we replaced the last linear layer of ResNet-50 used in the dynamic sparse training regime via RigL with our proposed SNACK layer for the forward pass and backpropagation. The layer replacement was directly performed in the code provided by [23]. We ran all the experiments five times and report the mean and standard deviation in Table 8. It is worth highlighting that the sparse matrix multiplication using SNACK produces identical gradients to sparse-via-mask, which are used in the training phase alongside RigL. This should be reflected in exact results in terms of accuracy between both implementations. However, as can be seen in Table 8, there are slight differences, due to float operations being handled differently. Furthermore, we investigate their energy consumption (see Table 9). We observed an improvement in energy consumption of RigL(SNACK) versus RigL(via mask) for sparsities above 90% (i.e 90%, 95%, and 99%). Of course, the small reduction in energy consumption is well calibrated with the expected impact of sparsifying only the last layer computation over the total training requirements in ResNet-50.

D

Performance benchmarks using various SpM sizes

D.1

Effects of different hidden layer sizes

In Figures 5, 6, 7, 8, 9, we show more granular memory, speed, and energy results similar to Figure 3 in the main text. We see the findings that SNACK is faster, more memory optimized, and more energy efficient hold across various input sizes. Our experiments showcase the efficiency of SNACK across various input size dimensions and a large hidden layer size (10000). Figures 5–9 report on the comparison between SNACK, Dense+Mask, and Dense Only layers on energy consumption, speed, and memory utilization. SNACK achieves drastic improvements scaling with sparsity levels. For example in Table 10, at extreme sparsity level (σ = 0.99 and input_size = 12500) SNACK achieves execution times as low as 45.3 µs (batch size 1), representing a 13.6× speedup over Dense (Only) (618.1 µs) and a 30.0× speedup over Dense+Mask (1358.3 µs). Energy consumption also decreases when we use SNACK, from 4.36 mWh at 90% sparsity (Dense+Mask) to 0.42 mWh at 99% sparsity using SNACK, representing 73% 15

less energy than Dense (Only) and 90% less energy than Dense+Mask at the same sparsity level. Memory usage scales favorably, with SNACK using only 839 MB at 99% sparsity for 12500 × 10000 dimensions, representing 89% memory reduction compared to Dense+Mask (7429 MB).

Table 10: Performance and environmental impact comparison of different implementations. All experiments use a 12500×10000 matrix with batch size of 1. Energy consumption (in mWh). Sparsity Params Memory CUDA τ Power Energy CO2 Type % (k) (MB) (µs) (W) (mWh) (mg) Dense+Mask SNACK (ours)

90 90

250,000 37,500

5995 1863

1358.3 163.1

102.8 100.2

4.36 0.39

1.96 0.18

Dense+Mask SNACK (ours)

95 95

250,000 18,750

6473 1023

1357.4 83.0

98.9 111.5

4.20 0.43

1.89 0.19

Dense+Mask SNACK (ours)

99 99

250,000 3,750

7429 839

1358.3 45.3

103.3 108.6

4.38 0.42

1.97 0.19

Dense (Only)

–

125,000

6951

618.6

100.4

1.55

0.70

Figure 5: SNACK performance when SNACK layer dimension is 2500x10000

Figure 6: SNACK performance when SNACK layer dimension is 5000x10000 16

Figure 7: SNACK performance when SNACK layer dimension is 7500x10000

Figure 8: SNACK performance when SNACK layer dimension is 10000x10000

Figure 9: SNACK performance when SNACK layer dimension is 12500x10000 17

D.2

Datacenter-class GPU sweep (NVIDIA A100)

To test whether the SNACK-COO advantage observed on the consumer-class RTX 3070 (Table 2) carries over to a datacenter-class accelerator, we re-ran the same SpMM micro-benchmark on an NVIDIA A100 across six leading-dimension sizes (500 through 10000, square matrices, 95% sparsity, batch 1, mean over 1,000 multiplications). Table 11 reports the resulting CUDA times. SNACKCOO remains the fastest sparse method at every size, with the constant-time SpMM_CSR(Cupy) path closing to within 8–13% at {8500, 10000}. At and above 5000×5000 SNACK-COO also beats the dense baseline outright, mirroring the RTX 3070 result. We exclude SpMM_BSR(Jax) because the JAX/CUDA toolchain on the A100 host could not be brought up against the same library versions used in Table 2; all other rows are matched to that table.

Table 11: A100 replication of Table 2. CUDA time (ms) for sparse matrix multiplication implementations on an NVIDIA A100, sparsity 95%, batch size 1, square matrices, mean of 1,000 multiplications. Best per-column value among sparse methods is bold; the DenseMM baseline is also bold where it remains the global best. † Simulates sparsity. Method

is sparse

SpMM_CSR(Torch) SpMM_CSR(Cupy) SpMM_COO(Torch) SNACK-COO (ours) SpMM_COO(Cupy) SpMM_CSR(cuSPARSE) SpMM_COO(cuSPARSE) SpMM_CSR(Sputnik DL-opt) SpMM_COO(Sputnik) FlashSparse DenseMM+Mask

Yes

DenseMM (baseline)

no

†

CUDA (ms) 500

1000

5000

7500

8500

10000

24.83 11.35 24.55 2.72 166.27 25.34 25.03 3.05 2.96 27.56 4.38

25.83 11.05 26.59 2.85 175.73 25.41 27.84 3.07 2.99 37.08 4.44

215.16 11.63 321.53 6.24 1153.05 215.40 305.23 8.10 8.01 329.91 46.68

443.18 11.45 693.69 8.12 2520.67 443.66 700.28 14.12 15.69 706.73 100.90

576.99 14.47 901.41 13.39 3410.45 577.62 890.00 17.57 17.52 835.93 130.43

811.52 15.68 1260.55 13.93 4893.86 812.43 1313.55 24.01 26.28 1076.32 174.30

2.20

2.30

23.19

51.96

66.45

89.99

Cross-GPU consistency at batch 1. The A100 numbers in Table 11 preserve the qualitative ordering established on the RTX 3070 in Table 2: SNACK-COO is the only sparse method that beats Dense+Mask across every size, the cuSPARSE / PyTorch / FlashSparse rows trail by an order of magnitude at small leading dimension, and the constant-time SpMM_CSR(Cupy) path remains a strong outlier at the largest size only. At 5000×5000 SNACK-COO is 7.5× faster than Dense+Mask on the A100 (vs. 7.0× on the RTX 3070) and 3.7× faster than the dense baseline (vs. 3.5× on the RTX 3070), confirming that the batch-to-SM mapping (Section 3.2) is not specific to consumer-class hardware. Batch 4: the regime crossover predicted by Section 3.2. Table 12 repeats the A100 sweep at batch 4 over the same method set as Table 11 (only SpMM_BSR(Jax) is omitted, for the JAX/CUDA toolchain reason given above). SNACK-COO retains its lead at the two smallest sizes (500 and 1000, where n = 4 ≪ m still holds strongly) but is overtaken from 5000×5000 upward by Sputnik, with SpMM_COO(Sputnik) winning at {5000, 8500} and SpMM_CSR(Sputnik DL-opt) winning at {7500, 10000}; the two Sputnik variants finish within ±10% of each other at every size (mirroring their behaviour on the RTX 3070 in Table 2), and the per-column-best at 10000×10000 is 21.44 ms versus SNACK-COO’s 65.54 ms. This is exactly the crossover predicted by the parallel-efficiency analysis in Appendix F: the SNACK-COO advantage is ϵSNACK-COO > ϵstandard for n ≪ m, and as the p p batch grows or the matrix shrinks Sputnik’s structured-sparse SpMM, which is tuned for moderatebatch / moderate-sparsity, becomes the better choice. This is the operational reason the SNACK layer keeps the SpMM backend swappable (Section 3.5) and why Appendix G trains GAMLP through Sputnik but serves it through SNACK-COO. Note that even at this less-favorable batch size, SNACK-COO still beats Dense+Mask on every size and beats the dense baseline at sizes ≥ 5000, so the kernel remains a valid choice whenever Sputnik is unavailable. 18

Table 12: A100 SpMM benchmark at batch size 4. CUDA time (ms) for sparse matrix multiplication implementations on an NVIDIA A100, sparsity 95%, batch size 4, square matrices, mean of 1,000 multiplications. Same method set as Table 11; SpMM_BSR(Jax) is omitted because the JAX/CUDA toolchain on the A100 host could not be brought up at the matched library versions. Best per-column value among sparse methods is bold; the DenseMM baseline is also bold where it remains the global best. † Simulates sparsity. Method

SpMM_CSR(Torch) SpMM_CSR(Cupy) SpMM_COO(Torch) SNACK-COO (ours) SpMM_COO(Cupy) SpMM_CSR(cuSPARSE) SpMM_COO(cuSPARSE) SpMM_CSR(Sputnik DL-opt) SpMM_COO(Sputnik) FlashSparse DenseMM+Mask DenseMM (baseline)

E

is sparse

Yes

†

no

CUDA (ms) 500

1000

5000

7500

8500

10000

27.46 12.15 27.37 2.98 180.59 27.49 27.27 3.36 3.26 30.04 4.66

28.25 12.19 30.03 3.05 185.85 28.28 29.93 3.33 3.24 39.86 4.68

248.35 12.81 342.17 21.12 1267.72 248.62 341.93 8.24 8.15 369.50 50.20

528.04 21.32 767.56 36.90 2800.83 528.62 767.45 12.29 13.66 774.58 110.71

675.79 27.07 979.77 52.17 3618.27 676.53 981.85 16.08 16.03 922.03 139.38

950.01 37.01 1392.60 65.54 5160.39 951.06 1391.79 21.44 23.46 1140.90 194.47

2.34

2.35

25.23

55.63

70.04

97.72

Extended Related Works

This appendix expands the discussion in Section 2 with material that we omit from the main text for brevity. Early pruning, lottery tickets, and pruning-at-init. Early sparsity work focused on post-training pruning, notably Optimal Brain Damage (OBD) [24] and Optimal Brain Surgeon [18], which removed weights with minimal impact on performance. The Lottery Ticket Hypothesis [9] showed that dense models contain sparse sub-networks capable of learning independently when trained from their original initialization. A complementary line removes weights before training: SNIP [25] prunes by single-shot connection sensitivity, and SynFlow [47] iteratively conserves synaptic flow without using any data; both establish a sparse mask once and then train dense-on-mask, in contrast to the prune-and-regrow loop of DST. Modern pruning techniques [10, 60, 27] compress models for resource-constrained deployment and typically apply binary masks over dense weights to simulate sparsity [60]. SparseGPT [10] enables one-shot unstructured pruning of large language models with minimal performance degradation, without retraining. Hardware-friendly transformations such as SPDF [48] combine unstructured sparse pre-training with dense fine-tuning, achieving a 2.5× FLOP reduction. Sparse-IFT [49] proposes drop-in transformations that preserve dense FLOP counts while improving accuracy, but require specialized hardware kernels. The structured/unstructured distinction [52] is central: structured patterns are regular and hardware-friendly; unstructured sparsity is more flexible but harder to accelerate. Dynamic sparse training history. Beyond SET [33] and RigL [8], two earlier works are central to the DST lineage that SNACK plugs into. Deep rewiring [4] formalised topology evolution as a sampling process governed by a posterior over connectivities and showed that very sparse networks could be trained from scratch on small benchmarks. Sparse-from-Scratch [7] introduced momentumbased regrowth and demonstrated, on the available hardware of the time, that DST can match dense baselines without any dense pretraining stage. SNACK is agnostic to the specific prune-and-regrow rule: any of SET, RigL, deep rewiring, or sparse-from-scratch can drive its uniform connection initializer. Memory-efficient LLM training. Two lines of work attack the LLM training memory wall along axes orthogonal to a truly sparse weight tensor. [55] (MEST) keeps the weight tensor dense at compute time but stores optimizer states only for active parameters, plus a soft memory-economic schedule on the regrowth set; the storage savings stack with SNACK’s, since SNACK additionally removes the dense weight and gradient tensors themselves. [59] (GaLore) projects gradients onto a 19

low-rank subspace before applying Adam, reducing optimizer-state memory at the cost of a periodic SVD; this is orthogonal to SNACK because the savings come from the gradient/optimizer side rather than the weight side, and the two could be combined (a sparse weight tensor with low-rank optimizer states) without architectural conflict. At LLM-pretraining scale, [1] pretrains high-sparsity Llama variants and pairs them with a tuned sparse-deployment stack; their kernel choice (CSR with structured 2:4 fallbacks) targets the moderate-batch deployment regime where Sputnik dominates, whereas SNACK-COO targets the small-batch regime they leave for further work. Sparse formats and earlier GPU SpMM kernels. The CSR format dominates GPU-based SpMM implementations and leverages warp-level threads to distribute row-wise computation [3]. Block Sparse Row (BSR) extends CSR by grouping non-zeros into n × n blocks: it behaves similarly to CSR for small block sizes and trades memory for cache efficiency at larger block sizes [32]. Customized kernel fusion and load-balancing techniques further improved SpMM on GPUs [54]. Energy efficiency in SpMM/SpMV is closely tied to reducing computational load: the less time required per multiplication, the less energy is required to power the GPU/TPU/CPU. Block-sparse Transformer pretraining (BLaST). BLaST [39] couples a fused Triton blockSpMM kernel (BCSC layout, Tensor-Core MMA-aligned blk_M×blk_N tiles) with a block pruneand-grow training loop and a second-order oLLM pruner inspired by Optimal Brain Surgeon [18]. It reaches up to 95% MLP sparsity at a median ≤ 2.25% accuracy drop and reports a 2.2× end-to-end inference speedup for Llama-3.2 across 16 GH200 GPUs, with up to 4.45× inference-memory reduction. BLaST is therefore complementary to SNACK on two axes. First, it requires structured (block) sparsity: the kernel’s speedup hinges on the BCSC block size matching MMA fragment shapes, whereas SNACK targets the unstructured patterns produced by SET/RigL-style DST. Second, BLaST is not truly sparse end-to-end: during training the dense weight and gradient tensors are kept intact and only the pruned forward/backward path runs through BCSC, so the dense-tensor materialization that SNACK explicitly avoids is still paid at training time (and amortized only as si grows). The two systems thus occupy adjacent regimes – BLaST for block-sparse, multi-GPU LLM pretraining at moderate batch sizes, SNACK for unstructured, small-batch, high-sparsity training and inference on commodity GPUs. Truly sparse implementations on CPU and IPU. Using sparse data structures in SciPy and Cython, [29] implemented and trained with DST a truly sparse Multi-Layer Perceptron (SET-MLP) with one million neurons on a single CPU core, with single-core execution being the practical limitation of their approach. [6] parallelized this approach across multiple CPU cores but reported a hard dependence on CPU. [2] adjusted the approach to autoencoders and showed that a truly sparse implementation on one CPU core can outperform the dense GPU baseline in run-time and energy efficiency, owing to DST’s quadratic parameter reduction. PopSparse [26] leverages IPU hardware to deliver 90%-sparsity speedups. SNACK is, in contrast, the first such truly sparse PyTorch building block on standard GPUs, with no hardware-specific sparsity support required. Sparse accelerators and adjacent kernels. Several theoretical sparse accelerators have been proposed, including SpArch’s fixed outer-product dataflow and Flexagon’s reconfigurable multidataflow architecture [58, 36]. Gamma [56] leverages Gustavson’s algorithm with specialized fiber cache and merger hardware to accelerate sparse matrix multiplication. [43] demonstrate custom CUDA kernels for uniform sparse matrix multiplication in extreme classification, where each output neuron receives the same number of inputs, enabling memory-efficient training with 500K labels on commodity GPUs. These directions are complementary to SNACK, which targets unstructured sparsity on commodity hardware without architectural changes.

F

SNACK Algorithmic Details

This appendix is the technical companion to Section 3: a parallel-efficiency proof for SNACKCOO (referenced from Section 3.2); the geometric-distribution analysis of the uniform initializer’s expected cost (referenced from Algorithm 2); the full chain-rule derivation of SNACK’s sparse forward and backward, including the mini-batch generalization (referenced from Section 3.4); the CUDA/autograd implementation plumbing (referenced from Section 3.5); and a closing summary of how end-to-end true sparsity is preserved across the entire pipeline. 20

Parallel-efficiency analysis for SNACK-COO (Section 3.2). The standard COO scheme creates m parallel row tasks with the batch dimension as an outer loop. When the layer is much wider than the batch, m > |SMs| · warpsSM , many SMs sit idle. SNACK-COO replaces this with a 2D block grid G = {(Nb , n)} where Nb = ⌈nnz/T ⌉, so the task pool is Nb · n blocks (i.e., nnz · n thread tasks). The work distribution per SM becomes ⌈nnz/T ⌉ · n Nb · n = , |SMs| |SMs| which scales with the non-zero workload rather than with the matrix height m. Within a block, thread τ in block (i, b) reads the COO triple at index k = iT + τ (k < nnz) and atomically updates CAr [k],b += Av [k] · BAc [k],b . The time complexity remains O(nnz · n), but the parallel efficiency Tseq ϵp = , ϵSNACK-COO > ϵstandard p p |SMs| · Tpar SNACK-COO standard improves because Tpar < Tpar at small n. Atomic adds are required when multiple non-zeros write to the same output row, but the batch-wise block organization confines contention to one output column per block, reducing it relative to the row-wise scheme. Threads within a block process consecutive non-zeros for the same column b, which also improves cache locality on B.

Uniform Initializer cost (Algorithm 2). Algorithm 2 draws nnz = (1 − η) mk unique pairs by rejection: at step i, the probability of drawing a fresh pair is pi = 1 − i/(mk), giving a geometric expected number of trials of E[1/pi ] = mk/(mk − i). Summing over the nnz steps, Z (1−η)mk (1−η)mk X mk mk ≈ di = −mk ln η, η ∈ (0, 1), E[trials] = mk − i mk −i 0 i=0 Pb where the harmonic approximation i=a 1/i ≈ ln(b/a) is used. The initializer is therefore essentially linear in nnz, and the rejection penalty ζ = E[trials]/nnz = − ln η/(1 − η) is bounded: ζ → 1 as η → 1 (very sparse, almost no collisions), and ζ grows only logarithmically as η shrinks toward the dense regime. Full gradient derivation (Section 3.4). For a SNACK layer with sparse W ∈ Rm×n , input X ∈ Rn×b , and pre-activation A = W X ∈ Rm×b , the chain rule gives  ∂A ⊤∂L ∂L ∂L ∂A ∂L = = ∇A · X ⊤ , = = W ⊤ ∇A, ∂W ∂A ∂W ∂X ∂X ∂A where ∇A = ∂L/∂A ∈ Rm×b . The input-gradient term W ⊤ ∇A is sparse-times-dense and reuses SNACK-COO directly: W ⊤ inherits W ’s sparsity by swapping row and column index arrays, introducing no new non-zeros. Expanding the weight-gradient product entry-wise, b  ∂L  X = ∇Ai,t Xj,t , ∂W i,j t=1 Pb which is a sum of b rank-1 outer products: ∂L/∂W = t=1 ∇A(t) (X (t) )⊤ . A dense backward evaluates all mn entries of this matrix; SNACK evaluates only (i, j) ∈ supp(W ), since entries outside the support are zeroed by the sparsity mask and contribute nothing to the SGD update. This drops the per-step cost from O(mnb) to O(nnz · b). Activation functions act element-wise on A and are unaffected by W ’s sparsity [17]. Implementation plumbing (Section 3.5). The SNACK-COO kernel parallelizes over sparse-matrix cells rather than over batch entries, which saturates the SMs from the non-zero workload alone and is what makes the kernel favorable at small n. Within a block we issue __syncthreads() before atomic writes to avoid races, and we use T = 256 threads per block as a default that balances occupancy against scheduling overhead. The CUDA launch is wrapped in a C++/Torch extension that converts the COO triple from PyTorch tensors to flat float/int arrays and back. The sparse outer-product gradient is a separate one-dimensional P CUDA kernel that loops over the nnz entries of supp(W ) and writes back the accumulated t ∇Ai,t Xj,t for each entry. The PyTorch nn.Module wraps both kernels behind a custom torch.autograd.Function, accepts a pluggable Initializer (default: the C++ uniform initializer of Algorithm 2), and stores weights as the COO triple (Av , Ar , Ac ) on-device for the entire training run. 21

End-to-end sparsity. Across the full pipeline, SNACK preserves true sparsity at every step: (i) the forward pass runs SpMM (SNACK-COO by default, or any pluggable backend) directly on the COO triple and never expands W to dense; (ii) the backward pass uses the sparse outer product over supp(W ) and never instantiates the m × n weight-gradient matrix; (iii) DST topology updates use the uniform initializer of Algorithm 2, which draws non-zeros directly with no dense m × k probability tensor. Consequently no step of training or inference allocates a dense m × n weight or gradient tensor, which is the property the memory and energy measurements of Sections 4–5 exploit.

G

GAMLP Graph-Workload Training Setup

This appendix details the training recipe behind the GAMLP-shaped sparse-inference results of Table 6 (Section 5). Inference is reported at batch = 1 on the trained sparse weights; training itself is run at the natural mini-batch size for graph workloads, b = 4096, where the leading dimension is no longer “small” relative to the layer width and SNACK-COO’s batch-to-SM mapping ceases to be the right fit. We therefore train through the SNACK layer with a Sputnik SpMM backend, which dominates at moderate sparsity and large batches [14], and only swap to SNACK-COO at inference time. This is exactly the backend-agnostic story of Section 3: the SNACK layer is fixed (sparse outer-product gradient, uniform initializer, COO weights), and the SpMM backend is chosen per regime. Model and data. We use R_GAMLP_RLU [57] on the OGB ogbn-products node-classification benchmark [20], with hidden dimension 512, 5 feature-propagation hops, and 2/2/2 layers in the three GAMLP stages (default R_GAMLP_RLU configuration). Pre-processing, residual connections, and batch norm are enabled; dropout, attention dropout, and label dropout are 0.5, with α = 0.5 for the residual mix. The sparsifiable substrate is the 5000×5000 feature-head projection used at inference; embeddings, label propagation, and the classification head are kept dense. Optimization. Training uses the Adam optimizer at learning rate 10−3 with cross-entropy loss for 40 epochs at batch size 4096 on a single GPU. We use the same dense GAMLP hyperparameters as the public reference and only swap the feature-head projection from a dense nn.Linear to a SNACK module with the Sputnik backend; everything else (input pipeline, scheduler, weight decay, seed) follows the upstream GAMLP defaults. Dense and Dense+Mask baselines for the same backbone are trained under an identical recipe, the only difference being the feature-head implementation. Sparsity schedule and DST. Each run is initialized at its target sparsity η ∈ {0.9, 0.95, 0.99} and trained with a magnitude-based prune-and-grow update: every 5 epochs (until epoch 25), we prune the ζ = 5% smallest-magnitude non-zeros and regrow an equal number of unique (i, j) coordinates with the uniform connection initializer of Algorithm 2. After epoch 25 the topology is frozen and only weight values continue to update, which gives the optimizer a stable support to converge on. The inference results in Table 6 use the resulting frozen SNACK checkpoints, loaded directly into the SNACK-COO inference path with no re-pruning. Why Sputnik for training and SNACK-COO for inference. SNACK-COO’s parallel-efficiency analysis (Appendix F) shows the kernel’s advantage relies on n ≪ m, i.e. batch much smaller than the layer width. At training-time batch 4096 that condition is reversed and Sputnik’s structured-sparse SpMM is the better fit; at inference-time batch 1 the condition is restored and SNACK-COO is the better fit. Both paths use the same SNACK layer, the same COO weights, and the same outer-product gradient; only the SpMM backend changes. This is the concrete realization of the backend-agnostic claim made in Section 3.5 and explains why the same module can deliver Table 4’s training memory savings under Sputnik and Table 6’s inference speedups under SNACK-COO. The full training script and reproducibility commands are available at https://github.com/JafarBadour/Snack-SNN under experiments/gamlp_snack_benchmark/. Empirical validation of the per-regime backend choice (Table 13). The GAMLP training measurements directly back the argument above: at the natural graph-training batch 4096, the layerwidth-to-batch ratio violates SNACK-COO’s n ≪ m assumption, and SNACK-COO accordingly pays a ≈ 9.4× per-step latency tax (160.39 vs. 17.11 ms) and a ≈ 10× per-step energy tax (17,177 vs. 1,700 mJ) over Dense. This is precisely the regime the paragraph above flags as Sputnik’s 22

Table 13: GAMLP-style dynamic sparse training on ogbn-products at sparsity η = 0.9, batch size 4096, 40 epochs on a single GPU. We report best and final validation accuracy, mean per-step training latency (ms), per-step energy (mJ), and peak GPU allocation (MB). The SNACK-COO (ours) row reflects training executed through the SNACK layer with the SNACK-COO backend, included to expose the n ≪ m regime mismatch at b = 4096 predicted in Appendix F; under the recommended per-regime backend choice the SNACK layer is trained with the Sputnik backend and only swapped to SNACK-COO at inference time (Table 6). Best per-column value is in bold. Method Dense Dense+Mask SNACK-COO (ours)

Best val acc

Final val acc

Latency_step_ms

Energy_step_mJ

Peak Mem (MB)

0.9232 0.9174 0.9172

0.9191 0.9140 0.9101

17.11 17.34 160.39

1700.3 1694.3 17177.1

6823 6840 6784

territory, which is why the recommended per-regime configuration trains the SNACK layer through the Sputnik backend at b = 4096 (whose moderate-sparsity, large-batch dominance is documented in [14] and reproduced in our A100 sweeps in Appendix D.2, where the Sputnik variants overtake SNACK-COO from 5000×5000 upward at batch 4) and only swaps to the SNACK-COO backend at the inference-time batch 1 regime, which is the configuration whose memory and latency we report in Table 6. Peak training memory is already minimally smaller for SNACK-COO (6784 vs. 6823 MB for Dense); the modest gap is consistent with the architectural ceiling noted in Section 5: only the 5000×5000 feature-head projection is sparsified while the rest of the GAMLP backbone (multi-hop feature propagation buffers, embeddings, classifier head) remains dense. Validation accuracy is essentially matched: SNACK-COO reaches 0.9101 final / 0.9172 best vs. Dense+Mask’s 0.9140 final / 0.9174 best, a sub-half-point gap that is independent of the SpMM backend, since both SNACK-COO and SNACK(Sputnik) share the same uniform sparse initializer, sparse outer-product gradient, and magnitude prune-and-grow schedule, so any kernel-level swap is mathematically a reordering and will not move accuracy materially. The remaining accuracy headroom is a DST-recipe question (initialization, prune-and-grow schedule, dropout/regularization), orthogonal to the kernel contribution of this paper.

23

Related documents

• Beyond Masked Sparsity: SNACK Enables Truly Sparse Neural Networks on GPU#1199367
• Penrose: Privacy-Preserving Performance Profiling of In-The-Wild GPUs#1062797
• Comparison between Supervised and Unsupervised Learning in Deep Unfolded Sparse Signal Recovery#160946
• When More Memory Fails: Hidden Feasibility Holes in Dynamic Tensor Rematerialization#1325265
• One Global Beam Across Many GPUs: High-Throughput Beam Search at Billion-Record Frontier Scale#1217565