Repository: rapidsai/kvikio (language: python)

Issue:
[FEATURE] Support for nvCOMP batch API
This feature adds support for [nvCOMP batch/low-level](https://github.com/NVIDIA/nvcomp/blob/main/doc/lowlevel_c_quickstart.md) API which allows to process multiple chunks in parallel.

The proposed implementation provides an easy way to use the API via well-known [numcodecs Codec](https://numcodecs.readthedocs.io/en/stable/abc.html) API. Using `numcodecs` also enables seamless integration with libraries such as `zarr` that use `numcodecs` internally.

Additionally, using nvCOMP batch API enables interoperability between existing codecs and nvCOMP batch codec. For example, the data can be compressed on CPU using default `LZ4` codec and then decompressed on GPU using proposed nvCOMP batch codec.

To support batch mode, Codec interface was extended with functions, `encode_batch` and `decode_batch`, which implement batch mode.

Note that the current version of `zarr` does not support chunk-parallel functionality, but there is a [proposal](https://github.com/zarr-developers/zarr-python/issues/1398) for this feature.

Currently the following compression/decompression algorithms are supported:
* LZ4
* Gdeflate
* zstd
* Snappy

nvCOMP also supports other algorithms which can be relatively easily added to kvikio.

Example of usage:

* Simple use of Codec batch API:
```python
import numcodecs
import numpy as np

# Get the codec from the numcodecs registry.
codec = numcodecs.registry.get_codec(dict(id="nvcomp_batch", algorithm="lz4"))

# Creater 2 chunks. The chunks do not have to be the same size.
shape = (4, 8)
chunk1, chunk2 = np.random.randn(2, *shape).astype(np.float32)

# Compress data.
data_comp = codec.encode_batch([chunk1, chunk2])

# Decompress.
data_decomp = codec.decode_batch(data_comp)

# Verify.
np.testing.assert_equal(data_decomp[0].view(np.float32).reshape(shape), chunk1)
np.testing.assert_equal(data_decomp[1].view(np.float32).reshape(shape), chunk2)
```

* Using with `zarr` (no parallel chunking yet - see the note above).
```python
import numcodecs
import numpy as np
import zarr

# Get the codec from the numcodecs registry.
codec = numcodecs.registry.get_codec(dict(id="nvcomp_batch", algorithm="lz4"))
shape = (16, 16)
chunks = (8, 8)

# Create data and compress.
data = np.random.randn(*shape).astype(np.float32)
z1 = zarr.array(data, chunks=chunks, compressor=codec)

# Store in compressed format.
zarr_store = zarr.MemoryStore()
zarr.save_array(zarr_store, z1, compressor=codec)

# Read back/decompress.
z2 = zarr.open_array(zarr_store)

np.testing.assert_equal(z1[:], z2[:])
```

If desired, the API can also be used directly, without having to use numcodecs API.



Discussion / hints:

Fix patch:
diff --git a/notebooks/nvcomp_batch_codec.ipynb b/notebooks/nvcomp_batch_codec.ipynb
new file mode 100644
index 0000000000..f4f4689f3c
--- /dev/null
+++ b/notebooks/nvcomp_batch_codec.ipynb
@@ -0,0 +1,357 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b543ae63",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "\n",
+    "import numcodecs\n",
+    "\n",
+    "import numpy as np\n",
+    "\n",
+    "import zarr\n",
+    "\n",
+    "from IPython.display import display\n",
+    "\n",
+    "np.set_printoptions(precision=4, suppress=True)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "1a8e55d5",
+   "metadata": {},
+   "source": [
+    "### Basic usage\n",
+    "\n",
+    "Get nvCOMP codec from numcodecs registry:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "75524650",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "NvCompBatchCodec(algorithm='lz4', options={})"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "NVCOMP_CODEC_ID = \"nvcomp_batch\"\n",
+    "\n",
+    "# Currently supported algorithms.\n",
+    "LZ4_ALGO = \"LZ4\"\n",
+    "GDEFLATE_ALGO = \"Gdeflate\"\n",
+    "SNAPPY_ALGO = \"snappy\"\n",
+    "ZSTD_ALGO = \"zstd\"\n",
+    "\n",
+    "codec = numcodecs.registry.get_codec(dict(id=NVCOMP_CODEC_ID, algorithm=LZ4_ALGO))\n",
+    "# To pass algorithm-specific options, use options parameter:\n",
+    "# codec = numcodecs.registry.get_codec(dict(id=NVCOMP_CODEC_ID, algo=LZ4_ALGO, options={\"data_type\": 1}))\n",
+    "\n",
+    "display(codec)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "46641ccb",
+   "metadata": {},
+   "source": [
+    "Create data:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "12a4fffd",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([[ 1.6243, -0.6118, -0.5282, ...,  0.0436, -0.62  ,  0.698 ],\n",
+       "       [-0.4471,  1.2245,  0.4035, ...,  0.4203,  0.811 ,  1.0444],\n",
+       "       [-0.4009,  0.824 , -0.5623, ...,  0.7848, -0.9554,  0.5859],\n",
+       "       ...,\n",
+       "       [ 1.3797,  0.1387,  1.2255, ...,  1.8051,  0.3722,  0.1253],\n",
+       "       [ 0.7348, -0.7115, -0.1248, ..., -1.9533, -0.7684, -0.5345],\n",
+       "       [ 0.2183, -0.8654,  0.8886, ..., -1.0141, -0.0627, -1.4379]],\n",
+       "      dtype=float32)"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<table class=\"zarr-info\"><tbody><tr><th style=\"text-align: left\">Type</th><td style=\"text-align: left\">zarr.core.Array</td></tr><tr><th style=\"text-align: left\">Data type</th><td style=\"text-align: left\">float32</td></tr><tr><th style=\"text-align: left\">Shape</th><td style=\"text-align: left\">(100, 100)</td></tr><tr><th style=\"text-align: left\">Chunk shape</th><td style=\"text-align: left\">(10, 10)</td></tr><tr><th style=\"text-align: left\">Order</th><td style=\"text-align: left\">C</td></tr><tr><th style=\"text-align: left\">Read-only</th><td style=\"text-align: left\">False</td></tr><tr><th style=\"text-align: left\">Compressor</th><td style=\"text-align: left\">NvCompBatchCodec(algorithm='lz4', options={})</td></tr><tr><th style=\"text-align: left\">Store type</th><td style=\"text-align: left\">zarr.storage.KVStore</td></tr><tr><th style=\"text-align: left\">No. bytes</th><td style=\"text-align: left\">40000 (39.1K)</td></tr><tr><th style=\"text-align: left\">No. bytes stored</th><td style=\"text-align: left\">41006 (40.0K)</td></tr><tr><th style=\"text-align: left\">Storage ratio</th><td style=\"text-align: left\">1.0</td></tr><tr><th style=\"text-align: left\">Chunks initialized</th><td style=\"text-align: left\">100/100</td></tr></tbody></table>"
+      ],
+      "text/plain": [
+       "Type               : zarr.core.Array\n",
+       "Data type          : float32\n",
+       "Shape              : (100, 100)\n",
+       "Chunk shape        : (10, 10)\n",
+       "Order              : C\n",
+       "Read-only          : False\n",
+       "Compressor         : NvCompBatchCodec(algorithm='lz4', options={})\n",
+       "Store type         : zarr.storage.KVStore\n",
+       "No. bytes          : 40000 (39.1K)\n",
+       "No. bytes stored   : 41006 (40.0K)\n",
+       "Storage ratio      : 1.0\n",
+       "Chunks initialized : 100/100"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "shape = (100, 100)\n",
+    "chunks = (10, 10)\n",
+    "\n",
+    "np.random.seed(1)\n",
+    "\n",
+    "x = zarr.array(np.random.randn(*shape).astype(np.float32), chunks=chunks, compressor=codec)\n",
+    "display(x[:])\n",
+    "display(x.info)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "c15cbdff",
+   "metadata": {},
+   "source": [
+    "Store and load back the data:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "730cde85",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'chunks': [10, 10],\n",
+       " 'compressor': {'algorithm': 'lz4', 'id': 'nvcomp_batch', 'options': {}},\n",
+       " 'dtype': '<f4',\n",
+       " 'fill_value': 0.0,\n",
+       " 'filters': None,\n",
+       " 'order': 'C',\n",
+       " 'shape': [100, 100],\n",
+       " 'zarr_format': 2}"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# Use simple dictionary store, see zarr documentation for other options.\n",
+    "zarr_store = {}\n",
+    "\n",
+    "zarr.save_array(zarr_store, x, compressor=codec)\n",
+    "\n",
+    "# Check stored metadata.\n",
+    "meta_info = json.loads(zarr_store[\".zarray\"])\n",
+    "display(meta_info)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "11338a99",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<table class=\"zarr-info\"><tbody><tr><th style=\"text-align: left\">Type</th><td style=\"text-align: left\">zarr.core.Array</td></tr><tr><th style=\"text-align: left\">Data type</th><td style=\"text-align: left\">float32</td></tr><tr><th style=\"text-align: left\">Shape</th><td style=\"text-align: left\">(100, 100)</td></tr><tr><th style=\"text-align: left\">Chunk shape</th><td style=\"text-align: left\">(10, 10)</td></tr><tr><th style=\"text-align: left\">Order</th><td style=\"text-align: left\">C</td></tr><tr><th style=\"text-align: left\">Read-only</th><td style=\"text-align: left\">False</td></tr><tr><th style=\"text-align: left\">Compressor</th><td style=\"text-align: left\">NvCompBatchCodec(algorithm='lz4', options={})</td></tr><tr><th style=\"text-align: left\">Store type</th><td style=\"text-align: left\">zarr.storage.KVStore</td></tr><tr><th style=\"text-align: left\">No. bytes</th><td style=\"text-align: left\">40000 (39.1K)</td></tr><tr><th style=\"text-align: left\">No. bytes stored</th><td style=\"text-align: left\">41006 (40.0K)</td></tr><tr><th style=\"text-align: left\">Storage ratio</th><td style=\"text-align: left\">1.0</td></tr><tr><th style=\"text-align: left\">Chunks initialized</th><td style=\"text-align: left\">100/100</td></tr></tbody></table>"
+      ],
+      "text/plain": [
+       "Type               : zarr.core.Array\n",
+       "Data type          : float32\n",
+       "Shape              : (100, 100)\n",
+       "Chunk shape        : (10, 10)\n",
+       "Order              : C\n",
+       "Read-only          : False\n",
+       "Compressor         : NvCompBatchCodec(algorithm='lz4', options={})\n",
+       "Store type         : zarr.storage.KVStore\n",
+       "No. bytes          : 40000 (39.1K)\n",
+       "No. bytes stored   : 41006 (40.0K)\n",
+       "Storage ratio      : 1.0\n",
+       "Chunks initialized : 100/100"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "y = zarr.open_array(zarr_store)\n",
+    "display(y.info)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "5b6cc2ca",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Test the roundtrip.\n",
+    "np.testing.assert_equal(y[:], x[:])"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "1a8eea79",
+   "metadata": {},
+   "source": [
+    "### CPU compression / GPU decompression\n",
+    "\n",
+    "Some algorithms, such as LZ4, can be used interchangeably on CPU and GPU. For example, the data might be created using CPU LZ4 codec and then decompressed using GPU version of LZ4 codec."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "87d25b76",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'chunks': [10, 10],\n",
+       " 'compressor': {'acceleration': 1, 'id': 'lz4'},\n",
+       " 'dtype': '<f4',\n",
+       " 'fill_value': 0.0,\n",
+       " 'filters': None,\n",
+       " 'order': 'C',\n",
+       " 'shape': [100, 100],\n",
+       " 'zarr_format': 2}"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<table class=\"zarr-info\"><tbody><tr><th style=\"text-align: left\">Type</th><td style=\"text-align: left\">zarr.core.Array</td></tr><tr><th style=\"text-align: left\">Data type</th><td style=\"text-align: left\">float32</td></tr><tr><th style=\"text-align: left\">Shape</th><td style=\"text-align: left\">(100, 100)</td></tr><tr><th style=\"text-align: left\">Chunk shape</th><td style=\"text-align: left\">(10, 10)</td></tr><tr><th style=\"text-align: left\">Order</th><td style=\"text-align: left\">C</td></tr><tr><th style=\"text-align: left\">Read-only</th><td style=\"text-align: left\">False</td></tr><tr><th style=\"text-align: left\">Compressor</th><td style=\"text-align: left\">LZ4(acceleration=1)</td></tr><tr><th style=\"text-align: left\">Store type</th><td style=\"text-align: left\">zarr.storage.KVStore</td></tr><tr><th style=\"text-align: left\">No. bytes</th><td style=\"text-align: left\">40000 (39.1K)</td></tr><tr><th style=\"text-align: left\">No. bytes stored</th><td style=\"text-align: left\">40973 (40.0K)</td></tr><tr><th style=\"text-align: left\">Storage ratio</th><td style=\"text-align: left\">1.0</td></tr><tr><th style=\"text-align: left\">Chunks initialized</th><td style=\"text-align: left\">100/100</td></tr></tbody></table>"
+      ],
+      "text/plain": [
+       "Type               : zarr.core.Array\n",
+       "Data type          : float32\n",
+       "Shape              : (100, 100)\n",
+       "Chunk shape        : (10, 10)\n",
+       "Order              : C\n",
+       "Read-only          : False\n",
+       "Compressor         : LZ4(acceleration=1)\n",
+       "Store type         : zarr.storage.KVStore\n",
+       "No. bytes          : 40000 (39.1K)\n",
+       "No. bytes stored   : 40973 (40.0K)\n",
+       "Storage ratio      : 1.0\n",
+       "Chunks initialized : 100/100"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<table class=\"zarr-info\"><tbody><tr><th style=\"text-align: left\">Type</th><td style=\"text-align: left\">zarr.core.Array</td></tr><tr><th style=\"text-align: left\">Data type</th><td style=\"text-align: left\">float32</td></tr><tr><th style=\"text-align: left\">Shape</th><td style=\"text-align: left\">(100, 100)</td></tr><tr><th style=\"text-align: left\">Chunk shape</th><td style=\"text-align: left\">(10, 10)</td></tr><tr><th style=\"text-align: left\">Order</th><td style=\"text-align: left\">C</td></tr><tr><th style=\"text-align: left\">Read-only</th><td style=\"text-align: left\">False</td></tr><tr><th style=\"text-align: left\">Compressor</th><td style=\"text-align: left\">NvCompBatchCodec(algorithm='lz4', options={})</td></tr><tr><th style=\"text-align: left\">Store type</th><td style=\"text-align: left\">zarr.storage.KVStore</td></tr><tr><th style=\"text-align: left\">No. bytes</th><td style=\"text-align: left\">40000 (39.1K)</td></tr><tr><th style=\"text-align: left\">No. bytes stored</th><td style=\"text-align: left\">40883 (39.9K)</td></tr><tr><th style=\"text-align: left\">Storage ratio</th><td style=\"text-align: left\">1.0</td></tr><tr><th style=\"text-align: left\">Chunks initialized</th><td style=\"text-align: left\">100/100</td></tr></tbody></table>"
+      ],
+      "text/plain": [
+       "Type               : zarr.core.Array\n",
+       "Data type          : float32\n",
+       "Shape              : (100, 100)\n",
+       "Chunk shape        : (10, 10)\n",
+       "Order              : C\n",
+       "Read-only          : False\n",
+       "Compressor         : NvCompBatchCodec(algorithm='lz4', options={})\n",
+       "Store type         : zarr.storage.KVStore\n",
+       "No. bytes          : 40000 (39.1K)\n",
+       "No. bytes stored   : 40883 (39.9K)\n",
+       "Storage ratio      : 1.0\n",
+       "Chunks initialized : 100/100"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# Get default (CPU) implementation of LZ4 codec.\n",
+    "cpu_codec = numcodecs.registry.get_codec({\"id\": \"lz4\"})\n",
+    "\n",
+    "x = zarr.array(np.random.randn(*shape).astype(np.float32), chunks=chunks, compressor=cpu_codec)\n",
+    "# Define a simple, dictionary-based store. In real scenarios this can be a filesystem or some other persistent store.\n",
+    "store = {}\n",
+    "zarr.save_array(store, x, compressor=cpu_codec)\n",
+    "\n",
+    "# Check that the data was written by the expected codec.\n",
+    "meta = json.loads(store[\".zarray\"])\n",
+    "display(meta)\n",
+    "assert meta[\"compressor\"][\"id\"] == \"lz4\"\n",
+    "\n",
+    "# Change codec to GPU/nvCOMP-based.\n",
+    "meta[\"compressor\"] = {\"id\": NVCOMP_CODEC_ID, \"algorithm\": LZ4_ALGO}\n",
+    "store[\".zarray\"] = json.dumps(meta).encode()\n",
+    "\n",
+    "y = zarr.open_array(store, compressor=codec)\n",
+    "\n",
+    "display(x.info)\n",
+    "display(y.info)\n",
+    "\n",
+    "np.testing.assert_equal(x[:], y[:])\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b9294992",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/python/kvikio/_lib/CMakeLists.txt b/python/kvikio/_lib/CMakeLists.txt
index d80e4a8991..f88c0126b0 100644
--- a/python/kvikio/_lib/CMakeLists.txt
+++ b/python/kvikio/_lib/CMakeLists.txt
@@ -13,9 +13,10 @@
 # =============================================================================
 
 # Set the list of Cython files to build
-set(cython_modules arr libnvcomp libkvikio)
+set(cython_modules arr libnvcomp libnvcomp_ll libkvikio)
 
 # Build all of the Cython targets
 add_cython_modules("${cython_modules}")
 
 target_link_libraries(libnvcomp nvcomp::nvcomp)
+target_link_libraries(libnvcomp_ll nvcomp::nvcomp)
diff --git a/python/kvikio/_lib/libnvcomp_ll.pyx b/python/kvikio/_lib/libnvcomp_ll.pyx
new file mode 100644
index 0000000000..1794d29819
--- /dev/null
+++ b/python/kvikio/_lib/libnvcomp_ll.pyx
@@ -0,0 +1,771 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+from abc import ABC, abstractmethod
+from enum import IntEnum
+
+from libc.stdint cimport uintptr_t
+
+from kvikio._lib.nvcomp_ll_cxx_api cimport cudaStream_t, nvcompStatus_t, nvcompType_t
+
+
+class nvCompStatus(IntEnum):
+    Success = nvcompStatus_t.nvcompSuccess,
+    ErrorInvalidValue = nvcompStatus_t.nvcompErrorInvalidValue,
+    ErrorNotSupported = nvcompStatus_t.nvcompErrorNotSupported,
+    ErrorCannotDecompress = nvcompStatus_t.nvcompErrorCannotDecompress,
+    ErrorBadChecksum = nvcompStatus_t.nvcompErrorBadChecksum,
+    ErrorCannotVerifyChecksums = nvcompStatus_t.nvcompErrorCannotVerifyChecksums,
+    ErrorCudaError = nvcompStatus_t.nvcompErrorCudaError,
+    ErrorInternal = nvcompStatus_t.nvcompErrorInternal,
+
+
+class nvCompType(IntEnum):
+    CHAR = nvcompType_t.NVCOMP_TYPE_CHAR
+    UCHAR = nvcompType_t.NVCOMP_TYPE_UCHAR
+    SHORT = nvcompType_t.NVCOMP_TYPE_SHORT
+    USHORT = nvcompType_t.NVCOMP_TYPE_USHORT
+    INT = nvcompType_t.NVCOMP_TYPE_INT
+    UINT = nvcompType_t.NVCOMP_TYPE_UINT
+    LONGLONG = nvcompType_t.NVCOMP_TYPE_LONGLONG
+    ULONGLONG = nvcompType_t.NVCOMP_TYPE_ULONGLONG
+    BITS = nvcompType_t.NVCOMP_TYPE_BITS
+
+
+class nvCompBatchAlgorithm(ABC):
+    """Abstract class that provides interface to nvCOMP batched algorithms."""
+
+    # TODO(akamenev): it might be possible to have a simpler implementation that
+    # eilminates the need to have a separate implementation class for each algorithm,
+    # potentially using fused types in Cython (similar to C++ templates),
+    # but I could not figure out how to do that (e.g. each algorithm API set has
+    # a different type for the options and so on).
+
+    def get_compress_temp_size(
+        self,
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+    ):
+        """Get temporary space required for compression.
+
+        Parameters
+        ----------
+        batch_size: int
+            The number of items in the batch.
+        max_uncompressed_chunk_bytes: int
+            The maximum size in bytes of a chunk in the batch.
+
+        Returns
+        -------
+        int
+            The size in bytes of the required GPU workspace for compression.
+        """
+        err, temp_size = self._get_comp_temp_size(
+            batch_size,
+            max_uncompressed_chunk_bytes
+        )
+        if err != nvcompStatus_t.nvcompSuccess:
+            raise RuntimeError(
+                f"Could not get compress temp buffer size, "
+                f"error: {nvCompStatus(err)!r}."
+            )
+        return temp_size
+
+    @abstractmethod
+    def _get_comp_temp_size(
+        self,
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+    ) -> (nvcompStatus_t, size_t):
+        """Algorithm-specific implementation."""
+        ...
+
+    def get_compress_chunk_size(self, size_t max_uncompressed_chunk_bytes):
+        """Get the maximum size any chunk could compress to in the batch.
+
+        Parameters
+        ----------
+        max_uncompressed_chunk_bytes: int
+            The maximum size in bytes of a chunk in the batch.
+
+        Returns
+        -------
+        int
+            The maximum compressed size in bytes of the largest chunk. That is,
+            the minimum amount of output memory required to be given to
+            the corresponding *CompressAsync function.
+        """
+        err, comp_chunk_size = self._get_comp_chunk_size(max_uncompressed_chunk_bytes)
+        if err != nvcompStatus_t.nvcompSuccess:
+            raise RuntimeError(
+                f"Could not get output buffer size, "
+                f"error: {nvCompStatus(err)!r}."
+            )
+        return comp_chunk_size
+
+    @abstractmethod
+    def _get_comp_chunk_size(self, size_t max_uncompressed_chunk_bytes):
+        """Algorithm-specific implementation."""
+        ...
+
+    def compress(
+        self,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        size_t max_uncomp_chunk_bytes,
+        size_t batch_size,
+        temp_buf,
+        comp_chunks,
+        comp_chunk_sizes,
+        stream,
+    ):
+        """Perform compression.
+
+        Parameters
+        ----------
+        uncomp_chunks: cp.ndarray
+            The pointers on the GPU, to uncompressed batched items.
+        uncomp_chunk_sizes: cp.ndarray
+            The size in bytes of each uncompressed batch item on the GPU.
+        max_uncomp_chunk_bytes: int
+            The maximum size in bytes of the largest chunk in the batch.
+        batch_size: int
+            The number of chunks to compress.
+        temp_buf: cp.ndarray
+            The temporary GPU workspace.
+        comp_chunks: cp.ndarray
+            (output) The pointers on the GPU, to the output location for each
+            compressed batch item.
+        comp_chunk_sizes: cp.ndarray
+            (output) The compressed size in bytes of each chunk on the GPU.
+        stream: cp.cuda.Stream
+            CUDA stream.
+        """
+        err = self._compress(
+            uncomp_chunks,
+            uncomp_chunk_sizes,
+            max_uncomp_chunk_bytes,
+            batch_size,
+            temp_buf,
+            comp_chunks,
+            comp_chunk_sizes,
+            stream,
+        )
+        if err != nvcompStatus_t.nvcompSuccess:
+            raise RuntimeError(f"Compression failed, error: {nvCompStatus(err)!r}.")
+
+    @abstractmethod
+    def _compress(
+        self,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        size_t max_uncomp_chunk_bytes,
+        size_t batch_size,
+        temp_buf,
+        comp_chunks,
+        comp_chunk_sizes,
+        stream
+    ):
+        """Algorithm-specific implementation."""
+        ...
+
+    def get_decompress_temp_size(
+        self,
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+    ):
+        """Get the amount of temp space required on the GPU for decompression.
+
+        Parameters
+        ----------
+        batch_size: int
+            The number of items in the batch.
+        max_uncompressed_chunk_bytes: int
+            The size in bytes of the largest chunk when uncompressed.
+
+        Returns
+        -------
+        int
+            The amount of temporary GPU space in bytes that will be
+            required to decompress.
+        """
+        err, temp_size = self._get_decomp_temp_size(
+            batch_size,
+            max_uncompressed_chunk_bytes
+        )
+        if err != nvcompStatus_t.nvcompSuccess:
+            raise RuntimeError(
+                f"Could not get decompress temp buffer size, "
+                f"error: {nvCompStatus(err)!r}."
+            )
+
+        return temp_size
+
+    @abstractmethod
+    def _get_decomp_temp_size(
+        self,
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+    ):
+        """Algorithm-specific implementation."""
+        ...
+
+    def decompress(
+        self,
+        comp_chunks,
+        comp_chunk_sizes,
+        size_t batch_size,
+        temp_buf,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        actual_uncomp_chunk_sizes,
+        statuses,
+        stream,
+    ):
+        """Perform decompression.
+
+        Parameters
+        ----------
+        comp_chunks: cp.ndarray
+            The pointers on the GPU, to compressed batched items.
+        comp_chunk_sizes: cp.ndarray
+            The size in bytes of each compressed batch item on the GPU.
+        batch_size: int
+            The number of chunks to decompress.
+        temp_buf: cp.ndarray
+            The temporary GPU workspace.
+        uncomp_chunks: cp.ndarray
+            (output) The pointers on the GPU, to the output location for each
+            decompressed batch item.
+        uncomp_chunk_sizes: cp.ndarray
+            The size in bytes of each decompress chunk location on the GPU.
+        actual_uncomp_chunk_sizes: cp.ndarray
+            (output) The actual decompressed size in bytes of each chunk on the GPU.
+        statuses: cp.ndarray
+            (output) The status for each chunk of whether it was decompressed or not.
+        stream: cp.cuda.Stream
+            CUDA stream.
+        """
+        err = self._decompress(
+            comp_chunks,
+            comp_chunk_sizes,
+            batch_size,
+            temp_buf,
+            uncomp_chunks,
+            uncomp_chunk_sizes,
+            actual_uncomp_chunk_sizes,
+            statuses,
+            stream,
+        )
+        if err != nvcompStatus_t.nvcompSuccess:
+            raise RuntimeError(f"Decompression failed, error: {nvCompStatus(err)!r}.")
+
+    @abstractmethod
+    def _decompress(
+        self,
+        comp_chunks,
+        comp_chunk_sizes,
+        size_t batch_size,
+        temp_buf,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        actual_uncomp_chunk_sizes,
+        statuses,
+        stream,
+    ):
+        """Algorithm-specific implementation."""
+        ...
+
+
+cdef uintptr_t to_ptr(buf):
+    return buf.data.ptr
+
+
+cdef cudaStream_t to_stream(stream):
+    return <cudaStream_t><size_t>stream.ptr
+
+
+#
+# LZ4 algorithm.
+#
+from kvikio._lib.nvcomp_ll_cxx_api cimport (
+    nvcompBatchedLZ4CompressAsync,
+    nvcompBatchedLZ4CompressGetMaxOutputChunkSize,
+    nvcompBatchedLZ4CompressGetTempSize,
+    nvcompBatchedLZ4DecompressAsync,
+    nvcompBatchedLZ4DecompressGetTempSize,
+    nvcompBatchedLZ4Opts_t,
+)
+
+
+class nvCompBatchAlgorithmLZ4(nvCompBatchAlgorithm):
+    """LZ4 algorithm implementation."""
+
+    algo_id: str = "lz4"
+
+    options: nvcompBatchedLZ4Opts_t
+
+    def __init__(self, data_type: int = 0):
+        self.options = nvcompBatchedLZ4Opts_t(data_type)
+
+    def _get_comp_temp_size(
+        self,
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+    ) -> (nvcompStatus_t, size_t):
+        cdef size_t temp_bytes = 0
+
+        err = nvcompBatchedLZ4CompressGetTempSize(
+            batch_size,
+            max_uncompressed_chunk_bytes,
+            self.options,
+            &temp_bytes
+        )
+
+        return (err, temp_bytes)
+
+    def _get_comp_chunk_size(self, size_t max_uncompressed_chunk_bytes):
+        cdef size_t max_compressed_bytes = 0
+
+        err = nvcompBatchedLZ4CompressGetMaxOutputChunkSize(
+            max_uncompressed_chunk_bytes,
+            self.options,
+            &max_compressed_bytes
+        )
+
+        return (err, max_compressed_bytes)
+
+    def _compress(
+        self,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        size_t max_uncomp_chunk_bytes,
+        size_t batch_size,
+        temp_buf,
+        comp_chunks,
+        comp_chunk_sizes,
+        stream
+    ):
+        # Cast buffer pointers that have Python int type to appropriate C types
+        # suitable for passing to nvCOMP API.
+        return nvcompBatchedLZ4CompressAsync(
+            <const void* const*>to_ptr(uncomp_chunks),
+            <const size_t*>to_ptr(uncomp_chunk_sizes),
+            max_uncomp_chunk_bytes,
+            batch_size,
+            <void*>to_ptr(temp_buf),
+            <size_t>temp_buf.nbytes,
+            <void* const*>to_ptr(comp_chunks),
+            <size_t*>to_ptr(comp_chunk_sizes),
+            self.options,
+            to_stream(stream),
+        )
+
+    def _get_decomp_temp_size(
+        self,
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+    ):
+        cdef size_t temp_bytes = 0
+
+        err = nvcompBatchedLZ4DecompressGetTempSize(
+            batch_size,
+            max_uncompressed_chunk_bytes,
+            &temp_bytes
+        )
+
+        return (err, temp_bytes)
+
+    def _decompress(
+        self,
+        comp_chunks,
+        comp_chunk_sizes,
+        size_t batch_size,
+        temp_buf,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        actual_uncomp_chunk_sizes,
+        statuses,
+        stream,
+    ):
+        # Cast buffer pointers that have Python int type to appropriate C types
+        # suitable for passing to nvCOMP API.
+        return nvcompBatchedLZ4DecompressAsync(
+            <const void* const*>to_ptr(comp_chunks),
+            <const size_t*>to_ptr(comp_chunk_sizes),
+            <const size_t*>to_ptr(uncomp_chunk_sizes),
+            <size_t*>NULL,
+            batch_size,
+            <void* const>to_ptr(temp_buf),
+            <size_t>temp_buf.nbytes,
+            <void* const*>to_ptr(uncomp_chunks),
+            <nvcompStatus_t*>NULL,
+            to_stream(stream),
+        )
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(data_type={self.options['data_type']})"
+
+
+#
+# Gdeflate algorithm.
+#
+from kvikio._lib.nvcomp_ll_cxx_api cimport (
+    nvcompBatchedGdeflateCompressAsync,
+    nvcompBatchedGdeflateCompressGetMaxOutputChunkSize,
+    nvcompBatchedGdeflateCompressGetTempSize,
+    nvcompBatchedGdeflateDecompressAsync,
+    nvcompBatchedGdeflateDecompressGetTempSize,
+    nvcompBatchedGdeflateOpts_t,
+)
+
+
+class nvCompBatchAlgorithmGdeflate(nvCompBatchAlgorithm):
+    """Gdeflate algorithm implementation."""
+
+    algo_id: str = "gdeflate"
+
+    options: nvcompBatchedGdeflateOpts_t
+
+    def __init__(self, algo: int = 0):
+        self.options = nvcompBatchedGdeflateOpts_t(algo)
+
+    def _get_comp_temp_size(
+        self,
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+    ) -> (nvcompStatus_t, size_t):
+        cdef size_t temp_bytes = 0
+
+        err = nvcompBatchedGdeflateCompressGetTempSize(
+            batch_size,
+            max_uncompressed_chunk_bytes,
+            self.options,
+            &temp_bytes
+        )
+
+        return (err, temp_bytes)
+
+    def _get_comp_chunk_size(self, size_t max_uncompressed_chunk_bytes):
+        cdef size_t max_compressed_bytes = 0
+
+        err = nvcompBatchedGdeflateCompressGetMaxOutputChunkSize(
+            max_uncompressed_chunk_bytes,
+            self.options,
+            &max_compressed_bytes
+        )
+
+        return (err, max_compressed_bytes)
+
+    def _compress(
+        self,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        size_t max_uncomp_chunk_bytes,
+        size_t batch_size,
+        temp_buf,
+        comp_chunks,
+        comp_chunk_sizes,
+        stream
+    ):
+        return nvcompBatchedGdeflateCompressAsync(
+            <const void* const*>to_ptr(uncomp_chunks),
+            <const size_t*>to_ptr(uncomp_chunk_sizes),
+            max_uncomp_chunk_bytes,
+            batch_size,
+            <void*>to_ptr(temp_buf),
+            <size_t>temp_buf.nbytes,
+            <void* const*>to_ptr(comp_chunks),
+            <size_t*>to_ptr(comp_chunk_sizes),
+            self.options,
+            to_stream(stream),
+        )
+
+    def _get_decomp_temp_size(
+        self,
+        size_t num_chunks,
+        size_t max_uncompressed_chunk_bytes,
+    ):
+        cdef size_t temp_bytes = 0
+
+        err = nvcompBatchedGdeflateDecompressGetTempSize(
+            num_chunks,
+            max_uncompressed_chunk_bytes,
+            &temp_bytes
+        )
+
+        return (err, temp_bytes)
+
+    def _decompress(
+        self,
+        comp_chunks,
+        comp_chunk_sizes,
+        size_t batch_size,
+        temp_buf,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        actual_uncomp_chunk_sizes,
+        statuses,
+        stream,
+    ):
+        return nvcompBatchedGdeflateDecompressAsync(
+            <const void* const*>to_ptr(comp_chunks),
+            <const size_t*>to_ptr(comp_chunk_sizes),
+            <const size_t*>to_ptr(uncomp_chunk_sizes),
+            <size_t*>NULL,
+            batch_size,
+            <void* const>to_ptr(temp_buf),
+            <size_t>temp_buf.nbytes,
+            <void* const*>to_ptr(uncomp_chunks),
+            <nvcompStatus_t*>NULL,
+            to_stream(stream),
+        )
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(algo={self.options['algo']})"
+
+
+#
+# zstd algorithm.
+#
+from kvikio._lib.nvcomp_ll_cxx_api cimport (
+    nvcompBatchedZstdCompressAsync,
+    nvcompBatchedZstdCompressGetMaxOutputChunkSize,
+    nvcompBatchedZstdCompressGetTempSize,
+    nvcompBatchedZstdDecompressAsync,
+    nvcompBatchedZstdDecompressGetTempSize,
+    nvcompBatchedZstdOpts_t,
+)
+
+
+class nvCompBatchAlgorithmZstd(nvCompBatchAlgorithm):
+    """zstd algorithm implementation."""
+
+    algo_id: str = "zstd"
+
+    options: nvcompBatchedZstdOpts_t
+
+    def __init__(self):
+        self.options = nvcompBatchedZstdOpts_t(0)
+
+    def _get_comp_temp_size(
+        self,
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+    ) -> (nvcompStatus_t, size_t):
+        cdef size_t temp_bytes = 0
+
+        err = nvcompBatchedZstdCompressGetTempSize(
+            batch_size,
+            max_uncompressed_chunk_bytes,
+            self.options,
+            &temp_bytes
+        )
+
+        return (err, temp_bytes)
+
+    def _get_comp_chunk_size(self, size_t max_uncompressed_chunk_bytes):
+        cdef size_t max_compressed_bytes = 0
+
+        err = nvcompBatchedZstdCompressGetMaxOutputChunkSize(
+            max_uncompressed_chunk_bytes,
+            self.options,
+            &max_compressed_bytes
+        )
+
+        return (err, max_compressed_bytes)
+
+    def _compress(
+        self,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        size_t max_uncomp_chunk_bytes,
+        size_t batch_size,
+        temp_buf,
+        comp_chunks,
+        comp_chunk_sizes,
+        stream
+    ):
+        return nvcompBatchedZstdCompressAsync(
+            <const void* const*>to_ptr(uncomp_chunks),
+            <const size_t*>to_ptr(uncomp_chunk_sizes),
+            max_uncomp_chunk_bytes,
+            batch_size,
+            <void*>to_ptr(temp_buf),
+            <size_t>temp_buf.nbytes,
+            <void* const*>to_ptr(comp_chunks),
+            <size_t*>to_ptr(comp_chunk_sizes),
+            self.options,
+            to_stream(stream),
+        )
+
+    def _get_decomp_temp_size(
+        self,
+        size_t num_chunks,
+        size_t max_uncompressed_chunk_bytes,
+    ):
+        cdef size_t temp_bytes = 0
+
+        err = nvcompBatchedZstdDecompressGetTempSize(
+            num_chunks,
+            max_uncompressed_chunk_bytes,
+            &temp_bytes
+        )
+
+        return (err, temp_bytes)
+
+    def _decompress(
+        self,
+        comp_chunks,
+        comp_chunk_sizes,
+        size_t batch_size,
+        temp_buf,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        actual_uncomp_chunk_sizes,
+        statuses,
+        stream,
+    ):
+        return nvcompBatchedZstdDecompressAsync(
+            <const void* const*>to_ptr(comp_chunks),
+            <const size_t*>to_ptr(comp_chunk_sizes),
+            <const size_t*>to_ptr(uncomp_chunk_sizes),
+            <size_t*>to_ptr(actual_uncomp_chunk_sizes),
+            batch_size,
+            <void* const>to_ptr(temp_buf),
+            <size_t>temp_buf.nbytes,
+            <void* const*>to_ptr(uncomp_chunks),
+            <nvcompStatus_t*>to_ptr(statuses),
+            to_stream(stream),
+        )
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}()"
+
+
+#
+# Snappy algorithm.
+#
+from kvikio._lib.nvcomp_ll_cxx_api cimport (
+    nvcompBatchedSnappyCompressAsync,
+    nvcompBatchedSnappyCompressGetMaxOutputChunkSize,
+    nvcompBatchedSnappyCompressGetTempSize,
+    nvcompBatchedSnappyDecompressAsync,
+    nvcompBatchedSnappyDecompressGetTempSize,
+    nvcompBatchedSnappyOpts_t,
+)
+
+
+class nvCompBatchAlgorithmSnappy(nvCompBatchAlgorithm):
+    """Snappy algorithm implementation."""
+
+    algo_id: str = "snappy"
+
+    options: nvcompBatchedSnappyOpts_t
+
+    def __init__(self):
+        self.options = nvcompBatchedSnappyOpts_t(0)
+
+    def _get_comp_temp_size(
+        self,
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+    ) -> (nvcompStatus_t, size_t):
+        cdef size_t temp_bytes = 0
+
+        err = nvcompBatchedSnappyCompressGetTempSize(
+            batch_size,
+            max_uncompressed_chunk_bytes,
+            self.options,
+            &temp_bytes
+        )
+
+        return (err, temp_bytes)
+
+    def _get_comp_chunk_size(self, size_t max_uncompressed_chunk_bytes):
+        cdef size_t max_compressed_bytes = 0
+
+        err = nvcompBatchedSnappyCompressGetMaxOutputChunkSize(
+            max_uncompressed_chunk_bytes,
+            self.options,
+            &max_compressed_bytes
+        )
+
+        return (err, max_compressed_bytes)
+
+    def _compress(
+        self,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        size_t max_uncomp_chunk_bytes,
+        size_t batch_size,
+        temp_buf,
+        comp_chunks,
+        comp_chunk_sizes,
+        stream
+    ):
+        return nvcompBatchedSnappyCompressAsync(
+            <const void* const*>to_ptr(uncomp_chunks),
+            <const size_t*>to_ptr(uncomp_chunk_sizes),
+            max_uncomp_chunk_bytes,
+            batch_size,
+            <void*>to_ptr(temp_buf),
+            <size_t>temp_buf.nbytes,
+            <void* const*>to_ptr(comp_chunks),
+            <size_t*>to_ptr(comp_chunk_sizes),
+            self.options,
+            to_stream(stream),
+        )
+
+    def _get_decomp_temp_size(
+        self,
+        size_t num_chunks,
+        size_t max_uncompressed_chunk_bytes,
+    ):
+        cdef size_t temp_bytes = 0
+
+        err = nvcompBatchedSnappyDecompressGetTempSize(
+            num_chunks,
+            max_uncompressed_chunk_bytes,
+            &temp_bytes
+        )
+
+        return (err, temp_bytes)
+
+    def _decompress(
+        self,
+        comp_chunks,
+        comp_chunk_sizes,
+        size_t batch_size,
+        temp_buf,
+        uncomp_chunks,
+        uncomp_chunk_sizes,
+        actual_uncomp_chunk_sizes,
+        statuses,
+        stream,
+    ):
+        return nvcompBatchedSnappyDecompressAsync(
+            <const void* const*>to_ptr(comp_chunks),
+            <const size_t*>to_ptr(comp_chunk_sizes),
+            <const size_t*>to_ptr(uncomp_chunk_sizes),
+            <size_t*>NULL,
+            batch_size,
+            <void* const>to_ptr(temp_buf),
+            <size_t>temp_buf.nbytes,
+            <void* const*>to_ptr(uncomp_chunks),
+            <nvcompStatus_t*>NULL,
+            to_stream(stream),
+        )
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}()"
+
+
+SUPPORTED_ALGORITHMS = {
+    a.algo_id: a for a in [
+        nvCompBatchAlgorithmLZ4,
+        nvCompBatchAlgorithmGdeflate,
+        nvCompBatchAlgorithmZstd,
+        nvCompBatchAlgorithmSnappy,
+    ]
+}
diff --git a/python/kvikio/_lib/nvcomp_ll_cxx_api.pxd b/python/kvikio/_lib/nvcomp_ll_cxx_api.pxd
new file mode 100644
index 0000000000..2ceb3037d1
--- /dev/null
+++ b/python/kvikio/_lib/nvcomp_ll_cxx_api.pxd
@@ -0,0 +1,250 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+# distutils: language = c++
+# cython: language_level=3
+
+cdef extern from "cuda_runtime.h":
+    ctypedef void* cudaStream_t
+
+cdef extern from "nvcomp.h":
+    ctypedef enum nvcompType_t:
+        NVCOMP_TYPE_CHAR = 0,       # 1B
+        NVCOMP_TYPE_UCHAR = 1,      # 1B
+        NVCOMP_TYPE_SHORT = 2,      # 2B
+        NVCOMP_TYPE_USHORT = 3,     # 2B
+        NVCOMP_TYPE_INT = 4,        # 4B
+        NVCOMP_TYPE_UINT = 5,       # 4B
+        NVCOMP_TYPE_LONGLONG = 6,   # 8B
+        NVCOMP_TYPE_ULONGLONG = 7,  # 8B
+        NVCOMP_TYPE_BITS = 0xff     # 1b
+
+cdef extern from "nvcomp/shared_types.h":
+    ctypedef enum nvcompStatus_t:
+        nvcompSuccess = 0,
+        nvcompErrorInvalidValue = 10,
+        nvcompErrorNotSupported = 11,
+        nvcompErrorCannotDecompress = 12,
+        nvcompErrorBadChecksum = 13,
+        nvcompErrorCannotVerifyChecksums = 14,
+        nvcompErrorCudaError = 1000,
+        nvcompErrorInternal = 10000,
+
+# nvCOMP Low-Level Interface.
+# https://github.com/NVIDIA/nvcomp/blob/main/doc/lowlevel_c_quickstart.md
+
+#
+# LZ4 batch compression/decompression API.
+#
+cdef extern from "nvcomp/lz4.h" nogil:
+    ctypedef struct nvcompBatchedLZ4Opts_t:
+        nvcompType_t data_type
+
+    # Compression API.
+    cdef nvcompStatus_t nvcompBatchedLZ4CompressGetTempSize(
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+        nvcompBatchedLZ4Opts_t format_opts,
+        size_t* temp_bytes
+    )
+
+    cdef nvcompStatus_t nvcompBatchedLZ4CompressGetMaxOutputChunkSize(
+        size_t max_uncompressed_chunk_bytes,
+        nvcompBatchedLZ4Opts_t format_opts,
+        size_t* max_compressed_bytes
+    )
+
+    cdef nvcompStatus_t nvcompBatchedLZ4CompressAsync(
+        const void* const* device_uncompressed_ptrs,
+        const size_t* device_uncompressed_bytes,
+        size_t max_uncompressed_chunk_bytes,
+        size_t batch_size,
+        void* device_temp_ptr,
+        size_t temp_bytes,
+        void* const* device_compressed_ptrs,
+        size_t* device_compressed_bytes,
+        nvcompBatchedLZ4Opts_t format_opts,
+        cudaStream_t stream
+    )
+
+    # Decompression API.
+    cdef nvcompStatus_t nvcompBatchedLZ4DecompressGetTempSize(
+        size_t num_chunks,
+        size_t max_uncompressed_chunk_bytes,
+        size_t* temp_bytes
+    )
+
+    nvcompStatus_t nvcompBatchedLZ4DecompressAsync(
+        const void* const* device_compressed_ptrs,
+        const size_t* device_compressed_bytes,
+        const size_t* device_uncompressed_bytes,
+        size_t* device_actual_uncompressed_bytes,
+        size_t batch_size,
+        void* const device_temp_ptr,
+        size_t temp_bytes,
+        void* const* device_uncompressed_ptrs,
+        nvcompStatus_t* device_statuses,
+        cudaStream_t stream
+    )
+
+#
+# Gdeflate batch compression/decompression API.
+#
+cdef extern from "nvcomp/gdeflate.h" nogil:
+    ctypedef struct nvcompBatchedGdeflateOpts_t:
+        int algo
+
+    # Compression API.
+    cdef nvcompStatus_t nvcompBatchedGdeflateCompressGetTempSize(
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+        nvcompBatchedGdeflateOpts_t format_opts,
+        size_t* temp_bytes
+    )
+
+    cdef nvcompStatus_t nvcompBatchedGdeflateCompressGetMaxOutputChunkSize(
+        size_t max_uncompressed_chunk_bytes,
+        nvcompBatchedGdeflateOpts_t format_opts,
+        size_t* max_compressed_bytes
+    )
+
+    cdef nvcompStatus_t nvcompBatchedGdeflateCompressAsync(
+        const void* const* device_uncompressed_ptrs,
+        const size_t* device_uncompressed_bytes,
+        size_t max_uncompressed_chunk_bytes,
+        size_t batch_size,
+        void* device_temp_ptr,
+        size_t temp_bytes,
+        void* const* device_compressed_ptrs,
+        size_t* device_compressed_bytes,
+        nvcompBatchedGdeflateOpts_t format_opts,
+        cudaStream_t stream
+    )
+
+    # Decompression API.
+    cdef nvcompStatus_t nvcompBatchedGdeflateDecompressGetTempSize(
+        size_t num_chunks,
+        size_t max_uncompressed_chunk_bytes,
+        size_t* temp_bytes
+    )
+
+    nvcompStatus_t nvcompBatchedGdeflateDecompressAsync(
+        const void* const* device_compressed_ptrs,
+        const size_t* device_compressed_bytes,
+        const size_t* device_uncompressed_bytes,
+        size_t* device_actual_uncompressed_bytes,
+        size_t batch_size,
+        void* const device_temp_ptr,
+        size_t temp_bytes,
+        void* const* device_uncompressed_ptrs,
+        nvcompStatus_t* device_statuses,
+        cudaStream_t stream
+    )
+
+#
+# zstd batch compression/decompression API.
+#
+cdef extern from "nvcomp/zstd.h" nogil:
+    ctypedef struct nvcompBatchedZstdOpts_t:
+        int reserved
+
+    # Compression API.
+    cdef nvcompStatus_t nvcompBatchedZstdCompressGetTempSize(
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+        nvcompBatchedZstdOpts_t format_opts,
+        size_t* temp_bytes
+    )
+
+    cdef nvcompStatus_t nvcompBatchedZstdCompressGetMaxOutputChunkSize(
+        size_t max_uncompressed_chunk_bytes,
+        nvcompBatchedZstdOpts_t format_opts,
+        size_t* max_compressed_bytes
+    )
+
+    cdef nvcompStatus_t nvcompBatchedZstdCompressAsync(
+        const void* const* device_uncompressed_ptrs,
+        const size_t* device_uncompressed_bytes,
+        size_t max_uncompressed_chunk_bytes,
+        size_t batch_size,
+        void* device_temp_ptr,
+        size_t temp_bytes,
+        void* const* device_compressed_ptrs,
+        size_t* device_compressed_bytes,
+        nvcompBatchedZstdOpts_t format_opts,
+        cudaStream_t stream
+    )
+
+    # Decompression API.
+    cdef nvcompStatus_t nvcompBatchedZstdDecompressGetTempSize(
+        size_t num_chunks,
+        size_t max_uncompressed_chunk_bytes,
+        size_t* temp_bytes
+    )
+
+    nvcompStatus_t nvcompBatchedZstdDecompressAsync(
+        const void* const* device_compressed_ptrs,
+        const size_t* device_compressed_bytes,
+        const size_t* device_uncompressed_bytes,
+        size_t* device_actual_uncompressed_bytes,
+        size_t batch_size,
+        void* const device_temp_ptr,
+        size_t temp_bytes,
+        void* const* device_uncompressed_ptrs,
+        nvcompStatus_t* device_statuses,
+        cudaStream_t stream
+    )
+
+#
+# Snappy batch compression/decompression API.
+#
+cdef extern from "nvcomp/snappy.h" nogil:
+    ctypedef struct nvcompBatchedSnappyOpts_t:
+        int reserved
+
+    # Compression API.
+    cdef nvcompStatus_t nvcompBatchedSnappyCompressGetTempSize(
+        size_t batch_size,
+        size_t max_uncompressed_chunk_bytes,
+        nvcompBatchedSnappyOpts_t format_opts,
+        size_t* temp_bytes
+    )
+
+    cdef nvcompStatus_t nvcompBatchedSnappyCompressGetMaxOutputChunkSize(
+        size_t max_uncompressed_chunk_bytes,
+        nvcompBatchedSnappyOpts_t format_opts,
+        size_t* max_compressed_bytes
+    )
+
+    cdef nvcompStatus_t nvcompBatchedSnappyCompressAsync(
+        const void* const* device_uncompressed_ptrs,
+        const size_t* device_uncompressed_bytes,
+        size_t max_uncompressed_chunk_bytes,
+        size_t batch_size,
+        void* device_temp_ptr,
+        size_t temp_bytes,
+        void* const* device_compressed_ptrs,
+        size_t* device_compressed_bytes,
+        nvcompBatchedSnappyOpts_t format_opts,
+        cudaStream_t stream
+    )
+
+    # Decompression API.
+    cdef nvcompStatus_t nvcompBatchedSnappyDecompressGetTempSize(
+        size_t num_chunks,
+        size_t max_uncompressed_chunk_bytes,
+        size_t* temp_bytes
+    )
+
+    nvcompStatus_t nvcompBatchedSnappyDecompressAsync(
+        const void* const* device_compressed_ptrs,
+        const size_t* device_compressed_bytes,
+        const size_t* device_uncompressed_bytes,
+        size_t* device_actual_uncompressed_bytes,
+        size_t batch_size,
+        void* const device_temp_ptr,
+        size_t temp_bytes,
+        void* const* device_uncompressed_ptrs,
+        nvcompStatus_t* device_statuses,
+        cudaStream_t stream
+    )
diff --git a/python/kvikio/nvcomp_codec.py b/python/kvikio/nvcomp_codec.py
new file mode 100644
index 0000000000..8f3b73dd79
--- /dev/null
+++ b/python/kvikio/nvcomp_codec.py
@@ -0,0 +1,272 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+from typing import Any, List, Mapping, Optional
+
+import cupy as cp
+import numpy as np
+from numcodecs.abc import Codec
+from numcodecs.compat import ensure_contiguous_ndarray_like
+
+import kvikio._lib.libnvcomp_ll as _ll
+
+
+class NvCompBatchCodec(Codec):
+    """Codec that uses batch algorithms from nvCOMP library.
+
+    An algorithm is selected using `algorithm` parameter.
+    If the algorithm takes additional options, they can be
+    passed to the algorithm using `options` dictionary.
+    """
+
+    # Header stores original uncompressed size. This is required to enable
+    # data compatibility between existing numcodecs codecs and NvCompBatchCodec.
+    HEADER_SIZE_BYTES: int = 4
+
+    codec_id: str = "nvcomp_batch"
+    algorithm: str
+    options: Mapping[str, Any]
+
+    def __init__(
+        self,
+        algorithm: str,
+        options: Optional[Mapping[str, Any]] = None,
+        stream: Optional[cp.cuda.Stream] = None,
+    ) -> None:
+        algo_id = algorithm.lower()
+        algo_t = _ll.SUPPORTED_ALGORITHMS.get(algo_id, None)
+        if algo_t is None:
+            raise ValueError(
+                f"{algorithm} is not supported. "
+                f"Must be one of: {list(_ll.SUPPORTED_ALGORITHMS.keys())}"
+            )
+
+        self.algorithm = algo_id
+        self.options = dict(options) if options is not None else {}
+
+        # Create an algorithm.
+        self._algo = algo_t(**self.options)
+        # Use default stream, if needed.
+        self._stream = stream if stream is not None else cp.cuda.Stream.ptds
+
+    def encode(self, buf):
+        """Encode data in `buf` using nvCOMP.
+
+        Parameters
+        ----------
+        buf : buffer-like
+            Data to be encoded. May be any object supporting the new-style
+            buffer protocol.
+
+        Returns
+        -------
+        enc : buffer-like
+            Encoded data. May be any object supporting the new-style buffer
+            protocol.
+        """
+        return self.encode_batch([buf])[0]
+
+    def encode_batch(self, bufs: List[Any]) -> List[Any]:
+        """Encode data in `bufs` using nvCOMP.
+
+        Parameters
+        ----------
+        bufs : List[buffer-like].
+            Data to be encoded. Each buffer in the list may be any object
+            supporting the new-style buffer protocol.
+
+        Returns
+        -------
+        enc : List[buffer-like]
+            List of encoded buffers. Each buffer may be any object supporting
+            the new-style buffer protocol.
+        """
+        num_chunks = len(bufs)
+        if num_chunks == 0:
+            return []
+
+        bufs = [cp.asarray(ensure_contiguous_ndarray_like(b)) for b in bufs]
+        buf_sizes = [b.size * b.itemsize for b in bufs]
+
+        max_chunk_size = max(buf_sizes)
+
+        # Get temp and output buffer sizes.
+        temp_size = self._algo.get_compress_temp_size(num_chunks, max_chunk_size)
+        comp_chunk_size = self._algo.get_compress_chunk_size(max_chunk_size)
+
+        # Prepare data and size buffers.
+        uncomp_chunks = cp.array(
+            [b.data.ptr for b in bufs],
+            dtype=cp.uint64,
+        )
+        uncomp_chunk_sizes = cp.array(buf_sizes, dtype=cp.uint64)
+
+        temp_buf = cp.empty(temp_size, dtype=cp.uint8)
+
+        # Includes header with the original buffer size,
+        # same as in numcodecs codec. This enables data compatibility between
+        # numcodecs default codecs and this nvCOMP batch codec.
+        # TODO(akamenev): probably should use contiguous buffer which stores all chunks?
+        comp_chunks_header = [
+            cp.empty(self.HEADER_SIZE_BYTES + comp_chunk_size, dtype=cp.uint8)
+            for _ in range(num_chunks)
+        ]
+        # comp_chunks is used as a container that stores pointers to actual chunks.
+        # nvCOMP requires this container to be in GPU memory.
+        comp_chunks = cp.array(
+            [c.data.ptr + self.HEADER_SIZE_BYTES for c in comp_chunks_header],
+            dtype=cp.uint64,
+        )
+        # Similar to comp_chunks, comp_chunk_sizes is an array that contains
+        # chunk sizes and is required by nvCOMP to be in GPU memory.
+        comp_chunk_sizes = cp.empty(num_chunks, dtype=cp.uint64)
+
+        self._algo.compress(
+            uncomp_chunks,
+            uncomp_chunk_sizes,
+            max_chunk_size,
+            num_chunks,
+            temp_buf,
+            comp_chunks,
+            comp_chunk_sizes,
+            self._stream,
+        )
+
+        # Write output buffers, each with the header.
+        res = []
+        for i in range(num_chunks):
+            comp_chunk = comp_chunks_header[i]
+            header = comp_chunk[:4].view(dtype=cp.uint32)
+            header[:] = buf_sizes[i]
+
+            res.append(
+                comp_chunk[: self.HEADER_SIZE_BYTES + comp_chunk_sizes[0]].tobytes()
+            )
+
+        return res
+
+    def decode(self, buf, out=None):
+        """Decode data in `buf` using nvCOMP.
+
+        Parameters
+        ----------
+        buf : buffer-like
+            Encoded data. May be any object supporting the new-style buffer
+            protocol.
+        out : buffer-like, optional
+            Writeable buffer to store decoded data. N.B. if provided, this buffer must
+            be exactly the right size to store the decoded data.
+
+        Returns
+        -------
+        dec : buffer-like
+            Decoded data. May be any object supporting the new-style
+            buffer protocol.
+        """
+        return self.decode_batch([buf], [out])[0]
+
+    def decode_batch(
+        self, bufs: List[Any], out: Optional[List[Any]] = None
+    ) -> List[Any]:
+        """Decode data in `bufs` using nvCOMP.
+
+        Parameters
+        ----------
+        bufs : List[buffer-like]
+            Encoded data. Each buffer in the list may be any object
+            supporting the new-style buffer protocol.
+        out : List[buffer-like], optional
+            List of writeable buffers to store decoded data.
+            N.B. if provided, each buffer must be exactly the right size
+            to store the decoded data.
+
+        Returns
+        -------
+        dec : List[buffer-like]
+            List of decoded buffers. Each buffer may be any object supporting
+            the new-style buffer protocol.
+        """
+        num_chunks = len(bufs)
+        if num_chunks == 0:
+            return []
+
+        # TODO(akamenev): check only first buffer, assuming they are all
+        # of the same kind.
+        is_host_buffer = not hasattr(bufs[0], "__cuda_array_interface__")
+        if is_host_buffer:
+            bufs = [cp.asarray(ensure_contiguous_ndarray_like(b)) for b in bufs]
+
+        # Get uncompressed chunk sizes from the header.
+        uncomp_chunk_sizes = [
+            int(b[: self.HEADER_SIZE_BYTES].view(dtype=cp.uint32)[0]) for b in bufs
+        ]
+        max_chunk_size = max(uncomp_chunk_sizes)
+
+        # Get temp buffer size.
+        temp_size = self._algo.get_decompress_temp_size(num_chunks, max_chunk_size)
+
+        # Prepare compressed chunks buffers.
+        comp_chunks = cp.array(
+            [b.data.ptr + self.HEADER_SIZE_BYTES for b in bufs],
+            dtype=cp.uint64,
+        )
+        comp_chunk_sizes = cp.array(
+            [b.size - self.HEADER_SIZE_BYTES for b in bufs],
+            dtype=cp.uint64,
+        )
+
+        temp_buf = cp.empty(temp_size, dtype=cp.uint8)
+
+        # Prepare uncompressed chunks buffers.
+        # First, allocate chunks of appropriate sizes and then
+        # copy the pointers to a pointer array in GPU memory as required by nvCOMP.
+        # TODO(akamenev): probably can allocate single contiguous buffer.
+        uncomp_chunks = [cp.empty(size, dtype=cp.uint8) for size in uncomp_chunk_sizes]
+        uncomp_chunk_ptrs = cp.array(
+            [c.data.ptr for c in uncomp_chunks], dtype=cp.uint64
+        )
+        # Sizes array must be in GPU memory.
+        uncomp_chunk_sizes = cp.array(uncomp_chunk_sizes, dtype=cp.uint64)
+
+        # TODO(akamenev): currently we provide the following 2 buffers to decompress()
+        # but do not check/use them afterwards since some of the algos
+        # (e.g. LZ4 and Gdeflate) do not require it and run faster
+        # without those arguments passed, while other algos (e.g. zstd) require
+        # these buffers to be valid.
+        actual_uncomp_chunk_sizes = cp.empty(num_chunks, dtype=cp.uint64)
+        statuses = cp.empty(num_chunks, dtype=cp.int32)
+
+        self._algo.decompress(
+            comp_chunks,
+            comp_chunk_sizes,
+            num_chunks,
+            temp_buf,
+            uncomp_chunk_ptrs,
+            uncomp_chunk_sizes,
+            actual_uncomp_chunk_sizes,
+            statuses,
+            self._stream,
+        )
+
+        res = []
+        for i in range(num_chunks):
+            ret = uncomp_chunks[i]
+            if out is not None and out[i] is not None:
+                o = ensure_contiguous_ndarray_like(out[i])
+                if hasattr(o, "__cuda_array_interface__"):
+                    cp.copyto(o, ret.view(dtype=o.dtype), casting="no")
+                else:
+                    np.copyto(o, cp.asnumpy(ret.view(dtype=o.dtype)), casting="no")
+                res.append(o)
+            elif is_host_buffer:
+                res.append(cp.asnumpy(ret))
+            else:
+                res.append(ret)
+
+        return res
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}"
+            f"(algorithm={self.algorithm!r}, options={self.options!r})"
+        )
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 5f484d8ee6..41d5d64947 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -105,3 +105,6 @@ skip = [
 
 [tool.mypy]
 ignore_missing_imports = true
+
+[project.entry-points."numcodecs.codecs"]
+nvcomp_batch = "kvikio.nvcomp_codec:NvCompBatchCodec"


Test patch:
diff --git a/python/tests/test_nvcomp_codec.py b/python/tests/test_nvcomp_codec.py
new file mode 100644
index 0000000000..99a90834c0
--- /dev/null
+++ b/python/tests/test_nvcomp_codec.py
@@ -0,0 +1,187 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+import itertools as it
+import json
+
+import numcodecs
+import numpy as np
+import pytest
+import zarr
+from numpy.testing import assert_equal
+
+from kvikio.nvcomp_codec import NvCompBatchCodec
+
+NVCOMP_CODEC_ID = "nvcomp_batch"
+
+LZ4_ALGO = "LZ4"
+GDEFLATE_ALGO = "Gdeflate"
+SNAPPY_ALGO = "snappy"
+ZSTD_ALGO = "zstd"
+
+SUPPORTED_CODECS = [LZ4_ALGO, GDEFLATE_ALGO, SNAPPY_ALGO, ZSTD_ALGO]
+
+
+def _get_codec(algo: str):
+    codec_args = {"id": NVCOMP_CODEC_ID, "algorithm": algo}
+    return numcodecs.registry.get_codec(codec_args)
+
+
+@pytest.fixture(params=[(16,), (8, 16), (16, 16)])
+def shape(request):
+    return request.param
+
+
+# Separate fixture for combinations of shapes and chunks, since
+# chunks array must have the same rank as data array.
+@pytest.fixture(
+    params=it.chain(
+        it.product([(32,)], [(16,), (32,), (40,)]),
+        it.product([(16, 8), (16, 16)], [(8, 16), (16, 16), (40, 12)]),
+    )
+)
+def shape_chunks(request):
+    return request.param
+
+
+@pytest.mark.parametrize("algo", SUPPORTED_CODECS)
+def test_codec_registry(algo: str):
+    codec = _get_codec(algo)
+    assert isinstance(codec, numcodecs.abc.Codec)
+
+
+@pytest.mark.parametrize("algo", SUPPORTED_CODECS)
+def test_basic(algo: str, shape):
+    codec = NvCompBatchCodec(algo)
+
+    # Create data.
+    dtype = np.float32
+    data = np.ones(shape, dtype=dtype)
+    # Do roundtrip.
+    comp_data = codec.encode(data)
+    # Decompress and cast to original data type/shape.
+    decomp_data = codec.decode(comp_data).view(dtype).reshape(shape)
+
+    assert_equal(decomp_data, data)
+
+
+@pytest.mark.parametrize("algo", SUPPORTED_CODECS)
+def test_basic_zarr(algo: str, shape_chunks):
+    shape, chunks = shape_chunks
+
+    codec = NvCompBatchCodec(algo)
+
+    data = np.ones(shape, dtype=np.float32)
+
+    # This will do the compression.
+    z = zarr.array(data, chunks=chunks, compressor=codec)
+
+    # Test the decompression.
+    assert_equal(z[:], data[:])
+
+
+@pytest.mark.parametrize("algo", SUPPORTED_CODECS)
+def test_batch_comp_decomp(algo: str):
+    codec = _get_codec(algo)
+
+    np.random.seed(1)
+
+    dtype = np.float32
+    # 2 equal-sized chunks.
+    chunks = list(np.random.randn(2, 100).astype(dtype))
+
+    comp_chunks = codec.encode_batch([c.tobytes() for c in chunks])
+    assert len(comp_chunks) == 2
+
+    decomp_chunks = codec.decode_batch(comp_chunks)
+    assert len(decomp_chunks) == 2
+
+    decomp_chunks = [c.view(dtype=dtype) for c in decomp_chunks]
+    assert_equal(decomp_chunks[0], chunks[0])
+    assert_equal(decomp_chunks[1], chunks[1])
+
+
+@pytest.mark.parametrize("algo", SUPPORTED_CODECS)
+def test_comp_decomp(algo: str, shape_chunks):
+    shape, chunks = shape_chunks
+
+    codec = _get_codec(algo)
+
+    np.random.seed(1)
+
+    data = np.random.randn(*shape).astype(np.float32)
+
+    z1 = zarr.array(data, chunks=chunks, compressor=codec)
+
+    zarr_store = zarr.MemoryStore()
+    zarr.save_array(zarr_store, z1, compressor=codec)
+    # Check the store.
+    meta = json.loads(zarr_store[".zarray"])
+    assert meta["compressor"]["id"] == NVCOMP_CODEC_ID
+    assert meta["compressor"]["algorithm"] == algo.lower()
+
+    # Read back/decompress.
+    z2 = zarr.open_array(zarr_store)
+
+    assert_equal(z1[:], z2[:])
+
+
+@pytest.mark.parametrize(
+    "algo, options",
+    [
+        ("lz4", {"data_type": 4}),  # NVCOMP_TYPE_INT data type.
+        ("gdeflate", {"algo": 1}),  # low-throughput, high compression ratio algo
+    ],
+)
+def test_codec_options(algo, options):
+    codec = NvCompBatchCodec(algo, options)
+
+    shape = (16, 16)
+    chunks = (8, 8)
+
+    data = np.ones(shape, dtype=np.float32)
+
+    z = zarr.array(data, chunks=chunks, compressor=codec)
+
+    assert_equal(z[:], data[:])
+
+
+def test_codec_invalid_options():
+    # There are currently only 3 supported algos in Gdeflate
+    codec = NvCompBatchCodec(GDEFLATE_ALGO, options={"algo": 10})
+
+    data = np.ones((16, 16), dtype=np.float32)
+
+    with pytest.raises(RuntimeError):
+        zarr.array(data, compressor=codec)
+
+
+def test_lz4_cpu_comp_gpu_decomp():
+    cpu_codec = numcodecs.registry.get_codec({"id": "lz4"})
+    gpu_codec = _get_codec(LZ4_ALGO)
+
+    shape = (16, 16)
+    chunks = (8, 8)
+
+    data = np.ones(shape, dtype=np.float32)
+
+    z1 = zarr.array(data, chunks=chunks)
+    store = {}
+    zarr.save_array(store, z1, compressor=cpu_codec)
+
+    meta = json.loads(store[".zarray"])
+    assert meta["compressor"]["id"] == "lz4"
+
+    meta["compressor"] = {"id": NVCOMP_CODEC_ID, "algorithm": LZ4_ALGO}
+    store[".zarray"] = json.dumps(meta).encode()
+
+    z2 = zarr.open_array(store, compressor=gpu_codec)
+
+    assert_equal(z1[:], z2[:])
+
+
+def test_empty_batch():
+    codec = _get_codec(LZ4_ALGO)
+
+    assert len(codec.encode_batch([])) == 0
+    assert len(codec.decode_batch([])) == 0