From c73d218f0078548bd9645c51d5769503d3d75c35 Mon Sep 17 00:00:00 2001 From: Songhao Jia Date: Thu, 12 Nov 2026 20:31:43 -0800 Subject: [PATCH] Lower grouped GEMM on sm_80 and above (#25017) Summary: Inductor offers its grouped GEMM Triton template only on sm_90 and above, and this backend autotunes with Triton only, so on an A100 (sm_80) a model using `aten._grouped_mm` fails to export with `NoValidChoicesError: No choices to select` (pytorch/executorch#24902). The template uses TMA only on sm_90 and above; its remaining path works on sm_80. During AOTInductor compilation, a `_CompilePatch` now reports `has_grouped_mm_triton_support()` as true when the current device is sm_80 or newer, following the existing small-GPU override for Inductor's 68-SM floor. The override is scoped to the threads inside the compile, is skipped on ROCm, and leaves devices below sm_80 refused. Reviewed By: shoumikhin Differential Revision: D131591208 --- backends/cuda/cuda_backend.py | 40 ++++++++++++++++ backends/cuda/tests/test_cuda_export.py | 61 +++++++++++++++++++++++++ 2 files changed, 101 insertions(+) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 64f0866a642..7b1ffae6b11 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -200,6 +200,45 @@ def remove() -> None: _SMALL_GPU_OVERRIDE = _CompilePatch(_install_small_gpu_override) +def _install_grouped_mm_override(active): + from torch._inductor.kernel import mm_grouped + + original = mm_grouped.has_grouped_mm_triton_support + + def has_grouped_mm_triton_support() -> bool: + if original(): + return True + if not active(): + return False + # The grouped GEMM Triton template uses TMA only on sm_90 and above; + # its remaining path works on sm_80. Inductor still offers the + # template only on sm_90 and above, and this backend allows no ATen + # choice, so below sm_90 a grouped GEMM has nothing to lower to. + return torch.cuda.is_available() and torch.cuda.get_device_capability() >= ( + 8, + 0, + ) + + mm_grouped.has_grouped_mm_triton_support = has_grouped_mm_triton_support + + def remove() -> None: + mm_grouped.has_grouped_mm_triton_support = original + + return remove + + +_GROUPED_MM_OVERRIDE = _CompilePatch(_install_grouped_mm_override) + + +def _allow_grouped_mm_triton_template() -> typing.ContextManager[None]: + """Let Inductor offer its grouped GEMM Triton template on sm_80 and above. + + Only threads inside this context get the override; other threads get the + original answer. + """ + return _GROUPED_MM_OVERRIDE.scope() + + def _allow_triton_templates_on_small_gpus() -> typing.ContextManager[None]: """Let Inductor use Triton templates on GPUs below its 68-SM floor. This backend autotunes with Triton only, so otherwise a matmul or convolution @@ -1464,6 +1503,7 @@ def _combined(): # 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()) + stack.enter_context(_allow_grouped_mm_triton_template()) 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. diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 5861112dacc..20a32256d0f 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -175,6 +175,35 @@ def test_triton_templates_allowed_on_small_gpus_during_compile(self): ) self.assertIs(inductor_utils.is_big_gpu, is_big_gpu) + def test_grouped_mm_triton_template_allowed_from_sm80_during_compile(self): + from torch._inductor.kernel import mm_grouped + + original = mm_grouped.has_grouped_mm_triton_support + with mock.patch.object( + torch.cuda, "is_available", return_value=True + ), mock.patch.object( + torch.cuda, "get_device_capability", return_value=(8, 0) + ), mock.patch.object( + torch.version, "hip", None + ): + self.assertFalse(mm_grouped.has_grouped_mm_triton_support()) + with CudaBackend.get_extra_aoti_compile_context_manager([]): + self.assertTrue(mm_grouped.has_grouped_mm_triton_support()) + self.assertIs(mm_grouped.has_grouped_mm_triton_support, original) + + def test_grouped_mm_triton_template_stays_refused_below_sm80(self): + from torch._inductor.kernel import mm_grouped + + with mock.patch.object( + torch.cuda, "is_available", return_value=True + ), mock.patch.object( + torch.cuda, "get_device_capability", return_value=(7, 5) + ), mock.patch.object( + torch.version, "hip", None + ): + with CudaBackend.get_extra_aoti_compile_context_manager([]): + self.assertFalse(mm_grouped.has_grouped_mm_triton_support()) + def test_small_gpu_override_is_skipped_on_rocm(self): from torch._inductor import utils as inductor_utils @@ -443,6 +472,38 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: edge_program_manager = self._export_to_cuda_with_lower(module, inputs) self.assertIsNotNone(edge_program_manager, "Simple add operation export failed") + def test_grouped_mm_export(self): + """A grouped GEMM lowers below sm_90, where Inductor offers its + Triton template only from sm_90 and this backend allows no ATen + choice (pytorch/executorch#24902).""" + + class GroupedMm(torch.nn.Module): + def forward(self, x, weight, offs): + return torch._grouped_mm(x, weight, offs) + + torch.manual_seed(0) + experts, tokens_per_expert, in_dim, out_dim = 4, 32, 64, 64 + x = torch.randn( + experts * tokens_per_expert, + in_dim, + dtype=torch.bfloat16, + device="cuda", + ) + weight = torch.randn( + experts, in_dim, out_dim, dtype=torch.bfloat16, device="cuda" + ) + offs = torch.arange( + tokens_per_expert, + (experts + 1) * tokens_per_expert, + tokens_per_expert, + dtype=torch.int32, + device="cuda", + ) + edge_program_manager = self._export_to_cuda_with_lower( + GroupedMm(), (x, weight, offs) + ) + self.assertIsNotNone(edge_program_manager, "Grouped MM operation export failed") + def test_detach(self): """Edge turns detach into detach_copy, which has no CUDA kernel of its own."""