Repository navigation
CUDA backend: no fallback kernels for nonzero and masked_select #24767
Description
What happens
A program that calls torch.nonzero or torch.masked_select directly cannot be lowered with the CUDA backend. Export stops with:
RuntimeError: missing fallback kernels (1 total): - aoti_torch_cuda_nonzero Please add them to the AOTI backend.
and the same with aoti_torch_cuda_masked_select. It fails loudly, never with wrong results. Plain boolean-mask indexing (x[x > 0]) exports fine, so only the direct calls hit it. Tools that trace models into explicit ops emit these direct calls for data-dependent shapes.
Repro
import torch
from executorch.backends.cuda.cuda_backend import CudaBackend
from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
class M(torch.nn.Module):
def forward(self, x):
return torch.nonzero(x > 0).sum() # or torch.masked_select(x, x > 0).sum()
x = torch.tensor([[1.0, -2.0, 3.0], [-4.0, 5.0, -6.0]], device="cuda")
ep = torch.export.export(M().eval().cuda(), (x,))
to_edge_transform_and_lower(
ep,
partitioner=[CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")])],
compile_config=EdgeCompileConfig(_check_ir_validity=False),
).to_executorch()Tested with executorch 1.6.0+cu132 and torch 2.15.0+cu132 on an H100.
Ask
Add aoti_torch_cuda_nonzero and aoti_torch_cuda_masked_select shims to the CUDA backend, like the existing fallback shims.
Related: #24684.
Metadata
Metadata
Assignees
Labels
No labels
Activity
Gasoonjia commentedon Nov 9, 2026
Thansk for sharing!
#24787
#24788
These two are trying to address the issue.