Let the CUDA backend autotune matmuls on GPUs with fewer than 68 SMs (pytorch#23352)

## What is wrong today

The CUDA backend compiles a model with AOTInductor and asks it to
autotune every matrix multiply and convolution with Triton only. It does
not allow the ATen fallback, because that would pull libtorch operators
into the compiled library.

Inductor only offers its Triton matmul and convolution templates on GPUs
it considers big enough: 68 or more streaming multiprocessors (SMs). On
a smaller GPU the Triton templates are skipped. A matmul then has no
choices left, because nothing else is allowed:

```
NoValidChoicesError: No choices to select. Provided reason: No choices exist for backend.
  target: aten.addmm.default
```

A convolution falls back to ATen instead, which this backend then
rejects as a missing fallback kernel (`aoti_torch_cuda_convolution`).

So on any CUDA GPU with fewer than 68 SMs, even a model with a single
`nn.Linear` cannot be lowered with `CudaPartitioner`. Today the only way
around it is to patch `torch._inductor.utils.is_big_gpu` in your own
export script.

## What this change does

While the backend's own AOTInductor compile runs, inductor treats the
GPU as big enough, so it considers the Triton templates it would
otherwise skip. The autotuner then picks the fastest Triton choice as
usual.

The override follows the pattern the file already uses for its other
inductor tweaks:

- It is a context manager, `_allow_triton_templates_on_small_gpus`,
entered in `get_extra_aoti_compile_context_manager` next to the existing
ones.
- It keeps the original `is_big_gpu` in a closure and puts it back when
the compile ends, even if the compile fails.
- It answers only on the thread doing the compile. Other threads get the
original answer.
- It is skipped on ROCm, where `is_big_gpu` checks the GPU architecture
instead of the SM count.

GPUs with 68 or more SMs already pass the check, so nothing changes for
them. The Triton-only backend list stays as it is.

## What was tested

- New unit tests in `test_cuda_export.py`, which need no small GPU: with
`is_big_gpu` reporting a small GPU, Inductor's own template check says
yes inside the backend's compile context and no outside it; the override
is skipped on ROCm; the original function and the thread's state are
back after a failed compile; a nested compile keeps the override until
the outer one ends; and another thread gets the original answer. Each
one fails when the part it covers is removed.
- On a CUDA GPU with 20 SMs (Linux aarch64), a module with one
`nn.Linear` lowered with `CudaPartitioner`, with no patch in the caller:
before, `NoValidChoicesError`; after, the `.pte` and `.ptd` files are
written.
- On the same GPU, `test_cuda_export.py`, `test_cuda_partitioner.py` and
`test_cuda_weight_metadata.py`: before, 6 export tests fail (linear, a
multi-operation module and SDPA with Triton kernels off with
`NoValidChoicesError`; conv1d, conv2d and a ResNet block with the
missing convolution kernel); after, they pass. A lowered `nn.Linear` and
`Conv2d` also give the same results as eager PyTorch.

Author: shoumikhin

backends/cuda/cuda_backend.py

@@ -113,6 +113,40 @@ def _has_rmask(_self) -> bool:
         IndexingOptions.has_rmask = orig_has_rmask
 
 
+_SMALL_GPU_GUARD = threading.local()
+
+
+@contextlib.contextmanager
+def _allow_triton_templates_on_small_gpus():
+    """Let Inductor use Triton templates on GPUs below its 68-SM floor. This
+    backend autotunes with Triton only, so otherwise a matmul or convolution
+    has no Triton choices left.
+
+    Only threads inside this context get the override; other threads get the
+    original answer.
+
+    TODO: remove once Inductor skips the floor when ATen is not an allowed
+    backend (https://github.com/pytorch/pytorch/issues/141690).
+    """
+    from torch._inductor import utils as inductor_utils
+
+    orig_is_big_gpu = inductor_utils.is_big_gpu
+
+    def _is_big_gpu(index_or_device=0) -> bool:
+        return getattr(_SMALL_GPU_GUARD, "active", False) or orig_is_big_gpu(
+            index_or_device
+        )
+
+    prev_active = getattr(_SMALL_GPU_GUARD, "active", False)
+    inductor_utils.is_big_gpu = _is_big_gpu
+    _SMALL_GPU_GUARD.active = True
+    try:
+        yield
+    finally:
+        _SMALL_GPU_GUARD.active = prev_active
+        inductor_utils.is_big_gpu = orig_is_big_gpu
+
+
 def _full_zeros_preserving_strides(x: torch.Tensor, device) -> torch.Tensor:
     """Allocate a zero-filled tensor matching ``x``'s size/stride/dtype on ``device``.
 
@@ -904,8 +938,10 @@ def get_extra_aoti_compile_context_manager(
         compilation for the CUDA backend. Each manager is documented at
         its own `enter_context` call site below.
 
-        The optional shared-memory target applies to every CUDA compilation.
-        The low-memory export monkey-patch (CPU clones for mutated buffers)
+        The optional shared-memory target applies to every CUDA compilation,
+        and so does the small-GPU override that lets Inductor use Triton
+        templates below 68 SMs (skipped on ROCm). The low-memory export
+        monkey-patch (CPU clones for mutated buffers)
         is gated on the ``low_memory_mode`` compile spec — only models that
         explicitly opt in (currently Qwen3.5 MoE) get it. Other models go
         through the unmodified AOTI codepath, which avoids regressions in
@@ -932,6 +968,9 @@ def _combined():
                 # `ReplaceEdgeOpWithTritonOpPass` are unaffected.
                 stack.enter_context(torch.nn.attention.sdpa_kernel([SDPBackend.MATH]))
                 stack.enter_context(target_smem_context())
+                # On ROCm, is_big_gpu gates on architecture, not SM count.
+                if torch.version.hip is None:
+                    stack.enter_context(_allow_triton_templates_on_small_gpus())
                 if low_memory_mode == "ON":
                     # Force AOTI's mutated-buffer clones onto CPU during
                     # compile so we stay under tight GPU memory caps (e.g.

backends/cuda/tests/test_cuda_export.py

@@ -6,6 +6,7 @@
 
 import contextlib
 import os
+import threading
 import unittest
 from typing import Tuple
 from unittest import mock
@@ -133,6 +134,81 @@ def test_target_smem_context_only_patches_exact_triton_limit(self):
 
                 self.assertIs(triton_compiler.max_shared_mem, local_max_shared_mem)
 
+    def test_triton_templates_allowed_on_small_gpus_during_compile(self):
+        from torch._inductor import utils as inductor_utils
+        from torch._inductor.ir import FixedLayout
+
+        layout = FixedLayout(torch.device("cuda", 0), torch.float32, [64, 64])
+        with mock.patch.object(
+            inductor_utils, "is_big_gpu", return_value=False
+        ) as is_big_gpu, mock.patch.object(torch.version, "hip", None):
+            self.assertFalse(
+                inductor_utils._use_template_for_gpu(layout, [torch.float32])
+            )
+            with CudaBackend.get_extra_aoti_compile_context_manager([]):
+                self.assertTrue(
+                    inductor_utils._use_template_for_gpu(layout, [torch.float32])
+                )
+            self.assertIs(inductor_utils.is_big_gpu, is_big_gpu)
+
+    def test_small_gpu_override_is_skipped_on_rocm(self):
+        from torch._inductor import utils as inductor_utils
+
+        with mock.patch.object(
+            inductor_utils, "is_big_gpu", return_value=False
+        ) as is_big_gpu, mock.patch.object(torch.version, "hip", "6.4.0"):
+            with CudaBackend.get_extra_aoti_compile_context_manager([]):
+                self.assertIs(inductor_utils.is_big_gpu, is_big_gpu)
+
+    def test_small_gpu_override_is_restored_when_a_compile_fails(self):
+        from executorch.backends.cuda.cuda_backend import (
+            _allow_triton_templates_on_small_gpus,
+        )
+        from torch._inductor import utils as inductor_utils
+
+        with mock.patch.object(
+            inductor_utils, "is_big_gpu", return_value=False
+        ) as is_big_gpu:
+            with self.assertRaises(RuntimeError):
+                with _allow_triton_templates_on_small_gpus():
+                    raise RuntimeError("compile failed")
+            self.assertIs(inductor_utils.is_big_gpu, is_big_gpu)
+            with _allow_triton_templates_on_small_gpus():
+                held = inductor_utils.is_big_gpu
+            self.assertFalse(held(0))
+
+    def test_small_gpu_override_survives_a_nested_compile(self):
+        from executorch.backends.cuda.cuda_backend import (
+            _allow_triton_templates_on_small_gpus,
+        )
+        from torch._inductor import utils as inductor_utils
+
+        with mock.patch.object(
+            inductor_utils, "is_big_gpu", return_value=False
+        ) as is_big_gpu:
+            with _allow_triton_templates_on_small_gpus():
+                with _allow_triton_templates_on_small_gpus():
+                    pass
+                self.assertTrue(inductor_utils.is_big_gpu(0))
+            self.assertIs(inductor_utils.is_big_gpu, is_big_gpu)
+
+    def test_small_gpu_override_answers_only_its_own_thread(self):
+        from executorch.backends.cuda.cuda_backend import (
+            _allow_triton_templates_on_small_gpus,
+        )
+        from torch._inductor import utils as inductor_utils
+
+        seen = {}
+        with mock.patch.object(inductor_utils, "is_big_gpu", return_value=False):
+            with _allow_triton_templates_on_small_gpus():
+                other = threading.Thread(
+                    target=lambda: seen.update(other=inductor_utils.is_big_gpu(0))
+                )
+                other.start()
+                other.join(timeout=10)
+                seen["inside"] = inductor_utils.is_big_gpu(0)
+        self.assertEqual({"inside": True, "other": False}, seen)
+
     def test_cuda_include_ptx_compile_spec(self):
         with mock.patch.object(
             CudaBackend, "_setup_cuda_environment_for_fatbin", return_value=True