Silently bypass bmm autotuning with out_dtype argument

PaulZhang12 · PaulZhang12 · commit b002365a38c0 · 2025-10-28T15:55:46.000-07:00
ghstack-source-id: b37202c Pull Request resolved: #166457
diff --git a/torch/_inductor/kernel/bmm.py b/torch/_inductor/kernel/bmm.py
@@ -239,9 +239,8 @@ def _to_dtype(x):
         templates_to_use.append(aten_handler)
         kwarg_overrides[aten_handler.uid] = aten_extra_kwargs
 
-    if use_triton_template(layout, check_max_autotune=False):
+    if use_triton_template(layout, check_max_autotune=False) and (out_dtype is None or out_dtype == mat1.get_dtype()):
         # TODO: add out_dtype support for Triton Template
-        assert out_dtype is None, "out_dtype is not supported for Triton"
         templates_to_use.append(bmm_template)
 
     # Single unified call for all templates
