$ show vllm-project/compressed-tensors
compressed-tensors is a production-grade Python library for compressing and quantizing large neural network models, primarily targeting LLMs. It extends the safetensors format to store and manage compressed tensors, supporting a wide range of quantization schemes (INT4/INT8/FP8/NVFP4/MXFP4/MXFP8), sparsity patterns, and rotation-based transforms (Hadamard, random matrix). The library is tightly integrated with the vLLM inference engine and Hugging Face transformers, providing a seamless path from model compression to deployment. Its architecture is modular: a registry pattern for compressors and transforms, a quantization lifecycle with calibration and forward hooks, a flexible offloading subsystem (CPU/disk/distributed) for running models beyond VRAM limits, and checkpoint conversion pipelines for formats like AutoAWQ and ModelOpt NVFP4. While primarily a software library, it also serves as a research platform for novel quantization techniques, with benchmarks and experimental features like MoE expert pruning and KV cache quantization. The codebase is well-structured with extensive tests, CI/CD pipelines, and active maintenance, though it exhibits some coupling to vLLM internals and experimental components that may not be stable.
The LLM identified 5 files worth highlighting for reuse:
Implements a generic RegistryMixin class that allows subclasses to register themselves by name with automatic name standardization, aliasing, and subclass validation. Supports module-level import for lazy registration. This pattern is used throughout the codebase for compressors, transforms, and configs, and can be extracted as a standalone reusable registry for any plugin-based system.
Defines the OffloadCache abstract base class, a MutableMapping that transparently offloads tensors to a backing store (CPU, disk, device) on set and onloads on get. Includes context managers for disabling offloading/onloading. This pattern is reusable for any system needing automatic memory management of large objects, independent of the specific storage backend.
Provides the ImplBackend class that implements a priority-based dispatch system for selecting optimized implementations (e.g., Triton vs. CPU fallback). It allows registering implementations with a priority and automatically selects the highest-priority available backend. This is a clean, reusable pattern for abstracting hardware-specific code paths.
Contains functions like match_named_modules, match_named_parameters, and match_quantizable_tensors that provide flexible module targeting using regex patterns, class names, and fused module mappings. Supports narrow matching and internal module exclusion. This is a reusable utility for any PyTorch model surgery task, independent of quantization.
Implements ParameterizedDefaultDict, a dict subclass that calls a factory function with the missing key to generate default values. This is a simple but powerful pattern for lazy initialization of per-key resources, used internally for caching transform weights. It can be extracted and reused in any Python project needing parameterized defaults.