From 9ff38e94af5cfbfa42458d375e092c3352ec7210 Mon Sep 17 00:00:00 2001 From: Alperen Konukbay Date: Fri, 24 Jul 2026 17:18:44 -0700 Subject: [PATCH 1/6] feat: fpz compressed pack format with batched GPU decode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fpz format: bf16 macroblocks stored as split byte-plane zstd (low/ mantissa plane raw — measured incompressible — high/sign-exponent plane compressed), 1.401x smaller on the real 38GB LTX transformer pack, byte-exact on read. v2 stores the high plane as many independent 16B- aligned zstd chunks — the shape a GPU decoder needs. Read paths: - CPU decode (thread-pooled, GIL-released zstd) for any device target. - Batched GPU decode (FLASHPACK_FPZ_GPU_DECODE=1 [+ FLASHPACK_FPZ_GPU_LL=1 for the batched C path]): ctypes bindings to libnvcomp's batched Zstd decompressor (flashpack/_nvcomp_ll.py; the pybind nvcomp wrapper never exposes the batched API and its per-chunk Python objects are a measured 7x floor). Per batch: one pinned int64 chunk table, one H2D, two on-device base-address adds, one foreign call. Measured: 38GB restored in 1.67-1.85s page-hot (H200 AND B200, byte parity everywhere), decode fully hidden under NVMe reads when cold. - fpz_gpu_warmup(): pays nvcomp's one-time CUDA kernel-load off the hot path (26.7s worst-case observed lazy -> ~1s residual with warmup+EAGER). - Optional fused interleave Triton kernel (default off; ~2% at best — HBM absorbs the strided writes — kept as dormant infrastructure). Streaming encoder writes packs without an uncompressed scratch file (38GB repack: 50min -> ~3.5min). Everything is opt-in via pack_to_file(compress="fpz-bf16"); plain packs are byte-for-byte unaffected. Built on the affinity-clamp + distributed-hardening branch; the combined tree is byte-identical to the gated speed/fpz-v1 head (f2f1bf0), on which every receipt in the PR body was measured. Full development history: https://github.com/fal-ai/flashpack/tree/speed/fpz-v1 Co-Authored-By: Claude Fable 5 --- pyproject.toml | 11 + src/flashpack/_interleave.py | 86 ++ src/flashpack/_nvcomp_ll.py | 236 ++++++ src/flashpack/constants.py | 34 + src/flashpack/deserialization.py | 1304 +++++++++++++++++++++++++++++- src/flashpack/serialization.py | 409 +++++++++- tests/test_fpz.py | 554 +++++++++++++ tests/test_fpz_gpu.py | 297 +++++++ tests/test_fpz_ll.py | 136 ++++ tests/test_guardrails.py | 10 +- tests/test_interleave.py | 21 + 11 files changed, 3094 insertions(+), 4 deletions(-) create mode 100644 src/flashpack/_interleave.py create mode 100644 src/flashpack/_nvcomp_ll.py create mode 100644 tests/test_fpz.py create mode 100644 tests/test_fpz_gpu.py create mode 100644 tests/test_fpz_ll.py create mode 100644 tests/test_interleave.py diff --git a/pyproject.toml b/pyproject.toml index 4952dd0..11af0fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,17 @@ dev = [ "pre-commit>=3.0.0", "setuptools-scm>=9.2.0", ] +fpz = [ + "zstandard>=0.22", +] +# GPU Zstd decode for the fpz read path (opt-in via FLASHPACK_FPZ_GPU_DECODE=1). +# NOTE: nvidia-nvcomp-cu12 is NVIDIA-proprietary redistributable software, not +# OSI-licensed -- a license review is required before this extra ships beyond +# the prototype. Kept out of the default and `fpz` deps so a plain install +# stays MIT-only. +fpz-gpu = [ + "nvidia-nvcomp-cu12", +] [tool.setuptools_scm] write_to = "src/flashpack/version.py" diff --git a/src/flashpack/_interleave.py b/src/flashpack/_interleave.py new file mode 100644 index 0000000..8a387a9 --- /dev/null +++ b/src/flashpack/_interleave.py @@ -0,0 +1,86 @@ +"""Fused byte-plane interleave kernel (Triton) for the fpz GPU read paths. + +The split-plane format stores a frame as separate low/high byte planes; the +reader must produce out[2i] = lo[i], out[2i+1] = hi[i]. The torch expression +of that -- two strided copies (``out[0::2] = lo; out[1::2] = hi``) -- makes +two passes of 2-byte-stride writes, a worst-case memory pattern that doubles +device traffic and dominates the GPU-side tail of the hot-tier load. + +The fused kernel makes one pass: read lo[i] and hi[i] once, write one +little-endian uint16 ``lo | hi << 8`` (stored through an int16 view -- same +bit pattern, and torch's int16 has full op support where uint16 does not). + +Triton ships inside the torch Linux wheels (pytorch-triton), so no extra +dependency; on hosts without it (or on any kernel failure) callers fall back +to the strided copies. Enable with FLASHPACK_FPZ_FUSED_INTERLEAVE=1. +""" + +from __future__ import annotations + +import threading + +import torch + +# 8 elements per thread (BLOCK / (32 * warps) == 8): two 8-byte loads feed one +# 16-byte store per thread, the measured-optimal shape for byte-plane joins on +# Hopper (dietgpu's FloatTypeInfo vectorization). Fixed config on +# purpose -- autotuning would compile extra variants on a fresh container's +# empty triton cache, which lands exactly on the cold-start path. +_BLOCK = 2048 +_NUM_WARPS = 8 + +_lock = threading.Lock() +_kernel_cache: tuple | None = None + + +def _get_kernel(): + """Build (or fetch) the JIT'd kernel; None when triton is unavailable.""" + global _kernel_cache + with _lock: + if _kernel_cache is not None: + return _kernel_cache[0] + try: + import triton + import triton.language as tl + + @triton.jit + def _interleave_u8(lo_ptr, hi_ptr, out_ptr, n_elem, BLOCK: tl.constexpr): + # int16 (not uint16) throughout: identical bit pattern, and + # triton's unsigned integer paths have known pointer-arith + # bugs (triton#6043) while int16 bitcast is the in-house + # precedent. .to(tl.int16) on a uint8 source zero-extends. + pid = tl.program_id(0).to(tl.int64) + offs = pid * BLOCK + tl.arange(0, BLOCK).to(tl.int64) + mask = offs < n_elem + lo = tl.load(lo_ptr + offs, mask=mask, other=0).to(tl.int16) + hi = tl.load(hi_ptr + offs, mask=mask, other=0).to(tl.int16) + tl.store(out_ptr + offs, lo | (hi << 8), mask=mask) + + def _launch(lo: torch.Tensor, hi: torch.Tensor, out_u8: torch.Tensor): + n = lo.numel() + out16 = out_u8.view(torch.int16) + grid = (triton.cdiv(n, _BLOCK),) + _interleave_u8[grid]( + lo, hi, out16, n, BLOCK=_BLOCK, num_warps=_NUM_WARPS + ) + + _kernel_cache = (_launch,) + except Exception: + _kernel_cache = (None,) + return _kernel_cache[0] + + +def fused_interleave_available() -> bool: + return _get_kernel() is not None + + +def interleave_into(out_u8: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> bool: + """Fused single-pass interleave of ``lo``/``hi`` uint8 planes into + ``out_u8`` (contiguous, even byte offset, ``2 * lo.numel()`` bytes) on + the CURRENT stream. Returns False (having written nothing) when the + kernel is unavailable so the caller can run the strided fallback.""" + launch = _get_kernel() + if launch is None: + return False + launch(lo, hi, out_u8) + return True diff --git a/src/flashpack/_nvcomp_ll.py b/src/flashpack/_nvcomp_ll.py new file mode 100644 index 0000000..e000bc9 --- /dev/null +++ b/src/flashpack/_nvcomp_ll.py @@ -0,0 +1,236 @@ +"""ctypes bindings to libnvcomp's batched Zstd decompress C API (v5 ABI). + +Why this exists: the ``nvidia.nvcomp`` pybind wrapper only exposes the +high-level Manager API (``Codec``/``Array``), which needs one Python Array +object per compressed chunk per decode call. At fpz chunk counts (tens of +thousands per pack) that per-object marshaling is the measured decode-wall +floor. The batched C API takes *device-resident tables* of chunk +pointers/sizes instead -- per batch, Python builds a few small tensors and +makes ONE foreign call, independent of chunk count. The pybind layer never +calls the batched C API at all, so going through ``libnvcomp.so`` directly +is the only route to it from Python. + +ABI notes (verified against the headers shipped in nvidia-libnvcomp-cu12 +5.3.0.16, the wheel the pybind package depends on): + +* every ``device_*`` argument is a buffer in device memory, including the + pointer tables themselves; only counts, the opts struct and the stream + pass by host value; +* ``nvcompBatchedZstdDecompressOpts_t`` is a 64-byte struct passed BY VALUE + (``backend`` enum + 60 reserved bytes; all-zero selects the default CUDA + backend); +* v5 renamed the temp-size query to ``...GetTempSizeAsync`` (the widely + documented ``...GetTempSizeEx`` is the v2-v4 name and does not exist in + this .so); +* standard zstd frames (one independent frame per chunk) are accepted -- + the batched decompressor reads them directly, no nvcomp framing; +* for Zstd the per-chunk ``actual sizes`` and ``statuses`` output arrays + must be real device buffers (unlike LZ4, where NULL is tolerated). + +Load failures (missing wheel, missing symbol, non-Linux) degrade to +``load()`` returning ``None``; callers keep the pybind wrapper path as the +fallback. +""" + +from __future__ import annotations + +import ctypes +import logging +import os +import threading + +logger = logging.getLogger(__name__) + +NVCOMP_SUCCESS = 0 + + +class _ZstdDecompressOpts(ctypes.Structure): + """``nvcompBatchedZstdDecompressOpts_t``: 64 bytes, passed by value.""" + + _fields_ = [("backend", ctypes.c_int), ("reserved", ctypes.c_char * 60)] + + +class _AlignmentRequirements(ctypes.Structure): + """``nvcompAlignmentRequirements_t``: input/output/temp minimums.""" + + _fields_ = [ + ("input", ctypes.c_size_t), + ("output", ctypes.c_size_t), + ("temp", ctypes.c_size_t), + ] + + +def _find_libnvcomp() -> str | None: + """Locate ``libnvcomp.so`` from the nvidia-libnvcomp wheel layout. + + The pip layout is ``site-packages/nvidia/libnvcomp/lib64/libnvcomp.so.5``; + resolve it from the package rather than the linker path so the binding + works without any LD_LIBRARY_PATH setup. + """ + override = os.environ.get("FLASHPACK_LIBNVCOMP_PATH") + if override: + return override if os.path.exists(override) else None + try: + import nvidia.libnvcomp as _libnvcomp_pkg + except ImportError: + return None + pkg_dir = os.path.dirname(_libnvcomp_pkg.__file__) + for sub in ("lib64", "lib"): + lib_dir = os.path.join(pkg_dir, sub) + if not os.path.isdir(lib_dir): + continue + for name in sorted(os.listdir(lib_dir)): + if name.startswith("libnvcomp.so"): + return os.path.join(lib_dir, name) + return None + + +class NvcompLL: + """Thin, stateless handle over the batched Zstd decompress entry points. + + All methods are thread-safe: the underlying C functions are stateless + launches, and ctypes releases the GIL for the duration of each call. + """ + + def __init__(self, lib: ctypes.CDLL): + self._lib = lib + self._opts = _ZstdDecompressOpts() # zero-filled = default backend + + self._get_alignments = lib.nvcompBatchedZstdDecompressGetRequiredAlignments + self._get_alignments.restype = ctypes.c_int + self._get_alignments.argtypes = [ + _ZstdDecompressOpts, + ctypes.POINTER(_AlignmentRequirements), + ] + + self._get_temp_size = lib.nvcompBatchedZstdDecompressGetTempSizeAsync + self._get_temp_size.restype = ctypes.c_int + self._get_temp_size.argtypes = [ + ctypes.c_size_t, # num_chunks + ctypes.c_size_t, # max_uncompressed_chunk_bytes + _ZstdDecompressOpts, + ctypes.POINTER(ctypes.c_size_t), # temp_bytes (host out) + ctypes.c_size_t, # max_total_uncompressed_bytes + ] + + self._decompress = lib.nvcompBatchedZstdDecompressAsync + self._decompress.restype = ctypes.c_int + self._decompress.argtypes = [ + ctypes.c_void_p, # device_compressed_chunk_ptrs (device void**) + ctypes.c_void_p, # device_compressed_chunk_bytes (device size_t*) + ctypes.c_void_p, # device_uncompressed_buffer_bytes (device size_t*) + ctypes.c_void_p, # device_uncompressed_chunk_bytes OUT (device size_t*) + ctypes.c_size_t, # num_chunks + ctypes.c_void_p, # device_temp_ptr + ctypes.c_size_t, # temp_bytes + ctypes.c_void_p, # device_uncompressed_chunk_ptrs (device void**) + _ZstdDecompressOpts, # by value + ctypes.c_void_p, # device_statuses (device nvcompStatus_t*) + ctypes.c_void_p, # cudaStream_t + ] + + self._status_string = lib.nvcompGetStatusString + self._status_string.restype = ctypes.c_char_p + self._status_string.argtypes = [ctypes.c_int] + + def status_string(self, status: int) -> str: + s = self._status_string(int(status)) + return s.decode() if s else f"nvcompStatus_t({status})" + + def _check(self, status: int, call: str) -> None: + if status != NVCOMP_SUCCESS: + raise RuntimeError(f"{call} failed: {self.status_string(status)}") + + def alignments(self) -> tuple[int, int, int]: + """Required (input, output, temp) buffer alignments for decompression.""" + reqs = _AlignmentRequirements() + self._check( + self._get_alignments(self._opts, ctypes.byref(reqs)), + "nvcompBatchedZstdDecompressGetRequiredAlignments", + ) + return int(reqs.input), int(reqs.output), int(reqs.temp) + + def temp_size( + self, num_chunks: int, max_chunk_bytes: int, max_total_bytes: int + ) -> int: + """Device scratch bytes needed for a decompress batch of this shape.""" + out = ctypes.c_size_t(0) + self._check( + self._get_temp_size( + num_chunks, + max_chunk_bytes, + self._opts, + ctypes.byref(out), + max_total_bytes, + ), + "nvcompBatchedZstdDecompressGetTempSizeAsync", + ) + return int(out.value) + + def decompress_async( + self, + src_ptrs_dev: int, + src_sizes_dev: int, + out_caps_dev: int, + actual_sizes_dev: int, + num_chunks: int, + temp_ptr_dev: int, + temp_bytes: int, + dst_ptrs_dev: int, + statuses_dev: int, + cuda_stream: int, + ) -> None: + """Enqueue one batched decompress; all ``*_dev`` args are raw device + addresses (``tensor.data_ptr()``). Raises on launch/validation errors; + per-chunk data errors land in ``statuses_dev`` (device int32 array) + for the caller to check stream-side.""" + self._check( + self._decompress( + src_ptrs_dev, + src_sizes_dev, + out_caps_dev, + actual_sizes_dev, + num_chunks, + temp_ptr_dev, + temp_bytes, + dst_ptrs_dev, + self._opts, + statuses_dev, + cuda_stream, + ), + "nvcompBatchedZstdDecompressAsync", + ) + + +_load_lock = threading.Lock() +_loaded: tuple[NvcompLL | None] | None = None + + +def load() -> NvcompLL | None: + """Load and bind libnvcomp once; ``None`` (with a single warning) on any + failure so callers can fall back to the pybind wrapper path.""" + global _loaded + with _load_lock: + if _loaded is not None: + return _loaded[0] + handle: NvcompLL | None = None + path = _find_libnvcomp() + if path is None: + logger.warning( + "flashpack: libnvcomp not found (pip install " + "nvidia-libnvcomp-cu12); batched GPU decode unavailable, " + "falling back to the nvcomp wrapper path." + ) + else: + try: + handle = NvcompLL(ctypes.CDLL(path)) + except (OSError, AttributeError) as e: + logger.warning( + "flashpack: could not bind batched nvcomp API from %s " + "(%s); falling back to the nvcomp wrapper path.", + path, + e, + ) + handle = None + _loaded = (handle,) + return handle diff --git a/src/flashpack/constants.py b/src/flashpack/constants.py index a6a5111..0a4d229 100644 --- a/src/flashpack/constants.py +++ b/src/flashpack/constants.py @@ -10,3 +10,37 @@ DEFAULT_NUM_WRITE_WORKERS = 32 DEFAULT_NUM_STREAMS = 4 DEFAULT_CHUNK_BYTES = 4 * 1024 * 1024 # 4 MiB + +# fpz split-plane zstd compression (V4-additive; opt-in via pack_to_file(compress=...)). +# bf16 macroblocks only: viewed as little-endian uint16, the low byte plane +# (mantissa LSB) is incompressible and stored raw, while the high byte plane +# (sign+exponent) compresses ~2.5x with zstd. A compressed macroblock's payload +# is a sequence of frames covering the block in FPZ_FRAME_UNCOMPRESSED_BYTES +# (uncompressed) steps; each frame is [lo raw bytes][hi zstd bytes] with the +# frame payload start 4096-byte aligned relative to the macroblock start. +FPZ_COMPRESS_BF16 = "fpz-bf16" +FPZ_CODEC_SPLITPLANE_V1 = "zstd-splitplane-v1" +# v2 chunks each frame's high plane into many small independent zstd frames so a +# GPU decoder (nvcomp) gets its native many-chunk batch shape. A single-frame +# (v1) high plane is ONE nvcomp chunk and decodes serially on the GPU (~0.6 +# GB/s); v2's chunks decode in parallel. In a v2 frame the payload is +# [lo raw bytes][hi zstd chunk 0][hi zstd chunk 1]... and the frame record +# carries "hi_chunks": the per-chunk compressed byte lengths (each chunk's +# uncompressed size is FPZ_HI_CHUNK_UNCOMPRESSED_BYTES except the frame's last). +FPZ_CODEC_SPLITPLANE_V2 = "zstd-splitplane-v2" +FPZ_FRAME_UNCOMPRESSED_BYTES = 64 * 1024 * 1024 # 64 MiB +FPZ_FRAME_ALIGN_BYTES = 4096 +# Per-chunk uncompressed size for the v2 high plane. Matches nvcomp's default +# uncomp_chunk_size (65536) and divides the 32 MiB half-frame evenly (512 +# chunks), so full frames have no odd-sized tail chunk. +FPZ_HI_CHUNK_UNCOMPRESSED_BYTES = 64 * 1024 +# Byte alignment of each v2 compressed chunk's START within the frame payload +# (chunks are padded to this; "hi_chunks" still records true zstd lengths and +# the reader recomputes padded offsets from the block's "hi_align" footer +# field, absent = 1 for pre-alignment packs). 16 covers nvcomp's batched +# decompressor input-alignment requirement -- its C API rejects unaligned +# device chunk pointers with nvcompErrorAlignment, and the reference callers +# align inputs to max(16, queried requirement). Cost: <= 15 pad bytes per +# chunk (~0.001% at 1 MiB chunks). +FPZ_HI_CHUNK_ALIGN_BYTES = 16 +DEFAULT_ZSTD_LEVEL = 3 diff --git a/src/flashpack/deserialization.py b/src/flashpack/deserialization.py index 8f5a84e..2e9d0e1 100644 --- a/src/flashpack/deserialization.py +++ b/src/flashpack/deserialization.py @@ -1,6 +1,9 @@ import json import math import os +import queue +import threading +import time import warnings from collections.abc import Iterator from dataclasses import dataclass @@ -11,11 +14,16 @@ import torch.distributed as dist import tqdm +from . import _interleave from .constants import ( DEFAULT_CHUNK_BYTES, DEFAULT_NUM_STREAMS, FILE_FORMAT_V3, FILE_FORMAT_V4, + FPZ_CODEC_SPLITPLANE_V1, + FPZ_CODEC_SPLITPLANE_V2, + FPZ_FRAME_UNCOMPRESSED_BYTES, + FPZ_HI_CHUNK_UNCOMPRESSED_BYTES, MAGIC, U64LE, ) @@ -24,11 +32,13 @@ parallel_read_supported, ) from .utils import ( + effective_read_threads, get_module_and_attribute, get_packing_dtype, human_num_elements, is_ignored_tensor_name, maybe_init_distributed, + require_zstandard, string_to_dtype, timer, torch_dtype_to_numpy_dtype, @@ -41,6 +51,15 @@ class MacroblockSpec: offset_bytes: int length_bytes: int length_elems: int + # For fpz-compressed blocks: the {"codec", "frames": [...]} record from the + # footer. None for plain (uncompressed) blocks. offset_bytes/length_bytes + # describe the compressed payload as stored; length_elems is always the + # logical (uncompressed) element count. + fpz: dict[str, Any] | None = None + + @property + def uncompressed_bytes(self) -> int: + return self.length_elems * torch.tensor([], dtype=self.dtype).element_size() @dataclass @@ -136,12 +155,18 @@ def _build_macroblock_specs(meta: dict[str, Any]) -> list[MacroblockSpec]: raise ValueError("Missing macroblock metadata for flashpack v4 file.") for block in macroblocks: dtype = string_to_dtype(block["dtype"]) + fpz = block.get("fpz") + if fpz is not None: + codec = fpz.get("codec") + if codec not in (FPZ_CODEC_SPLITPLANE_V1, FPZ_CODEC_SPLITPLANE_V2): + raise ValueError(f"Unsupported fpz codec: {codec!r}") specs.append( MacroblockSpec( dtype=dtype, offset_bytes=int(block["offset_bytes"]), length_bytes=int(block["length_bytes"]), length_elems=int(block["length_elems"]), + fpz=fpz, ) ) else: @@ -294,10 +319,15 @@ def _allocate_aligned_cpu_storage(specs: list[MacroblockSpec]) -> FlashTensorSto align = 4096 blocks: list[torch.Tensor] = [] for spec in specs: - raw = torch.empty(spec.length_bytes + align, dtype=torch.uint8) + # Size by the logical (uncompressed) byte count. This equals + # spec.length_bytes for plain blocks, but for fpz blocks length_bytes + # is the smaller compressed on-disk size, so the destination must be + # sized from the element count instead. + nbytes = spec.uncompressed_bytes + raw = torch.empty(nbytes + align, dtype=torch.uint8) off = (-raw.data_ptr()) % align packing_dtype = get_packing_dtype(spec.dtype) - block = raw.narrow(0, off, spec.length_bytes).view(packing_dtype) + block = raw.narrow(0, off, nbytes).view(packing_dtype) if spec.dtype != packing_dtype: block = block.view(spec.dtype) blocks.append(block) @@ -318,6 +348,1271 @@ def _broadcast_storage(storage: FlashTensorStorage, src: int) -> None: dist.broadcast(block.view(torch.uint8), src=src) +def _env_int(name: str, default: int) -> int: + try: + return int(os.environ.get(name, default)) + except ValueError: + return default + + +def _pread_into(fd: int, offset: int, mv: memoryview) -> None: + """Fill ``mv`` from ``fd`` at ``offset`` with ``preadv`` (reused reader + machinery). Raises ``IOError`` on a short read (e.g. a truncated file).""" + n = len(mv) + got = 0 + while got < n: + r = os.preadv(fd, [mv[got:]], offset + got) + if r <= 0: + raise IOError(f"short read: wanted {n} bytes at {offset}, got {got}") + got += r + + +def _fpz_frame_tasks(specs: list[MacroblockSpec]) -> list[tuple]: + """Build the per-block decode work list and validate frame coverage. + + Each item is ``("frame", block_idx, frame, out_pos)`` for an fpz frame or + ``("raw", block_idx, None, 0)`` for a plain block. Raises ``ValueError`` if + a block's frames do not exactly cover its uncompressed byte length -- the + same error surface the single-threaded decoder used to raise. + """ + tasks: list[tuple] = [] + for idx, spec in enumerate(specs): + if spec.fpz is None: + tasks.append(("raw", idx, None, 0)) + continue + total = spec.uncompressed_bytes + out_pos = 0 + for frame in spec.fpz["frames"]: + n_out = int(frame["n_out"]) + if out_pos + n_out > total: + raise ValueError("fpz frames exceed the macroblock size") + tasks.append(("frame", idx, frame, out_pos)) + out_pos += n_out + if out_pos != total: + raise ValueError(f"fpz frames cover {out_pos} bytes, expected {total}") + return tasks + + +def _align_up(n: int, align: int) -> int: + """Round ``n`` up to a multiple of ``align`` (``align`` >= 1).""" + return n + (-n % align) + + +def _fpz_read_frame_planes( + fd: int, + block_file_offset: int, + frame: dict[str, Any], + decompressor, + chunk_u: int = FPZ_HI_CHUNK_UNCOMPRESSED_BYTES, + hi_align: int = 1, +) -> tuple[np.ndarray, np.ndarray]: + """Read one fpz frame and return its ``(lo, hi)`` byte planes as uint8 + numpy arrays, each ``n_out // 2`` bytes. + + Shared CPU decode step for both read paths and both codec versions. v1 + frames store the high plane as a single zstd frame (``hi_len``); v2 frames + store it as many ``chunk_u``-uncompressed-byte chunks whose compressed + lengths are in ``hi_chunks`` (``chunk_u`` is the block's ``hi_chunk_usize``, + defaulting to the pre-parameterization 64 KiB when absent). zstd + ``decompress`` releases the GIL and takes the compressed input as a buffer, + so read targets are ``memoryview``s (no intermediate ``bytes`` copy) and N + threads scale near linearly. + """ + payload_off = int(frame["payload_off"]) + lo_len = int(frame["lo_len"]) + n_out = int(frame["n_out"]) + half = n_out - lo_len + + lo_raw = bytearray(lo_len) + _pread_into(fd, block_file_offset + payload_off, memoryview(lo_raw)) + lo = np.frombuffer(lo_raw, dtype=np.uint8) + # hi_align-packs pad after the lo plane and after each chunk so every + # chunk STARTS aligned (GPU batched decode needs aligned device chunk + # pointers); hi_chunks records true zstd lengths, offsets are padded. + hi_base = block_file_offset + payload_off + _align_up(lo_len, hi_align) + + if "hi_chunks" in frame: + # v2: decode each chunk (a standalone zstd frame) into its slice. + hi = np.empty(half, dtype=np.uint8) + uoff = 0 + src_off = hi_base + for clen in frame["hi_chunks"]: + clen = int(clen) + usize = min(chunk_u, half - uoff) + cbuf = bytearray(clen) + _pread_into(fd, src_off, memoryview(cbuf)) + dec = decompressor.decompress(memoryview(cbuf), max_output_size=usize) + if len(dec) != usize: + raise ValueError("fpz v2 chunk size mismatch") + hi[uoff : uoff + usize] = np.frombuffer(dec, dtype=np.uint8) + uoff += usize + src_off += _align_up(clen, hi_align) + if lo_len * 2 != n_out or uoff != half or lo.shape[0] != lo_len: + raise ValueError("fpz frame plane size mismatch") + return lo, hi + + # v1: the high plane is a single zstd frame. + hi_len = int(frame["hi_len"]) + hi_raw = bytearray(hi_len) + _pread_into(fd, hi_base, memoryview(hi_raw)) + # memoryview input avoids a GIL-held full copy of the compressed plane; + # decompress itself releases the GIL. + hi_bytes = decompressor.decompress(memoryview(hi_raw), max_output_size=half) + hi = np.frombuffer(hi_bytes, dtype=np.uint8) + if lo_len * 2 != n_out or lo.shape[0] != lo_len or hi.shape[0] != half: + raise ValueError("fpz frame plane size mismatch") + return lo, hi + + +def _fpz_read_into_cpu_storage( + path: str, specs: list[MacroblockSpec], blocks: list[torch.Tensor] +) -> None: + """Fill pre-allocated (uncompressed-sized) CPU ``blocks`` from an fpz file + with a pool of decode threads. + + Work is one item per fpz frame (or per plain block); frames write disjoint + destination byte ranges, so threads never collide. Each thread owns a file + descriptor and a zstd decompressor. The heavy step -- the zstd decode -- + releases the GIL, so throughput scales with ``FLASHPACK_READ_THREADS`` + (default 16) instead of running serially as it did before. + """ + zstandard = require_zstandard() + n_threads = effective_read_threads(_env_int("FLASHPACK_READ_THREADS", 16)) + + dst_u8 = [b.view(torch.uint8).numpy() for b in blocks] + tasks = _fpz_frame_tasks(specs) + n_threads = min(n_threads, max(1, len(tasks))) + + work: queue.SimpleQueue = queue.SimpleQueue() + for task in tasks: + work.put(task) + for _ in range(n_threads): + work.put(None) + + errors: list[BaseException] = [] + + def _reader() -> None: + try: + fd = os.open(path, os.O_RDONLY) + decompressor = zstandard.ZstdDecompressor() + try: + while True: + item = work.get() + if item is None: + break + kind, blk, frame, out_pos = item + if kind == "raw": + spec = specs[blk] + _pread_into(fd, spec.offset_bytes, memoryview(dst_u8[blk])) + continue + spec = specs[blk] + chunk_u = int( + (spec.fpz or {}).get( + "hi_chunk_usize", FPZ_HI_CHUNK_UNCOMPRESSED_BYTES + ) + ) + hi_align = int((spec.fpz or {}).get("hi_align", 1)) + lo, hi = _fpz_read_frame_planes( + fd, spec.offset_bytes, frame, decompressor, chunk_u, hi_align + ) + n_out = int(frame["n_out"]) + seg = dst_u8[blk][out_pos : out_pos + n_out] + # Disjoint destination ranges across threads; numpy releases + # the GIL for the strided byte copy. + seg[0::2] = lo + seg[1::2] = hi + finally: + os.close(fd) + except BaseException as e: + errors.append(e) + + threads = [threading.Thread(target=_reader, daemon=True) for _ in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + if errors: + raise errors[0] + + +_FPZ_CUDA_BUFFERS_PER_THREAD = 2 + + +def _fpz_read_into_cuda_storage( + path: str, + specs: list[MacroblockSpec], + blocks: list[torch.Tensor], + device: torch.device, +) -> None: + """Fill pre-allocated device ``blocks`` from an fpz file with a pool of + reader threads (mirrors ``parallel_read_into_storage``). + + Each reader owns a file descriptor, a CUDA stream, and a small ring of + double-buffered pinned/device staging slots. A work item is one fpz frame + (or a whole plain block); frames write disjoint destination segments so the + readers never collide. Per frame: ``preadv`` the low plane straight into a + pinned buffer, zstd-decode the high plane (GIL released) and copy it into a + pinned buffer, then enqueue on the stream the two H2Ds and the two strided + GPU copies (``dst_u8[0::2] = lo``, ``dst_u8[1::2] = hi``). A GPU decoder + (nvcomp) would slot in by replacing the decompress step. + + There is NO per-frame ``stream.synchronize()``: a CUDA event per staging + slot gates only buffer reuse, so a thread reads/decodes the next frame while + the GPU is still consuming the previous one. Removing that per-frame sync + (and the single-buffered staging) is the fix for the observed ~2.6 GB/s + stall -- the CPU decode now overlaps the H2D/copy instead of blocking on it. + + GPU-untested locally (no CUDA device); the frame read + decode is exercised + by the CPU tests via the shared ``_fpz_read_frame_planes`` helper. + """ + zstandard = require_zstandard() + n_threads = effective_read_threads(_env_int("FLASHPACK_READ_THREADS", 16)) + half_cap = FPZ_FRAME_UNCOMPRESSED_BYTES // 2 + n_slots = _FPZ_CUDA_BUFFERS_PER_THREAD + fused = _env_flag("FLASHPACK_FPZ_FUSED_INTERLEAVE") and ( + _interleave.fused_interleave_available() + ) + + byte_blocks = [b.view(torch.uint8) for b in blocks] + + # Order every reader stream after the destination allocation (same + # wait_event pattern as parallel_read_into_storage). + alloc_ready = torch.cuda.Event() + alloc_ready.record(torch.cuda.current_stream(device)) + + tasks = _fpz_frame_tasks(specs) + work: queue.SimpleQueue = queue.SimpleQueue() + for task in tasks: + work.put(task) + n_threads = min(n_threads, max(1, len(tasks))) + for _ in range(n_threads): + work.put(None) + + errors: list[BaseException] = [] + + def _reader() -> None: + try: + fd = os.open(path, os.O_RDONLY) + decompressor = zstandard.ZstdDecompressor() + stream = torch.cuda.Stream(device=device) + stream.wait_event(alloc_ready) + lo_pin = [ + torch.empty(half_cap, dtype=torch.uint8, pin_memory=True) + for _ in range(n_slots) + ] + hi_pin = [ + torch.empty(half_cap, dtype=torch.uint8, pin_memory=True) + for _ in range(n_slots) + ] + lo_dev = [ + torch.empty(half_cap, dtype=torch.uint8, device=device) + for _ in range(n_slots) + ] + hi_dev = [ + torch.empty(half_cap, dtype=torch.uint8, device=device) + for _ in range(n_slots) + ] + events = [torch.cuda.Event() for _ in range(n_slots)] + for ev in events: + ev.record(stream) + i = 0 + try: + while True: + item = work.get() + if item is None: + break + kind, blk, frame, out_pos = item + dst = byte_blocks[blk] + if kind == "raw": + spec = specs[blk] + buf = bytearray(spec.length_bytes) + _pread_into(fd, spec.offset_bytes, memoryview(buf)) + host = torch.frombuffer(buf, dtype=torch.uint8) + stream.synchronize() + with torch.cuda.stream(stream): + dst.copy_(host, non_blocking=False) + continue + + slot = i % n_slots + i += 1 + # The slot's previous H2D must be done before we overwrite + # its pinned buffers. + events[slot].synchronize() + + n_out = int(frame["n_out"]) + spec = specs[blk] + chunk_u = int( + (spec.fpz or {}).get( + "hi_chunk_usize", FPZ_HI_CHUNK_UNCOMPRESSED_BYTES + ) + ) + hi_align = int((spec.fpz or {}).get("hi_align", 1)) + # Shared CPU decode (handles both v1 single-frame and v2 + # chunked high planes), then copy both planes into pinned. + lo_np, hi_np = _fpz_read_frame_planes( + fd, spec.offset_bytes, frame, decompressor, chunk_u, hi_align + ) + half = int(hi_np.shape[0]) + lo_pin[slot].numpy()[:half] = lo_np + hi_pin[slot].numpy()[:half] = hi_np + + seg = dst.narrow(0, out_pos, n_out) + with torch.cuda.stream(stream): + lo_dev[slot][:half].copy_( + lo_pin[slot][:half], non_blocking=True + ) + hi_dev[slot][:half].copy_( + hi_pin[slot][:half], non_blocking=True + ) + if fused: + _interleave.interleave_into( + seg, lo_dev[slot][:half], hi_dev[slot][:half] + ) + else: + seg[0::2].copy_(lo_dev[slot][:half], non_blocking=True) + seg[1::2].copy_(hi_dev[slot][:half], non_blocking=True) + events[slot].record(stream) + finally: + os.close(fd) + stream.synchronize() + except BaseException as e: + errors.append(e) + + threads = [threading.Thread(target=_reader, daemon=True) for _ in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + torch.cuda.synchronize(device) + if errors: + raise errors[0] + + +# --------------------------------------------------------------------------- +# nvcomp GPU Zstd decode (prototype; opt-in via FLASHPACK_FPZ_GPU_DECODE=1). +# +# The CPU zstd decode caps fpz at ~9 GB/s logical (H200) versus ~21.6 GB/s for +# a page-hot raw pack, so decode -- not I/O -- is the fpz bottleneck. nvcomp's +# batched GPU Zstd decoder moves that work onto the device. +# +# API discovery -- nvidia-nvcomp-cu12 5.3.0 (pybind11 module ``nvidia.nvcomp``; +# signatures/docstrings read from the compiled nvcomp_impl .so, quoted below): +# +# Codec(algorithm="Zstd", device_id=, cuda_stream=, +# uncomp_chunk_size=65536, bitstream_kind=BitstreamKind.NVCOMP_NATIVE, +# checksum_policy=NO_COMPUTE_NO_VERIFY, decompress_backend=...) +# "Initialize codec." +# algorithm : name of the compression algorithm ("Zstd", "LZ4", ...). +# device_id : device to run on (default: current device). +# cuda_stream : cudaStream_t as a Python int (default: an internal +# stream). We pass each reader thread's own torch stream +# (``stream.cuda_stream``) so the decode is ordered on the +# SAME stream as our H2D copies and the interleave -- no +# cross-stream sync needed. +# bitstream_kind : BitstreamKind.{NVCOMP_NATIVE, RAW, WITH_UNCOMPRESSED_SIZE}. +# We use RAW: "Compresses input data as is, just using the +# underlying compression algorithm. Does not add a header +# with nvCOMP metadata." The fpz high plane is a standard +# single zstd frame written by python-zstandard, so it must +# be decoded as RAW (NVCOMP_NATIVE expects nvcomp's own +# chunked container and would reject a bare zstd frame). +# +# codec.decode(src, data_type="|u1", out=None, decompression_config=None) +# -> nvcomp.Array "Decode a single Array." +# codec.decode(srcs: list[Array], data_type=..., out=, +# decompression_config=None) -> list[Array] +# "Decode a batch of Arrays." +# out : "An optional writable buffer to store decoded data. ... If it is an +# externally-allocated buffer (e.g. cupy/numba array), its size is +# fixed and a ValueError is raised when it is too small." We pass a +# view over our pre-sized device high-plane tensor, so decode writes +# straight into it -- no extra device copy and no host round trip. +# data_type : output element type string; default "|u1" (uint8), which is +# exactly the byte plane we want, so we never pass it. +# decompression_config : when omitted, "decode internally calls +# configure_decompression on src, forcing a stream synchronization" +# on EVERY call -- that per-call sync serialized the whole pipeline +# and measured ~0.6 GB/s on the H200. We instead build a reusable +# DecompressConfig once per distinct batch shape via +# codec.decompression_config(srcs) ("reusable across multiple decode +# calls ... of the same uncompressed per-element shape") and pass it +# to decode, which is then sync-free. fpz packs have only a few +# shapes (full 64 MiB frames + one tail per block), so the one-time +# build sync is paid a handful of times per thread, not per decode. +# (A CompressConfig-derived config -- codec.decompression_config( +# codec.compression_config(size)) -- would skip even the build sync, +# but the docstring scopes that to same-process compress+decompress; +# our frames are compressed offline by python-zstandard, so we use +# the header-parsing overload that is proven against real frames.) +# +# nvcomp.as_array(src_object, cuda_stream=None) -> Array +# "Creates array from object with some standard interface." Zero-copy over +# any object exposing __cuda_array_interface__ / __dlpack__. A contiguous +# torch CUDA tensor qualifies, so we wrap the device staging tensors +# directly (no copy). nvcomp.from_dlpack(...) is the explicit-DLPack +# equivalent; as_array is sufficient here. +# +# nvcomp.set_device_allocator(allocator) -- "Sets a new allocator ... for +# future device allocations." allocator is +# ``allocator(nbytes: int, stream: nvcomp.Stream) -> obj`` where obj has an +# integer ``.ptr`` and frees on garbage collection. nvcomp grabs scratch +# from this for every decode; its default (cudaMalloc/cudaFree) syncs the +# device per call, so we install a torch-caching-allocator adapter (see +# _install_torch_nvcomp_allocator) to serve scratch pool-side with no sync. +# +# Compatibility, verified locally against the pack side (python-zstandard +# level-3, threads=-1, one-shot ``compress``): every high plane is a SINGLE +# standard zstd frame (magic 0xFD2FB528) with the content size embedded, a +# 2 MiB window (windowLog 21), no dictionary and no checksum -- all within +# nvcomp GPU Zstd's limits. nvcomp itself cannot be exercised without a device, +# so confirming decode correctness on a real frame is step 0 of the H200 run. +# The correctness-critical interleave (even byte = low plane, odd byte = high +# plane) is identical to the CPU path and is covered by the CPU tests. +# --------------------------------------------------------------------------- + +_FPZ_GPU_DECODE_WARNED = False + + +def _env_flag(name: str) -> bool: + return os.environ.get(name, "").strip().lower() in ("1", "true", "yes", "on") + + +def _fpz_gpu_decode_enabled() -> bool: + """Whether the opt-in nvcomp GPU decode path is requested (env-gated).""" + return _env_flag("FLASHPACK_FPZ_GPU_DECODE") + + +def _load_nvcomp(): + """Import the optional nvcomp module for GPU Zstd decode. + + Returns the module, or ``None`` if it is not importable -- in which case it + warns once (per process) so the caller can fall back to the CPU decode path + without spamming. nvcomp is NVIDIA-proprietary and never a hard dependency; + install it with ``pip install 'flashpack[fpz-gpu]'`` (see pyproject). + """ + global _FPZ_GPU_DECODE_WARNED + try: + from nvidia import nvcomp + + return nvcomp + except ImportError: + if not _FPZ_GPU_DECODE_WARNED: + _FPZ_GPU_DECODE_WARNED = True + warnings.warn( + "FLASHPACK_FPZ_GPU_DECODE=1 but the nvcomp package is not " + "importable; falling back to CPU zstd decode. Install the GPU " + "extra with: pip install 'flashpack[fpz-gpu]'.", + RuntimeWarning, + stacklevel=2, + ) + return None + + +def _env_flag_default(name: str, default: bool) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in ("1", "true", "yes", "on") + + +# nvcomp calls its device allocator once per decode for scratch. Its default +# allocator is cudaMalloc/cudaFree, and each of those synchronizes the device -- +# on the H200 that per-call sync (not the config sync) was the dominant fpz cost +# (~60ms/frame, tier-flat). Routing nvcomp's scratch through torch's stream-aware +# caching allocator serves it from an existing pool with no cudaMalloc/sync. +_fpz_nvcomp_alloc_tls = threading.local() +_FPZ_NVCOMP_ALLOC_INSTALLED = False +_FPZ_NVCOMP_ALLOC_WARNED = False + + +class _TorchNvcompDeviceBuffer: + """Adapter exposing a torch caching-allocator block to nvcomp's allocator + protocol: an object with an integer ``ptr`` that frees on ``__del__``. + + The allocation is tied to the calling reader thread's CUDA stream (stashed + in a thread-local by the reader) so torch's caching allocator won't hand the + block to another stream while nvcomp's decode -- which runs on that same + stream -- is still using it. + """ + + __slots__ = ("_ptr",) + + def __init__(self, nbytes: int, device_index: int, stream) -> None: + self._ptr = torch.cuda.caching_allocator_alloc(nbytes, device_index, stream) + + @property + def ptr(self) -> int: + return self._ptr + + def __del__(self) -> None: + try: + torch.cuda.caching_allocator_delete(self._ptr) + except Exception: + pass + + +def _install_torch_nvcomp_allocator(nvcomp, device: torch.device) -> bool: + """Route nvcomp's per-decode device scratch through torch's caching allocator. + + Global and idempotent. Guarded: any API mismatch or failure leaves nvcomp on + its default allocator (decode still works, just slower) and warns once. + + nvcomp API (from the wheel's ``set_device_allocator`` docstring): the + allocator is ``allocator(nbytes: int, stream: nvcomp.Stream) -> obj`` where + ``obj`` has an integer ``.ptr`` and releases its memory when garbage + collected. We ignore nvcomp's ``stream`` arg and instead read the reader + thread's torch stream from ``_fpz_nvcomp_alloc_tls`` (set per thread), which + is the stream nvcomp actually decodes on. + """ + global _FPZ_NVCOMP_ALLOC_INSTALLED, _FPZ_NVCOMP_ALLOC_WARNED + if _FPZ_NVCOMP_ALLOC_INSTALLED: + return True + dev_index = ( + device.index if device.index is not None else torch.cuda.current_device() + ) + + def _alloc(nbytes, stream=None): + return _TorchNvcompDeviceBuffer( + int(nbytes), dev_index, getattr(_fpz_nvcomp_alloc_tls, "stream", None) + ) + + try: + nvcomp.set_device_allocator(_alloc) + except Exception: + if not _FPZ_NVCOMP_ALLOC_WARNED: + _FPZ_NVCOMP_ALLOC_WARNED = True + warnings.warn( + "Could not install the torch caching allocator into nvcomp " + "(set_device_allocator failed); nvcomp keeps its default " + "cudaMalloc allocator. Set FLASHPACK_FPZ_GPU_TORCH_ALLOC=0 to " + "silence.", + RuntimeWarning, + stacklevel=2, + ) + return False + _FPZ_NVCOMP_ALLOC_INSTALLED = True + return True + + +def fpz_gpu_warmup(device: "str | torch.device" = "cuda") -> bool: + """Pay the batched GPU decoder's one-time init cost off the hot path. + + The first ``nvcompBatchedZstdDecompressAsync`` launch in a process pays + CUDA module loading for nvcomp's decompress kernels (measured 4-25s on + H200 under default lazy loading; ~1.6s residual with + ``CUDA_MODULE_LOADING=EAGER`` set before the first CUDA call, which is + the recommended companion setting). Apps can call this from ``setup()`` + -- e.g. while weights download -- so the first real load doesn't pay it. + + Decodes one tiny zstd chunk through the batched path end to end. Safe + no-op returning ``False`` when CUDA, libnvcomp, or zstandard is + unavailable; returns ``True`` only when the warmup decode round-tripped. + """ + try: + if not torch.cuda.is_available(): + return False + from . import _nvcomp_ll + + ll = _nvcomp_ll.load() + if ll is None: + return False + zstandard = require_zstandard() + dev = torch.device(device) + usize = 4096 + payload = zstandard.ZstdCompressor(level=1).compress(b"\x00" * usize) + src = torch.frombuffer(bytearray(payload), dtype=torch.uint8).to(dev) + dst = torch.empty(usize, dtype=torch.uint8, device=dev) + # Single-chunk device tables: [src ptr, src len, dst capacity, dst ptr] + tab = torch.tensor( + [src.data_ptr(), len(payload), usize, dst.data_ptr()], + dtype=torch.int64, + device=dev, + ) + actual = torch.zeros(1, dtype=torch.int64, device=dev) + statuses = torch.full((1,), -1, dtype=torch.int32, device=dev) + temp_bytes = ll.temp_size(1, usize, usize) + temp = torch.empty(max(1, temp_bytes), dtype=torch.uint8, device=dev) + stream = torch.cuda.current_stream(dev) + ll.decompress_async( + tab[0:1].data_ptr(), + tab[1:2].data_ptr(), + tab[2:3].data_ptr(), + actual.data_ptr(), + 1, + temp.data_ptr(), + temp_bytes, + tab[3:4].data_ptr(), + statuses.data_ptr(), + stream.cuda_stream, + ) + stream.synchronize() + return int(statuses.item()) == 0 and int(actual.item()) == usize + except Exception: + return False + + +# GPU-decode tuning knobs (env-overridable). +# +# Pipeline math (why these defaults). With v2 the GPU decode is fast and +# parallel, so the bottleneck moves to the file read: at the measured ~0.8 GB/s +# per-thread FUSE rate, and reading ~0.6x the logical bytes (the raw low plane +# plus the compressed high plane), hitting ~18 GB/s logical needs +# 0.6*18/0.8 ~= 14 read threads. So we restore the CPU path's read parallelism +# (~16 threads) instead of the 2 that the sync-free round used. Each thread +# reads AND decodes; preads (GIL released) overlap across threads and decodes +# overlap reads via the per-thread double-buffered slots. +# +# Memory: each slot holds the low / compressed-high / decompressed-high device +# planes (~3 x FPZ_FRAME_UNCOMPRESSED_BYTES/2 per frame) plus pinned host +# buffers (~2 x). Per-thread device ~= n_slots * batch_frames * 3 * 32 MiB and +# pinned ~= n_slots * batch_frames * 2 * 32 MiB; total scales by thread count. +# batch_frames=1 keeps per-thread memory small so we can afford ~16 threads +# (16*2*1*160 MiB ~= 5 GB device, ~2 GB pinned) -- and one 64 MiB frame already +# holds ~512 hi chunks, which is plenty of work for one nvcomp batched decode. +_FPZ_GPU_DEFAULT_THREADS = 16 +_FPZ_GPU_DEFAULT_BATCH_FRAMES = 1 +_FPZ_GPU_DEFAULT_BATCH_BYTES = 1024 * 1024 * 1024 # summed uncompressed per batch +_FPZ_GPU_DEFAULT_SLOTS = 2 # double-buffer depth per thread + + +def _fpz_hi_chunk_usizes(half: int, chunk_u: int) -> list[int]: + """Uncompressed sizes of a v2 frame's high-plane chunks (pure function). + + ``half`` bytes split into ``chunk_u``-sized pieces, the last holding the + remainder. Matches the encoder's chunking, and is the per-element shape the + GPU decode's DecompressConfig is keyed on. + """ + if half < 0 or chunk_u < 1: + raise ValueError("half must be >= 0 and chunk_u >= 1") + full, rem = divmod(half, chunk_u) + sizes = [chunk_u] * full + if rem: + sizes.append(rem) + return sizes + + +def _fpz_batch_signature(batch: list[tuple]) -> tuple[int, ...]: + """Config-cache key for a frame batch: the per-frame decompressed high-plane + sizes (``n_out // 2``), in order (pure function). + + An nvcomp ``DecompressConfig`` built from one batch is reusable for any other + batch with the same per-element uncompressed shape, so batches that share + this signature share a single config -- and the one-time + ``decompression_config`` stream sync is paid once per distinct signature + instead of once per decode call. + """ + return tuple(int(frame["n_out"]) // 2 for _, _blk, frame, _out_pos in batch) + + +def plan_fpz_gpu_batches( + frame_tasks: list[tuple], + max_batch_frames: int, + max_batch_bytes: int, +) -> list[list[tuple]]: + """Group fpz frame tasks into nvcomp batched-decode groups (pure function). + + ``frame_tasks`` are the ``("frame", block_idx, frame, out_pos)`` items from + :func:`_fpz_frame_tasks` (raw-block tasks are handled separately). Frames are + grouped in file order into batches bounded by BOTH a frame count + (``max_batch_frames``) and a summed uncompressed-output-byte budget + (``max_batch_bytes``). The byte budget matters because each fpz frame is up + to ``FPZ_FRAME_UNCOMPRESSED_BYTES`` (64 MiB) and every batched frame needs + device staging proportional to that size, so an unbounded batch would blow + the GPU memory budget. + + A single frame at or above the byte budget still forms its own size-1 batch + (the budget never drops a frame). An empty input yields an empty list. + """ + if max_batch_frames < 1: + raise ValueError("max_batch_frames must be >= 1") + if max_batch_bytes < 1: + raise ValueError("max_batch_bytes must be >= 1") + + batches: list[list[tuple]] = [] + current: list[tuple] = [] + current_bytes = 0 + for task in frame_tasks: + n_out = int(task[2]["n_out"]) + if current and ( + len(current) >= max_batch_frames or current_bytes + n_out > max_batch_bytes + ): + batches.append(current) + current = [] + current_bytes = 0 + current.append(task) + current_bytes += n_out + if current: + batches.append(current) + return batches + + +def _fpz_read_into_cuda_storage_gpu( + path: str, + specs: list[MacroblockSpec], + blocks: list[torch.Tensor], + device: torch.device, + nvcomp, +) -> None: + """GPU-decode variant of :func:`_fpz_read_into_cuda_storage` for v2 packs. + + v2 stores each frame's high plane as many small independent zstd chunks + (``FPZ_HI_CHUNK_UNCOMPRESSED_BYTES`` each). That is nvcomp's native shape: + the whole point of the GPU decoder is decoding MANY chunks in parallel. A v1 + single-frame high plane is one nvcomp chunk and decodes serially (~0.6 GB/s + measured, tier-flat), which is why the caller routes v1 to the CPU path. + + Per reader thread: read a frame's low plane and its whole compressed-high + blob into pinned staging, H2D both (moving the compressed high plane cuts + PCIe traffic ~2.4x), then submit ALL of the frame's high chunks as one + ``codec.decode`` batch (hundreds of Arrays), decoding straight into the + device high-plane staging; finally the same strided interleave + (``dst[0::2] = lo``, ``dst[1::2] = hi``). + + Two throughput levers, both load-bearing: + + * No per-decode sync. The naive path makes ``decode`` call + ``configure_decompression`` (a stream sync) every call. We build a + reusable ``DecompressConfig`` per distinct chunk-shape signature (one sync + each) and pass it to ``decode``; since v2 chunks are almost all a uniform + 64 KiB, that is ~1-2 configs total per thread and every steady-state decode + is sync-free. Reuse safety without the sync comes from an event-gated + double buffer (``n_slots`` staging sets; ``synchronize`` a slot's event + before reusing it, ``record`` it after decode+interleave). + * Read parallelism. Read (not decode) is now the bottleneck, so we use many + threads (see the tuning-knob pipeline math); preads overlap across threads + and decodes overlap reads via the slots. + + Each thread owns its fd, stream, Codec and config cache. Raw (uncompressed) + blocks take the same whole-block H2D as the CPU-decode path. + """ + half_cap = FPZ_FRAME_UNCOMPRESSED_BYTES // 2 + # A pack is written with one chunk size and one chunk-start alignment; + # read both from the first fpz block (absent for pre-parameterization v2 + # packs -> the 64 KiB default / packed layout). A frame whose block + # disagrees is caught by the chunk-count check in the loop. + chunk_u = FPZ_HI_CHUNK_UNCOMPRESSED_BYTES + hi_align = 1 + for spec in specs: + if spec.fpz is not None: + chunk_u = int(spec.fpz.get("hi_chunk_usize", chunk_u)) + hi_align = int(spec.fpz.get("hi_align", 1)) + break + max_chunks = (half_cap + chunk_u - 1) // chunk_u + # Upper bound on a frame's whole compressed-high blob: the zstd bound for + # half_cap uncompressed, plus per-chunk zstd frame-header overhead and + # chunk-start padding; aligned so per-frame staging bases (k * comp_cap) + # preserve the chunk-start alignment inside device staging. + comp_cap = half_cap + (half_cap // 255) + max_chunks * 80 + 4096 + comp_cap = _align_up(comp_cap, max(16, hi_align)) + + # Batched C-API decode (the "ll" path): one foreign call per batch over + # device-resident chunk tables instead of one pybind Array per chunk. + # Requires the pack's chunk starts to satisfy nvcomp's queried input + # alignment (hi_align-padded packs do; legacy packed layouts fall back). + # Fused single-pass interleave (Triton) vs the strided two-pass copies; + # opt-in while gating, falls back automatically when triton is missing. + fused = _env_flag("FLASHPACK_FPZ_FUSED_INTERLEAVE") and ( + _interleave.fused_interleave_available() + ) + + ll = None + if _env_flag("FLASHPACK_FPZ_GPU_LL"): + from . import _nvcomp_ll + + ll = _nvcomp_ll.load() + if ll is not None: + req_in, req_out, _req_temp = ll.alignments() + if hi_align % req_in != 0 or chunk_u % req_out != 0: + warnings.warn( + f"flashpack: fpz pack chunk alignment (hi_align={hi_align}" + f", chunk_u={chunk_u}) does not satisfy nvcomp's batched " + f"decode requirements (input={req_in}, output={req_out}); " + "using the wrapper decode path. Repack with current " + "flashpack for the batched path.", + RuntimeWarning, + stacklevel=2, + ) + ll = None + + n_threads = effective_read_threads( + _env_int("FLASHPACK_FPZ_GPU_DECODE_THREADS", _FPZ_GPU_DEFAULT_THREADS) + ) + batch_frames = max( + 1, _env_int("FLASHPACK_FPZ_GPU_BATCH_FRAMES", _FPZ_GPU_DEFAULT_BATCH_FRAMES) + ) + batch_bytes = max( + 1, _env_int("FLASHPACK_FPZ_GPU_BATCH_BYTES", _FPZ_GPU_DEFAULT_BATCH_BYTES) + ) + n_slots = max(1, _env_int("FLASHPACK_FPZ_GPU_SLOTS", _FPZ_GPU_DEFAULT_SLOTS)) + + byte_blocks = [b.view(torch.uint8) for b in blocks] + + # Order every reader stream after the destination allocation (same + # wait_event pattern as parallel_read_into_storage / the CPU-decode path). + alloc_ready = torch.cuda.Event() + alloc_ready.record(torch.cuda.current_stream(device)) + + tasks = _fpz_frame_tasks(specs) + raw_tasks = [task for task in tasks if task[0] == "raw"] + frame_tasks = [task for task in tasks if task[0] == "frame"] + frame_batches = plan_fpz_gpu_batches(frame_tasks, batch_frames, batch_bytes) + + # Work items: raw blocks (whole-block H2D) and frame batches (GPU decode). + work: queue.SimpleQueue = queue.SimpleQueue() + for task in raw_tasks: + work.put(("raw", task)) + for batch in frame_batches: + work.put(("batch", batch)) + n_threads = min(n_threads, max(1, work.qsize())) + for _ in range(n_threads): + work.put(None) + + errors: list[BaseException] = [] + trace_on = _env_flag("FLASHPACK_FPZ_GPU_TRACE") + traces: list[str] = [] + traces_lock = threading.Lock() + + # Route nvcomp's per-decode scratch through torch's caching allocator to kill + # the per-call cudaMalloc/cudaFree device sync (the round-2 bottleneck). + # Global + idempotent + guarded; disable with FLASHPACK_FPZ_GPU_TORCH_ALLOC=0. + # (Wrapper path only: the ll path manages its own torch-allocated scratch.) + if ll is None and _env_flag_default("FLASHPACK_FPZ_GPU_TORCH_ALLOC", True): + _install_torch_nvcomp_allocator(nvcomp, device) + + def _reader(thread_idx: int) -> None: + try: + fd = os.open(path, os.O_RDONLY) + stream = torch.cuda.Stream(device=device) + stream.wait_event(alloc_ready) + # nvcomp's device allocator (if installed) reads this thread's stream + # from the thread-local, so decode scratch is tied to the decode + # stream and torch won't reuse it out from under an in-flight decode. + _fpz_nvcomp_alloc_tls.stream = stream + # Resolve the concrete ordinal in this thread: an indexless "cuda" + # device places both the stream and the staging tensors on this + # thread's current device, and the Codec must match or nvcomp raises + # "Input array and Codec device id mismatched". + device_id = ( + device.index + if device.index is not None + else torch.cuda.current_device() + ) + codec = None + if ll is None: + codec = nvcomp.Codec( + algorithm="Zstd", + bitstream_kind=nvcomp.BitstreamKind.RAW, + device_id=device_id, + cuda_stream=stream.cuda_stream, + ) + # One contiguous staging set per slot; frame k lives at k*half_cap + # (low / decompressed-high) or k*comp_cap (compressed-high). Slices + # are contiguous, so nvcomp.as_array wraps them zero-copy. + lo_pin = [ + torch.empty(batch_frames * half_cap, dtype=torch.uint8, pin_memory=True) + for _ in range(n_slots) + ] + hiz_pin = [ + torch.empty(batch_frames * comp_cap, dtype=torch.uint8, pin_memory=True) + for _ in range(n_slots) + ] + lo_pin_view = [memoryview(b.numpy()) for b in lo_pin] + hiz_pin_view = [memoryview(b.numpy()) for b in hiz_pin] + lo_dev = [ + torch.empty(batch_frames * half_cap, dtype=torch.uint8, device=device) + for _ in range(n_slots) + ] + hiz_dev = [ + torch.empty(batch_frames * comp_cap, dtype=torch.uint8, device=device) + for _ in range(n_slots) + ] + hi_dev = [ + torch.empty(batch_frames * half_cap, dtype=torch.uint8, device=device) + for _ in range(n_slots) + ] + # Wrapper path: hoist the decode out= wrappers -- one nvcomp.Array + # per (slot, frame, chunk) over a fixed chunk_u slice at frame k's + # chunk j offset, built ONCE (layout is data-independent). decode + # writes the true (config-driven) size <= chunk_u into each, so the + # same wrappers serve every batch. Indexed [slot][k*max_chunks + j]. + # (The compressed-in src wrappers stay per-batch, sized to each + # chunk's exact compressed length, so nvcomp sees exactly one zstd + # frame per Array.) + out_wrap = None + if ll is None: + out_wrap = [ + [ + nvcomp.as_array( + hi_dev[s].narrow(0, k * half_cap + j * chunk_u, chunk_u) + ) + for k in range(batch_frames) + for j in range(max_chunks) + ] + for s in range(n_slots) + ] + else: + # ll path: no per-chunk Python objects at all. Per batch we + # fill ONE pinned int64 table -- rows (src rel offset, src + # size, dst rel offset, dst capacity) -- with vectorized numpy + # over the footer's chunk lengths, H2D it, add the staging + # base addresses on-device, and make one foreign call. The + # pinned table is per SLOT (host reuse is gated by the slot + # event, like the other pinned staging); the device table, + # scratch, and result arrays are per thread (reuse is + # stream-ordered). + cap_chunks = batch_frames * max_chunks + ll_temp_bytes = ll.temp_size( + cap_chunks, chunk_u, batch_frames * half_cap + ) + ll_temp = torch.empty( + max(1, ll_temp_bytes), dtype=torch.uint8, device=device + ) + ll_tab_pin = [ + torch.empty((4, cap_chunks), dtype=torch.int64, pin_memory=True) + for _ in range(n_slots) + ] + ll_tab_np = [t.numpy() for t in ll_tab_pin] + ll_tab_dev = torch.empty( + (4, cap_chunks), dtype=torch.int64, device=device + ) + ll_actual = torch.empty(cap_chunks, dtype=torch.int64, device=device) + ll_statuses = torch.empty(cap_chunks, dtype=torch.int32, device=device) + # Stream-side correctness accumulators: per-chunk statuses and + # actual-size mismatches fold into two scalars ON the decode + # stream (no syncs); read once after the final synchronize. + ll_status_max = torch.zeros((), dtype=torch.int32, device=device) + ll_size_bad = torch.zeros((), dtype=torch.bool, device=device) + ll_usizes_cache: dict[int, np.ndarray] = {} + ll_dst_rel_cache: dict[tuple[int, int], np.ndarray] = {} + # Recorded now so the first synchronize on any slot is a no-op. + events = [torch.cuda.Event() for _ in range(n_slots)] + for ev in events: + ev.record(stream) + # Per-thread cache: batch shape signature -> reusable DecompressConfig. + configs: dict[tuple[int, ...], object] = {} + batch_idx = 0 + t_pread = t_h2d = t_decode = t_interleave = t_evsync = t_final = 0.0 + t_wrap = 0.0 + n_batches = n_frames_done = n_cfg = 0 + try: + while True: + item = work.get() + if item is None: + break + kind, payload = item + if kind == "raw": + _, blk, _frame, _out_pos = payload + spec = specs[blk] + buf = bytearray(spec.length_bytes) + _pread_into(fd, spec.offset_bytes, memoryview(buf)) + host = torch.frombuffer(buf, dtype=torch.uint8) + stream.synchronize() + with torch.cuda.stream(stream): + byte_blocks[blk].copy_(host, non_blocking=False) + continue + + batch = payload + slot = batch_idx % n_slots + batch_idx += 1 + # Wait for this slot's previous batch (its interleave) before + # overwriting its pinned/device buffers -- decode no longer + # synchronizes, so this event is what keeps reuse safe. + _t = time.perf_counter() if trace_on else 0.0 + events[slot].synchronize() + if trace_on: + t_evsync += time.perf_counter() - _t + + halves: list[int] = [] + srcs: list = [] + outs: list = [] + sig_parts: list[int] = [] + n_ll = 0 # chunks staged into the ll table this batch + # Read every frame's planes into this slot's pinned staging, + # H2D them, and stage the per-chunk dispatch (wrapper: one + # src/out Array per chunk; ll: rows of the batch table). + for k, (_, blk, frame, _out_pos) in enumerate(batch): + payload_off = int(frame["payload_off"]) + lo_len = int(frame["lo_len"]) + n_out = int(frame["n_out"]) + half = n_out - lo_len + hi_chunks = frame.get("hi_chunks") + if hi_chunks is None: + raise ValueError( + "GPU fpz decode requires a v2 (chunked) pack" + ) + if lo_len != half or lo_len * 2 != n_out: + raise ValueError("fpz frame plane size mismatch") + clens = np.asarray(hi_chunks, dtype=np.int64) + m = int(clens.shape[0]) + # Chunk starts are hi_align-padded in the payload (and + # therefore in staging); hi_chunks holds true lengths. + aligned = clens + (-clens) % hi_align + hi_len_total = int(aligned.sum()) + if hi_len_total > comp_cap: + raise ValueError( + f"fpz compressed frame ({hi_len_total} bytes) exceeds " + f"staging capacity ({comp_cap} bytes)" + ) + usizes = _fpz_hi_chunk_usizes(half, chunk_u) + if len(usizes) != m: + raise ValueError("fpz v2 chunk count mismatch") + base = specs[blk].offset_bytes + payload_off + lo_off = k * half_cap + hiz_off = k * comp_cap + _t = time.perf_counter() if trace_on else 0.0 + _pread_into( + fd, base, lo_pin_view[slot][lo_off : lo_off + lo_len] + ) + _pread_into( + fd, + base + _align_up(lo_len, hi_align), + hiz_pin_view[slot][hiz_off : hiz_off + hi_len_total], + ) + if trace_on: + t_pread += time.perf_counter() - _t + halves.append(half) + _t = time.perf_counter() if trace_on else 0.0 + with torch.cuda.stream(stream): + lo_dev[slot].narrow(0, lo_off, half).copy_( + lo_pin[slot].narrow(0, lo_off, half), non_blocking=True + ) + hiz_dev[slot].narrow(0, hiz_off, hi_len_total).copy_( + hiz_pin[slot].narrow(0, hiz_off, hi_len_total), + non_blocking=True, + ) + if trace_on: + t_h2d += time.perf_counter() - _t + _t = time.perf_counter() if trace_on else 0.0 + if ll is None: + # One src Array per compressed chunk (exact length) + # and its hoisted out wrapper; chunk usizes drive + # the config. This per-chunk wrapper building is + # GIL-bound Python and is the dominant residual + # cost at small chunk sizes -- its own trace + # bucket so its share is visible. + coff = hiz_off + for j, (clen, alen) in enumerate( + zip(clens.tolist(), aligned.tolist()) + ): + srcs.append( + nvcomp.as_array( + hiz_dev[slot].narrow(0, coff, int(clen)) + ) + ) + outs.append(out_wrap[slot][k * max_chunks + j]) + coff += int(alen) + sig_parts.extend(usizes) + else: + # Vectorized table rows for this frame's chunks: + # (0) src offset within hiz staging = padded prefix + # sums, (1) true compressed length, (2) dst offset + # within hi staging, (3) expected uncompressed size. + tab = ll_tab_np[slot] + starts = np.empty(m, dtype=np.int64) + starts[0] = 0 + np.cumsum(aligned[: m - 1], out=starts[1:]) + tab[0, n_ll : n_ll + m] = hiz_off + starts + tab[1, n_ll : n_ll + m] = clens + dst_rel = ll_dst_rel_cache.get((k, m)) + if dst_rel is None: + dst_rel = k * half_cap + ( + np.arange(m, dtype=np.int64) * chunk_u + ) + ll_dst_rel_cache[(k, m)] = dst_rel + tab[2, n_ll : n_ll + m] = dst_rel + caps = ll_usizes_cache.get(half) + if caps is None: + caps = np.asarray(usizes, dtype=np.int64) + ll_usizes_cache[half] = caps + tab[3, n_ll : n_ll + m] = caps + n_ll += m + if trace_on: + t_wrap += time.perf_counter() - _t + + _t = time.perf_counter() if trace_on else 0.0 + if ll is None: + # Reusable config per chunk-shape signature: build once + # (one sync, waits on the H2D above), then decode + # sync-free here and on every later batch that shares + # the shape. + sig = tuple(sig_parts) + cfg = configs.get(sig) + if cfg is None: + cfg = codec.decompression_config(srcs) + configs[sig] = cfg + n_cfg += 1 + codec.decode(srcs, out=outs, decompression_config=cfg) + else: + # One H2D of the table, two on-device base-address + # adds, ONE foreign call for the whole batch -- Python + # cost is independent of the chunk count. The add + # outputs are fresh stream-local tensors; nvcomp reads + # them during the (stream-ordered) decode, so dropping + # the Python refs afterwards is safe. + with torch.cuda.stream(stream): + ll_tab_dev[:, :n_ll].copy_( + ll_tab_pin[slot][:, :n_ll], non_blocking=True + ) + src_ptrs = ll_tab_dev[0, :n_ll] + hiz_dev[slot].data_ptr() + dst_ptrs = ll_tab_dev[2, :n_ll] + hi_dev[slot].data_ptr() + ll.decompress_async( + src_ptrs.data_ptr(), + ll_tab_dev[1].data_ptr(), + ll_tab_dev[3].data_ptr(), + ll_actual.data_ptr(), + n_ll, + ll_temp.data_ptr(), + ll_temp_bytes, + dst_ptrs.data_ptr(), + ll_statuses.data_ptr(), + stream.cuda_stream, + ) + with torch.cuda.stream(stream): + torch.maximum( + ll_status_max, + ll_statuses[:n_ll].max(), + out=ll_status_max, + ) + torch.logical_or( + ll_size_bad, + (ll_actual[:n_ll] != ll_tab_dev[3, :n_ll]).any(), + out=ll_size_bad, + ) + if trace_on: + t_decode += time.perf_counter() - _t + + # Interleave per frame (same invariant as the CPU path): + # even bytes low plane, odd bytes high plane. Fused = one + # kernel pass; strided = two 2-byte-stride copy passes. + _t = time.perf_counter() if trace_on else 0.0 + for k, (_, blk, frame, out_pos) in enumerate(batch): + half = halves[k] + n_out = int(frame["n_out"]) + seg = byte_blocks[blk].narrow(0, out_pos, n_out) + with torch.cuda.stream(stream): + if fused: + _interleave.interleave_into( + seg, + lo_dev[slot].narrow(0, k * half_cap, half), + hi_dev[slot].narrow(0, k * half_cap, half), + ) + else: + seg[0::2].copy_( + lo_dev[slot].narrow(0, k * half_cap, half), + non_blocking=True, + ) + seg[1::2].copy_( + hi_dev[slot].narrow(0, k * half_cap, half), + non_blocking=True, + ) + events[slot].record(stream) + if trace_on: + t_interleave += time.perf_counter() - _t + n_batches += 1 + n_frames_done += len(batch) + finally: + os.close(fd) + _t = time.perf_counter() if trace_on else 0.0 + stream.synchronize() + _fpz_nvcomp_alloc_tls.stream = None + if ll is not None: + # Deferred per-chunk verification: both scalars were folded on + # the decode stream per batch, so this is the only D2H. + status_val = int(ll_status_max.item()) + if status_val != 0: + raise RuntimeError( + "fpz batched GPU decode reported a per-chunk error: " + + ll.status_string(status_val) + ) + if bool(ll_size_bad.item()): + raise ValueError( + "fpz batched GPU decode produced a chunk size mismatch" + ) + if trace_on: + t_final += time.perf_counter() - _t + # Enqueue phases (h2d, interleave) are async so their wall is + # small; a large `decode` wall means the decode CALL itself + # blocks (internal sync / scratch alloc), while a large + # `evsync`/`final` means the pipeline is GPU-bound waiting on + # decode+interleave to finish. In ll mode `wrap` is the numpy + # table fill and `decode` is the table H2D + foreign call. + mode = "ll" if ll is not None else "wrapper" + line = ( + f"[fpz-gpu-trace] thread={thread_idx} mode={mode} " + f"frames={n_frames_done} " + f"batches={n_batches} cfg_builds={n_cfg} " + f"pread={t_pread:.3f}s h2d_enq={t_h2d:.3f}s " + f"wrap={t_wrap:.3f}s decode={t_decode:.3f}s " + f"interleave_enq={t_interleave:.3f}s " + f"evsync={t_evsync:.3f}s final_sync={t_final:.3f}s" + ) + with traces_lock: + traces.append(line) + except BaseException as e: + errors.append(e) + + threads = [ + threading.Thread(target=_reader, args=(i,), daemon=True) + for i in range(n_threads) + ] + for t in threads: + t.start() + for t in threads: + t.join() + torch.cuda.synchronize(device) + if trace_on: + for line in traces: + print(line) + if errors: + raise errors[0] + + +_FPZ_V1_GPU_WARNED = False + + +def _fpz_specs_all_v2(specs: list[MacroblockSpec]) -> bool: + """True if every fpz block uses the v2 (chunked) codec -- the only format + the GPU decoder can decompress in parallel.""" + for spec in specs: + if spec.fpz is not None and spec.fpz.get("codec") != FPZ_CODEC_SPLITPLANE_V2: + return False + return True + + +def _read_fpz_into_storage( + path: str, specs: list[MacroblockSpec], device: torch.device +) -> FlashTensorStorage: + """Materialize an fpz (partially compressed) pack into ``device`` storage.""" + global _FPZ_V1_GPU_WARNED + if device.type == "cpu": + storage = _allocate_aligned_cpu_storage(specs) + _fpz_read_into_cpu_storage(path, specs, storage.blocks) + return storage + if device.type == "cuda": + storage = _allocate_empty_storage(specs, device) + nvcomp = _load_nvcomp() if _fpz_gpu_decode_enabled() else None + # The GPU decoder only helps v2 (chunked) packs; a v1 pack is one nvcomp + # chunk per frame and decodes serially, so fall back to the threaded CPU + # decode for it (still correct, and faster than serial GPU decode). + if nvcomp is not None and not _fpz_specs_all_v2(specs): + nvcomp = None + if not _FPZ_V1_GPU_WARNED: + _FPZ_V1_GPU_WARNED = True + warnings.warn( + "FLASHPACK_FPZ_GPU_DECODE=1 but this pack uses the v1 fpz " + "codec, which cannot be GPU-decoded in parallel; using the " + "CPU decode path. Repack with the v2 encoder for GPU decode.", + RuntimeWarning, + stacklevel=2, + ) + if nvcomp is not None: + _fpz_read_into_cuda_storage_gpu(path, specs, storage.blocks, device, nvcomp) + else: + _fpz_read_into_cuda_storage(path, specs, storage.blocks, device) + return storage + raise ValueError(f"Unsupported device: {device}") + + def read_flashpack_file( path: str, device: str | torch.device = "cpu", @@ -335,6 +1630,11 @@ def read_flashpack_file( specs = _build_macroblock_specs(meta) device = torch.device(device) if isinstance(device, str) else device + if any(spec.fpz is not None for spec in specs): + with timer("read_fpz", silent): + storage = _read_fpz_into_storage(path, specs, device) + return storage, meta + if device.type == "cpu": if parallel_read_supported(device): # Opt-in eager path (FLASHPACK_CPU_PARALLEL_READ=1): materialize diff --git a/src/flashpack/serialization.py b/src/flashpack/serialization.py index 5e1262c..b81ea99 100644 --- a/src/flashpack/serialization.py +++ b/src/flashpack/serialization.py @@ -11,12 +11,26 @@ from .constants import ( DEFAULT_ALIGN_BYTES, DEFAULT_NUM_WRITE_WORKERS, + DEFAULT_ZSTD_LEVEL, FILE_FORMAT_V3, FILE_FORMAT_V4, + FPZ_CODEC_SPLITPLANE_V1, + FPZ_CODEC_SPLITPLANE_V2, + FPZ_COMPRESS_BF16, + FPZ_FRAME_ALIGN_BYTES, + FPZ_FRAME_UNCOMPRESSED_BYTES, + FPZ_HI_CHUNK_ALIGN_BYTES, + FPZ_HI_CHUNK_UNCOMPRESSED_BYTES, MAGIC, U64LE, ) -from .utils import dtype_to_string, get_packing_dtype, timer, torch_dtype_to_numpy_dtype +from .utils import ( + dtype_to_string, + get_packing_dtype, + require_zstandard, + timer, + torch_dtype_to_numpy_dtype, +) @dataclass @@ -38,6 +52,35 @@ class MacroblockPlan: tensors: list[TensorIndexRecord] +def _resolve_hi_chunk_bytes(hi_chunk_bytes: int | None) -> int: + """Resolve and validate the v2 high-plane chunk size. + + Precedence: explicit ``hi_chunk_bytes`` arg > ``FLASHPACK_FPZ_CHUNK_BYTES`` + env (for the converter) > ``FPZ_HI_CHUNK_UNCOMPRESSED_BYTES`` default. An + explicitly requested value must be a positive multiple of + ``FPZ_FRAME_ALIGN_BYTES`` and no larger than a frame's high plane + (``FPZ_FRAME_UNCOMPRESSED_BYTES // 2``). The default is clamped to the frame + high plane rather than rejected (a chunk >= the high plane just yields one + chunk per frame -- the case tests hit by shrinking the frame size). + """ + half_frame = FPZ_FRAME_UNCOMPRESSED_BYTES // 2 + env = os.environ.get("FLASHPACK_FPZ_CHUNK_BYTES") + if hi_chunk_bytes is None and not env: + return max(1, min(FPZ_HI_CHUNK_UNCOMPRESSED_BYTES, half_frame)) + requested = hi_chunk_bytes if hi_chunk_bytes is not None else int(env) + if requested < FPZ_FRAME_ALIGN_BYTES or requested % FPZ_FRAME_ALIGN_BYTES: + raise ValueError( + f"hi_chunk_bytes must be a positive multiple of " + f"{FPZ_FRAME_ALIGN_BYTES} (got {requested})" + ) + if requested > half_frame: + raise ValueError( + f"hi_chunk_bytes ({requested}) exceeds the frame high-plane " + f"size ({half_frame})" + ) + return requested + + def pack_to_file( state_dict_or_model: dict[str, torch.Tensor] | torch.nn.Module, destination_path: str, @@ -46,10 +89,27 @@ def pack_to_file( align_bytes: int = DEFAULT_ALIGN_BYTES, silent: bool = True, num_workers: int = DEFAULT_NUM_WRITE_WORKERS, + compress: str | None = None, + hi_chunk_bytes: int | None = None, ) -> None: """ Pack the state dictionary or model to a flashpack file. + + ``compress="fpz-bf16"`` enables split-plane zstd compression for bf16 + macroblocks only (see ``constants.py``); every other dtype is stored + uncompressed, and the file falls back to the plain uncompressed format + when no bf16 macroblock is present. Requires the optional ``zstandard`` + package. ``hi_chunk_bytes`` overrides the v2 high-plane chunk size (default + ``FPZ_HI_CHUNK_UNCOMPRESSED_BYTES``, or ``FLASHPACK_FPZ_CHUNK_BYTES``); + larger chunks mean fewer per-chunk wrapper objects for the GPU decoder. """ + if compress is not None and compress != FPZ_COMPRESS_BF16: + raise ValueError( + f"Unsupported compress option: {compress!r} " + f"(expected None or {FPZ_COMPRESS_BF16!r})" + ) + resolved_hi_chunk_bytes = _resolve_hi_chunk_bytes(hi_chunk_bytes) + if isinstance(state_dict_or_model, torch.nn.Module): state_dict = state_dict_or_model.state_dict() else: @@ -154,6 +214,32 @@ def _lcm(a: int, b: int) -> int: dest_dir = os.path.dirname(os.path.abspath(destination_path)) or "." os.makedirs(dest_dir, exist_ok=True) + + # fpz path: bf16 macroblocks are split-plane zstd-compressed, every other + # dtype is stored verbatim. When no block is eligible the file is identical + # to the uncompressed pack, so fall through to the plain path below. + compress_flags = [ + compress == FPZ_COMPRESS_BF16 and block.dtype is torch.bfloat16 + for block in macroblocks + ] + if any(compress_flags): + # Single pass, no uncompressed scratch: convert each tensor to CPU + # bytes and stream them through a rolling frame buffer directly into + # the final compressed file. + with timer("fpz_stream_write", silent): + _write_fpz_pack_streaming( + state_dict=state_dict, + macroblocks=macroblocks, + index=index, + align_bytes=align_bytes, + compress_flags=compress_flags, + destination_path=destination_path, + dest_dir=dest_dir, + silent=silent, + hi_chunk_bytes=resolved_hi_chunk_bytes, + ) + return + fd_tmp = None tmp_path = None @@ -345,3 +431,324 @@ def copy_one(rec: TensorIndexRecord) -> None: os.remove(tmp_path) except OSError: pass + + +def _write_zeros(f, n: int) -> None: + """Write ``n`` zero bytes to ``f`` in bounded chunks.""" + while n > 0: + take = min(n, 1 << 20) + f.write(b"\x00" * take) + n -= take + + +def _iter_block_uncompressed_chunks( + block: MacroblockPlan, + state_dict: dict[str, torch.Tensor], + progress: "tqdm.tqdm | None" = None, +): + """Yield a macroblock's uncompressed payload in order, reproducing the + memmap layout exactly. + + Emits ``("zeros", nbytes)`` for the inter-tensor element-alignment gaps + (zero-filled, as a fresh memmap is) and ``("bytes", uint8_ndarray)`` for + each tensor -- the target-dtype CPU reinterpretation (packing view) that + the uncompressed copy loop writes. ``.to(device="cpu")`` handles the D2H + transfer for GPU-source tensors. + """ + elem_size = torch.tensor([], dtype=block.dtype).element_size() + packing_dtype = get_packing_dtype(block.dtype) + cursor_elems = 0 + for rec in block.tensors: + if rec.offset > cursor_elems: + yield ("zeros", (rec.offset - cursor_elems) * elem_size) + cursor_elems = rec.offset + src = state_dict[rec.name] + src_cpu = src.view(-1).to(dtype=block.dtype, device="cpu") + if block.dtype != packing_dtype: + src_cpu = src_cpu.view(packing_dtype) + raw = src_cpu.contiguous().view(torch.uint8).numpy() + yield ("bytes", raw) + cursor_elems += rec.length + if progress is not None: + progress.update(1) + if block.total_elems > cursor_elems: + yield ("zeros", (block.total_elems - cursor_elems) * elem_size) + + +def _fpz_encode_frame(f, block_start: int, frame_u8: np.ndarray, compressor) -> dict: + """Encode one split-plane zstd frame from ``frame_u8`` (the uncompressed + bytes of a single frame) and write it to ``f``. + + bf16 elements are little-endian, so even bytes are the low (mantissa-LSB) + plane -- kept raw -- and odd bytes are the high (sign+exponent) plane -- + zstd-compressed. The frame payload start is padded to a 4096-byte boundary + relative to the macroblock start. + """ + n_out = int(frame_u8.shape[0]) + lo = np.ascontiguousarray(frame_u8[0::2]) + hi = np.ascontiguousarray(frame_u8[1::2]) + lo_bytes = lo.tobytes() + hi_z = compressor.compress(hi.tobytes()) + + payload_off = f.tell() - block_start + pad = (-payload_off) % FPZ_FRAME_ALIGN_BYTES + if pad: + f.write(b"\x00" * pad) + payload_off += pad + + f.write(lo_bytes) + f.write(hi_z) + return { + "payload_off": int(payload_off), + "lo_len": int(len(lo_bytes)), + "hi_len": int(len(hi_z)), + "n_out": int(n_out), + } + + +# Default fpz codec version written by the streaming encoder. v2 (chunked high +# plane) is the GPU-decodable format; tests set this to 1 to exercise the +# v1-still-reads backward-compatibility path. +_DEFAULT_FPZ_VERSION = 2 + + +def _fpz_encode_frame_v2( + f, block_start: int, frame_u8: np.ndarray, chunk_compressor, chunk: int +) -> dict: + """Encode one split-plane frame with a CHUNKED high plane (codec v2). + + Same low/high split as v1, but the high plane is compressed as a sequence of + independent zstd frames of ``chunk`` uncompressed bytes each (the frame's + last chunk holds the remainder). Many small chunks are what a GPU decoder + needs to decompress in parallel; the frame record lists each chunk's + compressed length so the reader locates them by prefix sum. Larger chunks + mean fewer per-chunk wrapper objects for the GPU decoder to build (the read + bottleneck once decode is parallel) at the cost of slightly less parallelism + and a hair less ratio. Ratio drops slightly versus v1 because each chunk + compresses without the neighbouring chunks' context. + """ + n_out = int(frame_u8.shape[0]) + lo = np.ascontiguousarray(frame_u8[0::2]) + hi = np.ascontiguousarray(frame_u8[1::2]) + lo_bytes = lo.tobytes() + half = int(hi.shape[0]) + hi_z_chunks = [ + chunk_compressor.compress(hi[off : off + chunk].tobytes()) + for off in range(0, half, chunk) + ] + + payload_off = f.tell() - block_start + pad = (-payload_off) % FPZ_FRAME_ALIGN_BYTES + if pad: + f.write(b"\x00" * pad) + payload_off += pad + + f.write(lo_bytes) + # Pad after the lo plane and after every chunk so each chunk STARTS + # hi_align-aligned within the (FPZ_FRAME_ALIGN_BYTES-aligned) payload: + # batched GPU decode requires aligned device chunk pointers. A full + # frame's lo plane (32 MiB) is already aligned, but the tail frame's + # arbitrary half-length is not. "hi_chunks" records TRUE zstd lengths; + # the reader recomputes padded offsets from the footer's "hi_align". + pad = (-len(lo_bytes)) % FPZ_HI_CHUNK_ALIGN_BYTES + if pad: + f.write(b"\x00" * pad) + hi_chunks: list[int] = [] + for z in hi_z_chunks: + f.write(z) + hi_chunks.append(int(len(z))) + pad = (-len(z)) % FPZ_HI_CHUNK_ALIGN_BYTES + if pad: + f.write(b"\x00" * pad) + return { + "payload_off": int(payload_off), + "lo_len": int(len(lo_bytes)), + "n_out": int(n_out), + "hi_chunks": hi_chunks, + } + + +def _fpz_stream_compress_block( + f, + block_start: int, + block: MacroblockPlan, + state_dict: dict[str, torch.Tensor], + encode_frame, + progress: "tqdm.tqdm | None", +) -> list[dict]: + """Stream one bf16 macroblock through a rolling ``FPZ_FRAME_UNCOMPRESSED_BYTES`` + buffer, emitting a split-plane frame (via ``encode_frame``) each time it + fills (and once more for the tail). Peak extra memory is one frame buffer + plus one source tensor.""" + frame_bytes = FPZ_FRAME_UNCOMPRESSED_BYTES + buf = np.empty(frame_bytes, dtype=np.uint8) + fill = 0 + frames: list[dict] = [] + + for kind, data in _iter_block_uncompressed_chunks(block, state_dict, progress): + if kind == "zeros": + remaining = data + while remaining > 0: + take = min(remaining, frame_bytes - fill) + buf[fill : fill + take] = 0 + fill += take + remaining -= take + if fill == frame_bytes: + frames.append(encode_frame(f, block_start, buf)) + fill = 0 + else: + arr = data + pos = 0 + n = int(arr.shape[0]) + while pos < n: + take = min(n - pos, frame_bytes - fill) + buf[fill : fill + take] = arr[pos : pos + take] + fill += take + pos += take + if fill == frame_bytes: + frames.append(encode_frame(f, block_start, buf)) + fill = 0 + + if fill > 0: + frames.append(encode_frame(f, block_start, buf[:fill])) + return frames + + +def _write_fpz_pack_streaming( + state_dict: dict[str, torch.Tensor], + macroblocks: list[MacroblockPlan], + index: list[TensorIndexRecord], + align_bytes: int, + compress_flags: list[bool], + destination_path: str, + dest_dir: str, + silent: bool, + hi_chunk_bytes: int, +) -> None: + """Write a compressed (fpz) pack to ``destination_path`` atomically in a + single pass -- no uncompressed scratch file. + + Each macroblock's payload is produced on the fly from ``state_dict`` (same + dtype conversion and inter-tensor alignment as the uncompressed planner) + and either streamed through split-plane zstd frames (bf16 blocks, per + ``compress_flags``) or written verbatim. The footer/frame format is + byte-compatible with the read path. ``hi_chunk_bytes`` is the v2 high-plane + chunk size, recorded per fpz block so the reader reproduces the chunking. + """ + zstandard = require_zstandard() + version = _DEFAULT_FPZ_VERSION + if version == 2: + # v2 compresses each high-plane chunk (hi_chunk_bytes uncompressed) as + # its own zstd frame. threads=-1 (one worker per core) does nothing for + # a small input and only adds per-call overhead, so use a single-threaded + # compressor; parallelism at repack time now comes from the many chunks, + # not from one big multithreaded compress. (Chunks are compressed + # serially here; a chunk-level thread pool is a repack-speed follow-up.) + chunk_compressor = zstandard.ZstdCompressor(level=DEFAULT_ZSTD_LEVEL) + + def encode_frame(f_, block_start_, frame_u8_): + return _fpz_encode_frame_v2( + f_, block_start_, frame_u8_, chunk_compressor, hi_chunk_bytes + ) + + codec_name = FPZ_CODEC_SPLITPLANE_V2 + else: + # threads=-1 = one worker per core: a ~19GB high plane at single-threaded + # zstd-3 (~0.4 GB/s) would take ~45 min per repack; multithreaded frames + # keep converter jobs in minutes. Frame outputs are byte-compatible. + compressor = zstandard.ZstdCompressor(level=DEFAULT_ZSTD_LEVEL, threads=-1) + + def encode_frame(f_, block_start_, frame_u8_): + return _fpz_encode_frame(f_, block_start_, frame_u8_, compressor) + + codec_name = FPZ_CODEC_SPLITPLANE_V1 + + fd_tmp, tmp_path = tempfile.mkstemp(dir=dest_dir, prefix=".packtmp_") + os.close(fd_tmp) + progress = None + if not silent: + progress = tqdm.tqdm(desc="Packing (fpz)", total=len(index)) + try: + macroblock_records: list[dict] = [] + with open(tmp_path, "wb") as f: + for block_id, block in enumerate(macroblocks): + elem_size = torch.tensor([], dtype=block.dtype).element_size() + block_alignment = ( + math.lcm(align_bytes, elem_size) if align_bytes else elem_size + ) + if block_alignment: + pad = (-f.tell()) % block_alignment + if pad: + f.write(b"\x00" * pad) + block_offset = f.tell() + + record = { + "dtype": dtype_to_string(block.dtype), + "offset_bytes": int(block_offset), + "length_elems": int(block.total_elems), + } + if compress_flags[block_id]: + frames = _fpz_stream_compress_block( + f, block_offset, block, state_dict, encode_frame, progress + ) + record["length_bytes"] = int(f.tell() - block_offset) + fpz_record: dict = {"codec": codec_name, "frames": frames} + if version == 2: + # Record the chunk size so the reader reproduces the + # chunking regardless of the current default, and the + # chunk-start alignment so it can recompute the padded + # offsets (absent = 1: pre-alignment packed layout). + fpz_record["hi_chunk_usize"] = int(hi_chunk_bytes) + fpz_record["hi_align"] = FPZ_HI_CHUNK_ALIGN_BYTES + record["fpz"] = fpz_record + else: + for kind, data in _iter_block_uncompressed_chunks( + block, state_dict, progress + ): + if kind == "zeros": + _write_zeros(f, data) + else: + f.write(data) + record["length_bytes"] = int(block.length_bytes) + macroblock_records.append(record) + + total_payload_bytes = f.tell() + meta_payload = { + "format": FILE_FORMAT_V4, + "align_bytes": int(align_bytes), + "total_payload_bytes": int(total_payload_bytes), + "total_elems": sum(block.total_elems for block in macroblocks), + "macroblocks": macroblock_records, + "index": [ + { + "name": r.name, + "shape": r.shape, + "offset": int(r.offset), + "length": int(r.length), + "macroblock": int(r.macroblock), + } + for r in index + ], + } + footer_json = json.dumps( + meta_payload, separators=(",", ":"), ensure_ascii=False + ).encode("utf-8") + f.write(footer_json) + f.write(U64LE.pack(len(footer_json))) + f.write(MAGIC) + f.flush() + try: + os.fsync(f.fileno()) + except OSError: + pass + + os.replace(tmp_path, destination_path) + tmp_path = None + finally: + if progress is not None: + progress.close() + if tmp_path and os.path.exists(tmp_path): + try: + os.remove(tmp_path) + except OSError: + pass diff --git a/tests/test_fpz.py b/tests/test_fpz.py new file mode 100644 index 0000000..b14da4b --- /dev/null +++ b/tests/test_fpz.py @@ -0,0 +1,554 @@ +"""Unit tests for fpz split-plane zstd compression (``compress="fpz-bf16"``). + +These exercise the CPU read path end to end: byte-identical round trips, the +mixed-dtype policy (only bf16 blocks compress), the multi-frame/alignment +logic, and the error surfaces. The CUDA decoder shares the frame read + decode ++ interleave helpers with the CPU path, so covering them here also covers the +GPU decoder's correctness-critical core (the GPU path itself needs a device and +is not run in CI). +""" + +import os +import sys +import time + +import numpy as np +import pytest +import torch +from flashpack import serialization +from flashpack.constants import ( + FILE_FORMAT_V3, + FILE_FORMAT_V4, + FPZ_CODEC_SPLITPLANE_V1, + FPZ_CODEC_SPLITPLANE_V2, + FPZ_FRAME_ALIGN_BYTES, +) +from flashpack.deserialization import ( + assign_from_file, + get_flashpack_file_metadata, + iterate_from_flash_tensor, + read_flashpack_file, + revert_from_file, +) +from flashpack.serialization import pack_to_file +from flashpack.utils import require_zstandard + + +def _bf16_state_dict() -> dict[str, torch.Tensor]: + generator = torch.Generator().manual_seed(0) + # A genuinely high-entropy tensor (randn) and a low-entropy one (a smooth + # ramp -- exponents change slowly, so the high byte plane compresses hard). + random = torch.randn(1024, 512, generator=generator).to(torch.bfloat16) + ramp = (torch.arange(1024 * 512, dtype=torch.float32) * 0.01).reshape(1024, 512) + return {"block.random": random, "block.ramp": ramp.to(torch.bfloat16)} + + +def _uint16_view(t: torch.Tensor) -> torch.Tensor: + return t.contiguous().view(torch.uint16) + + +def _raw_bytes(t: torch.Tensor) -> torch.Tensor: + # dtype-agnostic byte view for bit-exact (NaN-safe) comparison. + return t.contiguous().view(torch.uint8) + + +def _pack(tmp_path, state_dict, name: str, **kwargs) -> str: + path = str(tmp_path / name) + kwargs.setdefault("target_dtype", None) + pack_to_file(state_dict, path, **kwargs) + return path + + +def test_bf16_roundtrip_is_bit_identical(tmp_path) -> None: + source = _bf16_state_dict() + plain = _pack(tmp_path, source, "plain.flashpack") + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + + storage_p, meta_p = read_flashpack_file(plain, device="cpu") + storage_c, meta_c = read_flashpack_file(comp, device="cpu") + plain_tensors = dict(iterate_from_flash_tensor(storage_p, meta_p)) + comp_tensors = dict(iterate_from_flash_tensor(storage_c, meta_c)) + + assert set(comp_tensors) == set(source) + for name, original in source.items(): + # bit-exact against both the uncompressed pack and the source bytes + assert torch.equal(_uint16_view(comp_tensors[name]), _uint16_view(original)) + assert torch.equal( + _uint16_view(comp_tensors[name]), _uint16_view(plain_tensors[name]) + ) + + +def test_compressed_file_is_v4_with_fpz_record(tmp_path) -> None: + comp = _pack(tmp_path, _bf16_state_dict(), "comp.flashpack", compress="fpz-bf16") + meta = get_flashpack_file_metadata(comp) + assert meta["format"] == FILE_FORMAT_V4 + (block,) = meta["macroblocks"] + assert block["fpz"]["codec"] == FPZ_CODEC_SPLITPLANE_V2 + assert len(block["fpz"]["frames"]) >= 1 + # v2 frames carry per-chunk compressed lengths instead of a single hi_len. + assert block["fpz"]["frames"][0]["hi_chunks"] + # length_elems stays logical; length_bytes is the smaller on-disk payload. + assert block["length_bytes"] < block["length_elems"] * 2 + + +def test_low_entropy_tensor_shrinks_file(tmp_path, capsys) -> None: + ramp = (torch.arange(2048 * 1024, dtype=torch.float32) * 0.01).reshape(2048, 1024) + source = {"w": ramp.to(torch.bfloat16)} + plain = _pack(tmp_path, source, "plain.flashpack") + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + + plain_size = os.path.getsize(plain) + comp_size = os.path.getsize(comp) + ratio = plain_size / comp_size + with capsys.disabled(): + print( + f"\n[fpz] low-entropy ramp: plain={plain_size} comp={comp_size} " + f"ratio={ratio:.3f}x" + ) + assert comp_size < plain_size + assert ratio > 1.5 + + +def test_mixed_dtype_only_bf16_compressed(tmp_path) -> None: + source = { + "bf16.big": torch.randn(512, 512).to(torch.bfloat16), + "fp32.big": torch.randn(512, 512), + } + plain = _pack(tmp_path, source, "plain.flashpack") + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + + meta = get_flashpack_file_metadata(comp) + assert meta["format"] == FILE_FORMAT_V4 + by_dtype = {b["dtype"]: b for b in meta["macroblocks"]} + assert "fpz" in by_dtype["bfloat16"] + assert "fpz" not in by_dtype["float32"] + + # Only bf16 shrank, but the whole file must still round-trip bit-exactly. + assert os.path.getsize(comp) < os.path.getsize(plain) + storage, m = read_flashpack_file(comp, device="cpu") + tensors = dict(iterate_from_flash_tensor(storage, m)) + assert torch.equal(tensors["fp32.big"], source["fp32.big"]) + assert torch.equal( + _uint16_view(tensors["bf16.big"]), _uint16_view(source["bf16.big"]) + ) + + +def test_fp32_only_compress_is_noop(tmp_path) -> None: + # No bf16 block -> nothing to compress -> identical to the plain pack. + source = {"a": torch.randn(64), "b": torch.randn(32)} + plain = _pack(tmp_path, source, "plain.flashpack") + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + + assert get_flashpack_file_metadata(comp)["format"] == FILE_FORMAT_V3 + with open(plain, "rb") as f: + plain_bytes = f.read() + with open(comp, "rb") as f: + comp_bytes = f.read() + assert plain_bytes == comp_bytes + + +def test_multi_frame_roundtrip_and_alignment(tmp_path, monkeypatch) -> None: + # Shrink the frame step so a modest tensor spans several frames, exercising + # the padding/alignment and multi-frame cover logic. + monkeypatch.setattr( + "flashpack.serialization.FPZ_FRAME_UNCOMPRESSED_BYTES", 8192, raising=True + ) + ramp = (torch.arange(8192 * 6, dtype=torch.float32) * 0.01).to(torch.bfloat16) + source = {"w": ramp} + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + + frames = get_flashpack_file_metadata(comp)["macroblocks"][0]["fpz"]["frames"] + assert len(frames) >= 2 + for frame in frames: + assert frame["payload_off"] % FPZ_FRAME_ALIGN_BYTES == 0 + assert frame["lo_len"] * 2 == frame["n_out"] + + storage, meta = read_flashpack_file(comp, device="cpu") + (w,) = [t for _, t in iterate_from_flash_tensor(storage, meta)] + assert torch.equal(_uint16_view(w), _uint16_view(ramp)) + + +def test_iterate_views_match_plain(tmp_path) -> None: + source = _bf16_state_dict() + plain = _pack(tmp_path, source, "plain.flashpack") + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + + sp, mp = read_flashpack_file(plain, device="cpu") + sc, mc = read_flashpack_file(comp, device="cpu") + plain_items = list(iterate_from_flash_tensor(sp, mp)) + comp_items = list(iterate_from_flash_tensor(sc, mc)) + assert [n for n, _ in plain_items] == [n for n, _ in comp_items] + for (_, tp), (_, tc) in zip(plain_items, comp_items): + assert tp.shape == tc.shape + assert torch.equal(_uint16_view(tp), _uint16_view(tc)) + + +def test_revert_from_compressed_matches_source(tmp_path) -> None: + source = _bf16_state_dict() + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + reverted = revert_from_file(comp) + assert set(reverted) == set(source) + for name, original in source.items(): + assert reverted[name].dtype is torch.bfloat16 + assert torch.equal(_uint16_view(reverted[name]), _uint16_view(original)) + + +def test_assign_from_compressed_pack(tmp_path) -> None: + class Net(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(128, 128).to(torch.bfloat16) + + torch.manual_seed(0) + source = Net() + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + torch.manual_seed(1) + destination = Net() + assign_from_file(destination, comp, device="cpu") + assert torch.equal( + _uint16_view(destination.linear.weight.detach()), + _uint16_view(source.linear.weight.detach()), + ) + + +def test_unknown_compress_option_rejected(tmp_path) -> None: + with pytest.raises(ValueError, match="Unsupported compress option"): + pack_to_file( + {"a": torch.randn(4).to(torch.bfloat16)}, + str(tmp_path / "pack.flashpack"), + target_dtype=None, + compress="gzip", + ) + + +def test_unknown_codec_rejected(tmp_path) -> None: + comp = _pack(tmp_path, _bf16_state_dict(), "comp.flashpack", compress="fpz-bf16") + meta = get_flashpack_file_metadata(comp) + meta["macroblocks"][0]["fpz"]["codec"] = "bogus-codec-v9" + with pytest.raises(ValueError, match="Unsupported fpz codec"): + read_flashpack_file(comp, device="cpu", metadata=meta) + + +def test_truncated_frame_rejected(tmp_path) -> None: + source = {"w": torch.randn(2048, 1024).to(torch.bfloat16)} + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + # Capture the full metadata, then physically truncate the payload so the + # declared frame bytes run past EOF. + meta = get_flashpack_file_metadata(comp) + size = os.path.getsize(comp) + with open(comp, "r+b") as f: + f.truncate(size // 2) + with pytest.raises(IOError, match="short read"): + read_flashpack_file(comp, device="cpu", metadata=meta) + + +def test_corrupt_high_plane_rejected(tmp_path) -> None: + source = {"w": torch.randn(1024, 512).to(torch.bfloat16)} + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + block = get_flashpack_file_metadata(comp)["macroblocks"][0] + frame = block["fpz"]["frames"][0] + hi_start = block["offset_bytes"] + frame["payload_off"] + frame["lo_len"] + # Scribble over the first compressed high-plane chunk; zstd must reject it. + first_chunk_len = int(frame["hi_chunks"][0]) + with open(comp, "r+b") as f: + f.seek(hi_start) + f.write(b"\xff" * min(64, first_chunk_len)) + zstandard = require_zstandard() + with pytest.raises((zstandard.ZstdError, ValueError)): + read_flashpack_file(comp, device="cpu") + + +def test_require_zstandard_message_when_missing(monkeypatch) -> None: + # Simulate the optional dependency being absent. + monkeypatch.setitem(sys.modules, "zstandard", None) + with pytest.raises(ImportError, match="zstandard"): + require_zstandard() + + +def test_interleave_byte_order_is_little_endian(tmp_path) -> None: + # Confirm the plane split matches bf16's little-endian layout: even bytes + # are the low (mantissa-LSB) plane, odd bytes are the high (sign+exp) plane. + values = torch.tensor([1.5, -2.0, 0.0, 3.25], dtype=torch.bfloat16) + comp = _pack(tmp_path, {"w": values}, "comp.flashpack", compress="fpz-bf16") + storage, meta = read_flashpack_file(comp, device="cpu") + (w,) = [t for _, t in iterate_from_flash_tensor(storage, meta)] + raw = _uint16_view(w).numpy().view(np.uint8) + expected = _uint16_view(values).numpy().view(np.uint8) + assert np.array_equal(raw, expected) + + +def test_streaming_decode_matches_uncompressed_pack(tmp_path) -> None: + # Equivalence: the streaming compressed pack decodes bit-for-bit identically + # to the plain uncompressed pack of the same mixed state dict. The odd sizes + # force intra-block alignment padding (a zero gap between bf16.a and bf16.b), + # exercising the streaming gap-fill against the memmap layout. + source = { + "bf16.a": torch.randn(301, 400).to(torch.bfloat16), + "bf16.b": torch.randn(51).to(torch.bfloat16), + "fp32.c": torch.randn(128, 64), + } + plain = _pack(tmp_path, source, "plain.flashpack") + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + + sp, mp = read_flashpack_file(plain, device="cpu") + sc, mc = read_flashpack_file(comp, device="cpu") + plain_tensors = dict(iterate_from_flash_tensor(sp, mp)) + comp_tensors = dict(iterate_from_flash_tensor(sc, mc)) + assert set(plain_tensors) == set(comp_tensors) == set(source) + for name in source: + assert torch.equal( + _raw_bytes(plain_tensors[name]), _raw_bytes(comp_tensors[name]) + ) + assert torch.equal(_raw_bytes(comp_tensors[name]), _raw_bytes(source[name])) + + +def test_streaming_creates_no_uncompressed_scratch(tmp_path, monkeypatch) -> None: + # The whole point of streaming: never materialize the uncompressed payload. + # Prove it two ways -- np.memmap (the only uncompressed-scratch allocator in + # the write path) is never called, and exactly ONE tempfile is created (the + # final compressed pack), not a scratch + final pair like the old post-pass. + import flashpack.serialization as serialization + + # Small frames so the rolling-buffer multi-frame path runs under the guards. + monkeypatch.setattr( + serialization, "FPZ_FRAME_UNCOMPRESSED_BYTES", 1 << 16, raising=True + ) + + def no_memmap(*args, **kwargs): + raise AssertionError("uncompressed scratch memmap must not be created") + + monkeypatch.setattr(serialization.np, "memmap", no_memmap) + + tempfiles: list[str] = [] + real_mkstemp = serialization.tempfile.mkstemp + + def counting_mkstemp(*args, **kwargs): + result = real_mkstemp(*args, **kwargs) + tempfiles.append(result[1]) + return result + + monkeypatch.setattr(serialization.tempfile, "mkstemp", counting_mkstemp) + + ramp = (torch.arange(1 << 18, dtype=torch.float32) * 0.001).to(torch.bfloat16) + comp = str(tmp_path / "comp.flashpack") + pack_to_file({"w": ramp}, comp, target_dtype=None, compress="fpz-bf16") + + assert len(tempfiles) == 1 # only the final compressed pack, no scratch + frames = get_flashpack_file_metadata(comp)["macroblocks"][0]["fpz"]["frames"] + assert len(frames) >= 2 # multi-frame rolling-buffer path exercised + + storage, meta = read_flashpack_file(comp, device="cpu") + (w,) = [t for _, t in iterate_from_flash_tensor(storage, meta)] + assert torch.equal(_uint16_view(w), _uint16_view(ramp)) + + +def test_streaming_gap_straddling_frame_boundary(tmp_path, monkeypatch) -> None: + # Two bf16 tensors with alignment padding between them, and a frame step + # small enough that the zero gap straddles a frame cut -- the case where + # the rolling buffer must carry a partial gap across the flush boundary. + import flashpack.serialization as serialization + + monkeypatch.setattr(serialization, "FPZ_FRAME_UNCOMPRESSED_BYTES", 64, raising=True) + source = { + "a": torch.randn(40).to(torch.bfloat16), + "b": torch.randn(40).to(torch.bfloat16), + } + plain = _pack(tmp_path, source, "plain.flashpack", align_bytes=128) + comp = _pack( + tmp_path, source, "comp.flashpack", align_bytes=128, compress="fpz-bf16" + ) + + frames = get_flashpack_file_metadata(comp)["macroblocks"][0]["fpz"]["frames"] + assert len(frames) >= 2 + + sp, mp = read_flashpack_file(plain, device="cpu") + sc, mc = read_flashpack_file(comp, device="cpu") + plain_tensors = dict(iterate_from_flash_tensor(sp, mp)) + comp_tensors = dict(iterate_from_flash_tensor(sc, mc)) + for name in source: + assert torch.equal( + _raw_bytes(plain_tensors[name]), _raw_bytes(comp_tensors[name]) + ) + assert torch.equal(_raw_bytes(comp_tensors[name]), _raw_bytes(source[name])) + + +@pytest.mark.skipif( + (os.cpu_count() or 1) < 4, reason="thread-scaling proof needs >=4 CPUs" +) +def test_cpu_decode_scales_with_threads(tmp_path, monkeypatch) -> None: + # Regression guard for the serialized CPU decode: the fpz read path must + # parallelize across FLASHPACK_READ_THREADS. zstd decompress releases the + # GIL, so 8 threads must materially beat 1. Generous bound (<=0.6x wall) + # with best-of-3 timing so CI jitter can't flake it. + import flashpack.serialization as serialization + + monkeypatch.setattr( + serialization, "FPZ_FRAME_UNCOMPRESSED_BYTES", 1 << 20, raising=True + ) + # ~64 MB uncompressed, many frames; entropy high enough that decode (not + # I/O from the warm page cache) dominates the wall. + data = (torch.randn(32 * 1024 * 1024) * 0.05).to(torch.bfloat16) + comp = str(tmp_path / "comp.flashpack") + pack_to_file({"w": data}, comp, target_dtype=None, compress="fpz-bf16") + + frames = get_flashpack_file_metadata(comp)["macroblocks"][0]["fpz"]["frames"] + assert len(frames) >= 16 # enough work to spread over 8 threads + + def best_wall(threads: int, reps: int = 3) -> float: + monkeypatch.setenv("FLASHPACK_READ_THREADS", str(threads)) + best = float("inf") + for _ in range(reps): + t0 = time.perf_counter() + read_flashpack_file(comp, device="cpu") + best = min(best, time.perf_counter() - t0) + return best + + read_flashpack_file(comp, device="cpu") # warm the page cache + single = best_wall(1) + multi = best_wall(8) + + # Sanity: decode is still correct under many threads. + monkeypatch.setenv("FLASHPACK_READ_THREADS", "8") + storage, meta = read_flashpack_file(comp, device="cpu") + (w,) = [t for _, t in iterate_from_flash_tensor(storage, meta)] + assert torch.equal(_uint16_view(w), _uint16_view(data)) + + assert multi <= 0.6 * single, ( + f"fpz CPU decode did not scale: 1-thread={single * 1e3:.1f} ms, " + f"8-thread={multi * 1e3:.1f} ms (expected 8-thread <= 0.6x)" + ) + + +# -------------------------------------------------------------------------- +# v2 (chunked high plane) format + v1 backward compatibility +# -------------------------------------------------------------------------- + + +def test_v2_frame_splits_high_plane_into_chunks(tmp_path) -> None: + # A 1024x512 bf16 tensor has a 512 KiB high plane -> several 64 KiB chunks. + source = {"w": torch.randn(1024, 512).to(torch.bfloat16)} + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + block = get_flashpack_file_metadata(comp)["macroblocks"][0] + assert block["fpz"]["codec"] == FPZ_CODEC_SPLITPLANE_V2 + (frame,) = block["fpz"]["frames"] + assert len(frame["hi_chunks"]) >= 2 # multiple parallel-decodable chunks + # Round-trips bit-exactly through the chunked decode. + storage, meta = read_flashpack_file(comp, device="cpu") + (w,) = [t for _, t in iterate_from_flash_tensor(storage, meta)] + assert torch.equal(_uint16_view(w), _uint16_view(source["w"])) + + +def test_v1_pack_still_reads(tmp_path, monkeypatch) -> None: + # Backward compatibility: a pack written by the v1 encoder must still decode + # bit-exactly (the reader supports both codecs). + monkeypatch.setattr(serialization, "_DEFAULT_FPZ_VERSION", 1) + source = _bf16_state_dict() + comp = _pack(tmp_path, source, "comp_v1.flashpack", compress="fpz-bf16") + block = get_flashpack_file_metadata(comp)["macroblocks"][0] + assert block["fpz"]["codec"] == FPZ_CODEC_SPLITPLANE_V1 + assert "hi_len" in block["fpz"]["frames"][0] # single-frame high plane + + storage, meta = read_flashpack_file(comp, device="cpu") + tensors = dict(iterate_from_flash_tensor(storage, meta)) + for name, original in source.items(): + assert torch.equal(_uint16_view(tensors[name]), _uint16_view(original)) + + +def test_v2_ratio_close_to_v1(tmp_path, monkeypatch, capsys) -> None: + # v2 compresses each 64 KiB chunk independently, so it shrinks slightly less + # than v1's single-frame high plane. Measure both on a realistic low-entropy + # tensor; v2 must still compress and stay within a modest margin of v1. + ramp = (torch.arange(2048 * 1024, dtype=torch.float32) * 0.01).reshape(2048, 1024) + source = {"w": ramp.to(torch.bfloat16)} + plain = _pack(tmp_path, source, "plain.flashpack") + + monkeypatch.setattr(serialization, "_DEFAULT_FPZ_VERSION", 1) + v1 = _pack(tmp_path, source, "v1.flashpack", compress="fpz-bf16") + monkeypatch.setattr(serialization, "_DEFAULT_FPZ_VERSION", 2) + v2 = _pack(tmp_path, source, "v2.flashpack", compress="fpz-bf16") + + plain_sz = os.path.getsize(plain) + v1_sz = os.path.getsize(v1) + v2_sz = os.path.getsize(v2) + with capsys.disabled(): + print( + f"\n[fpz] ratio plain={plain_sz} " + f"v1={v1_sz} ({plain_sz / v1_sz:.3f}x) " + f"v2={v2_sz} ({plain_sz / v2_sz:.3f}x) v2/v1={v2_sz / v1_sz:.3f}" + ) + assert v2_sz < plain_sz # v2 still compresses + assert v2_sz <= v1_sz * 1.25 # within a modest margin of v1 + + +# -------------------------------------------------------------------------- +# v2 parameterized chunk size +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("chunk_bytes", [64 * 1024, 256 * 1024, 1024 * 1024]) +def test_v2_roundtrip_at_various_chunk_sizes(tmp_path, chunk_bytes) -> None: + # A 2048x1024 bf16 tensor has a 2 MiB high plane -> several chunks at each + # size; every size must round-trip bit-exactly and record its chunk size. + ramp = (torch.arange(2048 * 1024, dtype=torch.float32) * 0.01).reshape(2048, 1024) + source = {"w": ramp.to(torch.bfloat16)} + comp = _pack( + tmp_path, + source, + f"c{chunk_bytes}.flashpack", + compress="fpz-bf16", + hi_chunk_bytes=chunk_bytes, + ) + block = get_flashpack_file_metadata(comp)["macroblocks"][0] + assert block["fpz"]["hi_chunk_usize"] == chunk_bytes # footer field present + assert len(block["fpz"]["frames"][0]["hi_chunks"]) >= 2 + + storage, meta = read_flashpack_file(comp, device="cpu") + (w,) = [t for _, t in iterate_from_flash_tensor(storage, meta)] + assert torch.equal(_uint16_view(w), _uint16_view(source["w"])) + + +def test_env_sets_v2_chunk_size(tmp_path, monkeypatch) -> None: + monkeypatch.setenv("FLASHPACK_FPZ_CHUNK_BYTES", str(256 * 1024)) + source = {"w": torch.randn(1024, 512).to(torch.bfloat16)} + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + block = get_flashpack_file_metadata(comp)["macroblocks"][0] + assert block["fpz"]["hi_chunk_usize"] == 256 * 1024 + + +def test_v2_reader_defaults_chunk_size_when_field_absent(tmp_path) -> None: + # Back-compat: a v2 pack written before hi_chunk_usize existed (field + # absent) must decode with the 64 KiB default. The default pack uses 64 KiB, + # so dropping the field and re-reading must still be bit-exact. + source = {"w": torch.randn(1024, 512).to(torch.bfloat16)} + comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") + meta = get_flashpack_file_metadata(comp) + del meta["macroblocks"][0]["fpz"]["hi_chunk_usize"] + + storage, m = read_flashpack_file(comp, device="cpu", metadata=meta) + (w,) = [t for _, t in iterate_from_flash_tensor(storage, m)] + assert torch.equal(_uint16_view(w), _uint16_view(source["w"])) + + +@pytest.mark.parametrize("bad", [4097, 1000, 64 * 1024 + 1]) +def test_hi_chunk_bytes_must_be_multiple_of_align(tmp_path, bad) -> None: + with pytest.raises(ValueError, match="multiple"): + pack_to_file( + {"w": torch.randn(64).to(torch.bfloat16)}, + str(tmp_path / "p.flashpack"), + target_dtype=None, + compress="fpz-bf16", + hi_chunk_bytes=bad, + ) + + +def test_hi_chunk_bytes_too_large_rejected(tmp_path) -> None: + from flashpack.constants import FPZ_FRAME_UNCOMPRESSED_BYTES + + with pytest.raises(ValueError, match="exceeds"): + pack_to_file( + {"w": torch.randn(64).to(torch.bfloat16)}, + str(tmp_path / "p.flashpack"), + target_dtype=None, + compress="fpz-bf16", + hi_chunk_bytes=FPZ_FRAME_UNCOMPRESSED_BYTES, + ) diff --git a/tests/test_fpz_gpu.py b/tests/test_fpz_gpu.py new file mode 100644 index 0000000..4d1348e --- /dev/null +++ b/tests/test_fpz_gpu.py @@ -0,0 +1,297 @@ +"""CPU-testable surface of the nvcomp GPU Zstd decode path for fpz. + +The GPU decode itself needs a CUDA device and nvcomp, so it is exercised on the +H200, not in CI. What IS testable without a device -- and what these cover -- is +everything around it: the pure frame-batch planner, the env gating, and the +guarded-import fallback (including warn-once), plus a guard that turning the +flag on never disturbs the CPU decode path. +""" + +import sys + +import pytest +import torch +from flashpack import deserialization +from flashpack.deserialization import ( + _env_flag, + _env_flag_default, + _fpz_batch_signature, + _fpz_gpu_decode_enabled, + _fpz_hi_chunk_usizes, + _load_nvcomp, + iterate_from_flash_tensor, + plan_fpz_gpu_batches, + read_flashpack_file, +) +from flashpack.serialization import pack_to_file + + +def _frame_task(block_idx: int, n_out: int, out_pos: int = 0) -> tuple: + # Matches _fpz_frame_tasks' ("frame", block_idx, frame, out_pos) shape. + return ("frame", block_idx, {"n_out": n_out, "payload_off": 0}, out_pos) + + +# -------------------------------------------------------------------------- +# plan_fpz_gpu_batches -- pure function +# -------------------------------------------------------------------------- + + +def test_plan_empty_input_is_empty() -> None: + assert plan_fpz_gpu_batches([], max_batch_frames=4, max_batch_bytes=1 << 30) == [] + + +def test_plan_single_frame_is_one_batch() -> None: + tasks = [_frame_task(0, 100)] + assert plan_fpz_gpu_batches(tasks, 4, 1 << 30) == [tasks] + + +def test_plan_bounded_by_frame_count() -> None: + tasks = [_frame_task(0, 10) for _ in range(10)] + batches = plan_fpz_gpu_batches(tasks, max_batch_frames=3, max_batch_bytes=1 << 30) + assert [len(b) for b in batches] == [3, 3, 3, 1] + # Every frame appears exactly once, order preserved. + assert [t for b in batches for t in b] == tasks + + +def test_plan_bounded_by_byte_budget() -> None: + tasks = [_frame_task(0, 100) for _ in range(5)] + # Budget of 250 bytes -> 2 frames per batch before the count cap bites. + batches = plan_fpz_gpu_batches(tasks, max_batch_frames=99, max_batch_bytes=250) + assert [len(b) for b in batches] == [2, 2, 1] + for b in batches: + assert sum(int(t[2]["n_out"]) for t in b) <= 250 + + +def test_plan_oversized_frame_gets_its_own_batch() -> None: + # A frame larger than the whole budget must not be dropped or merged away. + tasks = [_frame_task(0, 50), _frame_task(0, 500), _frame_task(0, 50)] + batches = plan_fpz_gpu_batches(tasks, max_batch_frames=99, max_batch_bytes=100) + assert [[int(t[2]["n_out"]) for t in b] for b in batches] == [[50], [500], [50]] + + +def test_plan_preserves_frames_across_block_boundaries() -> None: + # Frames from different macroblocks are grouped purely by order/size; the + # per-frame block index and out_pos ride along untouched. + tasks = [_frame_task(0, 10, 0), _frame_task(0, 10, 10), _frame_task(1, 10, 0)] + batches = plan_fpz_gpu_batches(tasks, max_batch_frames=2, max_batch_bytes=1 << 30) + assert [len(b) for b in batches] == [2, 1] + assert [t for b in batches for t in b] == tasks + + +@pytest.mark.parametrize("bad", [0, -1]) +def test_plan_rejects_nonpositive_frame_cap(bad: int) -> None: + with pytest.raises(ValueError, match="max_batch_frames"): + plan_fpz_gpu_batches( + [_frame_task(0, 1)], max_batch_frames=bad, max_batch_bytes=1 + ) + + +@pytest.mark.parametrize("bad", [0, -1]) +def test_plan_rejects_nonpositive_byte_budget(bad: int) -> None: + with pytest.raises(ValueError, match="max_batch_bytes"): + plan_fpz_gpu_batches( + [_frame_task(0, 1)], max_batch_frames=1, max_batch_bytes=bad + ) + + +# -------------------------------------------------------------------------- +# _fpz_batch_signature -- pure config-cache key +# -------------------------------------------------------------------------- + + +def test_signature_is_per_frame_half_sizes() -> None: + batch = [_frame_task(0, 100), _frame_task(0, 40)] + assert _fpz_batch_signature(batch) == (50, 20) + + +def test_signature_matches_for_same_shape_batches() -> None: + # Two batches of identical full-frame shapes hash to the same key -> they + # share one cached DecompressConfig (the whole point of the cache). + a = [_frame_task(0, 128), _frame_task(0, 128)] + b = [_frame_task(1, 128), _frame_task(2, 128)] + assert _fpz_batch_signature(a) == _fpz_batch_signature(b) == (64, 64) + + +def test_signature_differs_when_a_tail_frame_changes_shape() -> None: + full = [_frame_task(0, 128), _frame_task(0, 128)] + with_tail = [_frame_task(0, 128), _frame_task(0, 40)] + assert _fpz_batch_signature(full) != _fpz_batch_signature(with_tail) + + +# -------------------------------------------------------------------------- +# _fpz_hi_chunk_usizes -- pure v2 chunk sizing (matches the encoder) +# -------------------------------------------------------------------------- + + +def test_hi_chunk_usizes_exact_multiple_has_no_tail() -> None: + # 4 * chunk -> 4 equal chunks, no remainder (the full-frame case). + assert _fpz_hi_chunk_usizes(4 * 64, 64) == [64, 64, 64, 64] + + +def test_hi_chunk_usizes_has_remainder_tail() -> None: + assert _fpz_hi_chunk_usizes(200, 64) == [64, 64, 64, 8] + + +def test_hi_chunk_usizes_smaller_than_one_chunk() -> None: + assert _fpz_hi_chunk_usizes(40, 64) == [40] + + +def test_hi_chunk_usizes_zero_is_empty() -> None: + assert _fpz_hi_chunk_usizes(0, 64) == [] + + +def test_hi_chunk_usizes_sum_equals_half() -> None: + for half in (1, 63, 64, 65, 1000, 1 << 20): + assert sum(_fpz_hi_chunk_usizes(half, 64)) == half + + +@pytest.mark.parametrize("bad", [0, -1]) +def test_hi_chunk_usizes_rejects_bad_chunk(bad: int) -> None: + with pytest.raises(ValueError): + _fpz_hi_chunk_usizes(100, bad) + + +# -------------------------------------------------------------------------- +# env gating +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", " On "]) +def test_env_flag_truthy(monkeypatch, value: str) -> None: + monkeypatch.setenv("FLASHPACK_FPZ_GPU_DECODE", value) + assert _env_flag("FLASHPACK_FPZ_GPU_DECODE") is True + assert _fpz_gpu_decode_enabled() is True + + +@pytest.mark.parametrize("value", ["0", "false", "no", "", "maybe"]) +def test_env_flag_falsy(monkeypatch, value: str) -> None: + monkeypatch.setenv("FLASHPACK_FPZ_GPU_DECODE", value) + assert _env_flag("FLASHPACK_FPZ_GPU_DECODE") is False + assert _fpz_gpu_decode_enabled() is False + + +def test_env_flag_unset_is_false(monkeypatch) -> None: + monkeypatch.delenv("FLASHPACK_FPZ_GPU_DECODE", raising=False) + assert _fpz_gpu_decode_enabled() is False + + +def test_env_flag_default_respects_default_when_unset(monkeypatch) -> None: + monkeypatch.delenv("FLASHPACK_FPZ_GPU_TORCH_ALLOC", raising=False) + assert _env_flag_default("FLASHPACK_FPZ_GPU_TORCH_ALLOC", True) is True + assert _env_flag_default("FLASHPACK_FPZ_GPU_TORCH_ALLOC", False) is False + + +@pytest.mark.parametrize( + "value,expected", [("0", False), ("false", False), ("1", True), ("on", True)] +) +def test_env_flag_default_env_overrides(monkeypatch, value, expected) -> None: + monkeypatch.setenv("FLASHPACK_FPZ_GPU_TORCH_ALLOC", value) + # Env always wins over the default, in both directions. + assert _env_flag_default("FLASHPACK_FPZ_GPU_TORCH_ALLOC", True) is expected + assert _env_flag_default("FLASHPACK_FPZ_GPU_TORCH_ALLOC", False) is expected + + +# -------------------------------------------------------------------------- +# torch caching allocator wiring into nvcomp (guarded, idempotent) +# -------------------------------------------------------------------------- + + +class _FakeNvcompAllocOK: + def __init__(self) -> None: + self.installed = None + + def set_device_allocator(self, allocator) -> None: + self.installed = allocator + + +class _FakeNvcompAllocRaises: + def set_device_allocator(self, allocator) -> None: + raise RuntimeError("no allocator hook here") + + +def _reset_alloc_latches(monkeypatch) -> None: + monkeypatch.setattr(deserialization, "_FPZ_NVCOMP_ALLOC_INSTALLED", False) + monkeypatch.setattr(deserialization, "_FPZ_NVCOMP_ALLOC_WARNED", False) + + +def test_install_allocator_success_registers_a_callable(monkeypatch) -> None: + _reset_alloc_latches(monkeypatch) + fake = _FakeNvcompAllocOK() + ok = deserialization._install_torch_nvcomp_allocator(fake, torch.device("cuda:0")) + assert ok is True + assert callable(fake.installed) + + +def test_install_allocator_is_idempotent(monkeypatch) -> None: + _reset_alloc_latches(monkeypatch) + first = _FakeNvcompAllocOK() + assert deserialization._install_torch_nvcomp_allocator( + first, torch.device("cuda:0") + ) + # Already installed globally: a second call is a no-op and does not + # re-register on another (fake) module. + second = _FakeNvcompAllocOK() + assert deserialization._install_torch_nvcomp_allocator( + second, torch.device("cuda:0") + ) + assert second.installed is None + + +def test_install_allocator_failure_is_guarded_and_warns(monkeypatch) -> None: + _reset_alloc_latches(monkeypatch) + with pytest.warns(RuntimeWarning, match="set_device_allocator"): + ok = deserialization._install_torch_nvcomp_allocator( + _FakeNvcompAllocRaises(), torch.device("cuda:0") + ) + assert ok is False + + +# -------------------------------------------------------------------------- +# guarded import + warn-once fallback +# -------------------------------------------------------------------------- + + +def test_load_nvcomp_missing_returns_none_and_warns_once(monkeypatch) -> None: + # Force the import to fail regardless of what's installed, and reset the + # process-level warn-once latch so the assertion is deterministic. + monkeypatch.setitem(sys.modules, "nvidia", None) + monkeypatch.setattr(deserialization, "_FPZ_GPU_DECODE_WARNED", False) + + with pytest.warns(RuntimeWarning, match="nvcomp"): + assert _load_nvcomp() is None + + # Second call: still None, but no second warning (warn-once). + import warnings as _warnings + + with _warnings.catch_warnings(): + _warnings.simplefilter("error") # any warning would raise + assert _load_nvcomp() is None + + +# -------------------------------------------------------------------------- +# the flag must never disturb the CPU decode path +# -------------------------------------------------------------------------- + + +def test_gpu_flag_on_leaves_cpu_path_bit_identical(tmp_path, monkeypatch) -> None: + # device="cpu" never touches the GPU branch, so enabling the flag (with no + # nvcomp available) must produce exactly the same bytes as with it off. + monkeypatch.setattr(deserialization, "_FPZ_GPU_DECODE_WARNED", False) + source = {"w": torch.randn(1024, 512).to(torch.bfloat16)} + comp = str(tmp_path / "comp.flashpack") + pack_to_file(source, comp, target_dtype=None, compress="fpz-bf16") + + monkeypatch.delenv("FLASHPACK_FPZ_GPU_DECODE", raising=False) + s_off, m_off = read_flashpack_file(comp, device="cpu") + off = dict(iterate_from_flash_tensor(s_off, m_off)) + + monkeypatch.setenv("FLASHPACK_FPZ_GPU_DECODE", "1") + s_on, m_on = read_flashpack_file(comp, device="cpu") + on = dict(iterate_from_flash_tensor(s_on, m_on)) + + assert set(on) == set(off) == set(source) + for name in source: + assert torch.equal( + on[name].contiguous().view(torch.uint16), + off[name].contiguous().view(torch.uint16), + ) diff --git a/tests/test_fpz_ll.py b/tests/test_fpz_ll.py new file mode 100644 index 0000000..240d1ba --- /dev/null +++ b/tests/test_fpz_ll.py @@ -0,0 +1,136 @@ +"""Unit tests for the batched-nvcomp ("ll") groundwork that is testable +without a GPU: the v2 chunk-start alignment format change (padding, footer +field, back-compat) and the ctypes binding's loader/ABI guards. The foreign +calls themselves need a device and libnvcomp and are exercised by the H200 +validation app, not CI. +""" + +import ctypes + +import torch +from flashpack import _nvcomp_ll, serialization +from flashpack.constants import FPZ_HI_CHUNK_ALIGN_BYTES +from flashpack.deserialization import ( + _align_up, + get_flashpack_file_metadata, + iterate_from_flash_tensor, + read_flashpack_file, +) +from flashpack.serialization import pack_to_file + + +def _state_with_unaligned_tail() -> dict[str, torch.Tensor]: + # An odd bf16 element count makes the (single, partial) frame's half-plane + # length odd -- NOT a multiple of the chunk alignment -- so the test + # covers the padded-lo-plane case, not just full 32 MiB frames. + generator = torch.Generator().manual_seed(7) + return { + "w": torch.randn(100_001, generator=generator).to(torch.bfloat16), + } + + +def _fpz_blocks(meta: dict) -> list[dict]: + return [b for b in meta["macroblocks"] if b.get("fpz")] + + +def test_v2_footer_records_hi_align(tmp_path) -> None: + dest = str(tmp_path / "pack.flashpack") + pack_to_file(_state_with_unaligned_tail(), dest, None, compress="fpz-bf16") + blocks = _fpz_blocks(get_flashpack_file_metadata(dest)) + assert blocks, "expected at least one fpz block" + for block in blocks: + assert block["fpz"]["hi_align"] == FPZ_HI_CHUNK_ALIGN_BYTES + + +def test_v2_chunk_starts_are_aligned_on_disk(tmp_path) -> None: + """Every compressed chunk's absolute file offset must be a multiple of + the recorded alignment (this is what the batched GPU decoder's device + pointer table relies on), including after an unaligned tail lo plane.""" + dest = str(tmp_path / "pack.flashpack") + pack_to_file(_state_with_unaligned_tail(), dest, None, compress="fpz-bf16") + meta = get_flashpack_file_metadata(dest) + saw_unaligned_lo = False + for block in _fpz_blocks(meta): + align = int(block["fpz"]["hi_align"]) + assert align > 1 + for frame in block["fpz"]["frames"]: + lo_len = int(frame["lo_len"]) + saw_unaligned_lo |= lo_len % align != 0 + start = ( + int(block["offset_bytes"]) + + int(frame["payload_off"]) + + _align_up(lo_len, align) + ) + off = start + for clen in frame["hi_chunks"]: + assert off % align == 0 + off += _align_up(int(clen), align) + assert saw_unaligned_lo, "test state must produce an unaligned lo plane" + + +def test_v2_padded_pack_roundtrips_exactly(tmp_path) -> None: + dest = str(tmp_path / "pack.flashpack") + state = _state_with_unaligned_tail() + pack_to_file(state, dest, None, compress="fpz-bf16") + storage, meta = read_flashpack_file(dest, device="cpu") + out = dict(iterate_from_flash_tensor(storage, meta)) + assert set(out) == set(state) + for name, tensor in state.items(): + assert torch.equal(out[name], tensor), name + + +def test_v2_unpadded_layout_backcompat(tmp_path, monkeypatch) -> None: + """A pack written with hi_align=1 (the pre-alignment packed layout, as + older readers/writers produced) must still read byte-exactly: the reader + takes the alignment from the footer, never from the current constant.""" + monkeypatch.setattr(serialization, "FPZ_HI_CHUNK_ALIGN_BYTES", 1) + dest = str(tmp_path / "pack_old.flashpack") + state = _state_with_unaligned_tail() + pack_to_file(state, dest, None, compress="fpz-bf16") + for block in _fpz_blocks(get_flashpack_file_metadata(dest)): + assert block["fpz"]["hi_align"] == 1 + storage, meta = read_flashpack_file(dest, device="cpu") + out = dict(iterate_from_flash_tensor(storage, meta)) + for name, tensor in state.items(): + assert torch.equal(out[name], tensor), name + + +def test_align_up() -> None: + assert _align_up(0, 16) == 0 + assert _align_up(1, 16) == 16 + assert _align_up(16, 16) == 16 + assert _align_up(17, 16) == 32 + assert _align_up(123, 1) == 123 + + +def test_nvcomp_ll_abi_struct_sizes() -> None: + """The C structs are passed BY VALUE -- their byte sizes are load-bearing + ABI facts (64-byte opts struct, three size_t alignment requirements).""" + assert ctypes.sizeof(_nvcomp_ll._ZstdDecompressOpts) == 64 + assert ctypes.sizeof(_nvcomp_ll._AlignmentRequirements) == 3 * ctypes.sizeof( + ctypes.c_size_t + ) + + +def test_nvcomp_ll_load_degrades_to_none(monkeypatch) -> None: + """Without the libnvcomp wheel (or with a bad override path), load() + returns None instead of raising, and the result is cached.""" + monkeypatch.setattr(_nvcomp_ll, "_loaded", None) + monkeypatch.setenv("FLASHPACK_LIBNVCOMP_PATH", "/nonexistent/libnvcomp.so") + assert _nvcomp_ll.load() is None + # Cached: a second call must not re-probe (flip the env to a still-bad + # value and confirm the cached None is returned without error). + monkeypatch.delenv("FLASHPACK_LIBNVCOMP_PATH") + assert _nvcomp_ll.load() is None + monkeypatch.setattr(_nvcomp_ll, "_loaded", None) + + +def test_nvcomp_ll_find_library_env_override(tmp_path, monkeypatch) -> None: + fake = tmp_path / "libnvcomp.so.5" + fake.write_bytes(b"not a real library") + monkeypatch.setenv("FLASHPACK_LIBNVCOMP_PATH", str(fake)) + assert _nvcomp_ll._find_libnvcomp() == str(fake) + # A real path that is not a loadable library must degrade to None too. + monkeypatch.setattr(_nvcomp_ll, "_loaded", None) + assert _nvcomp_ll.load() is None + monkeypatch.setattr(_nvcomp_ll, "_loaded", None) diff --git a/tests/test_guardrails.py b/tests/test_guardrails.py index 1698da1..4e83d4d 100644 --- a/tests/test_guardrails.py +++ b/tests/test_guardrails.py @@ -1,8 +1,10 @@ -"""Unit tests for the reader-thread affinity clamp.""" +"""Unit tests for the resource guardrails: the reader-thread affinity clamp +and the GPU-decoder warmup helper's graceful degradation.""" import os from flashpack import utils +from flashpack.deserialization import fpz_gpu_warmup from flashpack.utils import effective_read_threads @@ -41,3 +43,9 @@ def test_affinity_unavailable_falls_back_to_cpu_count(monkeypatch) -> None: monkeypatch.delattr(os, "sched_getaffinity", raising=False) monkeypatch.setattr(utils.os, "cpu_count", lambda: 8) assert effective_read_threads(16) == 8 + + +def test_fpz_gpu_warmup_degrades_without_gpu() -> None: + """On a CUDA-less host the warmup must be a safe no-op returning False, + never an exception (it is called from app setup paths).""" + assert fpz_gpu_warmup() is False diff --git a/tests/test_interleave.py b/tests/test_interleave.py new file mode 100644 index 0000000..a2b4da4 --- /dev/null +++ b/tests/test_interleave.py @@ -0,0 +1,21 @@ +"""Unit tests for the fused-interleave module's degradation behavior (the +kernel itself needs a GPU and is validated by the H200 probe/gate: numeric +parity against the strided path plus the pack checksum gate).""" + +import torch +from flashpack import _interleave + + +def test_available_is_bool_and_cached() -> None: + first = _interleave.fused_interleave_available() + assert isinstance(first, bool) + assert _interleave.fused_interleave_available() == first + + +def test_interleave_into_degrades_without_triton() -> None: + if _interleave.fused_interleave_available(): + return # covered by the GPU gate on hosts that have triton+CUDA + out = torch.empty(8, dtype=torch.uint8) + lo = torch.zeros(4, dtype=torch.uint8) + hi = torch.ones(4, dtype=torch.uint8) + assert _interleave.interleave_into(out, lo, hi) is False From 52062654636e95724686e6ffae543c18d3de3be9 Mon Sep 17 00:00:00 2001 From: Alperen Konukbay Date: Sun, 26 Jul 2026 22:48:36 -0700 Subject: [PATCH 2/6] refactor: single GPU decode path, drop bench-only scaffolding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review cleanup — no functional additions: - The batched libnvcomp path is now THE GPU decode; the pybind-wrapper path (7x slower, only ever a comparison baseline) is removed along with its per-chunk Array objects, DecompressConfig cache, torch-allocator adapter and the FLASHPACK_FPZ_GPU_LL selection gate. Path selection is automatic: FLASHPACK_FPZ_GPU_DECODE=1 uses the batched decoder when libnvcomp is present, the pack is v2 and its chunk alignment satisfies the decompressor; otherwise it falls back to the threaded CPU decode with a warn-once explaining why. - The optional fused-interleave Triton kernel is removed (measured ~2%; the strided copies run at HBM speed). - Per-thread timing trace now goes through logging.debug (was print) with compacted buckets; FLASHPACK_FPZ_GPU_TORCH_ALLOC and the now-dead _env_flag_default helper are gone. - Tests updated accordingly: wrapper/allocator tests replaced by coverage of the new automatic fallback; docstrings rewritten as contracts. - fpz-gpu extra now depends only on nvidia-libnvcomp-cu12 (the pybind wheel is no longer used). Co-Authored-By: Claude Fable 5 --- pyproject.toml | 10 +- src/flashpack/_interleave.py | 86 ---- src/flashpack/_nvcomp_ll.py | 9 +- src/flashpack/deserialization.py | 781 ++++++++++--------------------- tests/test_fpz_gpu.py | 129 +---- tests/test_interleave.py | 21 - 6 files changed, 261 insertions(+), 775 deletions(-) delete mode 100644 src/flashpack/_interleave.py delete mode 100644 tests/test_interleave.py diff --git a/pyproject.toml b/pyproject.toml index 11af0fe..85f9e6b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,12 +48,12 @@ fpz = [ "zstandard>=0.22", ] # GPU Zstd decode for the fpz read path (opt-in via FLASHPACK_FPZ_GPU_DECODE=1). -# NOTE: nvidia-nvcomp-cu12 is NVIDIA-proprietary redistributable software, not -# OSI-licensed -- a license review is required before this extra ships beyond -# the prototype. Kept out of the default and `fpz` deps so a plain install -# stays MIT-only. +# NOTE: nvidia-libnvcomp-cu12 is NVIDIA-proprietary redistributable software, +# not OSI-licensed -- a license review is required before this extra ships. +# Kept out of the default and `fpz` deps so a plain install stays MIT-only. fpz-gpu = [ - "nvidia-nvcomp-cu12", + "zstandard>=0.22", + "nvidia-libnvcomp-cu12", ] [tool.setuptools_scm] diff --git a/src/flashpack/_interleave.py b/src/flashpack/_interleave.py deleted file mode 100644 index 8a387a9..0000000 --- a/src/flashpack/_interleave.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Fused byte-plane interleave kernel (Triton) for the fpz GPU read paths. - -The split-plane format stores a frame as separate low/high byte planes; the -reader must produce out[2i] = lo[i], out[2i+1] = hi[i]. The torch expression -of that -- two strided copies (``out[0::2] = lo; out[1::2] = hi``) -- makes -two passes of 2-byte-stride writes, a worst-case memory pattern that doubles -device traffic and dominates the GPU-side tail of the hot-tier load. - -The fused kernel makes one pass: read lo[i] and hi[i] once, write one -little-endian uint16 ``lo | hi << 8`` (stored through an int16 view -- same -bit pattern, and torch's int16 has full op support where uint16 does not). - -Triton ships inside the torch Linux wheels (pytorch-triton), so no extra -dependency; on hosts without it (or on any kernel failure) callers fall back -to the strided copies. Enable with FLASHPACK_FPZ_FUSED_INTERLEAVE=1. -""" - -from __future__ import annotations - -import threading - -import torch - -# 8 elements per thread (BLOCK / (32 * warps) == 8): two 8-byte loads feed one -# 16-byte store per thread, the measured-optimal shape for byte-plane joins on -# Hopper (dietgpu's FloatTypeInfo vectorization). Fixed config on -# purpose -- autotuning would compile extra variants on a fresh container's -# empty triton cache, which lands exactly on the cold-start path. -_BLOCK = 2048 -_NUM_WARPS = 8 - -_lock = threading.Lock() -_kernel_cache: tuple | None = None - - -def _get_kernel(): - """Build (or fetch) the JIT'd kernel; None when triton is unavailable.""" - global _kernel_cache - with _lock: - if _kernel_cache is not None: - return _kernel_cache[0] - try: - import triton - import triton.language as tl - - @triton.jit - def _interleave_u8(lo_ptr, hi_ptr, out_ptr, n_elem, BLOCK: tl.constexpr): - # int16 (not uint16) throughout: identical bit pattern, and - # triton's unsigned integer paths have known pointer-arith - # bugs (triton#6043) while int16 bitcast is the in-house - # precedent. .to(tl.int16) on a uint8 source zero-extends. - pid = tl.program_id(0).to(tl.int64) - offs = pid * BLOCK + tl.arange(0, BLOCK).to(tl.int64) - mask = offs < n_elem - lo = tl.load(lo_ptr + offs, mask=mask, other=0).to(tl.int16) - hi = tl.load(hi_ptr + offs, mask=mask, other=0).to(tl.int16) - tl.store(out_ptr + offs, lo | (hi << 8), mask=mask) - - def _launch(lo: torch.Tensor, hi: torch.Tensor, out_u8: torch.Tensor): - n = lo.numel() - out16 = out_u8.view(torch.int16) - grid = (triton.cdiv(n, _BLOCK),) - _interleave_u8[grid]( - lo, hi, out16, n, BLOCK=_BLOCK, num_warps=_NUM_WARPS - ) - - _kernel_cache = (_launch,) - except Exception: - _kernel_cache = (None,) - return _kernel_cache[0] - - -def fused_interleave_available() -> bool: - return _get_kernel() is not None - - -def interleave_into(out_u8: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> bool: - """Fused single-pass interleave of ``lo``/``hi`` uint8 planes into - ``out_u8`` (contiguous, even byte offset, ``2 * lo.numel()`` bytes) on - the CURRENT stream. Returns False (having written nothing) when the - kernel is unavailable so the caller can run the strided fallback.""" - launch = _get_kernel() - if launch is None: - return False - launch(lo, hi, out_u8) - return True diff --git a/src/flashpack/_nvcomp_ll.py b/src/flashpack/_nvcomp_ll.py index e000bc9..3493231 100644 --- a/src/flashpack/_nvcomp_ll.py +++ b/src/flashpack/_nvcomp_ll.py @@ -28,8 +28,7 @@ must be real device buffers (unlike LZ4, where NULL is tolerated). Load failures (missing wheel, missing symbol, non-Linux) degrade to -``load()`` returning ``None``; callers keep the pybind wrapper path as the -fallback. +``load()`` returning ``None``; callers fall back to the CPU decode path. """ from __future__ import annotations @@ -208,7 +207,7 @@ def decompress_async( def load() -> NvcompLL | None: """Load and bind libnvcomp once; ``None`` (with a single warning) on any - failure so callers can fall back to the pybind wrapper path.""" + failure so callers can fall back to the CPU decode path.""" global _loaded with _load_lock: if _loaded is not None: @@ -219,7 +218,7 @@ def load() -> NvcompLL | None: logger.warning( "flashpack: libnvcomp not found (pip install " "nvidia-libnvcomp-cu12); batched GPU decode unavailable, " - "falling back to the nvcomp wrapper path." + "falling back to the CPU decode path." ) else: try: @@ -227,7 +226,7 @@ def load() -> NvcompLL | None: except (OSError, AttributeError) as e: logger.warning( "flashpack: could not bind batched nvcomp API from %s " - "(%s); falling back to the nvcomp wrapper path.", + "(%s); falling back to the CPU decode path.", path, e, ) diff --git a/src/flashpack/deserialization.py b/src/flashpack/deserialization.py index 2e9d0e1..35c3733 100644 --- a/src/flashpack/deserialization.py +++ b/src/flashpack/deserialization.py @@ -1,4 +1,5 @@ import json +import logging import math import os import queue @@ -14,7 +15,6 @@ import torch.distributed as dist import tqdm -from . import _interleave from .constants import ( DEFAULT_CHUNK_BYTES, DEFAULT_NUM_STREAMS, @@ -44,6 +44,8 @@ torch_dtype_to_numpy_dtype, ) +logger = logging.getLogger(__name__) + @dataclass class MacroblockSpec: @@ -569,9 +571,6 @@ def _fpz_read_into_cuda_storage( n_threads = effective_read_threads(_env_int("FLASHPACK_READ_THREADS", 16)) half_cap = FPZ_FRAME_UNCOMPRESSED_BYTES // 2 n_slots = _FPZ_CUDA_BUFFERS_PER_THREAD - fused = _env_flag("FLASHPACK_FPZ_FUSED_INTERLEAVE") and ( - _interleave.fused_interleave_available() - ) byte_blocks = [b.view(torch.uint8) for b in blocks] @@ -664,13 +663,8 @@ def _reader() -> None: hi_dev[slot][:half].copy_( hi_pin[slot][:half], non_blocking=True ) - if fused: - _interleave.interleave_into( - seg, lo_dev[slot][:half], hi_dev[slot][:half] - ) - else: - seg[0::2].copy_(lo_dev[slot][:half], non_blocking=True) - seg[1::2].copy_(hi_dev[slot][:half], non_blocking=True) + seg[0::2].copy_(lo_dev[slot][:half], non_blocking=True) + seg[1::2].copy_(hi_dev[slot][:half], non_blocking=True) events[slot].record(stream) finally: os.close(fd) @@ -688,86 +682,15 @@ def _reader() -> None: raise errors[0] -# --------------------------------------------------------------------------- -# nvcomp GPU Zstd decode (prototype; opt-in via FLASHPACK_FPZ_GPU_DECODE=1). -# -# The CPU zstd decode caps fpz at ~9 GB/s logical (H200) versus ~21.6 GB/s for -# a page-hot raw pack, so decode -- not I/O -- is the fpz bottleneck. nvcomp's -# batched GPU Zstd decoder moves that work onto the device. -# -# API discovery -- nvidia-nvcomp-cu12 5.3.0 (pybind11 module ``nvidia.nvcomp``; -# signatures/docstrings read from the compiled nvcomp_impl .so, quoted below): -# -# Codec(algorithm="Zstd", device_id=, cuda_stream=, -# uncomp_chunk_size=65536, bitstream_kind=BitstreamKind.NVCOMP_NATIVE, -# checksum_policy=NO_COMPUTE_NO_VERIFY, decompress_backend=...) -# "Initialize codec." -# algorithm : name of the compression algorithm ("Zstd", "LZ4", ...). -# device_id : device to run on (default: current device). -# cuda_stream : cudaStream_t as a Python int (default: an internal -# stream). We pass each reader thread's own torch stream -# (``stream.cuda_stream``) so the decode is ordered on the -# SAME stream as our H2D copies and the interleave -- no -# cross-stream sync needed. -# bitstream_kind : BitstreamKind.{NVCOMP_NATIVE, RAW, WITH_UNCOMPRESSED_SIZE}. -# We use RAW: "Compresses input data as is, just using the -# underlying compression algorithm. Does not add a header -# with nvCOMP metadata." The fpz high plane is a standard -# single zstd frame written by python-zstandard, so it must -# be decoded as RAW (NVCOMP_NATIVE expects nvcomp's own -# chunked container and would reject a bare zstd frame). -# -# codec.decode(src, data_type="|u1", out=None, decompression_config=None) -# -> nvcomp.Array "Decode a single Array." -# codec.decode(srcs: list[Array], data_type=..., out=, -# decompression_config=None) -> list[Array] -# "Decode a batch of Arrays." -# out : "An optional writable buffer to store decoded data. ... If it is an -# externally-allocated buffer (e.g. cupy/numba array), its size is -# fixed and a ValueError is raised when it is too small." We pass a -# view over our pre-sized device high-plane tensor, so decode writes -# straight into it -- no extra device copy and no host round trip. -# data_type : output element type string; default "|u1" (uint8), which is -# exactly the byte plane we want, so we never pass it. -# decompression_config : when omitted, "decode internally calls -# configure_decompression on src, forcing a stream synchronization" -# on EVERY call -- that per-call sync serialized the whole pipeline -# and measured ~0.6 GB/s on the H200. We instead build a reusable -# DecompressConfig once per distinct batch shape via -# codec.decompression_config(srcs) ("reusable across multiple decode -# calls ... of the same uncompressed per-element shape") and pass it -# to decode, which is then sync-free. fpz packs have only a few -# shapes (full 64 MiB frames + one tail per block), so the one-time -# build sync is paid a handful of times per thread, not per decode. -# (A CompressConfig-derived config -- codec.decompression_config( -# codec.compression_config(size)) -- would skip even the build sync, -# but the docstring scopes that to same-process compress+decompress; -# our frames are compressed offline by python-zstandard, so we use -# the header-parsing overload that is proven against real frames.) +# GPU decode for fpz v2 packs (opt-in via FLASHPACK_FPZ_GPU_DECODE=1). # -# nvcomp.as_array(src_object, cuda_stream=None) -> Array -# "Creates array from object with some standard interface." Zero-copy over -# any object exposing __cuda_array_interface__ / __dlpack__. A contiguous -# torch CUDA tensor qualifies, so we wrap the device staging tensors -# directly (no copy). nvcomp.from_dlpack(...) is the explicit-DLPack -# equivalent; as_array is sufficient here. -# -# nvcomp.set_device_allocator(allocator) -- "Sets a new allocator ... for -# future device allocations." allocator is -# ``allocator(nbytes: int, stream: nvcomp.Stream) -> obj`` where obj has an -# integer ``.ptr`` and frees on garbage collection. nvcomp grabs scratch -# from this for every decode; its default (cudaMalloc/cudaFree) syncs the -# device per call, so we install a torch-caching-allocator adapter (see -# _install_torch_nvcomp_allocator) to serve scratch pool-side with no sync. -# -# Compatibility, verified locally against the pack side (python-zstandard -# level-3, threads=-1, one-shot ``compress``): every high plane is a SINGLE -# standard zstd frame (magic 0xFD2FB528) with the content size embedded, a -# 2 MiB window (windowLog 21), no dictionary and no checksum -- all within -# nvcomp GPU Zstd's limits. nvcomp itself cannot be exercised without a device, -# so confirming decode correctness on a real frame is step 0 of the H200 run. -# The correctness-critical interleave (even byte = low plane, odd byte = high -# plane) is identical to the CPU path and is covered by the CPU tests. +# v2 stores each frame's high plane as many small independent zstd chunks -- +# the shape a batched GPU decompressor needs. The decode runs through ctypes +# bindings to libnvcomp's batched Zstd API (see ``_nvcomp_ll``); when +# libnvcomp is unavailable, or a pack predates the v2 chunk alignment, reads +# fall back to the threaded CPU decode path. The chunks are standard zstd +# frames (python-zstandard output: single frame, embedded content size, +# 2 MiB window, no dictionary), which the batched interface decodes directly. # --------------------------------------------------------------------------- _FPZ_GPU_DECODE_WARNED = False @@ -778,122 +701,10 @@ def _env_flag(name: str) -> bool: def _fpz_gpu_decode_enabled() -> bool: - """Whether the opt-in nvcomp GPU decode path is requested (env-gated).""" + """Whether the opt-in GPU decode path is requested (env-gated).""" return _env_flag("FLASHPACK_FPZ_GPU_DECODE") -def _load_nvcomp(): - """Import the optional nvcomp module for GPU Zstd decode. - - Returns the module, or ``None`` if it is not importable -- in which case it - warns once (per process) so the caller can fall back to the CPU decode path - without spamming. nvcomp is NVIDIA-proprietary and never a hard dependency; - install it with ``pip install 'flashpack[fpz-gpu]'`` (see pyproject). - """ - global _FPZ_GPU_DECODE_WARNED - try: - from nvidia import nvcomp - - return nvcomp - except ImportError: - if not _FPZ_GPU_DECODE_WARNED: - _FPZ_GPU_DECODE_WARNED = True - warnings.warn( - "FLASHPACK_FPZ_GPU_DECODE=1 but the nvcomp package is not " - "importable; falling back to CPU zstd decode. Install the GPU " - "extra with: pip install 'flashpack[fpz-gpu]'.", - RuntimeWarning, - stacklevel=2, - ) - return None - - -def _env_flag_default(name: str, default: bool) -> bool: - raw = os.environ.get(name) - if raw is None: - return default - return raw.strip().lower() in ("1", "true", "yes", "on") - - -# nvcomp calls its device allocator once per decode for scratch. Its default -# allocator is cudaMalloc/cudaFree, and each of those synchronizes the device -- -# on the H200 that per-call sync (not the config sync) was the dominant fpz cost -# (~60ms/frame, tier-flat). Routing nvcomp's scratch through torch's stream-aware -# caching allocator serves it from an existing pool with no cudaMalloc/sync. -_fpz_nvcomp_alloc_tls = threading.local() -_FPZ_NVCOMP_ALLOC_INSTALLED = False -_FPZ_NVCOMP_ALLOC_WARNED = False - - -class _TorchNvcompDeviceBuffer: - """Adapter exposing a torch caching-allocator block to nvcomp's allocator - protocol: an object with an integer ``ptr`` that frees on ``__del__``. - - The allocation is tied to the calling reader thread's CUDA stream (stashed - in a thread-local by the reader) so torch's caching allocator won't hand the - block to another stream while nvcomp's decode -- which runs on that same - stream -- is still using it. - """ - - __slots__ = ("_ptr",) - - def __init__(self, nbytes: int, device_index: int, stream) -> None: - self._ptr = torch.cuda.caching_allocator_alloc(nbytes, device_index, stream) - - @property - def ptr(self) -> int: - return self._ptr - - def __del__(self) -> None: - try: - torch.cuda.caching_allocator_delete(self._ptr) - except Exception: - pass - - -def _install_torch_nvcomp_allocator(nvcomp, device: torch.device) -> bool: - """Route nvcomp's per-decode device scratch through torch's caching allocator. - - Global and idempotent. Guarded: any API mismatch or failure leaves nvcomp on - its default allocator (decode still works, just slower) and warns once. - - nvcomp API (from the wheel's ``set_device_allocator`` docstring): the - allocator is ``allocator(nbytes: int, stream: nvcomp.Stream) -> obj`` where - ``obj`` has an integer ``.ptr`` and releases its memory when garbage - collected. We ignore nvcomp's ``stream`` arg and instead read the reader - thread's torch stream from ``_fpz_nvcomp_alloc_tls`` (set per thread), which - is the stream nvcomp actually decodes on. - """ - global _FPZ_NVCOMP_ALLOC_INSTALLED, _FPZ_NVCOMP_ALLOC_WARNED - if _FPZ_NVCOMP_ALLOC_INSTALLED: - return True - dev_index = ( - device.index if device.index is not None else torch.cuda.current_device() - ) - - def _alloc(nbytes, stream=None): - return _TorchNvcompDeviceBuffer( - int(nbytes), dev_index, getattr(_fpz_nvcomp_alloc_tls, "stream", None) - ) - - try: - nvcomp.set_device_allocator(_alloc) - except Exception: - if not _FPZ_NVCOMP_ALLOC_WARNED: - _FPZ_NVCOMP_ALLOC_WARNED = True - warnings.warn( - "Could not install the torch caching allocator into nvcomp " - "(set_device_allocator failed); nvcomp keeps its default " - "cudaMalloc allocator. Set FLASHPACK_FPZ_GPU_TORCH_ALLOC=0 to " - "silence.", - RuntimeWarning, - stacklevel=2, - ) - return False - _FPZ_NVCOMP_ALLOC_INSTALLED = True - return True - - def fpz_gpu_warmup(device: "str | torch.device" = "cuda") -> bool: """Pay the batched GPU decoder's one-time init cost off the hot path. @@ -991,19 +802,6 @@ def _fpz_hi_chunk_usizes(half: int, chunk_u: int) -> list[int]: return sizes -def _fpz_batch_signature(batch: list[tuple]) -> tuple[int, ...]: - """Config-cache key for a frame batch: the per-frame decompressed high-plane - sizes (``n_out // 2``), in order (pure function). - - An nvcomp ``DecompressConfig`` built from one batch is reusable for any other - batch with the same per-element uncompressed shape, so batches that share - this signature share a single config -- and the one-time - ``decompression_config`` stream sync is paid once per distinct signature - instead of once per decode call. - """ - return tuple(int(frame["n_out"]) // 2 for _, _blk, frame, _out_pos in batch) - - def plan_fpz_gpu_batches( frame_tasks: list[tuple], max_batch_frames: int, @@ -1046,94 +844,67 @@ def plan_fpz_gpu_batches( return batches +def _fpz_pack_chunk_layout(specs: list[MacroblockSpec]) -> tuple[int, int]: + """(chunk_usize, chunk_alignment) recorded by the encoder. + + A pack is written with one chunk size and one chunk-start alignment; + both are read from the first fpz block (absent fields mean the 64 KiB + default / the unpadded pre-alignment layout). A frame whose block + disagrees is caught by the chunk-count check in the read loop. + """ + chunk_u = FPZ_HI_CHUNK_UNCOMPRESSED_BYTES + hi_align = 1 + for spec in specs: + if spec.fpz is not None: + chunk_u = int(spec.fpz.get("hi_chunk_usize", chunk_u)) + hi_align = int(spec.fpz.get("hi_align", 1)) + break + return chunk_u, hi_align + + def _fpz_read_into_cuda_storage_gpu( path: str, specs: list[MacroblockSpec], blocks: list[torch.Tensor], device: torch.device, - nvcomp, + ll, ) -> None: """GPU-decode variant of :func:`_fpz_read_into_cuda_storage` for v2 packs. - v2 stores each frame's high plane as many small independent zstd chunks - (``FPZ_HI_CHUNK_UNCOMPRESSED_BYTES`` each). That is nvcomp's native shape: - the whole point of the GPU decoder is decoding MANY chunks in parallel. A v1 - single-frame high plane is one nvcomp chunk and decodes serially (~0.6 GB/s - measured, tier-flat), which is why the caller routes v1 to the CPU path. - - Per reader thread: read a frame's low plane and its whole compressed-high - blob into pinned staging, H2D both (moving the compressed high plane cuts - PCIe traffic ~2.4x), then submit ALL of the frame's high chunks as one - ``codec.decode`` batch (hundreds of Arrays), decoding straight into the - device high-plane staging; finally the same strided interleave - (``dst[0::2] = lo``, ``dst[1::2] = hi``). - - Two throughput levers, both load-bearing: - - * No per-decode sync. The naive path makes ``decode`` call - ``configure_decompression`` (a stream sync) every call. We build a - reusable ``DecompressConfig`` per distinct chunk-shape signature (one sync - each) and pass it to ``decode``; since v2 chunks are almost all a uniform - 64 KiB, that is ~1-2 configs total per thread and every steady-state decode - is sync-free. Reuse safety without the sync comes from an event-gated - double buffer (``n_slots`` staging sets; ``synchronize`` a slot's event - before reusing it, ``record`` it after decode+interleave). - * Read parallelism. Read (not decode) is now the bottleneck, so we use many - threads (see the tuning-knob pipeline math); preads overlap across threads - and decodes overlap reads via the slots. - - Each thread owns its fd, stream, Codec and config cache. Raw (uncompressed) - blocks take the same whole-block H2D as the CPU-decode path. + v2 stores each frame's high plane as many small independent zstd chunks, + which the batched decompressor decodes in a single launch. Per reader + thread: pread a frame's low plane and compressed high plane into pinned + staging, H2D both (moving the high plane compressed cuts PCIe traffic by + the compression ratio), decode every chunk of the batch with ONE batched + call, then interleave the planes into the destination block (even bytes + low, odd bytes high -- the same invariant as the CPU path). + + Design notes: + + * Per batch, Python fills one pinned int64 chunk table with vectorized + numpy, issues one small H2D plus two on-device base-address adds, and + makes one foreign call -- cost independent of the chunk count. + * Decode never synchronizes: a slot's staging buffers are only reused + after its CUDA event (recorded after the interleave) has fired, and + per-chunk statuses/actual sizes fold into two on-stream scalars that + are checked once at the end. + * Reads dominate, so many threads overlap preads while decodes overlap + reads via the slots. + + ``ll`` is the loaded :mod:`flashpack._nvcomp_ll` binding; the caller + guarantees it is usable and that the pack's chunk alignment satisfies + the decompressor's requirements. """ half_cap = FPZ_FRAME_UNCOMPRESSED_BYTES // 2 - # A pack is written with one chunk size and one chunk-start alignment; - # read both from the first fpz block (absent for pre-parameterization v2 - # packs -> the 64 KiB default / packed layout). A frame whose block - # disagrees is caught by the chunk-count check in the loop. - chunk_u = FPZ_HI_CHUNK_UNCOMPRESSED_BYTES - hi_align = 1 - for spec in specs: - if spec.fpz is not None: - chunk_u = int(spec.fpz.get("hi_chunk_usize", chunk_u)) - hi_align = int(spec.fpz.get("hi_align", 1)) - break + chunk_u, hi_align = _fpz_pack_chunk_layout(specs) max_chunks = (half_cap + chunk_u - 1) // chunk_u - # Upper bound on a frame's whole compressed-high blob: the zstd bound for - # half_cap uncompressed, plus per-chunk zstd frame-header overhead and - # chunk-start padding; aligned so per-frame staging bases (k * comp_cap) - # preserve the chunk-start alignment inside device staging. + # Upper bound on a frame's compressed-high blob: the zstd bound plus + # per-chunk frame-header overhead and chunk-start padding; aligned so + # per-frame staging bases (k * comp_cap) preserve the chunk alignment + # inside device staging. comp_cap = half_cap + (half_cap // 255) + max_chunks * 80 + 4096 comp_cap = _align_up(comp_cap, max(16, hi_align)) - # Batched C-API decode (the "ll" path): one foreign call per batch over - # device-resident chunk tables instead of one pybind Array per chunk. - # Requires the pack's chunk starts to satisfy nvcomp's queried input - # alignment (hi_align-padded packs do; legacy packed layouts fall back). - # Fused single-pass interleave (Triton) vs the strided two-pass copies; - # opt-in while gating, falls back automatically when triton is missing. - fused = _env_flag("FLASHPACK_FPZ_FUSED_INTERLEAVE") and ( - _interleave.fused_interleave_available() - ) - - ll = None - if _env_flag("FLASHPACK_FPZ_GPU_LL"): - from . import _nvcomp_ll - - ll = _nvcomp_ll.load() - if ll is not None: - req_in, req_out, _req_temp = ll.alignments() - if hi_align % req_in != 0 or chunk_u % req_out != 0: - warnings.warn( - f"flashpack: fpz pack chunk alignment (hi_align={hi_align}" - f", chunk_u={chunk_u}) does not satisfy nvcomp's batched " - f"decode requirements (input={req_in}, output={req_out}); " - "using the wrapper decode path. Repack with current " - "flashpack for the batched path.", - RuntimeWarning, - stacklevel=2, - ) - ll = None - n_threads = effective_read_threads( _env_int("FLASHPACK_FPZ_GPU_DECODE_THREADS", _FPZ_GPU_DEFAULT_THREADS) ) @@ -1172,42 +943,14 @@ def _fpz_read_into_cuda_storage_gpu( traces: list[str] = [] traces_lock = threading.Lock() - # Route nvcomp's per-decode scratch through torch's caching allocator to kill - # the per-call cudaMalloc/cudaFree device sync (the round-2 bottleneck). - # Global + idempotent + guarded; disable with FLASHPACK_FPZ_GPU_TORCH_ALLOC=0. - # (Wrapper path only: the ll path manages its own torch-allocated scratch.) - if ll is None and _env_flag_default("FLASHPACK_FPZ_GPU_TORCH_ALLOC", True): - _install_torch_nvcomp_allocator(nvcomp, device) - def _reader(thread_idx: int) -> None: try: fd = os.open(path, os.O_RDONLY) stream = torch.cuda.Stream(device=device) stream.wait_event(alloc_ready) - # nvcomp's device allocator (if installed) reads this thread's stream - # from the thread-local, so decode scratch is tied to the decode - # stream and torch won't reuse it out from under an in-flight decode. - _fpz_nvcomp_alloc_tls.stream = stream - # Resolve the concrete ordinal in this thread: an indexless "cuda" - # device places both the stream and the staging tensors on this - # thread's current device, and the Codec must match or nvcomp raises - # "Input array and Codec device id mismatched". - device_id = ( - device.index - if device.index is not None - else torch.cuda.current_device() - ) - codec = None - if ll is None: - codec = nvcomp.Codec( - algorithm="Zstd", - bitstream_kind=nvcomp.BitstreamKind.RAW, - device_id=device_id, - cuda_stream=stream.cuda_stream, - ) - # One contiguous staging set per slot; frame k lives at k*half_cap - # (low / decompressed-high) or k*comp_cap (compressed-high). Slices - # are contiguous, so nvcomp.as_array wraps them zero-copy. + # One contiguous staging set per slot; frame k lives at + # k * half_cap (low / decompressed-high) or k * comp_cap + # (compressed-high). lo_pin = [ torch.empty(batch_frames * half_cap, dtype=torch.uint8, pin_memory=True) for _ in range(n_slots) @@ -1230,70 +973,40 @@ def _reader(thread_idx: int) -> None: torch.empty(batch_frames * half_cap, dtype=torch.uint8, device=device) for _ in range(n_slots) ] - # Wrapper path: hoist the decode out= wrappers -- one nvcomp.Array - # per (slot, frame, chunk) over a fixed chunk_u slice at frame k's - # chunk j offset, built ONCE (layout is data-independent). decode - # writes the true (config-driven) size <= chunk_u into each, so the - # same wrappers serve every batch. Indexed [slot][k*max_chunks + j]. - # (The compressed-in src wrappers stay per-batch, sized to each - # chunk's exact compressed length, so nvcomp sees exactly one zstd - # frame per Array.) - out_wrap = None - if ll is None: - out_wrap = [ - [ - nvcomp.as_array( - hi_dev[s].narrow(0, k * half_cap + j * chunk_u, chunk_u) - ) - for k in range(batch_frames) - for j in range(max_chunks) - ] - for s in range(n_slots) - ] - else: - # ll path: no per-chunk Python objects at all. Per batch we - # fill ONE pinned int64 table -- rows (src rel offset, src - # size, dst rel offset, dst capacity) -- with vectorized numpy - # over the footer's chunk lengths, H2D it, add the staging - # base addresses on-device, and make one foreign call. The - # pinned table is per SLOT (host reuse is gated by the slot - # event, like the other pinned staging); the device table, - # scratch, and result arrays are per thread (reuse is - # stream-ordered). - cap_chunks = batch_frames * max_chunks - ll_temp_bytes = ll.temp_size( - cap_chunks, chunk_u, batch_frames * half_cap - ) - ll_temp = torch.empty( - max(1, ll_temp_bytes), dtype=torch.uint8, device=device - ) - ll_tab_pin = [ - torch.empty((4, cap_chunks), dtype=torch.int64, pin_memory=True) - for _ in range(n_slots) - ] - ll_tab_np = [t.numpy() for t in ll_tab_pin] - ll_tab_dev = torch.empty( - (4, cap_chunks), dtype=torch.int64, device=device - ) - ll_actual = torch.empty(cap_chunks, dtype=torch.int64, device=device) - ll_statuses = torch.empty(cap_chunks, dtype=torch.int32, device=device) - # Stream-side correctness accumulators: per-chunk statuses and - # actual-size mismatches fold into two scalars ON the decode - # stream (no syncs); read once after the final synchronize. - ll_status_max = torch.zeros((), dtype=torch.int32, device=device) - ll_size_bad = torch.zeros((), dtype=torch.bool, device=device) - ll_usizes_cache: dict[int, np.ndarray] = {} - ll_dst_rel_cache: dict[tuple[int, int], np.ndarray] = {} + # Per-batch chunk table, rows: (0) src offset within hiz staging, + # (1) true compressed length, (2) dst offset within hi staging, + # (3) expected uncompressed size. The pinned copy is per SLOT + # (host reuse is gated by the slot event, like the other pinned + # staging); the device table, scratch and result arrays are per + # thread (their reuse is stream-ordered). + cap_chunks = batch_frames * max_chunks + ll_temp_bytes = ll.temp_size(cap_chunks, chunk_u, batch_frames * half_cap) + ll_temp = torch.empty( + max(1, ll_temp_bytes), dtype=torch.uint8, device=device + ) + ll_tab_pin = [ + torch.empty((4, cap_chunks), dtype=torch.int64, pin_memory=True) + for _ in range(n_slots) + ] + ll_tab_np = [t.numpy() for t in ll_tab_pin] + ll_tab_dev = torch.empty((4, cap_chunks), dtype=torch.int64, device=device) + ll_actual = torch.empty(cap_chunks, dtype=torch.int64, device=device) + ll_statuses = torch.empty(cap_chunks, dtype=torch.int32, device=device) + # Stream-side correctness accumulators: per-chunk statuses and + # actual-size mismatches fold into two scalars ON the decode + # stream (no syncs); read once after the final synchronize. + ll_status_max = torch.zeros((), dtype=torch.int32, device=device) + ll_size_bad = torch.zeros((), dtype=torch.bool, device=device) + ll_usizes_cache: dict[int, np.ndarray] = {} + ll_dst_rel_cache: dict[tuple[int, int], np.ndarray] = {} # Recorded now so the first synchronize on any slot is a no-op. events = [torch.cuda.Event() for _ in range(n_slots)] for ev in events: ev.record(stream) - # Per-thread cache: batch shape signature -> reusable DecompressConfig. - configs: dict[tuple[int, ...], object] = {} batch_idx = 0 - t_pread = t_h2d = t_decode = t_interleave = t_evsync = t_final = 0.0 - t_wrap = 0.0 - n_batches = n_frames_done = n_cfg = 0 + t_pread = t_h2d = t_table = t_decode = t_interleave = 0.0 + t_evsync = t_final = 0.0 + n_batches = n_frames_done = 0 try: while True: item = work.get() @@ -1314,22 +1027,16 @@ def _reader(thread_idx: int) -> None: batch = payload slot = batch_idx % n_slots batch_idx += 1 - # Wait for this slot's previous batch (its interleave) before - # overwriting its pinned/device buffers -- decode no longer - # synchronizes, so this event is what keeps reuse safe. + # Wait for this slot's previous batch (its interleave) + # before overwriting its pinned/device buffers -- decode + # does not synchronize, so this event keeps reuse safe. _t = time.perf_counter() if trace_on else 0.0 events[slot].synchronize() if trace_on: t_evsync += time.perf_counter() - _t halves: list[int] = [] - srcs: list = [] - outs: list = [] - sig_parts: list[int] = [] - n_ll = 0 # chunks staged into the ll table this batch - # Read every frame's planes into this slot's pinned staging, - # H2D them, and stage the per-chunk dispatch (wrapper: one - # src/out Array per chunk; ll: rows of the batch table). + n_chunks = 0 for k, (_, blk, frame, _out_pos) in enumerate(batch): payload_off = int(frame["payload_off"]) lo_len = int(frame["lo_len"]) @@ -1350,8 +1057,8 @@ def _reader(thread_idx: int) -> None: hi_len_total = int(aligned.sum()) if hi_len_total > comp_cap: raise ValueError( - f"fpz compressed frame ({hi_len_total} bytes) exceeds " - f"staging capacity ({comp_cap} bytes)" + f"fpz compressed frame ({hi_len_total} bytes) " + f"exceeds staging capacity ({comp_cap} bytes)" ) usizes = _fpz_hi_chunk_usizes(half, chunk_u) if len(usizes) != m: @@ -1374,7 +1081,8 @@ def _reader(thread_idx: int) -> None: _t = time.perf_counter() if trace_on else 0.0 with torch.cuda.stream(stream): lo_dev[slot].narrow(0, lo_off, half).copy_( - lo_pin[slot].narrow(0, lo_off, half), non_blocking=True + lo_pin[slot].narrow(0, lo_off, half), + non_blocking=True, ) hiz_dev[slot].narrow(0, hiz_off, hi_len_total).copy_( hiz_pin[slot].narrow(0, hiz_off, hi_len_total), @@ -1382,129 +1090,84 @@ def _reader(thread_idx: int) -> None: ) if trace_on: t_h2d += time.perf_counter() - _t + # Vectorized table rows for this frame's chunks. _t = time.perf_counter() if trace_on else 0.0 - if ll is None: - # One src Array per compressed chunk (exact length) - # and its hoisted out wrapper; chunk usizes drive - # the config. This per-chunk wrapper building is - # GIL-bound Python and is the dominant residual - # cost at small chunk sizes -- its own trace - # bucket so its share is visible. - coff = hiz_off - for j, (clen, alen) in enumerate( - zip(clens.tolist(), aligned.tolist()) - ): - srcs.append( - nvcomp.as_array( - hiz_dev[slot].narrow(0, coff, int(clen)) - ) - ) - outs.append(out_wrap[slot][k * max_chunks + j]) - coff += int(alen) - sig_parts.extend(usizes) - else: - # Vectorized table rows for this frame's chunks: - # (0) src offset within hiz staging = padded prefix - # sums, (1) true compressed length, (2) dst offset - # within hi staging, (3) expected uncompressed size. - tab = ll_tab_np[slot] - starts = np.empty(m, dtype=np.int64) - starts[0] = 0 - np.cumsum(aligned[: m - 1], out=starts[1:]) - tab[0, n_ll : n_ll + m] = hiz_off + starts - tab[1, n_ll : n_ll + m] = clens - dst_rel = ll_dst_rel_cache.get((k, m)) - if dst_rel is None: - dst_rel = k * half_cap + ( - np.arange(m, dtype=np.int64) * chunk_u - ) - ll_dst_rel_cache[(k, m)] = dst_rel - tab[2, n_ll : n_ll + m] = dst_rel - caps = ll_usizes_cache.get(half) - if caps is None: - caps = np.asarray(usizes, dtype=np.int64) - ll_usizes_cache[half] = caps - tab[3, n_ll : n_ll + m] = caps - n_ll += m + tab = ll_tab_np[slot] + starts = np.empty(m, dtype=np.int64) + starts[0] = 0 + np.cumsum(aligned[: m - 1], out=starts[1:]) + tab[0, n_chunks : n_chunks + m] = hiz_off + starts + tab[1, n_chunks : n_chunks + m] = clens + dst_rel = ll_dst_rel_cache.get((k, m)) + if dst_rel is None: + dst_rel = k * half_cap + ( + np.arange(m, dtype=np.int64) * chunk_u + ) + ll_dst_rel_cache[(k, m)] = dst_rel + tab[2, n_chunks : n_chunks + m] = dst_rel + caps = ll_usizes_cache.get(half) + if caps is None: + caps = np.asarray(usizes, dtype=np.int64) + ll_usizes_cache[half] = caps + tab[3, n_chunks : n_chunks + m] = caps + n_chunks += m if trace_on: - t_wrap += time.perf_counter() - _t + t_table += time.perf_counter() - _t + # One H2D of the table, two on-device base-address adds, + # ONE foreign call for the whole batch. The add outputs + # are fresh stream-local tensors; the decompressor reads + # them during the (stream-ordered) decode, so dropping + # the Python refs afterwards is safe. _t = time.perf_counter() if trace_on else 0.0 - if ll is None: - # Reusable config per chunk-shape signature: build once - # (one sync, waits on the H2D above), then decode - # sync-free here and on every later batch that shares - # the shape. - sig = tuple(sig_parts) - cfg = configs.get(sig) - if cfg is None: - cfg = codec.decompression_config(srcs) - configs[sig] = cfg - n_cfg += 1 - codec.decode(srcs, out=outs, decompression_config=cfg) - else: - # One H2D of the table, two on-device base-address - # adds, ONE foreign call for the whole batch -- Python - # cost is independent of the chunk count. The add - # outputs are fresh stream-local tensors; nvcomp reads - # them during the (stream-ordered) decode, so dropping - # the Python refs afterwards is safe. - with torch.cuda.stream(stream): - ll_tab_dev[:, :n_ll].copy_( - ll_tab_pin[slot][:, :n_ll], non_blocking=True - ) - src_ptrs = ll_tab_dev[0, :n_ll] + hiz_dev[slot].data_ptr() - dst_ptrs = ll_tab_dev[2, :n_ll] + hi_dev[slot].data_ptr() - ll.decompress_async( - src_ptrs.data_ptr(), - ll_tab_dev[1].data_ptr(), - ll_tab_dev[3].data_ptr(), - ll_actual.data_ptr(), - n_ll, - ll_temp.data_ptr(), - ll_temp_bytes, - dst_ptrs.data_ptr(), - ll_statuses.data_ptr(), - stream.cuda_stream, + with torch.cuda.stream(stream): + ll_tab_dev[:, :n_chunks].copy_( + ll_tab_pin[slot][:, :n_chunks], non_blocking=True + ) + src_ptrs = ll_tab_dev[0, :n_chunks] + hiz_dev[slot].data_ptr() + dst_ptrs = ll_tab_dev[2, :n_chunks] + hi_dev[slot].data_ptr() + ll.decompress_async( + src_ptrs.data_ptr(), + ll_tab_dev[1].data_ptr(), + ll_tab_dev[3].data_ptr(), + ll_actual.data_ptr(), + n_chunks, + ll_temp.data_ptr(), + ll_temp_bytes, + dst_ptrs.data_ptr(), + ll_statuses.data_ptr(), + stream.cuda_stream, + ) + with torch.cuda.stream(stream): + torch.maximum( + ll_status_max, + ll_statuses[:n_chunks].max(), + out=ll_status_max, + ) + torch.logical_or( + ll_size_bad, + (ll_actual[:n_chunks] != ll_tab_dev[3, :n_chunks]).any(), + out=ll_size_bad, ) - with torch.cuda.stream(stream): - torch.maximum( - ll_status_max, - ll_statuses[:n_ll].max(), - out=ll_status_max, - ) - torch.logical_or( - ll_size_bad, - (ll_actual[:n_ll] != ll_tab_dev[3, :n_ll]).any(), - out=ll_size_bad, - ) if trace_on: t_decode += time.perf_counter() - _t # Interleave per frame (same invariant as the CPU path): - # even bytes low plane, odd bytes high plane. Fused = one - # kernel pass; strided = two 2-byte-stride copy passes. + # even bytes low plane, odd bytes high plane. _t = time.perf_counter() if trace_on else 0.0 for k, (_, blk, frame, out_pos) in enumerate(batch): half = halves[k] n_out = int(frame["n_out"]) seg = byte_blocks[blk].narrow(0, out_pos, n_out) with torch.cuda.stream(stream): - if fused: - _interleave.interleave_into( - seg, - lo_dev[slot].narrow(0, k * half_cap, half), - hi_dev[slot].narrow(0, k * half_cap, half), - ) - else: - seg[0::2].copy_( - lo_dev[slot].narrow(0, k * half_cap, half), - non_blocking=True, - ) - seg[1::2].copy_( - hi_dev[slot].narrow(0, k * half_cap, half), - non_blocking=True, - ) + seg[0::2].copy_( + lo_dev[slot].narrow(0, k * half_cap, half), + non_blocking=True, + ) + seg[1::2].copy_( + hi_dev[slot].narrow(0, k * half_cap, half), + non_blocking=True, + ) events[slot].record(stream) if trace_on: t_interleave += time.perf_counter() - _t @@ -1514,35 +1177,29 @@ def _reader(thread_idx: int) -> None: os.close(fd) _t = time.perf_counter() if trace_on else 0.0 stream.synchronize() - _fpz_nvcomp_alloc_tls.stream = None - if ll is not None: - # Deferred per-chunk verification: both scalars were folded on - # the decode stream per batch, so this is the only D2H. - status_val = int(ll_status_max.item()) - if status_val != 0: - raise RuntimeError( - "fpz batched GPU decode reported a per-chunk error: " - + ll.status_string(status_val) - ) - if bool(ll_size_bad.item()): - raise ValueError( - "fpz batched GPU decode produced a chunk size mismatch" - ) + # Deferred per-chunk verification: both scalars were folded on + # the decode stream per batch, so this is the only D2H. + status_val = int(ll_status_max.item()) + if status_val != 0: + raise RuntimeError( + "fpz batched GPU decode reported a per-chunk error: " + + ll.status_string(status_val) + ) + if bool(ll_size_bad.item()): + raise ValueError( + "fpz batched GPU decode produced a chunk size mismatch" + ) if trace_on: t_final += time.perf_counter() - _t # Enqueue phases (h2d, interleave) are async so their wall is - # small; a large `decode` wall means the decode CALL itself - # blocks (internal sync / scratch alloc), while a large - # `evsync`/`final` means the pipeline is GPU-bound waiting on - # decode+interleave to finish. In ll mode `wrap` is the numpy - # table fill and `decode` is the table H2D + foreign call. - mode = "ll" if ll is not None else "wrapper" + # small; a large `decode` means the foreign call itself + # blocks, while large `evsync`/`final` means the pipeline is + # GPU-bound waiting on decode + interleave. line = ( - f"[fpz-gpu-trace] thread={thread_idx} mode={mode} " - f"frames={n_frames_done} " - f"batches={n_batches} cfg_builds={n_cfg} " - f"pread={t_pread:.3f}s h2d_enq={t_h2d:.3f}s " - f"wrap={t_wrap:.3f}s decode={t_decode:.3f}s " + f"fpz-gpu thread={thread_idx} frames={n_frames_done} " + f"batches={n_batches} pread={t_pread:.3f}s " + f"h2d_enq={t_h2d:.3f}s table={t_table:.3f}s " + f"decode={t_decode:.3f}s " f"interleave_enq={t_interleave:.3f}s " f"evsync={t_evsync:.3f}s final_sync={t_final:.3f}s" ) @@ -1562,12 +1219,13 @@ def _reader(thread_idx: int) -> None: torch.cuda.synchronize(device) if trace_on: for line in traces: - print(line) + logger.debug(line) if errors: raise errors[0] _FPZ_V1_GPU_WARNED = False +_FPZ_ALIGN_WARNED = False def _fpz_specs_all_v2(specs: list[MacroblockSpec]) -> bool: @@ -1590,29 +1248,64 @@ def _read_fpz_into_storage( return storage if device.type == "cuda": storage = _allocate_empty_storage(specs, device) - nvcomp = _load_nvcomp() if _fpz_gpu_decode_enabled() else None - # The GPU decoder only helps v2 (chunked) packs; a v1 pack is one nvcomp - # chunk per frame and decodes serially, so fall back to the threaded CPU - # decode for it (still correct, and faster than serial GPU decode). - if nvcomp is not None and not _fpz_specs_all_v2(specs): - nvcomp = None - if not _FPZ_V1_GPU_WARNED: - _FPZ_V1_GPU_WARNED = True - warnings.warn( - "FLASHPACK_FPZ_GPU_DECODE=1 but this pack uses the v1 fpz " - "codec, which cannot be GPU-decoded in parallel; using the " - "CPU decode path. Repack with the v2 encoder for GPU decode.", - RuntimeWarning, - stacklevel=2, - ) - if nvcomp is not None: - _fpz_read_into_cuda_storage_gpu(path, specs, storage.blocks, device, nvcomp) + ll = _fpz_gpu_decoder(specs) if _fpz_gpu_decode_enabled() else None + if ll is not None: + _fpz_read_into_cuda_storage_gpu(path, specs, storage.blocks, device, ll) else: _fpz_read_into_cuda_storage(path, specs, storage.blocks, device) return storage raise ValueError(f"Unsupported device: {device}") +def _fpz_gpu_decoder(specs: list[MacroblockSpec]): + """Resolve the batched GPU decoder for this pack, or ``None`` to use the + CPU decode path. Warns once per reason: libnvcomp missing, a v1 pack (one + chunk per frame decodes serially -- the CPU path is faster), or a pack + whose chunk layout predates the alignment the decompressor requires.""" + global _FPZ_GPU_DECODE_WARNED, _FPZ_V1_GPU_WARNED, _FPZ_ALIGN_WARNED + from . import _nvcomp_ll + + ll = _nvcomp_ll.load() + if ll is None: + if not _FPZ_GPU_DECODE_WARNED: + _FPZ_GPU_DECODE_WARNED = True + warnings.warn( + "FLASHPACK_FPZ_GPU_DECODE=1 but libnvcomp is not available; " + "falling back to CPU zstd decode. Install the GPU extra " + "with: pip install 'flashpack[fpz-gpu]'.", + RuntimeWarning, + stacklevel=3, + ) + return None + if not _fpz_specs_all_v2(specs): + if not _FPZ_V1_GPU_WARNED: + _FPZ_V1_GPU_WARNED = True + warnings.warn( + "FLASHPACK_FPZ_GPU_DECODE=1 but this pack uses the v1 fpz " + "codec, which cannot be GPU-decoded in parallel; using the " + "CPU decode path. Repack with the v2 encoder for GPU decode.", + RuntimeWarning, + stacklevel=3, + ) + return None + chunk_u, hi_align = _fpz_pack_chunk_layout(specs) + req_in, req_out, _req_temp = ll.alignments() + if hi_align % req_in != 0 or chunk_u % req_out != 0: + if not _FPZ_ALIGN_WARNED: + _FPZ_ALIGN_WARNED = True + warnings.warn( + f"FLASHPACK_FPZ_GPU_DECODE=1 but this pack's chunk layout " + f"(hi_align={hi_align}, chunk_usize={chunk_u}) does not " + f"satisfy the decompressor's alignment requirements " + f"(input={req_in}, output={req_out}); using the CPU decode " + "path. Repack with the current encoder for GPU decode.", + RuntimeWarning, + stacklevel=3, + ) + return None + return ll + + def read_flashpack_file( path: str, device: str | torch.device = "cpu", diff --git a/tests/test_fpz_gpu.py b/tests/test_fpz_gpu.py index 4d1348e..6e7a918 100644 --- a/tests/test_fpz_gpu.py +++ b/tests/test_fpz_gpu.py @@ -1,24 +1,19 @@ -"""CPU-testable surface of the nvcomp GPU Zstd decode path for fpz. +"""CPU-testable surface of the fpz GPU decode path. -The GPU decode itself needs a CUDA device and nvcomp, so it is exercised on the -H200, not in CI. What IS testable without a device -- and what these cover -- is -everything around it: the pure frame-batch planner, the env gating, and the -guarded-import fallback (including warn-once), plus a guard that turning the -flag on never disturbs the CPU decode path. +The GPU decode itself needs a CUDA device and libnvcomp, so it is exercised +on GPU hosts, not in CI. What IS testable without a device -- and what these +cover -- is everything around it: the pure frame-batch planner, the env +gating, the CPU-fallback warnings, and a guard that turning the flag on +never disturbs the CPU decode path. """ -import sys - import pytest import torch from flashpack import deserialization from flashpack.deserialization import ( _env_flag, - _env_flag_default, - _fpz_batch_signature, _fpz_gpu_decode_enabled, _fpz_hi_chunk_usizes, - _load_nvcomp, iterate_from_flash_tensor, plan_fpz_gpu_batches, read_flashpack_file, @@ -94,30 +89,6 @@ def test_plan_rejects_nonpositive_byte_budget(bad: int) -> None: ) -# -------------------------------------------------------------------------- -# _fpz_batch_signature -- pure config-cache key -# -------------------------------------------------------------------------- - - -def test_signature_is_per_frame_half_sizes() -> None: - batch = [_frame_task(0, 100), _frame_task(0, 40)] - assert _fpz_batch_signature(batch) == (50, 20) - - -def test_signature_matches_for_same_shape_batches() -> None: - # Two batches of identical full-frame shapes hash to the same key -> they - # share one cached DecompressConfig (the whole point of the cache). - a = [_frame_task(0, 128), _frame_task(0, 128)] - b = [_frame_task(1, 128), _frame_task(2, 128)] - assert _fpz_batch_signature(a) == _fpz_batch_signature(b) == (64, 64) - - -def test_signature_differs_when_a_tail_frame_changes_shape() -> None: - full = [_frame_task(0, 128), _frame_task(0, 128)] - with_tail = [_frame_task(0, 128), _frame_task(0, 40)] - assert _fpz_batch_signature(full) != _fpz_batch_signature(with_tail) - - # -------------------------------------------------------------------------- # _fpz_hi_chunk_usizes -- pure v2 chunk sizing (matches the encoder) # -------------------------------------------------------------------------- @@ -175,97 +146,27 @@ def test_env_flag_unset_is_false(monkeypatch) -> None: assert _fpz_gpu_decode_enabled() is False -def test_env_flag_default_respects_default_when_unset(monkeypatch) -> None: - monkeypatch.delenv("FLASHPACK_FPZ_GPU_TORCH_ALLOC", raising=False) - assert _env_flag_default("FLASHPACK_FPZ_GPU_TORCH_ALLOC", True) is True - assert _env_flag_default("FLASHPACK_FPZ_GPU_TORCH_ALLOC", False) is False - - -@pytest.mark.parametrize( - "value,expected", [("0", False), ("false", False), ("1", True), ("on", True)] -) -def test_env_flag_default_env_overrides(monkeypatch, value, expected) -> None: - monkeypatch.setenv("FLASHPACK_FPZ_GPU_TORCH_ALLOC", value) - # Env always wins over the default, in both directions. - assert _env_flag_default("FLASHPACK_FPZ_GPU_TORCH_ALLOC", True) is expected - assert _env_flag_default("FLASHPACK_FPZ_GPU_TORCH_ALLOC", False) is expected - - # -------------------------------------------------------------------------- -# torch caching allocator wiring into nvcomp (guarded, idempotent) +# CPU-fallback warning (warn once) # -------------------------------------------------------------------------- -class _FakeNvcompAllocOK: - def __init__(self) -> None: - self.installed = None - - def set_device_allocator(self, allocator) -> None: - self.installed = allocator - - -class _FakeNvcompAllocRaises: - def set_device_allocator(self, allocator) -> None: - raise RuntimeError("no allocator hook here") - - -def _reset_alloc_latches(monkeypatch) -> None: - monkeypatch.setattr(deserialization, "_FPZ_NVCOMP_ALLOC_INSTALLED", False) - monkeypatch.setattr(deserialization, "_FPZ_NVCOMP_ALLOC_WARNED", False) - - -def test_install_allocator_success_registers_a_callable(monkeypatch) -> None: - _reset_alloc_latches(monkeypatch) - fake = _FakeNvcompAllocOK() - ok = deserialization._install_torch_nvcomp_allocator(fake, torch.device("cuda:0")) - assert ok is True - assert callable(fake.installed) - - -def test_install_allocator_is_idempotent(monkeypatch) -> None: - _reset_alloc_latches(monkeypatch) - first = _FakeNvcompAllocOK() - assert deserialization._install_torch_nvcomp_allocator( - first, torch.device("cuda:0") - ) - # Already installed globally: a second call is a no-op and does not - # re-register on another (fake) module. - second = _FakeNvcompAllocOK() - assert deserialization._install_torch_nvcomp_allocator( - second, torch.device("cuda:0") - ) - assert second.installed is None - - -def test_install_allocator_failure_is_guarded_and_warns(monkeypatch) -> None: - _reset_alloc_latches(monkeypatch) - with pytest.warns(RuntimeWarning, match="set_device_allocator"): - ok = deserialization._install_torch_nvcomp_allocator( - _FakeNvcompAllocRaises(), torch.device("cuda:0") - ) - assert ok is False - - -# -------------------------------------------------------------------------- -# guarded import + warn-once fallback -# -------------------------------------------------------------------------- - +def test_gpu_decoder_missing_libnvcomp_warns_once_and_falls_back( + monkeypatch, +) -> None: + from flashpack import _nvcomp_ll -def test_load_nvcomp_missing_returns_none_and_warns_once(monkeypatch) -> None: - # Force the import to fail regardless of what's installed, and reset the - # process-level warn-once latch so the assertion is deterministic. - monkeypatch.setitem(sys.modules, "nvidia", None) + monkeypatch.setattr(_nvcomp_ll, "_loaded", (None,)) monkeypatch.setattr(deserialization, "_FPZ_GPU_DECODE_WARNED", False) - with pytest.warns(RuntimeWarning, match="nvcomp"): - assert _load_nvcomp() is None + with pytest.warns(RuntimeWarning, match="libnvcomp"): + assert deserialization._fpz_gpu_decoder([]) is None - # Second call: still None, but no second warning (warn-once). import warnings as _warnings with _warnings.catch_warnings(): _warnings.simplefilter("error") # any warning would raise - assert _load_nvcomp() is None + assert deserialization._fpz_gpu_decoder([]) is None # -------------------------------------------------------------------------- diff --git a/tests/test_interleave.py b/tests/test_interleave.py deleted file mode 100644 index a2b4da4..0000000 --- a/tests/test_interleave.py +++ /dev/null @@ -1,21 +0,0 @@ -"""Unit tests for the fused-interleave module's degradation behavior (the -kernel itself needs a GPU and is validated by the H200 probe/gate: numeric -parity against the strided path plus the pack checksum gate).""" - -import torch -from flashpack import _interleave - - -def test_available_is_bool_and_cached() -> None: - first = _interleave.fused_interleave_available() - assert isinstance(first, bool) - assert _interleave.fused_interleave_available() == first - - -def test_interleave_into_degrades_without_triton() -> None: - if _interleave.fused_interleave_available(): - return # covered by the GPU gate on hosts that have triton+CUDA - out = torch.empty(8, dtype=torch.uint8) - lo = torch.zeros(4, dtype=torch.uint8) - hi = torch.ones(4, dtype=torch.uint8) - assert _interleave.interleave_into(out, lo, hi) is False From 13ff0ed318c0a9778ae1444694146b31363f12ec Mon Sep 17 00:00:00 2001 From: Alperen Konukbay Date: Sun, 26 Jul 2026 22:58:35 -0700 Subject: [PATCH 3/6] test: add zstandard to the dev extra (CI installs .[dev]; fpz tests need it) --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 85f9e6b..68736d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,8 @@ dev = [ "pytest>=8.4.1", "pre-commit>=3.0.0", "setuptools-scm>=9.2.0", + # fpz tests exercise the compressed format end to end + "zstandard>=0.22", ] fpz = [ "zstandard>=0.22", From f8ec563531f3f4dc818c5a2f416b8d3fe0b5d7bb Mon Sep 17 00:00:00 2001 From: Alperen Konukbay Date: Sun, 26 Jul 2026 23:04:25 -0700 Subject: [PATCH 4/6] test: skip fpz read tests where preadv is unavailable (Windows), same policy as the O_DIRECT reader tests --- tests/test_fpz.py | 8 ++++++++ tests/test_fpz_gpu.py | 10 ++++++++++ tests/test_fpz_ll.py | 10 ++++++++++ 3 files changed, 28 insertions(+) diff --git a/tests/test_fpz.py b/tests/test_fpz.py index b14da4b..958086e 100644 --- a/tests/test_fpz.py +++ b/tests/test_fpz.py @@ -33,6 +33,14 @@ from flashpack.serialization import pack_to_file from flashpack.utils import require_zstandard +# The fpz read paths pull bytes with os.preadv (POSIX-only), matching the +# repo's O_DIRECT reader; the format targets Linux GPU fleets. Encoder and +# reader are exercised on Linux/macOS. +pytestmark = pytest.mark.skipif( + not hasattr(os, "preadv"), + reason="fpz read paths require os.preadv (POSIX-only)", +) + def _bf16_state_dict() -> dict[str, torch.Tensor]: generator = torch.Generator().manual_seed(0) diff --git a/tests/test_fpz_gpu.py b/tests/test_fpz_gpu.py index 6e7a918..30bd63d 100644 --- a/tests/test_fpz_gpu.py +++ b/tests/test_fpz_gpu.py @@ -7,6 +7,8 @@ never disturbs the CPU decode path. """ +import os + import pytest import torch from flashpack import deserialization @@ -20,6 +22,14 @@ ) from flashpack.serialization import pack_to_file +# The fpz read paths pull bytes with os.preadv (POSIX-only), matching the +# repo's O_DIRECT reader; the format targets Linux GPU fleets. Encoder and +# reader are exercised on Linux/macOS. +pytestmark = pytest.mark.skipif( + not hasattr(os, "preadv"), + reason="fpz read paths require os.preadv (POSIX-only)", +) + def _frame_task(block_idx: int, n_out: int, out_pos: int = 0) -> tuple: # Matches _fpz_frame_tasks' ("frame", block_idx, frame, out_pos) shape. diff --git a/tests/test_fpz_ll.py b/tests/test_fpz_ll.py index 240d1ba..46404a5 100644 --- a/tests/test_fpz_ll.py +++ b/tests/test_fpz_ll.py @@ -6,7 +6,9 @@ """ import ctypes +import os +import pytest import torch from flashpack import _nvcomp_ll, serialization from flashpack.constants import FPZ_HI_CHUNK_ALIGN_BYTES @@ -18,6 +20,14 @@ ) from flashpack.serialization import pack_to_file +# The fpz read paths pull bytes with os.preadv (POSIX-only), matching the +# repo's O_DIRECT reader; the format targets Linux GPU fleets. Encoder and +# reader are exercised on Linux/macOS. +pytestmark = pytest.mark.skipif( + not hasattr(os, "preadv"), + reason="fpz read paths require os.preadv (POSIX-only)", +) + def _state_with_unaligned_tail() -> dict[str, torch.Tensor]: # An odd bf16 element count makes the (single, partial) frame's half-plane From 9601e5f4a9d88a6ee1f253eecf2d52782fa9b208 Mon Sep 17 00:00:00 2001 From: Alperen Konukbay Date: Sun, 26 Jul 2026 23:15:59 -0700 Subject: [PATCH 5/6] test: drop informational ratio prints from fpz size tests The assertions carry the contract; the capsys-disabled prints were bench instrumentation. Co-Authored-By: Claude Fable 5 --- tests/test_fpz.py | 18 +++--------------- 1 file changed, 3 insertions(+), 15 deletions(-) diff --git a/tests/test_fpz.py b/tests/test_fpz.py index 958086e..c5e2381 100644 --- a/tests/test_fpz.py +++ b/tests/test_fpz.py @@ -99,7 +99,7 @@ def test_compressed_file_is_v4_with_fpz_record(tmp_path) -> None: assert block["length_bytes"] < block["length_elems"] * 2 -def test_low_entropy_tensor_shrinks_file(tmp_path, capsys) -> None: +def test_low_entropy_tensor_shrinks_file(tmp_path) -> None: ramp = (torch.arange(2048 * 1024, dtype=torch.float32) * 0.01).reshape(2048, 1024) source = {"w": ramp.to(torch.bfloat16)} plain = _pack(tmp_path, source, "plain.flashpack") @@ -107,14 +107,8 @@ def test_low_entropy_tensor_shrinks_file(tmp_path, capsys) -> None: plain_size = os.path.getsize(plain) comp_size = os.path.getsize(comp) - ratio = plain_size / comp_size - with capsys.disabled(): - print( - f"\n[fpz] low-entropy ramp: plain={plain_size} comp={comp_size} " - f"ratio={ratio:.3f}x" - ) assert comp_size < plain_size - assert ratio > 1.5 + assert plain_size / comp_size > 1.5 def test_mixed_dtype_only_bf16_compressed(tmp_path) -> None: @@ -462,7 +456,7 @@ def test_v1_pack_still_reads(tmp_path, monkeypatch) -> None: assert torch.equal(_uint16_view(tensors[name]), _uint16_view(original)) -def test_v2_ratio_close_to_v1(tmp_path, monkeypatch, capsys) -> None: +def test_v2_ratio_close_to_v1(tmp_path, monkeypatch) -> None: # v2 compresses each 64 KiB chunk independently, so it shrinks slightly less # than v1's single-frame high plane. Measure both on a realistic low-entropy # tensor; v2 must still compress and stay within a modest margin of v1. @@ -478,12 +472,6 @@ def test_v2_ratio_close_to_v1(tmp_path, monkeypatch, capsys) -> None: plain_sz = os.path.getsize(plain) v1_sz = os.path.getsize(v1) v2_sz = os.path.getsize(v2) - with capsys.disabled(): - print( - f"\n[fpz] ratio plain={plain_sz} " - f"v1={v1_sz} ({plain_sz / v1_sz:.3f}x) " - f"v2={v2_sz} ({plain_sz / v2_sz:.3f}x) v2/v1={v2_sz / v1_sz:.3f}" - ) assert v2_sz < plain_sz # v2 still compresses assert v2_sz <= v1_sz * 1.25 # within a modest margin of v1 From c9e3a3318abcd0ca92eff8083da55679518a8e33 Mon Sep 17 00:00:00 2001 From: Alperen Konukbay Date: Fri, 31 Jul 2026 13:07:41 -0700 Subject: [PATCH 6/6] fix(fpz): default high-plane chunk size 64 KiB -> 1 MiB Pre-merge regression gate (same-boot interleaved A/B on a 38 GB pack): the shipped 64 KiB default -- chosen to match nvcomp's default chunk size -- decodes 2.6x slower than 1 MiB on the threaded CPU path (39.5 s vs 15.3 s hot; per-chunk overhead x 512 chunks per frame) for a compression-ratio difference under 1% (1.412x vs 1.401x). Every GPU batched-decode receipt (1.67-1.85 s per 38 GB) was earned on 1 MiB packs, so the default now matches the receipted configuration. Co-Authored-By: Claude Fable 5 --- src/flashpack/constants.py | 11 +++++++---- tests/test_fpz.py | 4 ++-- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/flashpack/constants.py b/src/flashpack/constants.py index 0a4d229..a226ce4 100644 --- a/src/flashpack/constants.py +++ b/src/flashpack/constants.py @@ -30,10 +30,13 @@ FPZ_CODEC_SPLITPLANE_V2 = "zstd-splitplane-v2" FPZ_FRAME_UNCOMPRESSED_BYTES = 64 * 1024 * 1024 # 64 MiB FPZ_FRAME_ALIGN_BYTES = 4096 -# Per-chunk uncompressed size for the v2 high plane. Matches nvcomp's default -# uncomp_chunk_size (65536) and divides the 32 MiB half-frame evenly (512 -# chunks), so full frames have no odd-sized tail chunk. -FPZ_HI_CHUNK_UNCOMPRESSED_BYTES = 64 * 1024 +# Per-chunk uncompressed size for the v2 high plane. Divides the 32 MiB +# half-frame evenly (32 chunks), so full frames have no odd-sized tail chunk. +# Measured on a 38 GB pack (same-boot interleaved A/B): 64 KiB chunks -- the +# nvcomp default this constant originally matched -- decode 2.6x slower than +# 1 MiB on the threaded CPU path (per-chunk overhead x 512 chunks per frame), +# for a compression-ratio difference under 1% (1.412x vs 1.401x). +FPZ_HI_CHUNK_UNCOMPRESSED_BYTES = 1024 * 1024 # Byte alignment of each v2 compressed chunk's START within the frame payload # (chunks are padded to this; "hi_chunks" still records true zstd lengths and # the reader recomputes padded offsets from the block's "hi_align" footer diff --git a/tests/test_fpz.py b/tests/test_fpz.py index c5e2381..212a0e0 100644 --- a/tests/test_fpz.py +++ b/tests/test_fpz.py @@ -427,8 +427,8 @@ def best_wall(threads: int, reps: int = 3) -> float: def test_v2_frame_splits_high_plane_into_chunks(tmp_path) -> None: - # A 1024x512 bf16 tensor has a 512 KiB high plane -> several 64 KiB chunks. - source = {"w": torch.randn(1024, 512).to(torch.bfloat16)} + # A 4096x1024 bf16 tensor has a 4 MiB high plane -> several 1 MiB chunks. + source = {"w": torch.randn(4096, 1024).to(torch.bfloat16)} comp = _pack(tmp_path, source, "comp.flashpack", compress="fpz-bf16") block = get_flashpack_file_metadata(comp)["macroblocks"][0] assert block["fpz"]["codec"] == FPZ_CODEC_SPLITPLANE_V2