Add TPU backend support (torch_tpu / Pallas kernels)#714
Open
tengomucho wants to merge 13 commits into
Open
Conversation
Make the Python loader recognize torch_tpu as a backend so a kernel built on torch_tpu._internal.pallas.jax_op can be published to / fetched from the Hub and dropped into kernelize(...) like any CUDA kernel. - Add TPU backend dataclass and parse_backend. - Probe torch_tpu via _get_torch_privateuse_backend_name() == "tpu", placed before the neuron/cuda/hip checks. - Add a tpu section (torch_tpu, jax, libtpu) to both python_depends JSON files so per-backend dependency validation resolves. - Add _TPURepos and 'tpu' to _validate_device_type's supported_devices so kernelize(model, device="tpu") works and resolves layer mappings. - Add has_tpu / tpu_only markers in conftest and a test_tpu test that loads kernels-test/relu-tpu on a torch_tpu interpreter.
Add Tpu to the canonical Backend enum in kernels-data so metadata-tpu.json and the Python-level kernels_data.Backend.TPU can be parsed and emitted. TPU is a noarch backend, so the noarch-only TpuGeneral block follows the shape of NeuronGeneral. - Add Backend::Tpu to mod.rs (and to Backend::all() / as_str / Display / FromStr). - Add tpu: Option<TpuGeneral> to the canonical General struct, and to the v3/v4/v5 per-edition schemas with matching From<*> / From<super::*> conversions both directions. - v2 migration seeds tpu: None. - Add PyBackend::Tpu (binding name "TPU") to the Python class, with both From directions and the __repr__ match arm; update the .pyi stub. - Add Backend.TPU member and update the Backend.from_str docstring in the .pyi stub. - Update the metadata.rs unit-test General literals to include tpu: None.
Add a 'tpu' arm to the templated get_backend() function, placed before the existing torch.cuda/hip/xpu probes and matching the canonical privateuse1 backend name check used in kernels/backends.py and kernels/tests/conftest.py. The three layers (_backend(), has_tpu, _kernel._ops.get_backend()) must agree so a noarch kernel built with backends=["tpu"] finds the matching torch_tpu ops namespace at import time.
A noarch Pallas-backed ReLU kernel that serves as the reference for
TPU kernel authors. Built on torch_tpu._internal.pallas.jax_op, so the
op namespace discovered at import is just the torch_tpu one registered
by the Torch XLA plugin rather than the kernel-builder _OPS_NAME
namespace (which is still emitted for symmetry but unused here).
- build.toml edition 5 with backends=["tpu"], [general.tpu]
python-depends = ["torch_tpu", "jax"], [general.hub] repo-id =
"kernels-test/relu-tpu", and no per-backend kernel sources
(torch-noarch).
- torch-ext/relu_tpu/__init__.py registers a plain `relu` callable via
jax_op("relu_tpu::relu", _jax_relu) and re-exports the layers submodule.
- torch-ext/relu_tpu/layers/__init__.py exposes ReLU as a torch.nn.Module
wrapping `relu` so kernelize(device="tpu") can swap a forward in.
- CARD.md mirrors the templated CARD.md committed under examples so
kernel-builder can publish without further edits.
test_basic.py and test_layer.py now use @pytest.mark.tpu_only to gate tests that require a torch_tpu interpreter. Register the marker so pytest will not warn about an unknown marker when those tests run on a non-TPU host and the tpu_only skip is exercised.
Add TPU as a first-class backend so a kernel repo declaring `backends = ["tpu"]` produces a `torch-tpu` noarch variant, dev shell, and CI/bundle outputs through the same generators as the other five backends. Plumbing: - lib/kernel-config.nix, lib/torch-version-utils.nix: add `tpu` to the backend init attrsets and an `isTpu` predicate + `backend` arm. - lib/mk-build-set.nix: add a `tpu` backendConfig and an empty backendOverlay arm (torch_tpu/jax/libtpu layer on the CPU torch wheel via pythonPackagesExtensions, not nixpkgs overlays). The tpu buildConfig sets allowUnfree = true because libtpu's wheel METADATA declares its license as "Google Cloud Platform Terms of Service" (unfree), unlike the Apache-2.0 torch_tpu. - lib/gen-flake-outputs.nix: add a `tpu` arm to buildConfigBackend and onePerFramework; backendCi/backendBundle pick it up automatically. - build-variants.json + scripts/gen_variants_markdown.py: register the `torch-tpu` noarch variant and the "TPU" platform name. - overlay.nix: wire jaxlib, libtpu, torch_tpu into pythonPackagesExtensions. - README.md: add a TPU row (Tier 3, experimental). Nix packages for the TPU Python wheels: - jax, jaxlib: PyPI wheels pinned to 0.10.2 (cp312). - libtpu 0.0.43, torch_tpu 0.1.1.dev20260707090224: wheels fetched directly from Google's Artifact Registry (gcloud), which requires an OAuth2 bearer token passed at fetch time via GCLOUD_ACCESS_TOKEN and turned into a netrc entry. Hashes are pinned for the cp312 wheel; scripts/helpers/get_torch_tpu_hash.sh prefetches the sha256 for a given ABI tag when refreshing versions. - lib/mk-build-set.nix: pin the noarch TPU extension env to python312 (the wheels are cp312-only at this point; a later commit derives the ABI tag from the package set's python instead).
Without a versions.nix record carrying tpu = true no TPU build set is ever instantiated, so kernels declaring backends = ["tpu"] resolved to zero applicable build sets. Add one torch 2.11 entry (torch_tpu pins torch>=2.11,<2.12), x86_64-linux only. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Mirrors the standard example flake (relu) so the kernel can be built with nix build directly, producing backendBundle.tpu and the redistributable torch-tpu variant. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The torch binary rewrap ships one wheel per (version, system, framework) keyed to the nixpkgs default python (currently 3.13), so pinning the TPU extension env to python312 forced an unusable cp313 torch into a 3.12 env. The registry publishes torch_tpu and libtpu wheels for cp311..cp314, so drop the python312 pin and derive the wheel ABI tag from the package set's python instead (hashes pinned for cp313). Also: - accept the tpu attribute in mk-build-set's buildConfig signature (build configs from versions.nix carry it) - declare torch_tpu's full wheel Requires-Dist set (torch, portpicker, tensorboard, frozendict, immutabledict; numpy/absl-py come via jax) - autoPatchelfHook with torch's lib dir on the search path, since the bundled extensions link libtorch_python.so/libc10.so/libtorch_cpu.so - get_torch_tpu_hash.sh takes the ABI tag as an optional argument Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
torch_tpu and jax had empty nix package lists in the dependency maps, so nix-builder's dependency resolution produced a check environment without them and the sandboxed get_kernel check could not detect the TPU backend. Point them at the torch_tpu/jax packages provided by the nix-builder overlay. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
torch_tpu only registers the torch.tpu namespace when TPU hardware is present (its loader is gated on the device count), so probing hasattr(torch, "tpu") made backend detection hardware-dependent — unlike every other backend, where detection reflects the installed torch stack (e.g. torch.version.cuda on GPU-less hosts). This broke nix-builder's sandboxed get_kernel check, which runs with torch_tpu installed but no TPU device. torch.backends.tpu is set unconditionally when torch_tpu is imported (torch's device-backend autoload triggers this on import torch), so probe that instead. Tests that need real hardware keep using the device-count-based has_tpu fixture. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
backend_python_depends only matched Cuda and Xpu, so a build.toml [general.tpu] python-depends section never reached the generated metadata.json and installed TPU kernels skipped dependency validation. (Neuron has the same pre-existing gap; left untouched.) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
torch_tpu's jax_op verifies the wrapped function's signature and rejects unannotated arguments, so the kernel failed to import. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Coverage report —
|
| Name | Stmts | Miss | Cover | Missing |
|---|---|---|---|---|
| src/kernels/__init__.py | 10 | 0 | 100% | |
| src/kernels/_system.py | 6 | 1 | 83% | 10 |
| src/kernels/_versions.py | 63 | 7 | 89% | 46, 49, 52-53, 56-57, 100 |
| src/kernels/backends.py | 212 | 62 | 71% | 40, 44, 48-51, 68, 90, 108, 117, 121, 125-127, 148, 157, 161, 165-167, 188, 199, 201, 208-211, 224, 228, 232-252, 260, 283-303 |
| src/kernels/compat.py | 8 | 1 | 88% | 5 |
| src/kernels/deps.py | 54 | 4 | 93% | 58-59, 95, 98 |
| src/kernels/layer/__init__.py | 6 | 0 | 100% | |
| src/kernels/layer/_interval_tree.py | 103 | 4 | 96% | 23, 52, 147, 150 |
| src/kernels/layer/device.py | 48 | 14 | 71% | 42, 47-49, 91, 96-98, 101, 149, 152, 155-157 |
| src/kernels/layer/func.py | 81 | 7 | 91% | 81, 111, 183, 301, 307, 320, 338 |
| src/kernels/layer/globals.py | 5 | 0 | 100% | |
| src/kernels/layer/kernelize.py | 74 | 8 | 89% | 255, 281, 289-290, 296, 300, 316-318 |
| src/kernels/layer/layer.py | 210 | 16 | 92% | 167, 210, 216, 229, 337, 417-418, 430, 439, 447, 458, 487, 491, 504, 557, 587 |
| src/kernels/layer/mode.py | 14 | 0 | 100% | |
| src/kernels/layer/repos.py | 144 | 42 | 71% | 27, 33, 36-43, 63-64, 70, 73-76, 90, 94, 103-104, 110, 113-116, 123-124, 130, 133-136, 143-144, 150, 153-156, 163-164, 170, 173-176, 257 |
| src/kernels/lockfile.py | 71 | 46 | 35% | 37-104, 108-131 |
| src/kernels/status.py | 49 | 2 | 96% | 23, 81 |
| src/kernels/utils.py | 301 | 55 | 82% | 65, 77-81, 87-88, 218, 222, 225, 287, 295, 334-335, 373, 404, 409, 444, 673, 676, 678, 684, 697-698, 719-731, 735-742, 750, 754-764, 768-775, 813, 817, 836, 838 |
| src/kernels/variants.py | 262 | 19 | 93% | 56, 87, 108, 138, 247-248, 289, 291, 371-378, 384-390, 421-427, 439-445, 534-536 |
| src/kernels/verify.py | 88 | 1 | 99% | 32 |
| TOTAL | 1809 | 289 | 84% |
Updated by the Test kernels workflow on commit 7126d26f2cdebfc7b1dd56d91f6a75dbd56015b9.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Important: Read before submitting
Closes #711.
Adds TPU as a first-class backend across the stack, so a kernel repo declaring
backends = ["tpu"]builds and loads atorch-tpuvariant (Pallas/JAX kernels running through torch_tpu):torch.backends.tpu,torch-tpunoarch variant resolution,kernelizedevice support,tpu_onlypytest marker.Backend::Tpuenum variant with serde/pyo3 bindings and TPUpython_dependsin backend metadata.tpuaccepted ininit --backends, TPU detection in the noarch_ops.pytemplate.torch-tpubuild variant. This part was made and tested with the help of Fable.relu-tpuexample kernel repository.TPU support is experimental (Tier 3): the libtpu/torch_tpu wheels are not public yet and are fetched from Google's Artifact Registry using a gcloud token (
GCLOUD_ACCESS_TOKEN), sorelu-tpuis excluded from the examples CI for now (until we have access to the TPUs for CI).Note:
KNOWN_BACKENDSinkernels.utilsintentionally does not gaintpu, any value change hits an API compatibility lint error in the CI.Checklist