Zstd Codec on the GPU by akshaysubr · Pull Request #2863 · zarr-developers/zarr-python

akshaysubr · 2025-02-25T16:15:20Z

This PR adds a Zstd codec that runs on the GPU using the nvCOMP 4.2 python APIs.

TODO:

Make fully async
Performance benchmarking
CPU-GPU roundtrip testing
Add unit tests and/or doctests in docstrings
Add docstrings and API docs for any new/modified user-facing classes and functions
New/modified features documented in docs/user-guide/*.rst
Changes documented as a new file in changes/
GitHub Actions have all passed
Test coverage is 100% (Codecov passes)

dstansby · 2025-02-28T09:09:22Z

Thanks for opening this PR! At the moment we do not have any codecs implemented in the zarr-python package, but instead store them in numcodecs. So although it looks like there are some zarr-python specific changes that are needed to support the new GPU codec, the actual codec should be implemented in numcodecs, and then imported in zarr-python.

akshaysubr · 2025-02-28T16:00:06Z

@dstansby My understanding was that numcodecs is the place that python bindings to native codec implementations live and that with v3, the Codec class itself lives in zarr-python. The GPU codecs and python bindings are implemented in nvCOMP and imported through the nvidia-nvcomp-cu12 python package so I'm not sure which part of this would need to go in numcodecs. What did you have in mind?

…hread

TomAugspurger

Looks nice overall. I think the async side of things ended up in a pretty good spot. The code itself is pretty easy to follow (a normal stream synchronize). Having to do that on another host thread is a bit unfortunate, but there's only one synchronize per batch so this should be fine.

I left comments on a few things to clean up that I can help out with if you want @akshaysubr.

TomAugspurger · 2025-07-07T11:38:18Z

+    checksum: bool = False
+
+    def __init__(self, *, level: int = 0, checksum: bool = False) -> None:
+        # TODO: Set CUDA device appropriately here and also set CUDA stream


Agreed with leaving devices / streams as a TODO for now.

I want to enable users to overlap host-to-device memcpys with compute operations (like decode, but their own compute operations as well), but I'm not sure yet what that API will look like.

If you have any thoughts on how best to do this I'd love to hear them, and write them up as an issue.

#3271 for planning on devices and streams.

TomAugspurger · 2025-07-07T11:41:03Z

+        chunks_and_specs: Iterable[tuple[Buffer | None, ArraySpec]],
+    ) -> Iterable[Buffer | None]:
+        return [
+            spec.prototype.buffer.from_array_like(cp.array(a, dtype=np.dtype("b"), copy=False))


This is one spot where @weiji14's idea to use dlpack in #2658 (comment) would help. If NDBuffer new how to consume objects implementing the dlpack protocol, we could (maybe) get rid of the cp.array

TomAugspurger · 2025-07-09T19:30:53Z

Fixed some merge conflicts and changed the _convert_{to,from}_nvcomp_arrays methods to be sync.

TomAugspurger · 2025-07-09T19:38:38Z

        if array_like.ndim != 1:
            raise ValueError("array_like: only 1-dim allowed")
-        if array_like.dtype != np.dtype("B"):
+        if array_like.dtype.itemsize != 1:


The new tests in test_nvcomp.py were failing without this change.

I'd like to get us to a point where we don't care as much about the details of the buffer passed in here. This is an OK start I think.

what exactly does this check for? It's not clear to me why any numpy array that can be viewed as bytes wouldn't be allowed in here

and same for the dimensionality check, since any N-dimensional numpy array can be viewed as a 1D array.

Yeah, I'm not really sure...

I agree that the actual data we store internally here needs to be a byte dtype. Just doing cp.asarray(input).view("b") seems pretty reasonable to me.

i'm not even convinced that we need Buffer / NDBuffer, when Buffer is just a special case of NDBuffer where there's 1 dimension and the data type is bytes

we could even express this formally by:

making NDBuffer generic with two type parameters (number of dimensions and dtype)

having APIs that insist on consuming a Buffer instead insist on consuming NDBuffer[Literal[1], np.uint8]

(super out of scope for this PR ofc)

It's not clear to me why any numpy array that can be viewed as bytes wouldn't be allowed in here

I think this is mainly because NDBuffer objects don't need to be contiguous, but Buffer objects must be contiguous in memory which might be important when we send those out to codecs that expect a contiguous memory slice.

But I agree that we can probably merge those two and make Buffer a specialization of NDBuffer.

codecov · 2025-07-09T19:59:31Z

Codecov Report

❌ Patch coverage is 47.25275% with 48 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.61%. Comparing base (4dc5b01) to head (92dbfb8).

Files with missing lines	Patch %	Lines
src/zarr/codecs/cuda.py	50.61%	40 Missing ⚠️
src/zarr/core/buffer/gpu.py	0.00%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2863      +/-   ##
==========================================
- Coverage   94.90%   94.61%   -0.30%     
==========================================
  Files          96       97       +1     
  Lines       14051    14129      +78     
==========================================
+ Hits        13335    13368      +33     
- Misses        716      761      +45

Files with missing lines	Coverage Δ
src/zarr/codecs/__init__.py	`100.00% <100.00%> (ø)`
src/zarr/core/array.py	`98.60% <100.00%> (-0.01%)`	⬇️
src/zarr/core/config.py	`100.00% <ø> (ø)`
src/zarr/core/buffer/gpu.py	`82.50% <0.00%> (-6.66%)`	⬇️
src/zarr/codecs/cuda.py	`50.61% <50.61%> (ø)`

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

…ed config - decode()/encode() created a cp.cuda.Event() but never called .record() before .synchronize(), so the "wait for GPU completion" was a no-op -- a race, not just wasted asyncio.to_thread overhead. Now records on the stream the nvcomp op was enqueued on before waiting. - _zstd_codec was a cached_property binding device/stream once, forever, even though codec/pipeline instances are cached and reused (e.g. by sharding's pipeline cache). Now built fresh per call from the caller's current device/stream. - level/checksum were parsed and round-tripped through to_dict/from_dict but never passed to nvcomp.Codec(...). checksum now maps to nvcomp's ChecksumPolicy; level has no nvCOMP equivalent, so a nonzero value now raises NotImplementedError instead of being silently dropped. - is_fixed_size was True for a compressor (encoded size is data-dependent, same as the CPU ZstdCodec's is_fixed_size = False); left is_fixed_size True made compute_encoded_size's NotImplementedError reachable by any caller that correctly trusts is_fixed_size before calling it.

# Conflicts: # docs/user-guide/gpu.md # pyproject.toml # src/zarr/codecs/__init__.py # src/zarr/codecs/gpu.py # src/zarr/core/buffer/gpu.py # tests/test_codecs/test_codecs.py # tests/test_codecs/test_nvcomp.py

# Conflicts: # src/zarr/core/buffer/gpu.py # tests/test_api.py

…in CI

…of cast Answers review question on why the cast was needed: get_codec_class returns type[Codec]; the registry lookup exists so enable_gpu() can swap in NvcompZstdCodec. A runtime isinstance check replaces the static cast.

…n gpu.NDBuffer (S101)

…Buffer object-dtype error with main

negin513 · 2026-10-14T19:49:29Z

@akshaysubr I've opened akshaysubr#4 against this PR's gpu-codecs branch: it merges current main (resolving the conflicts here), carries @TomAugspurger's follow-up commits that hadn't been pushed to this branch, and fixes a few NvcompZstdCodec issues (unrecorded sync event, cached device/stream, ignored level/checksum config, is_fixed_size). CPU tests/lint/mypy pass locally; GPU tests still need a run on GPU hardware. Merging it into gpu-codecs will update this PR.

…ailable in CI" This reverts commit 4d4463f.

Per review discussion on zarr-developers#2863 (JeppeKlitgaard, d-v-b, TomAugspurger): the extra installs CUDA-specific packages (cupy-cuda12x, nvidia-nvcomp-cu12), so name it after CUDA rather than GPUs in general, leaving room for other accelerator backends. The 'gpu' extra is removed (extras can't be deprecated). enable_gpu() and zarr.buffer.gpu are unchanged: they predate this PR and are public API.

Raising NotImplementedError in __init__ also fired from from_dict, so any array written by the CPU ZstdCodec with level != 0 could not be opened under enable_gpu(). level is an encode-only parameter; accept and preserve it in metadata as the CPU codec does. Restore the device/stream TODO.

nvCOMP only allows NO_COMPUTE_NO_VERIFY with a RAW bitstream, so mapping checksum to ChecksumPolicy made checksum=True fail on both encode and decode. The CUDA event is still recorded before synchronizing.

Sync gpu-codecs with main, rename the gpu extra to cuda, fix nvCOMP Zstd sync and options

negin513 · 2026-10-27T17:09:27Z

Hey @d-v-b and @TomAugspurger ; this should be ready to merged. Can you please take a look at this - so we can bring it in? 🙌

TomAugspurger · 2026-10-28T17:16:50Z

Sweet, I'll take a look in the next few days.

TomAugspurger

Looks good to me overall. I think the only things worth sorting out are

What do we want to call the new extra? If we're breaking install scripts by moving from gpu to cuda (a good call IMO) we should consider handling the cu12 vs. cu13 side too.
An approval from someone outside Nvidia, like @d-v-b would be worthwhile to ensure folks are still happy with the overall approach of developing accelerated / non-NumPy variants here in zarr-python.

TomAugspurger · 2026-10-31T13:37:24Z

+For example, after calling `zarr.config.enable_gpu()`, an nvcomp-based
+codec will be used:
+
+```python exec="true" session="config" source="above" result="ansi"


Does exec="true" actually execute this code block? I don't recall the behavior of enable_gpu() on a machine without a GPU, but I suspect it'll either error, or the codecs.zstd won't change so the print won't show what we want.

I think this is fine. exec="true" does execute the code block, but this "config" session is only interacting with the configuration, whereas the error happens during encoding/decoding.

TomAugspurger · 2026-10-31T14:22:44Z

+cuda = [
    "cupy-cuda12x; sys_platform != 'darwin'",
+    "nvidia-nvcomp-cu12; sys_platform != 'darwin'",
 ]


Since we're breaking this anyway, should we spit it into cu12 and cu13? And then zarr[cuda-12] can depend on cupy-cudad12x / nvidia-nvcomp-cu12?

There's a bunch of work going on in the packaging world to make this nicer, but I'm guessing that won't be ready for a year or two.

I think breaking this up is a good idea. I also don't see a harm in using a gpu alias for both environments to keep code from breaking.

TomAugspurger · 2026-10-31T16:42:55Z

+        # Record event for synchronization
+        event = cp.cuda.Event()
+        # Wait for decode to complete in a separate async thread
+        await asyncio.to_thread(event.synchronize)


I still think we want to work towards avoiding this stream synchronization. A useful exercise would be to write down all the information we promise to be "ready" when this function returns. I suspect that that's just the shape and dtype of the arrays, but not the contents of those arrays. Ideally those would be scheduled but running asynchronously in the background when we return.

maxrjones

thanks for your persistence with this!

I think the failing tests point to a real issue. I have a few other comments as well

maxrjones · 2026-11-13T04:52:37Z

+.. _user-guide-gpu:
+
+Using GPUs with Zarr
+====================
+
+Zarr can use GPUs to accelerate your workload by running
+:meth:`zarr.config.enable_gpu`.
+
+Reading data into device memory
+-------------------------------
+
+:meth:`zarr.config.enable_gpu` configures Zarr to use GPU memory for the data
+buffers used internally by Zarr.
+
+.. code-block:: python
+
+   >>> import zarr
+   >>> import cupy as cp  # doctest: +SKIP
+   >>> zarr.config.enable_gpu()  # doctest: +SKIP
+   >>> store = zarr.storage.MemoryStore()  # doctest: +SKIP
+   >>> z = zarr.create_array(  # doctest: +SKIP
+   ...     store=store, shape=(100, 100), chunks=(10, 10), dtype="float32",
+   ... )
+   >>> type(z[:10, :10])  # doctest: +SKIP
+   cupy.ndarray
+
+Note that the output type is a ``cupy.ndarray`` rather than a NumPy array.
+
+For supported codecs, data will be decoded using the GPU via the `nvcomp`_
+library. See :ref:`user-guide-config` for more. Isseus and feature requests
+for NVIDIA nvCOMP can be reported in the `nvcomp issue tracker`_.
+
+.. _nvcomp: https://docs.nvidia.com/cuda/nvcomp/samples/python_samples.html
+.. _nvcomp issue tracker: https://github.com/NVIDIA/CUDALibrarySamples/issues


Any new documentation should go in gpu.md. Please migrate the docs there and remove this file.

maxrjones · 2026-11-13T05:01:01Z

+For example, after calling `zarr.config.enable_gpu()`, an nvcomp-based
+codec will be used:
+
+```python exec="true" session="config" source="above" result="ansi"


I think this is fine. exec="true" does execute the code block, but this "config" session is only interacting with the configuration, whereas the error happens during encoding/decoding.

maxrjones · 2026-11-13T05:02:32Z

+cuda = [
    "cupy-cuda12x; sys_platform != 'darwin'",
+    "nvidia-nvcomp-cu12; sys_platform != 'darwin'",
 ]


I think breaking this up is a good idea. I also don't see a harm in using a gpu alias for both environments to keep code from breaking.

maxrjones · 2026-11-13T05:03:52Z

+instead of `gpu`, since it installs CUDA-specific packages (CuPy and nvCOMP). The `gpu` extra
+has been removed.


Suggested change

instead of `gpu`, since it installs CUDA-specific packages (CuPy and nvCOMP). The `gpu` extra

has been removed.

instead of `gpu`, since it installs CUDA-specific packages (CuPy and nvCOMP).

Please see comment below about keeping a gpu alias for backwards compatibility

maxrjones · 2026-11-13T05:19:20Z

+                    "batch_size": 65536,
+                },
+            }
        )


I'm curious if enable_gpu() should be left as-is (made vestigial and eventually deprecated), with this configuration feature added under a new enable_cuda since it's cude specific.

maxrjones · 2026-11-13T05:21:04Z

+            {
+                "buffer": "zarr.buffer.gpu.Buffer",
+                "ndbuffer": "zarr.buffer.gpu.NDBuffer",
+                "codecs": {"zstd": "zarr.codecs.cuda.NvcompZstdCodec"},
+                "codec_pipeline": {
+                    "path": "zarr.core.codec_pipeline.BatchedCodecPipeline",
+                    "batch_size": 65536,
+                },
+            }


Suggested change

{

"buffer": "zarr.buffer.gpu.Buffer",

"ndbuffer": "zarr.buffer.gpu.NDBuffer",

"codecs": {"zstd": "zarr.codecs.cuda.NvcompZstdCodec"},

"codec_pipeline": {

"path": "zarr.core.codec_pipeline.BatchedCodecPipeline",

"batch_size": 65536,

},

}

{

"buffer": "zarr.buffer.gpu.Buffer",

"ndbuffer": "zarr.buffer.gpu.NDBuffer",

"codecs.zstd": "zarr.codecs.cuda.NvcompZstdCodec",

"codec_pipeline.path": "zarr.core.codec_pipeline.BatchedCodecPipeline",

"codec_pipeline.batch_size": 65536,

}

We need to fix the config to be less fragile, but using a nested dict would replace the other codec mappings (e.g., codec_pipeline.max_workers). Please use dotted keys

maxrjones · 2026-11-13T05:23:00Z

+        level_parsed = _parse_zstd_level(level)
+        checksum_parsed = _parse_checksum(checksum)
+
+        object.__setattr__(self, "level", level_parsed)
+        object.__setattr__(self, "checksum", checksum_parsed)


should these be used somewhere?

maxrjones · 2026-11-13T05:24:27Z

+
+@dataclass(frozen=True)
+class NvcompZstdCodec(BytesBytesCodec):
+    is_fixed_size = True


will the difference between this default and the regular Zstd default of False cause an issue?

maxrjones · 2026-11-13T05:27:10Z

+try:
+    import cupy as cp
+except ImportError:  # pragma: no cover
+    cp = None
+
+try:
+    from nvidia import nvcomp
+except ImportError:  # pragma: no cover
+    nvcomp = None


should these be lazy imports? I'm wondering about whether nvcomp / cupy should be imported when enable_gpu (or preferably enable_cuda) is false

github-actions Bot added the needs release notes Automatically applied to PRs which haven't added release notes label Feb 25, 2025

TomAugspurger reviewed Feb 25, 2025

View reviewed changes

Comment thread src/zarr/codecs/gpu.py Resolved

akshaysubr added 2 commits February 25, 2025 21:30

First working version of Zstd codec on the GPU

49d5ee8

Adding nvcomp to the GPU dependency list

d548adc

akshaysubr force-pushed the gpu-codecs branch from d6608e7 to d548adc Compare February 25, 2025 21:30

Updating codec pipeline batch size for GPU codecs to enable parallelism

a8c0db3

madsbk mentioned this pull request Feb 28, 2025

Multi-GPU setup and python/zarr api rapidsai/kvikio#651

Closed

This was referenced Mar 3, 2025

Support zarr-python 3.x rapidsai/kvikio#646

Merged

Tracking upstream changes pangeo-data/ncar-hackathon-xarray-on-gpus#28

Open

weiji14 mentioned this pull request Mar 11, 2025

Kvikio backend entrypoint with Zarr v3 xarray-contrib/cupy-xarray#70

Draft

3 of 6 tasks

TomAugspurger reviewed Mar 18, 2025

View reviewed changes

Comment thread src/zarr/codecs/cuda.py Resolved

vyasr mentioned this pull request Apr 7, 2025

Evaluate proper home for Zarr bindings in kvikio rapidsai/kvikio#686

Open

weiji14 mentioned this pull request Apr 7, 2025

Request for packaging Python bindings conda-forge/nvcomp-feedstock#21

Closed

TomAugspurger reviewed Jun 9, 2025

View reviewed changes

Comment thread src/zarr/codecs/cuda.py Resolved

Comment thread src/zarr/codecs/gpu.py Resolved

Comment thread src/zarr/codecs/cuda.py Resolved

akshaysubr added 2 commits July 3, 2025 11:59

Making encode and decode async

69aa274

Removing custom awaitable in favor of event synchronize in an async t…

10e1bc9

…hread

TomAugspurger reviewed Jul 7, 2025

View reviewed changes

TomAugspurger added 2 commits July 9, 2025 12:28

Merge remote-tracking branch 'upstream/main' into gpu-codecs

771c0c1

Sync convert methods

ec07100

TomAugspurger added 2 commits July 9, 2025 12:36

test coverage

d1c37a3

loosen dtype restriction

69ea74e

TomAugspurger reviewed Jul 9, 2025

View reviewed changes

TomAugspurger marked this pull request as ready for review July 9, 2025 19:38

fixed Buffer.__add__

1b85fdc

TomAugspurger added 2 commits July 11, 2025 12:30

Added whatsnew

f5c7814

Merge remote-tracking branch 'upstream/main' into gpu-codecs

7671274

TomAugspurger and others added 17 commits August 25, 2026 15:58

Some buffer coverage

576ba5f

coverage

93d6010

update error message

314bc54

private

56b9e00

test fixup

ef2e2f7

doc fix

166c011

lint

fbefa5e

resolve ambiguous link

21dea73

cupy 14 compat

1ad457d

Apply suggestion from @TomAugspurger

58cee0e

Merge akshay/gpu-codecs (PR zarr-developers#2863 tip) into update branch

1acd153

# Conflicts: # docs/user-guide/gpu.md # pyproject.toml # src/zarr/codecs/__init__.py # src/zarr/codecs/gpu.py # src/zarr/core/buffer/gpu.py # tests/test_codecs/test_codecs.py # tests/test_codecs/test_nvcomp.py

Merge remote-tracking branch 'origin/main' into gpu-codecs-update-main

54a8430

# Conflicts: # src/zarr/core/buffer/gpu.py # tests/test_api.py

exclude GPU files from coverage because they require GPU unavailable …

4d4463f

…in CI

lint: drop unused ChunkCoords import, replace assert with TypeError i…

f2722ab

…n gpu.NDBuffer (S101)

docs: remove duplicated Codecs section from config.md merge; align ND…

8287247

…Buffer object-dtype error with main

negin513 and others added 7 commits October 14, 2026 14:46

Revert "exclude GPU files from coverage because they require GPU unav…

fba64c2

…ailable in CI" This reverts commit 4d4463f.

Drop explanatory comment on is_fixed_size

b6dc00d

Set NvcompZstdCodec.is_fixed_size = True

f87dce1

Restore cached _zstd_codec; drop checksum_policy mapping

05d97a1

nvCOMP only allows NO_COMPUTE_NO_VERIFY with a RAW bitstream, so mapping checksum to ChecksumPolicy made checksum=True fail on both encode and decode. The CUDA event is still recorded before synchronizing.

Merge pull request #4 from negin513/gpu-codecs-review-fixes

5f9fe39

Sync gpu-codecs with main, rename the gpu extra to cuda, fix nvCOMP Zstd sync and options

TomAugspurger approved these changes Oct 31, 2026

View reviewed changes

Merge branch 'main' into gpu-codecs

92dbfb8

maxrjones reviewed Nov 13, 2026

View reviewed changes

		instead of `gpu`, since it installs CUDA-specific packages (CuPy and nvCOMP). The `gpu` extra
		has been removed.

Conversation

akshaysubr commented Feb 25, 2025

Uh oh!

Uh oh!

dstansby commented Feb 28, 2025

Uh oh!

akshaysubr commented Feb 28, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

TomAugspurger left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

TomAugspurger commented Jul 9, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

negin513 commented Oct 14, 2026

Uh oh!

negin513 commented Oct 27, 2026

Uh oh!

TomAugspurger commented Oct 28, 2026

Uh oh!

TomAugspurger left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

maxrjones left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

codecov Bot commented Jul 9, 2025 •

edited

Loading