[MUSA] Use MUSA-optimized operators in piecewise CUDA graph (#23633)

Signed-off-by: popsiclexu <zhenxuexu@gmail.com>
2026-05-11 17:55:27 -07:00
parent 74d70af09a
commit bfc2eda42d
10 changed files with 105 additions and 12 deletions
@@ -56,7 +56,7 @@ def cache_once(fn):

@cache_once
 def is_arch_support_pdl() -> bool:
-    if bool(torch.version.hip):
+    if getattr(torch.version, "hip", None) or getattr(torch.version, "musa", None):
        return False
    try:
        device = torch.cuda.current_device()