Key CUDA-capability lookup by device index instead of process-wide lru_cache

[AmirF194]

Description

_find_capability() in kernels/src/kernels/layer/repos.py (used by both
_CUDARepos.repos and _ROCMRepos.repos to pick which LayerRepository a kernelize()
call should dispatch to) is a bare @functools.lru_cache function with no arguments:

@lru_cache
def _find_capability() -> int:
    import torch
    major, minor = torch.cuda.get_device_capability(device=None)
    return major * 10 + minor

Since it takes no arguments and is never invalidated, it memoizes whatever
torch.cuda.get_device_capability(device=None) returns on its first call for the
rest of the process, and device=None means "whatever CUDA device is current right
now". Any process that kernelizes modules living on more than one physical GPU of
different compute capability in one process, for example a model sharded across
heterogeneous GPUs with device_map="auto", or any code that calls
torch.cuda.set_device(n) between two kernelize() calls targeting different
devices, gets the wrong capability (and therefore the wrong LayerRepository) for
every device after the first one queried. Nothing errors or warns.

I traced this by reading repos.py / kernelize.py / layer.py and confirmed
kernelize() never passes a specific device index into the capability lookup either
(_find_device only derives a device type). This looks related to the intent behind
CUDAProperties.min_capability/max_capability that @danieldk described in #707
("the kernel mapping for kernelize supports CUDAProperties which can indicate the
capabilities supported by that kernel for that particular mapping") since that whole
mechanism silently breaks once more than one capability shows up in a process.

Steps to reproduce

I don't have access to a multi-GPU box, so I reproduced the caching behavior directly
against the real repos.py at HEAD (8ac69fc5) with torch.cuda.get_device_capability
mocked to return two different values on successive calls (standing in for
torch.cuda.set_device() switching between two physical GPUs):

import sys, types, importlib.util, importlib.machinery

fake_torch = types.ModuleType("torch")
fake_torch.__spec__ = importlib.machinery.ModuleSpec("torch", loader=None)
fake_cuda = types.ModuleType("torch.cuda")
state = {"capability": (7, 5)}
fake_cuda.get_device_capability = lambda device=None: state["capability"]
fake_torch.cuda = fake_cuda
sys.modules["torch"], sys.modules["torch.cuda"] = fake_torch, fake_cuda

# ... load the real kernels.layer.repos module ...

cuda_repos = repos._CUDARepos()
cuda_repos.insert(repos.Device(type="cuda", properties=repos.CUDAProperties(min_capability=70, max_capability=79)), {"variant": "sm75"})
cuda_repos.insert(repos.Device(type="cuda", properties=repos.CUDAProperties(min_capability=90, max_capability=99)), {"variant": "sm90"})

print(cuda_repos.repos)          # {'variant': 'sm75'}, correct for cc 7.5
state["capability"] = (9, 0)     # process now on a cc 9.0 device
print(cuda_repos.repos)          # still {'variant': 'sm75'} -- stale

Output:

{'variant': 'sm75'}
{'variant': 'sm75'}

Expected behavior

The second call should resolve to the sm90 repository, since torch.cuda's current
device capability changed. Instead DeviceRepos.repos keeps returning the first
device's repository for the rest of the process.

Environment

• kernels version: HEAD, 8ac69fc5 (2026-08-20)
• Python version: 3.12
• PyTorch version: n/a, mocked (verified on a CPU-only host, no multi-GPU hardware
  available to me to hit this with a real torch.cuda)
• CUDA version: n/a
• GPU model: n/a
• OS: Linux (Docker python:3.12-slim)

Additional context

I only verified the caching mechanism itself (shown above, against the real
_CUDARepos/_find_capability code), not a live heterogeneous-GPU run, since I don't
have that hardware. The fix likely needs a design call between two shapes: keying
_find_capability by torch.cuda.current_device() (cheap, still wrong if a module's
own parameters aren't on the "current" device, which is the common device_map="auto"
case), or threading a resolved device index from kernelize_layer down into
DeviceRepos.repos so capability comes from the actual module being kernelized. Happy
to send a PR once there's agreement on which of those you'd want; opening this first per
CONTRIBUTING.md.

This report and the repro above were put together with the help of an LLM coding agent;
I read and verified the reasoning and the repro myself before posting.

[sayakpaul]

This isn't supposed to be a user-facing thing anyway. But we could probably log in the following case:

Since it takes no arguments and is never invalidated, it memoizes whatever
torch.cuda.get_device_capability(device=None) returns on its first call for the
rest of the process, and device=None means "whatever CUDA device is current right
now". Any process that kernelizes modules living on more than one physical GPU of
different compute capability in one process, for example a model sharded across
heterogeneous GPUs with device_map="auto", or any code that calls
torch.cuda.set_device(n) between two kernelize() calls targeting different
devices, gets the wrong capability (and therefore the wrong LayerRepository) for
every device after the first one queried. Nothing errors or warns.

@danieldk thoughts?

[AmirF194]

Logging would surface it but wouldn't fix the dispatch itself, since the wrong LayerRepository still gets used either way. Between the two shapes in the issue, keying by torch.cuda.current_device() is the cheap patch but stays wrong under device_map="auto", where a module's parameters are not necessarily on the current device. Threading the resolved device index from kernelize_layer down into DeviceRepos.repos is the only one of the two that is correct for that case too.

[DsThakurRawat]

read through this against 8c59f17 because the two proposed shapes seemed to be pulling in different directions, and there's a third thing underneath both of them that i think decides which one is even possible.

Device has nowhere to put a device index. it's Device(type: str, properties: CUDAProperties | ROCMProperties | None), frozen, @strict, and exported in __all__. so "thread a resolved device index down into DeviceRepos.repos" isn't only plumbing, it needs somewhere to live first, either a new field on a public frozen dataclass or a second value carried alongside it through kernelize_layer.

related, and the reason the cheap patch is cheaper than it looks: kernelize() resolves the device exactly once, at lines 249-255, before the for _, module in model.named_modules() loop, and then hands that same device_type to every kernelize_layer call. so there is currently no per-module device anywhere in the call path to thread. _find_device also only looks at next(model.parameters()), the first parameter of the whole model, and keeps just param.device.type. under device_map="auto" that's one parameter's device standing in for every module in the model, which is the same failure as the lru_cache one layer up, just earlier.

the part i'd push back on re "this isn't supposed to be a user-facing thing anyway": the documented escape hatch is user-facing and it drops the index too. kernelize(model, device=...) accepts a torch.device, and line 255 does Device(device.type). i ran that path with a stubbed torch:

user passed      : device(type='cuda', index=1)
kernelize builds : Device(type='cuda', properties=None)
index survived?  : False

so someone who already knows about this and tries to be explicit about which GPU they mean gets the same wrong dispatch as someone who doesn't. that's what makes logging-only feel thin to me, it would report a condition the user has no supported way to avoid, since both the inferred and the explicit path discard the index before it reaches _find_capability.

if it helps narrow the call: keying by torch.cuda.current_device() fixes the set_device() case and nothing else, and it'd still be wrong under device_map="auto" where the current device has no particular relationship to where a given module's parameters live. the per-module resolution is the one that's actually correct, and it looks like it needs _find_device to move inside the module loop and return an index alongside the type, which is a bigger change than the issue title suggests but is contained to kernelize.py plus whatever carries the index.

same caveat as @AmirF194 here, no multi-GPU box, so the above is from reading the code and running the pure-python paths with torch stubbed, not from a real heterogeneous run.

[danieldk]

We have to stop caching the capability or maybe even pass through the device index so that the capability can be looked up for a particular device. I have to think a bit more about this and maybe also discuss with the transformers team.

We haven't encountered this much yet because most people are running homogeneous GPUs in multi-GPU scenarios, but I think it's worth supporting.