From e1d9cde3cc15f4bdd1cfb52dfe4404d4da5673cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Wed, 1 Jul 2026 14:45:43 +0000 Subject: [PATCH 1/3] kernel-builder: ensure that `HIP_ARCHITECTURES` gets properly cleared We want to clear `HIP_ARCHITECTURES`, on the shared library target, since we want to set archs per kernel component. However, conditioning this on `GPU_ARCHES` caused the property clearing to never happen. Since we do not use `GPU_ARCHES` anymore, rip it out. --- .../src/pyproject/templates/torch/preamble.cmake | 1 - .../pyproject/templates/torch/torch-extension.cmake | 1 - .../src/pyproject/templates/tvm_ffi/preamble.cmake | 1 - kernel-builder/src/pyproject/templates/utils.cmake | 13 +++---------- 4 files changed, 3 insertions(+), 13 deletions(-) diff --git a/kernel-builder/src/pyproject/templates/torch/preamble.cmake b/kernel-builder/src/pyproject/templates/torch/preamble.cmake index 5169693d..7f6b014d 100644 --- a/kernel-builder/src/pyproject/templates/torch/preamble.cmake +++ b/kernel-builder/src/pyproject/templates/torch/preamble.cmake @@ -198,7 +198,6 @@ elseif(GPU_LANG STREQUAL "SYCL") set(sycl_link_flags "-Wl,-z,noexecstack;-fsycl;--offload-compress;-fsycl-targets=spir64_gen,spir64;-Xs;-device pvc,xe-lpg,ats-m150 -options ' -cl-intel-enable-auto-large-GRF-mode -cl-poison-unsupported-fp64-kernels -cl-intel-greater-than-4GB-buffer-required';") set(sycl_flags "-fPIC;-fsycl;-fhonor-nans;-fhonor-infinities;-fno-associative-math;-fno-approx-func;-fno-sycl-instrument-device-code;--offload-compress;-fsycl-targets=spir64_gen,spir64;") set(GPU_FLAGS "${sycl_flags}") - set(GPU_ARCHES "") add_compile_definitions(XPU_KERNEL) diff --git a/kernel-builder/src/pyproject/templates/torch/torch-extension.cmake b/kernel-builder/src/pyproject/templates/torch/torch-extension.cmake index 7e4be0f5..361e1589 100644 --- a/kernel-builder/src/pyproject/templates/torch/torch-extension.cmake +++ b/kernel-builder/src/pyproject/templates/torch/torch-extension.cmake @@ -11,7 +11,6 @@ define_gpu_extension_target( LANGUAGE ${GPU_LANG} SOURCES ${SRC} COMPILE_FLAGS ${GPU_FLAGS} - ARCHITECTURES ${GPU_ARCHES} USE_SABI 3 WITH_SOABI) diff --git a/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake b/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake index 3f679765..2eefe654 100644 --- a/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake +++ b/kernel-builder/src/pyproject/templates/tvm_ffi/preamble.cmake @@ -138,7 +138,6 @@ elseif(GPU_LANG STREQUAL "SYCL") set(sycl_link_flags "-Wl,-z,noexecstack;-fsycl;--offload-compress;-fsycl-targets=spir64_gen,spir64;-Xs;-device pvc,xe-lpg,ats-m150 -options ' -cl-intel-enable-auto-large-GRF-mode -cl-poison-unsupported-fp64-kernels -cl-intel-greater-than-4GB-buffer-required';") set(sycl_flags "-fPIC;-fsycl;-fhonor-nans;-fhonor-infinities;-fno-associative-math;-fno-approx-func;-fno-sycl-instrument-device-code;--offload-compress;-fsycl-targets=spir64_gen,spir64;") set(GPU_FLAGS "${sycl_flags}") - set(GPU_ARCHES "") add_compile_definitions(XPU_KERNEL) endif() diff --git a/kernel-builder/src/pyproject/templates/utils.cmake b/kernel-builder/src/pyproject/templates/utils.cmake index eb96150d..4c19c65c 100644 --- a/kernel-builder/src/pyproject/templates/utils.cmake +++ b/kernel-builder/src/pyproject/templates/utils.cmake @@ -531,7 +531,7 @@ function (define_gpu_extension_target GPU_MOD_NAME) GPU "WITH_SOABI" "DESTINATION;LANGUAGE;USE_SABI" - "SOURCES;ARCHITECTURES;COMPILE_FLAGS;INCLUDE_DIRECTORIES;LIBRARIES") + "SOURCES;COMPILE_FLAGS;INCLUDE_DIRECTORIES;LIBRARIES") # Add hipify preprocessing step when building with HIP/ROCm. if (GPU_LANGUAGE STREQUAL "HIP") @@ -553,16 +553,9 @@ function (define_gpu_extension_target GPU_MOD_NAME) if (GPU_LANGUAGE STREQUAL "HIP") # Make this target dependent on the hipify preprocessor step. add_dependencies(${GPU_MOD_NAME} hipify${GPU_MOD_NAME}) - endif() - if (GPU_ARCHITECTURES) - if (GPU_LANGUAGE STREQUAL "HIP") - # Clear target architectures, we are passing arch flags per source file. - set_property(TARGET ${GPU_MOD_NAME} PROPERTY HIP_ARCHITECTURES off) - else() - set_target_properties(${GPU_MOD_NAME} PROPERTIES - ${GPU_LANGUAGE}_ARCHITECTURES "${GPU_ARCHITECTURES}") - endif() + # Clear target architectures, we are passing arch flags per source file. + set_property(TARGET ${GPU_MOD_NAME} PROPERTY HIP_ARCHITECTURES off) endif() if (TORCH_VERSION VERSION_LESS 2.12.0) From 73731337453e7645c605a71208bc7ec6446080b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Wed, 1 Jul 2026 14:48:25 +0000 Subject: [PATCH 2/3] nix-builder: remove unused variable that broke ROCm devenvs --- nix-builder/lib/build.nix | 1 - 1 file changed, 1 deletion(-) diff --git a/nix-builder/lib/build.nix b/nix-builder/lib/build.nix index fcc3bc88..b1a41ec5 100644 --- a/nix-builder/lib/build.nix +++ b/nix-builder/lib/build.nix @@ -447,7 +447,6 @@ rec { buildInputs = [ python ]; inputsFrom = [ extension ]; env = lib.optionalAttrs rocmSupport { - PYTORCH_ROCM_ARCH = lib.concatStringsSep ";" buildSet.torch.rocmArchs; HIP_PATH = pkgs.rocmPackages.clr; }; From c01b6bd78a2e6562d0d93b6fb492a4c8b415df91 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Thu, 2 Jul 2026 13:53:39 +0000 Subject: [PATCH 3/3] ci: add tests for architectures and architecture intersections --- examples/kernels/flake.nix | 109 ++++++++++++++++++ examples/kernels/relu-archs-subset/CARD.md | 61 ++++++++++ examples/kernels/relu-archs-subset/build.toml | 45 ++++++++ examples/kernels/relu-archs-subset/flake.nix | 17 +++ .../relu-archs-subset/relu_cuda/relu.cu | 47 ++++++++ .../torch-ext/relu_archs_subset/__init__.py | 15 +++ .../torch-ext/torch_binding.cpp | 19 +++ .../torch-ext/torch_binding.h | 5 + examples/kernels/relu/build.toml | 12 -- 9 files changed, 318 insertions(+), 12 deletions(-) create mode 100644 examples/kernels/relu-archs-subset/CARD.md create mode 100644 examples/kernels/relu-archs-subset/build.toml create mode 100644 examples/kernels/relu-archs-subset/flake.nix create mode 100644 examples/kernels/relu-archs-subset/relu_cuda/relu.cu create mode 100644 examples/kernels/relu-archs-subset/torch-ext/relu_archs_subset/__init__.py create mode 100644 examples/kernels/relu-archs-subset/torch-ext/torch_binding.cpp create mode 100644 examples/kernels/relu-archs-subset/torch-ext/torch_binding.h diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index e5baeaa0..c4e6b208 100644 --- a/examples/kernels/flake.nix +++ b/examples/kernels/flake.nix @@ -27,6 +27,12 @@ # - drv (system -> flakeOutputs -> derivation): the derivation for the given # system and flake outputs. # - torchVersions: optional override for the torchVersions argument + # - checkRocmArchs: optional list of ROCm architectures (e.g. "gfx942"). + # When set, the kernel dylib must contain exactly this set of + # architectures. + # - checkCudaCapabilities: optional list of CUDA capabilities (e.g. "9.0"). + # When set, the kernel dylib must contain exactly this set of + # capabilities. ciKernels = [ { name = "cpp20-symbols-kernel"; @@ -38,6 +44,28 @@ path = ./relu; drv = sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${cudaVersion}-${sys}"}; + checkCudaCapabilities = [ + "7.0" + "7.2" + "7.5" + "8.0" + "8.6" + "8.7" + "8.9" + "9.0" + ]; + } + { + # Check arch intersection. 5.0 is dropped because it is not supported. + name = "relu-archs-subset"; + path = ./relu-archs-subset; + drv = + sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${cudaVersion}-${sys}"}; + checkCudaCapabilities = [ + "7.5" + "8.0" + "9.0" + ]; } { name = "relu-torch-stable-abi-kernel"; @@ -175,6 +203,30 @@ path = ./relu; drv = sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${rocmVersion}-${sys}"}; + checkRocmArchs = [ + "gfx906" + "gfx908" + "gfx90a" + "gfx942" + "gfx950" + "gfx1030" + "gfx1100" + "gfx1101" + "gfx1200" + "gfx1201" + ]; + } + { + # Check arch intersection. gfx940 is dropped because it is not supported. + name = "relu-archs-subset"; + path = ./relu-archs-subset; + drv = + sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${rocmVersion}-${sys}"}; + checkRocmArchs = [ + "gfx90a" + "gfx942" + "gfx1100" + ]; } { name = "relu-compiler-flags"; @@ -279,6 +331,59 @@ let pkgs = nixpkgs.legacyPackages.${system}; + # Check that every *.so in the output of `drv` contains exactly the + # given set of architectures. `listArchs` is a shell command that + # prints the architectures of the kernel library `$so`, one per + # line. On success, the output is a symlink to `drv`. + checkKernelArchs = + drv: expectedArchs: listArchs: + pkgs.runCommand "${drv.name}-check-archs" + { + expected = lib.concatLines expectedArchs; + passAsFile = [ "expected" ]; + } + '' + kernelLibs=$(find ${drv}/ -name '*.so') + if [ -z "$kernelLibs" ]; then + echo "no *.so files found in ${drv}" >&2 + exit 1 + fi + + sort -u "$expectedPath" > expected + for so in $kernelLibs; do + ( + ${listArchs} + ) | sort -u > actual + if ! diff -u expected actual; then + echo "unexpected architectures in $so" >&2 + exit 1 + fi + done + + ln -s ${drv} $out + ''; + + checkRocmArchs = + drv: expectedArchs: + checkKernelArchs drv expectedArchs + # Use the ROCm LLVM that the kernel was built with. Stock + # llvm-readobj does not support `--offloading`. + '' + ${drv.torch.rocmPackages.rocm-llvm}/llvm/bin/llvm-readobj --offloading "$so" \ + | grep -oE 'gfx[0-9]+[a-z]?' + ''; + + checkCudaCapabilities = + drv: expectedCapabilities: + let + # Convert capabilities like "9.0a" to sm names like "sm_90a". + smArchs = map (cap: "sm_" + lib.replaceStrings [ "." ] [ "" ] cap) expectedCapabilities; + in + checkKernelArchs drv smArchs '' + ${drv.torch.cudaPackages.cuda_cuobjdump}/bin/cuobjdump "$so" \ + | grep -oE 'sm_[0-9]+[af]?' + ''; + resolveKernels = kernelOutputsList: map (kernel: { @@ -292,6 +397,10 @@ drv = baseDrv; expectedBuilderLogEntries = kernel.assertFailLogs or [ ]; } + else if kernel ? checkRocmArchs then + checkRocmArchs baseDrv kernel.checkRocmArchs + else if kernel ? checkCudaCapabilities then + checkCudaCapabilities baseDrv kernel.checkCudaCapabilities else baseDrv; }) kernelOutputsList; diff --git a/examples/kernels/relu-archs-subset/CARD.md b/examples/kernels/relu-archs-subset/CARD.md new file mode 100644 index 00000000..92ecc33d --- /dev/null +++ b/examples/kernels/relu-archs-subset/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 %} diff --git a/examples/kernels/relu-archs-subset/build.toml b/examples/kernels/relu-archs-subset/build.toml new file mode 100644 index 00000000..a86f45a4 --- /dev/null +++ b/examples/kernels/relu-archs-subset/build.toml @@ -0,0 +1,45 @@ +[general] +name = "relu-archs-subset" +version = 1 +license = "Apache-2.0" +backends = [ + "cuda", + "rocm", +] + +[general.hub] +repo-id = "kernels-test/relu-archs-subset" + +[torch] +src = [ + "torch-ext/torch_binding.cpp", + "torch-ext/torch_binding.h", +] + +# The capability/architecture sets below intentionally contain entries that +# are not supported by the toolchain versions used in CI (e.g. CUDA 10.0/12.0 +# capabilities require CUDA >= 12.8, gfx940 is not in the supported arch set). +# CI checks that the built kernels contain exactly the intersection with the +# supported sets. + +[kernel.relu_rocm] +backend = "rocm" +depends = ["torch"] +rocm-archs = [ + "gfx90a", + "gfx940", + "gfx942", + "gfx1100", +] +src = ["relu_cuda/relu.cu"] + +[kernel.relu] +backend = "cuda" +depends = ["torch"] +cuda-capabilities = [ + "5.0", + "7.5", + "8.0", + "9.0", +] +src = ["relu_cuda/relu.cu"] diff --git a/examples/kernels/relu-archs-subset/flake.nix b/examples/kernels/relu-archs-subset/flake.nix new file mode 100644 index 00000000..416b2df0 --- /dev/null +++ b/examples/kernels/relu-archs-subset/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for ReLU kernel with a subset of CUDA/ROCm archs"; + + inputs = { + kernel-builder.url = "path:../../.."; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/examples/kernels/relu-archs-subset/relu_cuda/relu.cu b/examples/kernels/relu-archs-subset/relu_cuda/relu.cu new file mode 100644 index 00000000..50b49f91 --- /dev/null +++ b/examples/kernels/relu-archs-subset/relu_cuda/relu.cu @@ -0,0 +1,47 @@ +#include +#include +#include + +#include + +__global__ void relu_kernel(float *__restrict__ out, + float const *__restrict__ input, const int d) { + const int64_t token_idx = blockIdx.x; + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + auto x = input[token_idx * d + idx]; + out[token_idx * d + idx] = x > 0.0f ? x : 0.0f; + } +} + +void relu(torch::Tensor &out, torch::Tensor const &input) { + TORCH_CHECK(input.device().is_cuda(), "input must be a CUDA tensor"); + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(input.scalar_type() == at::ScalarType::Float && + input.scalar_type() == at::ScalarType::Float, + "relu_kernel only supports float32"); + + TORCH_CHECK(input.sizes() == out.sizes(), + "Tensors must have the same shape. Got input shape: ", + input.sizes(), " and output shape: ", out.sizes()); + + TORCH_CHECK(input.scalar_type() == out.scalar_type(), + "Tensors must have the same data type. Got input dtype: ", + input.scalar_type(), " and output dtype: ", out.scalar_type()); + + TORCH_CHECK(input.device() == out.device(), + "Tensors must be on the same device. Got input device: ", + input.device(), " and output device: ", out.device()); + + if (input.numel() == 0) { + return; + } + + int d = input.size(-1); + int64_t num_tokens = input.numel() / d; + dim3 grid(num_tokens); + dim3 block(std::min(d, 1024)); + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + relu_kernel<<>>(out.data_ptr(), + input.data_ptr(), d); +} diff --git a/examples/kernels/relu-archs-subset/torch-ext/relu_archs_subset/__init__.py b/examples/kernels/relu-archs-subset/torch-ext/relu_archs_subset/__init__.py new file mode 100644 index 00000000..a844aae8 --- /dev/null +++ b/examples/kernels/relu-archs-subset/torch-ext/relu_archs_subset/__init__.py @@ -0,0 +1,15 @@ +from typing import Optional + +import torch + +from ._ops import ops + + +def relu(x: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor: + if out is None: + out = torch.empty_like(x) + ops.relu(out, x) + return out + + +__all__ = ["relu"] diff --git a/examples/kernels/relu-archs-subset/torch-ext/torch_binding.cpp b/examples/kernels/relu-archs-subset/torch-ext/torch_binding.cpp new file mode 100644 index 00000000..1765d92d --- /dev/null +++ b/examples/kernels/relu-archs-subset/torch-ext/torch_binding.cpp @@ -0,0 +1,19 @@ +#include + +#include "registration.h" +#include "torch_binding.h" + +TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + ops.def("relu(Tensor! out, Tensor input) -> ()"); +#if defined(CPU_KERNEL) + ops.impl("relu", torch::kCPU, &relu); +#elif defined(CUDA_KERNEL) || defined(ROCM_KERNEL) + ops.impl("relu", torch::kCUDA, &relu); +#elif defined(METAL_KERNEL) + ops.impl("relu", torch::kMPS, relu); +#elif defined(XPU_KERNEL) + ops.impl("relu", torch::kXPU, &relu); +#endif +} + +REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/examples/kernels/relu-archs-subset/torch-ext/torch_binding.h b/examples/kernels/relu-archs-subset/torch-ext/torch_binding.h new file mode 100644 index 00000000..3bcf2904 --- /dev/null +++ b/examples/kernels/relu-archs-subset/torch-ext/torch_binding.h @@ -0,0 +1,5 @@ +#pragma once + +#include + +void relu(torch::Tensor &out, torch::Tensor const &input); \ No newline at end of file diff --git a/examples/kernels/relu/build.toml b/examples/kernels/relu/build.toml index 6b9f4d09..6aa04405 100644 --- a/examples/kernels/relu/build.toml +++ b/examples/kernels/relu/build.toml @@ -36,18 +36,6 @@ src = [ [kernel.relu_rocm] backend = "rocm" depends = ["torch"] -rocm-archs = [ - "gfx906", - "gfx908", - "gfx90a", - "gfx942", - "gfx950", - "gfx1030", - "gfx1100", - "gfx1101", - "gfx1200", - "gfx1201", -] src = ["relu_cuda/relu.cu"] [kernel.relu_cpu]