torch-lattice 0.4.0

Sparse convolution library for modern PyTorch and CUDA.

Torch Lattice

torch-lattice is the Torch/CUDA training-side companion to mlx-lattice. It keeps the sparse model authoring and CUDA provenance workflow on the Torch side, then exports portable lattice MLIR artifacts for MLX/Metal deployment.

torch-lattice is a project-owned fork of MIT HAN Lab's TorchSparse. The public semantics are aligned to mlx-lattice and the lattice MLIR contract rather than to historical TorchSparse API quirks.

Documentation | MLX Lattice | PyPI | Acknowledgements

Install

For normal use, install the published package from PyPI:

uv add torch-lattice --torch-backend cu128

If you are installing into an existing environment instead of adding a project dependency, use:

uv pip install --torch-backend cu128 torch-lattice

The published wheel targets Python 3.14 and the PyTorch CUDA 12.8 wheel stack. At runtime, the important requirement is a Linux system with a compatible NVIDIA driver for the CUDA runtime provided by PyTorch. A local CUDA toolkit is only needed when building from source or developing the native extension.

For development from a checkout:

uv sync --all-packages --extra test

The repository also provides a CUDA Linux GitHub workflow that builds and smoke checks the native CUDA wheel on an Ubuntu runner.

Relationship to MLX Lattice

The two packages are intentionally split by runtime role:

• torch-lattice is the CUDA training and artifact-production side.
• mlx-lattice is the Apple Silicon inference and deployment side.
• lattice-contract defines the shared artifact constants and MLIR contract metadata used by both sides.

Portable artifacts use graph.mlir plus weights.safetensors. Torch-side exporters write those files; MLX-side artifact loading compiles them into an executable MLX program.

Convolution semantics

Convolution classes are explicit:

• torch_lattice.nn.Conv3d is forward support-generating sparse convolution and exports to lattice.conv3d, including stride=1. Calling the same module as conv(x, coordinates=target) exports target-aligned convolution without a separate module class.
• torch_lattice.nn.SubmConv3d is support-preserving submanifold convolution and exports to lattice.subm_conv3d.
• torch_lattice.nn.ConvTranspose3d exports to lattice.conv_transpose3d.
• torch_lattice.nn.GenerativeConvTranspose3d exports to lattice.generative_conv_transpose3d.

Artifact builders lower module identity directly. They do not infer submanifold semantics from stride, padding, or legacy indice-key conventions.

Tooling

After uv sync --all-packages, use the workspace scripts from the repository root:

uv run bench --preset smoke
uv run fuzz --cases 32 --device cuda --archive /tmp/torch_lattice_fuzz.tar.gz
uv run conformance fuzz --cases 32 --device cuda
uv run migration all --device cuda

The corresponding MLX-side replay command is:

uv run conformance replay /tmp/torch_lattice_fuzz.tar.gz \
  --report /tmp/torch_lattice_fuzz_report.json

Migration compatibility checks

Original TorchSparse and torch-lattice are not assumed to have identical class semantics. The supported migration rule is explicit:

• original torchsparse.nn.Conv3d(kernel_size > 1, stride = 1) maps to torch_lattice.nn.SubmConv3d;
• original pointwise Conv3d(kernel_size = 1) maps to torch_lattice.nn.Conv3d;
• original strided forward convolutions map to torch_lattice.nn.Conv3d with the same stride.

The migration CLI verifies the covered subset against a kept original TorchSparse package/worktree in separate subprocesses.

Documentation

The full documentation is hosted at torch-lattice.iki.moe:

• Installation
• Quickstart
• Concept references
• Tooling references
• API reference

The source for the documentation lives in docs/. Build it locally with:

uv sync --group docs --no-install-workspace
uv run --no-sync sphinx-build -W -b html docs docs/_build/html

Development

Common local checks:

uv run --all-packages --extra test pytest tests -q
uv run bench --list

Build CUDA Linux distributions locally with:

export CUDA_PATH=/usr/local/cuda-12.8
uv build \
  --sdist \
  --wheel \
  --config-setting=cmake.define.CMAKE_CUDA_COMPILER="$CUDA_PATH/bin/nvcc" \
  --config-setting=cmake.define.CUDAToolkit_ROOT="$CUDA_PATH"

Acknowledgements

torch-lattice is based on MIT HAN Lab's original TorchSparse project.

It is developed together with mlx-lattice, which provides the MLX/Metal deployment runtime for the same artifact contract.

License

Open sourced under the MIT license.