From f9139a56ad6213475a92270d7ed3c9ebbe2c4831 Mon Sep 17 00:00:00 2001
From: "zhiguo.qin" <zhiguo.qin@mthreads.com>
Date: Sun, 23 Aug 2026 15:24:56 +0800
Subject: [PATCH] compat: reuse CUDA Inductor heuristics for MUSA

---
 src/torchada/_patch.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/src/torchada/_patch.py b/src/torchada/_patch.py
index 1202009..f8b6a1e 100644
--- a/src/torchada/_patch.py
+++ b/src/torchada/_patch.py
@@ -108,6 +108,39 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
+@patch_function
+@requires_import("torch._inductor.template_heuristics.registry")
+def _patch_inductor_template_heuristics():
+    """Reuse CUDA Inductor template heuristics for CUDA-compatible MUSA templates."""
+    if not is_musa_platform():
+        return
+
+    import torch._inductor.template_heuristics.registry as registry
+
+    heuristic_registry = getattr(registry, "_TEMPLATE_HEURISTIC_REGISTRY", None)
+    if not isinstance(heuristic_registry, dict):
+        return
+
+    changed = False
+    for key, heuristic_class in list(heuristic_registry.items()):
+        if len(key) != 3:
+            continue
+        template_name, device_type, op_name = key
+        if device_type != "cuda":
+            continue
+        if not isinstance(template_name, str) or not template_name.startswith("triton::"):
+            continue
+        musa_key = (template_name, "musa", op_name)
+        if musa_key not in heuristic_registry:
+            heuristic_registry[musa_key] = heuristic_class
+            changed = True
+
+    if changed:
+        heuristic_cache = getattr(registry, "_HEURISTIC_CACHE", None)
+        if isinstance(heuristic_cache, dict):
+            heuristic_cache.clear()
+
+
 # Cache for translated device strings - avoids repeated string operations
 _device_str_cache = {}