nix-builder: fully enable TPU support

.github/workflows/build_kernel_tpu.yaml

@@ -0,0 +1,45 @@
+name: "Build and test kernel (TPU)"
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+    types: [opened, synchronize, reopened] # trigger on PRs
+    paths-ignore:
+      - "docs/**"
+      - "*.md"
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+permissions: {}
+
+jobs:
+  build:
+    permissions:
+      contents: read
+    name: Build kernels (TPU)
+    runs-on:
+      group: aws-highmemory-32-plus-nix
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      - uses: DeterminateSystems/nix-installer-action@3138316df39ed29be04236d7ffc686fa525866aa # v23
+        with:
+          extra-conf: |
+            max-jobs = 4
+            cores = 12
+            sandbox-fallback = false
+      - uses: cachix/cachix-action@1eb2ef646ac0255473d23a5907ad7b04ce94065c # v17
+        with:
+          name: huggingface
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+        env:
+          USER: runner
+      - name: Nix info
+        run: nix-shell -p nix-info --run "nix-info -m"
+      # For now we only test that there are no regressions in building TPU
+      # kernels. Also run tests once we have a TPU runner.
+      - name: Build all TPU example kernels
+        run: nix build -L ./examples/kernels#ci-build-tpu

README.md

@@ -81,7 +81,8 @@ different tiers:
 | Huawei NPU | ✓               | ✗                      | ✗                       | 3    |
 | Neuron     | ✓               | x                      | x                       | 3    |
 
-**Warning:** Neuron and TPU support are experimental and packages are not public yet.
+> [!WARNING] 
+> Neuron support is experimental and packages are not public yet. TPU support is is experimental.
 
 ## 📚 Documentation
 

docs/source/builder/build-variants.md

@@ -98,6 +98,10 @@ available. This list will be updated as new PyTorch versions are released.
 - `torch215-rocm100-x86_64-linux`
 - `torch215-rocm714-x86_64-linux`
 
+## TPU x86_64-linux
+
+- `torch-tpu`
+
 ## XPU x86_64-linux
 
 - `torch210-xpu20253-x86_64-linux`
@@ -116,4 +120,5 @@ one or more of the following variants:
 - `torch-cuda`
 - `torch-metal`
 - `torch-rocm`
+- `torch-tpu`
 - `torch-xpu`

examples/kernels/flake.nix

@@ -394,6 +394,15 @@
         }
       ];
 
+      # TPU kernels to build in CI.
+      ciTpuKernels = [
+        {
+          name = "relu-tpu-kernel";
+          path = ./relu-tpu;
+          drv = sys: out: out.packages.${sys}.redistributable.torch-tpu;
+        }
+      ];
+
       # Metal kernels to build in CI.
       ciMetalKernels = [
         {
@@ -477,6 +486,7 @@
       ];
 
       ciXpuKernelOutputs = mkKernelOutputs' ciXpuKernels;
+      ciTpuKernelOutputs = mkKernelOutputs' ciTpuKernels;
       ciMetalKernelOutputs = mkKernelOutputs' ciMetalKernels;
       ciCpuKernelOutputs = mkKernelOutputs' ciCpuKernels;
     in
@@ -617,6 +627,7 @@
           ci-build-cuda = mkCiBuild "ci-kernels-cuda" ciKernelOutputs;
           ci-build-rocm = mkCiBuild "ci-kernels-rocm" ciRocmKernelOutputs;
           ci-build-xpu = mkCiBuild "ci-kernels-xpu" ciXpuKernelOutputs;
+          ci-build-tpu = mkCiBuild "ci-kernels-tpu" ciTpuKernelOutputs;
           ci-build-metal = mkCiBuild "ci-kernels-metal" ciMetalKernelOutputs;
           ci-build-cpu = mkCiBuild "ci-kernels-cpu" ciCpuKernelOutputs;
         in
@@ -626,6 +637,7 @@
               ci-build-cuda
               ci-build-rocm
               ci-build-xpu
+              ci-build-tpu
               ci-build-metal
               ci-build-cpu
               ;

nix-builder/build-variants.json

@@ -90,6 +90,9 @@
       "torch215-rocm100-x86_64-linux",
       "torch215-rocm714-x86_64-linux"
     ],
+    "tpu": [
+      "torch-tpu"
+    ],
     "xpu": [
       "torch210-xpu20253-x86_64-linux",
       "torch211-xpu20253-x86_64-linux",

nix-builder/pkgs/python-modules/libtpu/default.nix

@@ -1,33 +1,11 @@
 {
   lib,
+  fetchPypi,
   buildPythonPackage,
   requireFile,
   python,
 }:
 
-# libtpu is served from Google's gated Artifact Registry, which requires
-# an OAuth2 bearer token. That token can't be handed to a build sandbox
-# without making the derivation impure, so instead the wheel is
-# prefetched out-of-band and handed to `requireFile`, which only checks
-# that a file matching `name` + `hash` already exists in the store:
-#
-#   GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
-#     nix-builder/scripts/helpers/prefetch-tpu-wheel.sh libtpu 0.0.43
-#
-# The script prints the `hash` line to paste in below. To bump the
-# version, update `version`, then re-run the script with the new
-# version.
-#
-# libtpu ships the `libtpu/libtpu.so` runtime that both torch_tpu and
-# jaxlib dlopen. Pinned to 0.0.x to match the jax 0.10.x / torch_tpu
-# 0.1.x stack (torch_tpu requires libtpu>=0.0.40; jax's "tpu" extra
-# pins libtpu==0.0.43.*), so bump it together with jaxlib/torch_tpu.
-#
-# NOTE on license: libtpu's wheel METADATA declares its license as
-# "Google Cloud Platform Terms of Service" — i.e. unfree, NOT
-# Apache-2.0 (unlike torch_tpu). The tpu buildSet therefore sets
-# allowUnfree = true (see lib/mk-build-set.nix).
-
 let
   # The wheel ships per-CPython-ABI builds (cp311..cp314); pick the tag
   # matching the python this package set is built for. The hash below is
@@ -37,19 +15,16 @@ let
 in
 buildPythonPackage rec {
   pname = "libtpu";
-  version = "0.0.48";
+  version = "0.0.50";
   format = "wheel";
 
-  src = requireFile {
-    name = "libtpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl";
-    hash = "sha256-QD0G/ssztiQYMmd/gwI/IXlYnwyZ4lwRocEf/7V7NBY=";
-    message = ''
-      libtpu is served from a gated Google Artifact Registry and cannot
-      be fetched by a pure Nix build. Fetch and register it with:
-
-        GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
-          nix-builder/scripts/helpers/prefetch-tpu-wheel.sh libtpu ${version} ${abi}
-    '';
+  src = fetchPypi {
+    inherit format pname version;
+    dist = abi;
+    python = abi;
+    inherit abi;
+    platform = "manylinux_2_31_x86_64";
+    hash = "sha256-vyMsa9hUVwlGJi3269OB5fcJqRdSxNZDNXj0I1lfz1w=";
   };
 
   pythonImportsCheck = [ "libtpu" ];

nix-builder/pkgs/python-modules/torch/binary/generic.nix

@@ -6,6 +6,8 @@
   symlinkJoin,
   buildPythonPackage,
   fetchFromGitHub,
+
+  fetchPypi,
   fetchurl,
   requireFile,
 
@@ -45,6 +47,7 @@
   numpy,
   portpicker,
   pyyaml,
+  pyzmq,
   pyzes,
   requests,
   setuptools,
@@ -157,23 +160,22 @@ let
   };
 
   # Once we require multiple versions, factor this out.
+  # The hash is for the cp314 wheel (the nixpkgs default python).
   torch-tpu =
     let
       abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}";
-      version = "0.1.1.dev20260916100444";
     in
-    requireFile {
-      name = "torch_tpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl";
-      hash = "sha256-A8Rw9/QkMCjj9jfb4s+/d04/m/xH5L9rJDa0+giessM=";
-      message = ''
-        torch_tpu is served from a gated Google Artifact Registry and
-        cannot be fetched by a pure Nix build. Fetch and register it with:
-
-          GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
-            nix-builder/scripts/helpers/prefetch-tpu-wheel.sh torch_tpu ${version} ${abi}
-      '';
+    fetchPypi {
+      format = "wheel";
+      # Wheel file names use an underscore.
+      pname = "torch_tpu";
+      version = "0.2.0";
+      dist = abi;
+      python = abi;
+      inherit abi;
+      platform = "manylinux_2_31_x86_64";
+      hash = "sha256-xdqU1lQfQchadDw0En1s2huCAqHafZR6OlLFE00Z3ak=";
     };
-
 in
 buildPythonPackage.override { stdenv = effectiveStdenv; } {
   pname = "torch";
@@ -293,6 +295,7 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } {
     frozendict
     immutabledict
     portpicker
+    pyzmq
     tensorboard
   ]
   ++ lib.optionals tritonSupport [

nix-builder/versions.nix

@@ -127,10 +127,7 @@
     torchVersion = "2.11";
     tpu = true;
     systems = [ "x86_64-linux" ];
-    # Excluded from the shared build cache: torch_tpu's autoPatchelfHook
-    # step produces an output that is itself a modified copy of Google's
-    # gated libtpu/torch_tpu wheels, so it can't be published publicly.
-    bundleBuild = false;
+    bundleBuild = true;
   }
 
   {