Add TPU backend support (torch_tpu / Pallas kernels)

Makefile

@@ -3,7 +3,7 @@
 
 export check_dirs := kernels/src kernels/tests kernels-data/bindings/python
 
-all: src/kernels/python_depends.json
+all: kernels/src/kernels/python_depends.json
 
 kernels/src/kernels/python_depends.json: kernels-data/src/python_dependencies.json
 	cp $< $@

docs/source/builder-cli.md

@@ -83,7 +83,7 @@ Initialize a new kernel project from template
 
   Default value: `Apache-2.0`
 * `--name <OWNER/REPO>` — Name of the kernel repo (e.g. `drbh/my-kernel`)
-* `--backends <BACKENDS>` — Backends to enable (`all`, `cpu`, `cuda`, `metal`, `neuron`, `rocm`, `xpu`)
+* `--backends <BACKENDS>` — Backends to enable (`all`, `cpu`, `cuda`, `metal`, `rocm`, `xpu`)
 
 * `--overwrite` — Overwrite existing scaffold files (preserves other files)
 

docs/source/kernel-requirements.md

@@ -149,7 +149,7 @@ The `backend` specifies a dictionary of the following form:
 ```
 
 The backend `type` must be one of `cann`, `cpu`, `cuda`, `metal`, `neuron`,
-`rocm`, or `xpu`. For CUDA and ROCm, the supported architectures must
+`rocm`, `tpu`, or `xpu`. For CUDA and ROCm, the supported architectures must
 be specified in the `archs` field.
 
 ### Python dependencies

examples/kernels/relu-tpu/CARD.md

@@ -0,0 +1,61 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Upstream
+
+The original source code for this kernel comes from {{ upstream }}.
+{% endif %}
+{% if source %}
+
+## Source
+
+The kernel-builder formatted source for this kernel is available at {{ source }}.
+{% endif %}

examples/kernels/relu-tpu/build.toml

@@ -0,0 +1,13 @@
+[general]
+name = "relu-tpu"
+version = 1
+edition = 5
+license = "Apache-2.0"
+backends = ["tpu"]
+
+[general.hub]
+repo-id = "kernels-test/relu-tpu"
+
+[torch-noarch]
+
+[kernel]

examples/kernels/relu-tpu/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for TPU ReLU kernel";
+
+  inputs = {
+    kernel-builder.url = "path:../../..";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py

@@ -0,0 +1,14 @@
+import jax
+import jax.numpy as jnp
+from torch_tpu._internal.pallas import jax_op
+
+
+def _jax_relu(x: jax.Array) -> jax.Array:
+    return jnp.maximum(x, 0)
+
+
+relu = jax_op("relu_tpu::relu", _jax_relu)
+
+from . import layers  # noqa: E402
+
+__all__ = ["layers", "relu"]

examples/kernels/relu-tpu/torch-ext/relu_tpu/layers/__init__.py

@@ -0,0 +1,9 @@
+import torch
+import torch.nn as nn
+
+from .. import relu
+
+
+class ReLU(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return relu(x)

kernel-builder/src/init.rs

@@ -44,7 +44,7 @@ pub struct InitArgs {
     #[arg(long, value_name = "OWNER/REPO")]
     pub name: Option<RepoInfo>,
 
-    /// Backends to enable (`all`, `cpu`, `cuda`, `metal`, `neuron`, `rocm`, `xpu`).
+    /// Backends to enable (`all`, `cpu`, `cuda`, `metal`, `rocm`, `xpu`).
     #[arg(long, num_args = 1.., default_values_t = default_init_backends())]
     pub backends: Vec<BackendSelection>,
 

kernel-builder/src/pyproject/templates/torch/noarch/_ops.py

@@ -4,7 +4,13 @@ def get_backend() -> str:
     """Detect the backend by inspecting torch."""
     import torch
 
-    if hasattr(torch, "neuron"):
+    if hasattr(torch.backends, "tpu"):
+        # torch_tpu sets torch.backends.tpu when it is imported (via
+        # torch's device-backend autoload), regardless of whether TPU
+        # hardware is present — analogous to torch.version.cuda being
+        # set on CUDA builds without a GPU.
+        return "tpu"
+    elif hasattr(torch, "neuron"):
         # Needs to be sorted before specific Torch builds, since Neuron
         # extension can be loaded into e.g. CUDA Torch builds.
         return "neuron"

kernels-data/bindings/python/kernels_data.pyi

@@ -32,6 +32,7 @@ class Backend(Enum):
     Metal = "Metal"
     Neuron = "Neuron"
     ROCm = "ROCm"
+    TPU = "TPU"
     XPU = "XPU"
 
     @staticmethod
@@ -40,7 +41,7 @@ class Backend(Enum):
 
         Args:
             s: One of `"cann"`, `"cpu"`, `"cuda"`, `"metal"`,
-               `"neuron"`, `"rocm"`, `"xpu"`.
+               `"neuron"`, `"rocm"`, `"tpu"`, `"xpu"`.
 
         Raises:
             ValueError: If the backend name is unknown.

kernels-data/bindings/python/src/lib.rs

@@ -88,6 +88,8 @@ enum PyBackend {
     Neuron,
     #[pyo3(name = "ROCm")]
     Rocm,
+    #[pyo3(name = "TPU")]
+    Tpu,
     #[pyo3(name = "XPU")]
     Xpu,
 }
@@ -101,6 +103,7 @@ impl From<Backend> for PyBackend {
             Backend::Metal => PyBackend::Metal,
             Backend::Neuron => PyBackend::Neuron,
             Backend::Rocm => PyBackend::Rocm,
+            Backend::Tpu => PyBackend::Tpu,
             Backend::Xpu => PyBackend::Xpu,
         }
     }
@@ -115,6 +118,7 @@ impl From<PyBackend> for Backend {
             PyBackend::Metal => Backend::Metal,
             PyBackend::Neuron => Backend::Neuron,
             PyBackend::Rocm => Backend::Rocm,
+            PyBackend::Tpu => Backend::Tpu,
             PyBackend::Xpu => Backend::Xpu,
         }
     }
@@ -123,7 +127,7 @@ impl From<PyBackend> for Backend {
 #[pymethods]
 impl PyBackend {
     /// Parse a backend name (`"cann"`, `"cpu"`, `"cuda"`, `"metal"`,
-    /// `"neuron"`, `"rocm"`, `"xpu"`).
+    /// `"neuron"`, `"rocm"`, `"tpu"`, `"xpu"`).
     #[staticmethod]
     #[pyo3(name = "from_str")]
     fn py_from_str(s: &str) -> PyResult<Self> {
@@ -144,6 +148,7 @@ impl PyBackend {
             PyBackend::Metal => "Metal",
             PyBackend::Neuron => "Neuron",
             PyBackend::Rocm => "ROCm",
+            PyBackend::Tpu => "TPU",
             PyBackend::Xpu => "XPU",
         };
         format!("Backend.{variant}")

kernels-data/bindings/python/tests/test_kernels_data.py

@@ -68,7 +68,7 @@ def test_backend_hash():
 
 def test_backend_unknown():
     with pytest.raises(ValueError):
-        Backend.from_str("tpu")
+        Backend.from_str("dsp")
 
 
 def test_backend_all_variants_and_casing():
@@ -80,8 +80,11 @@ def test_backend_all_variants_and_casing():
     assert repr(Backend.ROCm) == "Backend.ROCm"
     assert repr(Backend.XPU) == "Backend.XPU"
     assert repr(Backend.CANN) == "Backend.CANN"
+    assert str(Backend.TPU) == "tpu"
+    assert repr(Backend.TPU) == "Backend.TPU"
     assert Backend.from_str("cann") == Backend.CANN
     assert Backend.from_str("ROCM") == Backend.ROCm
+    assert Backend.from_str("TPU") == Backend.TPU
     assert Backend.from_str("metal") == Backend.Metal
 
 

kernels-data/src/config/mod.rs

@@ -113,6 +113,7 @@ pub struct General {
 
     pub cuda: Option<CudaGeneral>,
     pub neuron: Option<NeuronGeneral>,
+    pub tpu: Option<TpuGeneral>,
     pub xpu: Option<XpuGeneral>,
 }
 
@@ -144,6 +145,10 @@ impl General {
                 .cuda
                 .as_ref()
                 .and_then(|cuda| cuda.python_depends.as_ref()),
+            Backend::Tpu => self
+                .tpu
+                .as_ref()
+                .and_then(|tpu| tpu.python_depends.as_ref()),
             Backend::Xpu => self
                 .xpu
                 .as_ref()
@@ -191,6 +196,10 @@ pub struct NeuronGeneral {
     pub python_depends: Option<Vec<String>>,
 }
 
+pub struct TpuGeneral {
+    pub python_depends: Option<Vec<String>>,
+}
+
 pub struct Hub {
     pub repo_id: Option<String>,
     pub branch: Option<String>,
@@ -379,18 +388,20 @@ pub enum Backend {
     Metal,
     Neuron,
     Rocm,
+    Tpu,
     Xpu,
 }
 
 impl Backend {
-    pub const fn all() -> [Backend; 7] {
+    pub const fn all() -> [Backend; 8] {
         [
             Backend::Cann,
             Backend::Cpu,
             Backend::Cuda,
             Backend::Metal,
             Backend::Neuron,
             Backend::Rocm,
+            Backend::Tpu,
             Backend::Xpu,
         ]
     }
@@ -403,6 +414,7 @@ impl Backend {
             Backend::Metal => "metal",
             Backend::Neuron => "neuron",
             Backend::Rocm => "rocm",
+            Backend::Tpu => "tpu",
             Backend::Xpu => "xpu",
         }
     }
@@ -417,6 +429,7 @@ impl Display for Backend {
             Backend::Metal => write!(f, "metal"),
             Backend::Neuron => write!(f, "neuron"),
             Backend::Rocm => write!(f, "rocm"),
+            Backend::Tpu => write!(f, "tpu"),
             Backend::Xpu => write!(f, "xpu"),
         }
     }
@@ -433,6 +446,7 @@ impl FromStr for Backend {
             "metal" => Ok(Backend::Metal),
             "neuron" => Ok(Backend::Neuron),
             "rocm" => Ok(Backend::Rocm),
+            "tpu" => Ok(Backend::Tpu),
             "xpu" => Ok(Backend::Xpu),
             _ => Err(format!("Unknown backend: {s}")),
         }
@@ -444,3 +458,53 @@ pub enum ConfigError {
     #[error("Cannot migrate configuration: {reason:?}")]
     Migration { reason: String },
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn general_with_cuda_deps() -> General {
+        General {
+            name: KernelName::new("test-kernel").unwrap(),
+            version: 1,
+            license: "apache-2.0".to_string(),
+            upstream: None,
+            source: None,
+            backends: vec![Backend::Tpu],
+            hub: None,
+            python_depends: None,
+            cuda: Some(CudaGeneral {
+                minver: None,
+                maxver: None,
+                python_depends: Some(vec!["nvidia-cutlass-dsl".to_string()]),
+            }),
+            neuron: None,
+            tpu: None,
+            xpu: None,
+        }
+    }
+
+    #[test]
+    fn backend_python_depends_resolves_cuda() {
+        let general = general_with_cuda_deps();
+        let deps = general
+            .backend_python_depends(Backend::Cuda)
+            .map(|dep| dep.map(|(name, _)| name.to_string()))
+            .collect::<Result<Vec<_>>>()
+            .unwrap();
+
+        assert_eq!(deps, vec!["nvidia-cutlass-dsl".to_string()]);
+    }
+
+    #[test]
+    fn backend_python_depends_empty_for_backend_without_deps() {
+        let general = general_with_cuda_deps();
+
+        assert!(
+            general
+                .backend_python_depends(Backend::Cpu)
+                .next()
+                .is_none()
+        );
+    }
+}

kernels-data/src/config/v2.rs

@@ -179,6 +179,7 @@ impl General {
             hub: general.hub.map(Into::into),
             neuron: None,
             python_depends: None,
+            tpu: None,
             xpu: None,
         }
     }