Lower grouped GEMM on sm_80 and above (#25017)

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.

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."""