From 7405b691421a24245f7593f0bd18d0a1e3f106cb Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 2 Sep 2026 12:03:30 +0530 Subject: [PATCH 1/4] feat: implement user_agent in layers --- kernels/src/kernels/layer/layer.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index fa34b8c8..d3e3855e 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -58,6 +58,8 @@ class LayerRepository: Whether to allow loading kernels from untrusted organisations. A list of signing identities can be provided for future verification support; until then it warns and falls back to the default trust check. + user_agent (`Union[str, dict]`, *optional*): + The `user_agent` info to pass to `snapshot_download()` for internal telemetry. Example: ```python @@ -80,6 +82,7 @@ def __init__( revision: str | None = None, version: int | None = None, trust_remote_code: bool | list[str] = False, + user_agent: str | dict | None = None, ): if revision is not None and version is not None: raise ValueError("Either a revision or a version must be specified, not both.") @@ -89,6 +92,7 @@ def __init__( self._repo_id = repo_id self.layer_name = layer_name self._trust_remote_code = trust_remote_code + self._user_agent = user_agent # We are going to resolve these lazily, since we do not want # to do a network request for every registered LayerRepository. @@ -109,6 +113,7 @@ def load(self) -> Type["nn.Module"]: self._repo_id, revision=self._resolve_revision(), trust_remote_code=self._trust_remote_code, + user_agent=self._user_agent ) return _get_kernel_layer(self, kernel) From e79767526f2edd528194fb43cef67eec7a666fd9 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 2 Sep 2026 12:09:25 +0530 Subject: [PATCH 2/4] add test and docs. --- docs/source/layers.md | 18 ++++++++++++++++++ kernels/src/kernels/layer/layer.py | 4 ++-- kernels/tests/test_layer.py | 30 ++++++++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/docs/source/layers.md b/docs/source/layers.md index e1112041..b2e5d9dc 100644 --- a/docs/source/layers.md +++ b/docs/source/layers.md @@ -212,6 +212,24 @@ with use_kernel_mapping(kernel_layer_mapping): This ensures that the mapping is not active anymore outside the `with`-scope. +### Attributing Hub requests + +Libraries and applications can identify the Hub requests made while loading a +layer by passing optional user-agent metadata to [`~kernels.LayerRepository`]: + +```python +layer_repo = LayerRepository( + repo_id="kernels-community/activation", + layer_name="SiluAndMul", + version=1, + user_agent={"my-library": "1.0.0"}, +) +``` + +`user_agent` accepts either a string or a dictionary. If it is omitted, no +application-specific metadata is added. Setting `HF_HUB_DISABLE_TELEMETRY=1` +disables user-agent telemetry, including metadata supplied this way. + If the layer is stateless (it does not use member variables in its forward _or_ it was originally a function that was converted into a kernel layer with [`~kernels.use_kernel_func_from_hub`]), it can also be mapped to a kernel function: diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index d3e3855e..4a23b3d7 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -59,7 +59,7 @@ class LayerRepository: of signing identities can be provided for future verification support; until then it warns and falls back to the default trust check. user_agent (`Union[str, dict]`, *optional*): - The `user_agent` info to pass to `snapshot_download()` for internal telemetry. + Optional application metadata to include in the user-agent for Hub requests. Example: ```python @@ -113,7 +113,7 @@ def load(self) -> Type["nn.Module"]: self._repo_id, revision=self._resolve_revision(), trust_remote_code=self._trust_remote_code, - user_agent=self._user_agent + user_agent=self._user_agent, ) return _get_kernel_layer(self, kernel) diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py index 699c4a11..0e6239d6 100644 --- a/kernels/tests/test_layer.py +++ b/kernels/tests/test_layer.py @@ -1,5 +1,6 @@ import sys from contextlib import nullcontext +from types import SimpleNamespace import pytest import torch @@ -594,6 +595,35 @@ def test_layer_repository_requires_version_or_revision(): LayerRepository(repo_id="kernels-test/silu-and-mul", layer_name="SiluAndMul") +@pytest.mark.parametrize("user_agent", [None, "transformers/5.0.0", {"transformers": "5.0.0"}]) +def test_layer_repository_forwards_user_agent(monkeypatch, user_agent): + calls = [] + + def get_kernel(repo_id, **kwargs): + calls.append((repo_id, kwargs)) + return SimpleNamespace(layers=SimpleNamespace(SiluAndMul=SiluAndMul)) + + monkeypatch.setattr("kernels.layer.layer.get_kernel", get_kernel) + repo = LayerRepository( + repo_id="kernels-test/silu-and-mul", + layer_name="SiluAndMul", + revision="main", + user_agent=user_agent, + ) + + assert repo.load() is SiluAndMul + assert calls == [ + ( + "kernels-test/silu-and-mul", + { + "revision": "main", + "trust_remote_code": False, + "user_agent": user_agent, + }, + ) + ] + + def test_validate_kernel_layer(): class BadLayer(nn.Module): def __init__(self, *args, **kwargs): From 001b24543a5f19331913ccd7ff7f60de1759d7d3 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 3 Sep 2026 08:02:23 +0530 Subject: [PATCH 3/4] style --- kernels/tests/test_layer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py index 371b8688..f8fa41c5 100644 --- a/kernels/tests/test_layer.py +++ b/kernels/tests/test_layer.py @@ -622,6 +622,8 @@ def get_kernel(repo_id, **kwargs): }, ) ] + + def test_layer_repository_with_trust_remote_code_allowlist_is_hashable(): allowlist = ["untrusted-org/allowed-kernel"] repo = LayerRepository( From 7e4a0776c56c09b3499429fd708cda71a620a653 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 3 Sep 2026 08:14:13 +0530 Subject: [PATCH 4/4] nix formatting --- nix-builder/lib/extension/torch/arch.nix | 21 ++++++++++----------- nix-builder/lib/extension/tvm-ffi/arch.nix | 21 ++++++++++----------- 2 files changed, 20 insertions(+), 22 deletions(-) diff --git a/nix-builder/lib/extension/torch/arch.nix b/nix-builder/lib/extension/torch/arch.nix index f9f5955e..d2c3c82e 100644 --- a/nix-builder/lib/extension/torch/arch.nix +++ b/nix-builder/lib/extension/torch/arch.nix @@ -244,17 +244,16 @@ stdenv.mkDerivation (prevAttrs: { ] ++ extraDeps; - env = - { - inherit variant; - } - // lib.optionalAttrs cudaSupport { - CUDAToolkit_ROOT = "${lib.getDev cudaPackages.cuda_nvcc}"; - } - // lib.optionalAttrs xpuSupport { - MKLROOT = oneapi-torch-dev; - SYCL_ROOT = oneapi-torch-dev; - }; + env = { + inherit variant; + } + // lib.optionalAttrs cudaSupport { + CUDAToolkit_ROOT = "${lib.getDev cudaPackages.cuda_nvcc}"; + } + // lib.optionalAttrs xpuSupport { + MKLROOT = oneapi-torch-dev; + SYCL_ROOT = oneapi-torch-dev; + }; # If we use the default setup, CMAKE_CUDA_HOST_COMPILER gets set to nixpkgs g++. dontSetupCUDAToolkitCompilers = true; diff --git a/nix-builder/lib/extension/tvm-ffi/arch.nix b/nix-builder/lib/extension/tvm-ffi/arch.nix index de5a58dd..c4b39c6a 100644 --- a/nix-builder/lib/extension/tvm-ffi/arch.nix +++ b/nix-builder/lib/extension/tvm-ffi/arch.nix @@ -236,17 +236,16 @@ stdenv.mkDerivation (prevAttrs: { ] ++ extraDeps; - env = - { - inherit variant; - } - // lib.optionalAttrs cudaSupport { - CUDAToolkit_ROOT = "${lib.getDev cudaPackages.cuda_nvcc}"; - } - // lib.optionalAttrs xpuSupport { - MKLROOT = oneapi-torch-dev; - SYCL_ROOT = oneapi-torch-dev; - }; + env = { + inherit variant; + } + // lib.optionalAttrs cudaSupport { + CUDAToolkit_ROOT = "${lib.getDev cudaPackages.cuda_nvcc}"; + } + // lib.optionalAttrs xpuSupport { + MKLROOT = oneapi-torch-dev; + SYCL_ROOT = oneapi-torch-dev; + }; # If we use the default setup, CMAKE_CUDA_HOST_COMPILER gets set to nixpkgs g++. dontSetupCUDAToolkitCompilers = true;