entropack 0.1.0__py3-none-any.whl

entropack/__init__.py

@@ -0,0 +1,15 @@
+from importlib import metadata
+
+try:
+    __version__ = metadata.version("entropack")
+except metadata.PackageNotFoundError:
+    __version__ = "0+unknown"
+
+from .compression import CompressedTensor, compress, decompress
+from .linear import CompressedFP8Linear, CompressedINT8Linear, CompressedLinear
+from .schemes import CompressionConfig, DFloat11Config, LatticeRANSConfig, RawConfig, TileANSConfig
+
+__all__ = [
+    "CompressedFP8Linear", "CompressedINT8Linear", "CompressedLinear", "CompressedTensor", "CompressionConfig",
+    "DFloat11Config", "LatticeRANSConfig", "RawConfig", "TileANSConfig", "__version__", "compress", "decompress",
+]

entropack/backends/__init__.py

@@ -0,0 +1,19 @@
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+
+from . import cuda
+
+
+@dataclass(frozen=True)
+class Backend:
+    name: str
+    priority: int = 0
+    probe: Callable[[], str | None] | None = None
+    device: Callable[[], torch.device] | None = None
+
+
+BACKENDS: tuple[Backend, ...] = (
+    Backend("cuda", priority=100, probe=cuda.probe, device=cuda.runs_on), Backend("eager", priority=0),
+)

entropack/backends/cuda/__init__.py

@@ -0,0 +1,16 @@
+import torch
+
+
+def probe() -> str | None:
+    if not torch.cuda.is_available():
+        return "torch.cuda.is_available() is False (no NVIDIA device / driver)"
+    try:
+        import cupy
+    except ImportError as e:
+        return (f"cupy is not installed ({e}); install a matching build, e.g. "
+                "`pip install cupy-cuda13x` for CUDA 13 or `pip install cupy-cuda12x` for CUDA 12")
+    return None
+
+
+def runs_on() -> torch.device:
+    return torch.device("cuda", torch.cuda.current_device())

entropack/backends/cuda/device.py

@@ -0,0 +1,96 @@
+from dataclasses import dataclass
+
+import cupy
+import torch
+
+from .kernels import device_index
+
+_FALLBACK_WARP_SIZE = 32
+
+_DEFAULT_GRID_WAVES = 8
+
+__all__ = ["DeviceCaps", "caps", "resolve_threads", "validate_threads_per_block"]
+
+
+def _shared_optin(index: int, properties) -> int:
+    optin = getattr(properties, "shared_memory_per_block_optin", 0)
+    if optin:
+        return int(optin)
+    with cupy.cuda.Device(index):
+        attributes = cupy.cuda.Device(index).attributes
+    return int(attributes["MaxSharedMemoryPerBlockOptin"])
+
+
+@dataclass(frozen=True)
+class DeviceCaps:
+    index: int
+    name: str
+    compute_capability: tuple[int, int]
+    sm_count: int
+    warp_size: int
+    max_threads_per_block: int
+    threads_per_sm: int
+    regs_per_sm: int
+    shared_per_block: int
+    shared_optin: int
+    shared_per_sm: int
+    l2_bytes: int
+
+    def blocks_per_sm(self, threads_per_block: int, shared_per_block: int = 0, regs_per_thread: int = 0) -> int:
+        blocks = self.threads_per_sm // threads_per_block
+        if shared_per_block > 0:
+            blocks = min(blocks, self.shared_per_sm // shared_per_block)
+        if regs_per_thread > 0:
+            blocks = min(blocks, self.regs_per_sm // (regs_per_thread * threads_per_block))
+        return max(1, blocks)
+
+    def resident_blocks(self, threads_per_block: int, shared_per_block: int = 0, regs_per_thread: int = 0) -> int:
+        return self.sm_count * self.blocks_per_sm(threads_per_block, shared_per_block, regs_per_thread)
+
+    def grid(self, wanted: int, threads_per_block: int, shared_per_block: int = 0, waves: int = _DEFAULT_GRID_WAVES) -> int:
+        limit = self.resident_blocks(threads_per_block, shared_per_block) * waves
+        return max(1, min(wanted, limit))
+
+    def shared_limit(self, static_slack: int = 0) -> int:
+        return max(0, min(self.shared_optin, self.shared_per_sm - static_slack))
+
+    def threads_per_block(self, wanted: int) -> int:
+        usable = min(wanted, self.max_threads_per_block)
+        usable -= usable % self.warp_size
+        return max(self.warp_size, usable)
+
+
+_caps_cache: dict[int, DeviceCaps] = {}
+
+
+def caps(device=None) -> DeviceCaps:
+    index = device_index(device)
+    cached = _caps_cache.get(index)
+    if cached is not None:
+        return cached
+    properties = torch.cuda.get_device_properties(index)
+    queried = DeviceCaps(
+        index=index, name=properties.name, compute_capability=(properties.major, properties.minor),
+        sm_count=properties.multi_processor_count, warp_size=getattr(properties, "warp_size", 0) or _FALLBACK_WARP_SIZE,
+        max_threads_per_block=properties.max_threads_per_block, threads_per_sm=properties.max_threads_per_multi_processor,
+        regs_per_sm=properties.regs_per_multiprocessor, shared_per_block=properties.shared_memory_per_block,
+        shared_optin=_shared_optin(index, properties), shared_per_sm=properties.shared_memory_per_multiprocessor,
+        l2_bytes=getattr(properties, "L2_cache_size", 0),
+    )
+    _caps_cache[index] = queried
+    return queried
+
+
+def validate_threads_per_block(caps: DeviceCaps, threads_per_block: int) -> None:
+    if threads_per_block % caps.warp_size or threads_per_block > caps.max_threads_per_block:
+        raise ValueError(
+            f"threads_per_block={threads_per_block} is not launchable on {caps.name}: it must be "
+            f"a multiple of the {caps.warp_size}-thread warp size and at most {caps.max_threads_per_block}"
+        )
+
+
+def resolve_threads(caps: DeviceCaps, requested: int | None, default: int) -> int:
+    if requested is None:
+        return caps.threads_per_block(default)
+    validate_threads_per_block(caps, requested)
+    return requested

entropack/backends/cuda/kernels.py

@@ -0,0 +1,70 @@
+from collections.abc import Callable, Hashable, Sequence
+from pathlib import Path
+
+import cupy
+import numpy as np
+import torch
+
+_modules: dict[tuple, object] = {}
+_kernels: dict[tuple, object] = {}
+_streams: dict[tuple, object] = {}
+
+__all__ = ["KernelLibrary", "device_index", "ensure_dynamic_shared", "external_stream", "pointer"]
+
+
+def ensure_dynamic_shared(kernel, shared_bytes: int) -> None:
+    if shared_bytes > kernel.max_dynamic_shared_size_bytes:
+        kernel.max_dynamic_shared_size_bytes = shared_bytes
+
+
+class KernelLibrary:
+    def __init__(
+        self, key: str, source: Path, defines: Callable[[int, Hashable], Sequence[str]],
+        includes: Sequence[Path] = (), kernel_names: Sequence[str] | None = None,
+    ):
+        self.key = key
+        self.source = source
+        self.defines = defines
+        self.includes = tuple(includes)
+        self.kernel_names = None if kernel_names is None else frozenset(kernel_names)
+
+    def kernel(self, device_index: int, variant: Hashable, name: str):
+        if self.kernel_names is not None and name not in self.kernel_names:
+            raise KeyError(name)
+        module_key = (self.key, device_index, variant)
+        module = _modules.get(module_key)
+        if module is None:
+            options = ["--std=c++17"]
+            options += [f"-D{definition}" for definition in self.defines(device_index, variant)]
+            options += [f"-I{include}" for include in self.includes]
+            with cupy.cuda.Device(device_index):
+                module = cupy.RawModule(code=self.source.read_text(), options=tuple(options))
+            _modules[module_key] = module
+        kernel_key = (self.key, device_index, variant, name)
+        kernel = _kernels.get(kernel_key)
+        if kernel is None:
+            kernel = module.get_function(name)
+            _kernels[kernel_key] = kernel
+        return kernel
+
+
+def external_stream(stream: torch.cuda.Stream):
+    key = (device_index(stream.device), int(stream.cuda_stream))
+    wrapper = _streams.get(key)
+    if wrapper is None:
+        factory = getattr(cupy.cuda.Stream, "from_external", None)
+        wrapper = factory(stream) if factory else cupy.cuda.ExternalStream(key[1])
+        _streams[key] = wrapper
+    return wrapper
+
+
+def pointer(tensor: torch.Tensor) -> np.uint64:
+    return np.uint64(tensor.data_ptr())
+
+
+def device_index(target) -> int:
+    if isinstance(target, int):
+        return target
+    device = target.device if isinstance(target, torch.Tensor) else target
+    index = getattr(device, "index", None)
+    return torch.cuda.current_device() if index is None else index

entropack/compression/__init__.py

@@ -0,0 +1,4 @@
+from .api import CompressionFallbackWarning, compress, decompress
+from .compressed_tensor import CompressedTensor
+
+__all__ = ["CompressedTensor", "CompressionFallbackWarning", "compress", "decompress"]

entropack/compression/api.py

@@ -0,0 +1,76 @@
+import logging
+import warnings
+
+import torch
+
+from ..registry import DispatchError, get_scheme, name_for_config
+from ..schemes import CompressionConfig, RawConfig
+from ..schemes.config import validate_config
+from .compressed_tensor import CompressedTensor
+
+logger = logging.getLogger("entropack")
+
+
+class CompressionFallbackWarning(RuntimeWarning):
+    """Warning emitted when compression falls back to storing the tensor uncompressed."""
+
+
+def compress(tensor: torch.Tensor, config: CompressionConfig) -> CompressedTensor:
+    """Compress a tensor using the selected scheme.
+
+    Args:
+        tensor: input values. The container preserves the input shape, dtype, and device.
+        config: configuration selecting the compression scheme and execution backend.
+
+    Returns:
+        A :class:`CompressedTensor` containing the encoded data and metadata.
+
+    Encoding failures can return an uncompressed ``raw`` container with a
+    :class:`CompressionFallbackWarning`. Its header records the requested scheme and
+    failure reason. Invalid configurations and backend dispatch failures raise errors.
+    """
+    if isinstance(tensor, CompressedTensor):
+        raise TypeError("compress expects an uncompressed tensor; decompress the container before recompressing")
+    scheme = get_scheme(name_for_config(config))
+    validate_config(config)
+    try:
+        packed = scheme.encode(tensor, config)
+    except DispatchError:
+        raise
+    except Exception as error:
+        logger.warning(
+            "%s could not encode a %s %s tensor, storing it uncompressed: %s",
+            scheme.name, tuple(tensor.shape), tensor.dtype, error,
+        )
+        return _compress_raw(tensor, scheme.name, error)
+    return CompressedTensor(header={"compress_method": scheme.name}, buffers=packed, shape=tuple(tensor.shape), dtype=tensor.dtype)
+
+
+def _compress_raw(tensor: torch.Tensor, requested: str, error: Exception) -> CompressedTensor:
+    reason = f"{type(error).__name__}: {error}"
+    warnings.warn(
+        f"entropack stored a {tuple(tensor.shape)} {tensor.dtype} tensor uncompressed, at "
+        f"{tensor.element_size() * 8} bits per element: compress_method={requested!r} failed with {reason}",
+        CompressionFallbackWarning, stacklevel=3,
+    )
+    return CompressedTensor(
+        header={"compress_method": "raw", "requested": requested, "reason": reason},
+        buffers=get_scheme("raw").encode(tensor, RawConfig()), shape=tuple(tensor.shape), dtype=tensor.dtype,
+    )
+
+
+def decompress(compressed: CompressedTensor, config: CompressionConfig) -> torch.Tensor:
+    """Restore a tensor from its compressed representation.
+
+    Args:
+        compressed: container to decode. Its header identifies the compression scheme.
+        config: configuration for the same scheme. Decode settings select the backend and
+            execution options. Encode settings such as the target bitrate do not recompress
+            the stored data.
+
+    Returns:
+        A tensor with the container's shape and dtype, on the device holding its buffers.
+    """
+    validate_config(config)
+    restored = compressed.scheme.decode(compressed.buffers, shape=compressed.shape, dtype=compressed.encoded_dtype, config=config)
+    return restored.to(compressed.dtype)

entropack/compression/compressed_tensor.py

@@ -0,0 +1,205 @@
+import copy
+import json
+import math
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+from torch.utils._python_dispatch import return_and_correct_aliasing
+
+from ..registry import get_scheme
+from ..schemes import Scheme
+
+
+def _parse_dtype(name: str) -> torch.dtype:
+    if not name:
+        raise TypeError("CompressedTensor serialized dtype must be a non-empty string")
+    dtype = getattr(torch, name, None)
+    if not isinstance(dtype, torch.dtype):
+        raise ValueError(f"Unsupported serialized torch dtype '{name}'")
+    return dtype
+
+
+class CompressedTensor(torch.Tensor):
+    """Frozen compressed tensor with a logical output dtype.
+
+    ``encoded_dtype`` records the codec's input type. ``to(dtype=...)`` changes
+    the output dtype without casting packed buffers or re-encoding values.
+    """
+
+    header: dict[str, Any]
+    buffers: dict[str, torch.Tensor]
+
+    @staticmethod
+    def __new__(cls, header, buffers, shape, dtype, *, encoded_dtype=None):
+        device = next(iter(buffers.values())).device
+        return torch.Tensor._make_wrapper_subclass(cls, tuple(shape), dtype=dtype, device=device, requires_grad=False)
+
+    def __init__(self, header, buffers, shape, dtype, *, encoded_dtype=None):
+        self.header = copy.deepcopy(header)
+        self.buffers = dict(buffers)
+        self.encoded_dtype = dtype if encoded_dtype is None else encoded_dtype
+        self.validate()
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}(shape={tuple(self.shape)}, dtype={self.dtype}, device={self.device}, scheme={self.compress_method!r})"
+
+    def __tensor_flatten__(self):
+        names = tuple(self.buffers)
+        for name, value in self.buffers.items():
+            setattr(self, "_packed_" + name, value)
+        metadata = (names, copy.deepcopy(self.header), tuple(self.shape), self.dtype, self.encoded_dtype)
+        return ["_packed_" + name for name in names], metadata
+
+    @classmethod
+    def __tensor_unflatten__(cls, inner_tensors, metadata, outer_size, outer_stride):
+        names, header, shape, dtype, encoded_dtype = metadata
+        return cls(
+            header=header, buffers={name: inner_tensors["_packed_" + name] for name in names},
+            shape=shape, dtype=dtype, encoded_dtype=encoded_dtype,
+        )
+
+    def _map_buffers(self, fn, *, dtype=None) -> "CompressedTensor":
+        return type(self)(
+            header=self.header, buffers={name: fn(value) for name, value in self.buffers.items()},
+            shape=self.shape, dtype=self.dtype if dtype is None else dtype, encoded_dtype=self.encoded_dtype,
+        )
+
+    def to(self, *args, copy: bool = False, **kwargs) -> "CompressedTensor":
+        """Move buffers or change the logical dtype, with an optional keyword-only copy."""
+        device, dtype, non_blocking, memory_format = torch._C._nn._parse_to(*args, **kwargs)
+        result = super().to(device=device, dtype=dtype, non_blocking=non_blocking, copy=copy, memory_format=memory_format)
+        return result.clone() if copy and result.device == self.device else result
+
+    def __copy__(self) -> "CompressedTensor":
+        result = self._map_buffers(lambda value: value)
+        if getattr(self, "_is_param", False):
+            result._is_param = True
+        return result
+
+    def __deepcopy__(self, memo) -> "CompressedTensor":
+        if id(self) in memo:
+            return memo[id(self)]
+        result = self._map_buffers(lambda value: copy.deepcopy(value, memo))
+        if getattr(self, "_is_param", False):
+            result._is_param = True
+        memo[id(self)] = result
+        return result
+
+    def requires_grad_(self, requires_grad: bool = False):
+        """Compressed values are frozen; gradients may flow through their consumers."""
+        if requires_grad:
+            raise RuntimeError("CompressedTensor cannot require gradients; decompress it to train dense values")
+        return torch.Tensor.requires_grad_(self, False)
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
+        kwargs = kwargs or {}
+        source = args[0] if args else None
+        aten = torch.ops.aten
+
+        if func in (aten.detach.default, aten.alias.default):
+            result = source._map_buffers(torch.detach)
+            return return_and_correct_aliasing(func, args, kwargs, result)
+
+        if func in (aten.clone.default, aten._to_copy.default):
+            if kwargs.get("memory_format") not in (None, torch.preserve_format, torch.contiguous_format):
+                raise ValueError("CompressedTensor does not support that memory format")
+            if func == aten.clone.default:
+                return source._map_buffers(torch.clone)
+            if kwargs.get("layout", torch.strided) not in (None, torch.strided):
+                raise ValueError("CompressedTensor only supports strided storage")
+            options = {name: value for name, value in kwargs.items() if name in ("device", "non_blocking")}
+            return source._map_buffers(lambda value: value.to(**options), dtype=kwargs.get("dtype"))
+
+        result = func.decompose(*args, **kwargs)
+        if result is not NotImplemented:
+            return result
+        raise NotImplementedError(f"CompressedTensor does not implement {func}; decompress it before numerical operations")
+
+    @property
+    def compress_method(self) -> str:
+        """The scheme name the header carries."""
+        return self.header["compress_method"]
+
+    @property
+    def scheme(self) -> Scheme:
+        """The codec named by the header."""
+        return get_scheme(self.compress_method)
+
+    @property
+    def lossless(self) -> bool:
+        """Whether the encoding scheme is lossless, before any output dtype conversion."""
+        return self.scheme.lossless
+
+    @property
+    def actual_bpp(self) -> float:
+        """Bits per element of :attr:`shape`, serialized header included."""
+        return self.storage_nbytes() * 8 / math.prod(self.shape)
+
+    def validate(self) -> None:
+        """Check the scheme, dtype, and buffer names without inspecting contents."""
+        if any(isinstance(value, CompressedTensor) for value in self.buffers.values()):
+            raise TypeError("CompressedTensor buffers cannot contain another CompressedTensor")
+        scheme = self.scheme
+        if not scheme.supports(self.encoded_dtype):
+            raise ValueError(f"'{scheme.name}' does not support format {self.encoded_dtype}")
+        if set(self.buffers) != set(scheme.buffer_names):
+            raise ValueError(
+                f"CompressedTensor buffers for {self.compress_method} must be "
+                f"{list(scheme.buffer_names)}, got {sorted(self.buffers)}"
+            )
+
+    def _serialized_header_tensor(self) -> torch.Tensor:
+        metadata = {
+            "compress_method": self.compress_method, "dtype": str(self.encoded_dtype).removeprefix("torch."),
+            "shape": list(self.shape), "buffer_names": list(self.buffers), "header": self.header,
+        }
+        encoded = json.dumps(metadata, sort_keys=True, separators=(",", ":"), allow_nan=False).encode("utf-8")
+        return torch.frombuffer(bytearray(encoded), dtype=torch.uint8)
+
+    def storage_nbytes(self, include_header: bool = True) -> int:
+        """Bytes the buffers occupy, plus the serialized header unless ``include_header`` is false."""
+        total = sum(value.numel() * value.element_size() for value in self.buffers.values())
+        if include_header:
+            total += self._serialized_header_tensor().numel()
+        return total
+
+    def state_dict(self, prefix: str = "") -> dict[str, torch.Tensor]:
+        """The container as flat tensors under ``prefix``: one 1D uint8 header, then one entry per buffer."""
+        self.validate()
+        state = {f"{prefix}header": self._serialized_header_tensor()}
+        state.update({f"{prefix}buffers.{name}": value.detach() for name, value in self.buffers.items()})
+        return state
+
+    @classmethod
+    def from_state_dict(cls, state: Mapping[str, torch.Tensor], prefix: str = "") -> "CompressedTensor":
+        """Restore a container in its encoded dtype from a tensor-only checkpoint."""
+        header_key = f"{prefix}header"
+        if header_key not in state:
+            raise ValueError(f"CompressedTensor state is missing '{header_key}'")
+        header_tensor = state[header_key]
+        if header_tensor.dtype != torch.uint8 or header_tensor.ndim != 1:
+            raise TypeError("CompressedTensor serialized header must be a 1D uint8 tensor")
+        try:
+            metadata = json.loads(bytes(header_tensor.detach().cpu().tolist()).decode("utf-8"))
+        except (UnicodeDecodeError, json.JSONDecodeError) as error:
+            raise ValueError("CompressedTensor serialized header is invalid") from error
+        if not isinstance(metadata, dict):
+            raise ValueError("CompressedTensor serialized header must contain a JSON object")
+
+        required = {"compress_method", "dtype", "shape", "buffer_names", "header"}
+        missing = required - metadata.keys()
+        if missing:
+            raise ValueError(f"CompressedTensor serialized header is missing fields: {sorted(missing)}")
+        buffer_names = metadata["buffer_names"]
+        if not isinstance(buffer_names, list) or any(not isinstance(name, str) or not name for name in buffer_names):
+            raise TypeError("CompressedTensor buffer_names must be a list of strings")
+        buffers = {}
+        for name in buffer_names:
+            key = f"{prefix}buffers.{name}"
+            if key not in state:
+                raise ValueError(f"CompressedTensor state is missing '{key}'")
+            buffers[name] = state[key]
+
+        return cls(header=metadata["header"], buffers=buffers, shape=tuple(metadata["shape"]), dtype=_parse_dtype(metadata["dtype"]))

entropack/linear/__init__.py

@@ -0,0 +1,5 @@
+from .linear import CompressedFP8Linear, CompressedINT8Linear, CompressedLinear, QuantizedLinear
+
+__all__ = [
+    "CompressedFP8Linear", "CompressedINT8Linear", "CompressedLinear", "QuantizedLinear",
+]