Skip to content

AIDASLab/SPDP

Repository files navigation

SPDP

Unified Static–Dynamic Pruning for Efficient LLM Inference

PVLDB Vol. 19, No. 11 (VLDB 2026)

arXiv (coming soon) Project Page Code


📌 Overview

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.


⚙️ Installation

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.sh

This 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.


🚀 Kernel Benchmarks

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.sh

🧩 PyTorch Binding

The 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.


🤖 Model

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 ./results

Other 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.


📂 Layout

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

🙏 Attribution

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.

📖 Citation

@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}
}

About

SPDP: Unified Static-Dynamic Pruning for Efficient LLM Inference (PVLDB Vol.19 No.11)

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors