SPDP unifies unstructured static pruning (Wanda) with input-adaptive dynamic pruning (TEAL) behind a single compressed representation — the Tiled-Column-wise Bitmap Compressed (Tiled-CBC) format — served by two phase-specific GPU kernels:
| Phase | Kernel | Op |
|---|---|---|
Decode (N == 1) |
CUDA-core spMspV with Hybrid Activation-aware Dynamic Shared-Memory Bitmap Decoding (HAD-SMBD) | ops.spdp_v5 |
Prefill (N > 1) |
Tensor-Core SpMM | ops.spdp_v3 |
The sparse layer uses the decode kernel by default and routes to SpMM automatically when N > 1
(override with config.spdp_kernel). Baselines included: SpInfer, Flash-LLM, SparTA, Sputnik,
cuSPARSE, cuBLAS.
Requires Linux, an NVIDIA GPU with CUDA 12.x, Python 3.11, and access to meta-llama/Llama-2-7b-hf.
git clone --recurse-submodules https://github.com/AIDASLab/SPDP.git && cd SPDP
bash install_one_shot.sh
source .venv/bin/activate && source Init_spdp.shThis builds the uv venv, initialises the submodules (Sputnik, glog), builds libSpMM_API.so and the
kernel benchmarks, then installs extension-pytorch and transformers-spdp.
Re-enter an existing build with source activate_one_shot.sh.
cd extension-pytorch/csrc/spdp-kernel/kernel_benchmark
source test_env && bash make_all.sh
bash run_spdp_v5_only.sh # SPDP decode (spMspV)
bash run_all_main_prefill.sh # SPDP prefill (SpMM)
bash run_cusparse.sh; bash run_sputnik.sh; bash run_sparta.sh; bash run_flash_llm_only.shThe kernels are exposed as native PyTorch ops via spmm_package (torch.ops.spmm_package):
import torch, spmm_package
from spmm_package import ops
# Decode — spMspV, N == 1
y = ops.spdp_v5(compressed_val, tile_offsets, col_offset_packed, col_id_packed,
threshold, bitmap, x, M, K, split_k)
# Prefill — Tensor-Core SpMM, N > 1
y = ops.spdp_v3(compressed_val, tile_offsets, col_info, threshold,
bitmap, x, M, K, N, split_k, reduction_workspace)
# v5 column metadata, derived once from the Tiled-CBC buffers
col_offset_packed, col_id_packed = ops.build_colmeta_v4(tile_offsets, col_info)
# SpInfer baseline
y = ops.spinfer(A, compressed_val, tile_offsets, tile_offsets_median,
bitmap, max_nnz_intile, x, M, N, K, split_k, reduction_workspace)All ops are registered with fake/meta kernels, so they are torch.compile-traceable.
transformers-spdp is a HuggingFace Transformers fork whose LlamaSPDPLinear runs the kernels
above. The flow is prune → convert to Tiled-CBC → run.
1) Prune — static with wanda/, then collect TEAL activation histograms in TEAL/.
2) Convert a pruned checkpoint into the SPDP (Tiled-CBC) format:
cd transformers-spdp
python -m spdp_utils.scripts.run_edit_model --format spdp --cached \
--repo $SPDP_HOME/wanda/out/llama2_7b/unstructured/wanda0.45 \
--target_repo meta-llama/llama-2-7b-hf-wanda0.45 \
--out_dir $SPDP_HOME/transformers-spdp/spdp_utils/This rewrites config.json so the HF auto-classes resolve to LlamaSparseForCausalLM
(use --format spinfer for the SpInfer baseline).
3) Run the sparse model:
python spdp_utils/scripts/run_spdp_model_v5.py \
--model-id $SPDP_HOME/transformers-spdp/spdp_utils/meta-llama/llama-2-7b-hf-wanda0.45 \
--device cuda --dtype float16 --attn flash_attention_2 \
--input-tokens 32 --max-new-tokens 256 --warmup 5 --repeat 15 \
--dynamic-pruning True --dp-ratio 0.5 --sp-ratio 0.1 \
--histogram-path $SPDP_HOME/TEAL/models/llama2_7b/unstructured/wanda0.45/histogram/ \
--compression-format spdp --outdir ./resultsOther runners: run_spdp_model_v3.py (SpMM path), run_sparse_model.py (SpInfer),
run_dense_model.py (dense). TTFT/TPOT results are written as JSON to --outdir.
See transformers-spdp/spdp_utils/README.md for details.
extension-pytorch/ SPDP CUDA kernels + PyTorch bindings (spmm_package)
└ csrc/spdp-kernel/ spdp_kernel.cuh (spMspV + SpMM), SpMM_API.cu, utility.cuh
transformers-spdp/ HF Transformers fork with the SPDP execution path
└ spdp_utils/ checkpoint conversion + model runners
TEAL/ Dynamic pruning (activation histograms) + PPL evaluation
wanda/ Static (unstructured) pruning
See NOTICE for full details.
- FastGEMV (MIT) — SPDP's spMspV decode kernels are derived from its half-precision GEMV kernels; SPDP adds Tiled-CBC, HAD-SMBD and Split-K.
- SpInfer (Apache-2.0) — baseline, and the basis of the Tensor-Core data-movement primitives.
- Flash-LLM (Apache-2.0) — SpInfer's ancestor (itself extending NVIDIA CUTLASS); baseline.
- Also vendored: TEAL, Wanda, HuggingFace Transformers, Sputnik, glog.
@article{kim2026spdp,
author = {Jinhyeok Kim and Yejoon Lee and Jaeyoung Do},
title = {Unified Static--Dynamic Pruning for Efficient LLM Inference},
journal = {Proceedings of the VLDB Endowment},
volume = {19},
number = {11},
year = {2026}
}