Skip to content

Commit 6d4a2cd

Browse files
zhzhcookiezhzhcookie
committed
[HINT] Fix backend selection in hint manager
1 parent c1ea828 commit 6d4a2cd

2 files changed

Lines changed: 14 additions & 123 deletions

File tree

‎python/triton/compiler/hint_manager.py‎

Lines changed: 7 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
import sys
2424
import importlib
2525

26+
from triton._flagtree_backend import FLAGTREE_BACKEND
27+
2628

2729
class BaseHintHandler:
2830
# dynamicly find method
@@ -62,7 +64,7 @@ def __init__(self, backend_name):
6264
self.handler = self._load_handler(backend_name)
6365

6466
def _load_handler(self, backend):
65-
if backend == 'npu':
67+
if backend == 'ascend':
6668
try:
6769
module = importlib.import_module("triton.backends.ascend.ascend_hint_handler")
6870
return module.AscendHintHandler()
@@ -76,7 +78,7 @@ def _load_handler(self, backend):
7678
except ImportError as e:
7779
print(f"[FlagTree] Warning: Failed to load aipu Hint Handler: {e}", file=sys.stderr)
7880
return BaseHintHandler()
79-
elif backend == 'cuda':
81+
elif backend == 'nvidia':
8082
try:
8183
module = importlib.import_module("triton.backends.nvidia.nvidia_hint_handler")
8284
return module.NvidiaHintHandler()
@@ -94,66 +96,6 @@ def _load_handler(self, backend):
9496
return BaseHintHandler()
9597

9698

97-
# supported backend with matched version
98-
SUPPORTED_BACKENDS = ["aipu", "npu", "cuda", "sunrise"]
99-
100-
# TODO : npu will have conflicts if more backend involved
101-
# mapping name
102-
BACKEND_ALIASES = {
103-
"ascend": "npu",
104-
"huawei": "npu",
105-
"nvidia": "cuda",
106-
# sunrise: GPUTarget backend name is "tang", torch device type is "ptpu".
107-
"tang": "sunrise",
108-
"ptpu": "sunrise",
109-
}
110-
111-
112-
def normalize_backend_name(name: str) -> str:
113-
if not name:
114-
return ""
115-
name = name.lower()
116-
return BACKEND_ALIASES.get(name, name)
117-
118-
119-
def hint_get_flagtree_backend() -> str:
120-
detected_backend = ""
121-
122-
# Priority 1: Triton Driver
123-
try:
124-
import torch
125-
from triton.runtime import driver
126-
if hasattr(driver, 'active') and hasattr(driver.active, 'get_active_torch_device'):
127-
device = driver.active.get_active_torch_device()
128-
if isinstance(device, torch.device):
129-
detected_backend = device.type
130-
# unimplemented support
131-
elif isinstance(device, str):
132-
detected_backend = device
133-
except ImportError:
134-
return ""
135-
136-
# TODO : some backend may not support priority 1, so keep priority 2 is necessary
137-
# Priority 2: Torch Global State
138-
if not detected_backend:
139-
check_priority = ["aipu", "npu", "cuda"]
140-
141-
# 3. parse according to benefit
142-
for candidate in check_priority:
143-
module = getattr(torch, candidate, None)
144-
if module and hasattr(module, "is_available") and module.is_available():
145-
detected_backend = candidate
146-
break
147-
148-
# (Normalization and Validation)
149-
canonical_backend = normalize_backend_name(detected_backend)
150-
151-
if not canonical_backend or canonical_backend not in SUPPORTED_BACKENDS:
152-
return ""
153-
154-
return canonical_backend
155-
156-
15799
# lazy load after first call hint trigger
158100
_global_hint_manager = None
159101

@@ -162,5 +104,7 @@ def hint_trigger(hook_name, *args, **kwargs):
162104
global _global_hint_manager
163105

164106
if _global_hint_manager is None:
165-
_global_hint_manager = HintManager(hint_get_flagtree_backend())
107+
# NVIDIA builds have no FlagTree backend marker.
108+
backend_name = FLAGTREE_BACKEND or "nvidia"
109+
_global_hint_manager = HintManager(backend_name)
166110
return _global_hint_manager.handler.trigger(hook_name, *args, **kwargs)

‎third_party/xpu/python/triton/compiler/hint_manager.py‎

Lines changed: 7 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import sys
22
import importlib
33

4+
from triton._flagtree_backend import FLAGTREE_BACKEND
5+
46

57
class BaseHintHandler:
68
# dynamicly find method
@@ -40,7 +42,7 @@ def __init__(self, backend_name):
4042
self.handler = self._load_handler(backend_name)
4143

4244
def _load_handler(self, backend):
43-
if backend == 'npu':
45+
if backend == 'ascend':
4446
try:
4547
module = importlib.import_module("triton.backends.ascend.ascend_hint_handler")
4648
return module.AscendHintHandler()
@@ -54,7 +56,7 @@ def _load_handler(self, backend):
5456
except ImportError as e:
5557
print(f"[FlagTree] Warning: Failed to load aipu Hint Handler: {e}", file=sys.stderr)
5658
return BaseHintHandler()
57-
elif backend == 'cuda':
59+
elif backend == 'nvidia':
5860
try:
5961
module = importlib.import_module("triton.backends.nvidia.nvidia_hint_handler")
6062
return module.NvidiaHintHandler()
@@ -65,63 +67,6 @@ def _load_handler(self, backend):
6567
return BaseHintHandler()
6668

6769

68-
# supported backend with matched version
69-
SUPPORTED_BACKENDS = ["aipu", "npu", "cuda"]
70-
71-
# TODO : npu will have conflicts if more backend involved
72-
# mapping name
73-
BACKEND_ALIASES = {
74-
"ascend": "npu",
75-
"huawei": "npu",
76-
"nvidia": "cuda",
77-
}
78-
79-
80-
def normalize_backend_name(name: str) -> str:
81-
if not name:
82-
return ""
83-
name = name.lower()
84-
return BACKEND_ALIASES.get(name, name)
85-
86-
87-
def hint_get_flagtree_backend() -> str:
88-
detected_backend = ""
89-
90-
# Priority 1: Triton Driver
91-
try:
92-
import torch
93-
from triton.runtime import driver
94-
if hasattr(driver, 'active') and hasattr(driver.active, 'get_active_torch_device'):
95-
device = driver.active.get_active_torch_device()
96-
if isinstance(device, torch.device):
97-
detected_backend = device.type
98-
# unimplemented support
99-
elif isinstance(device, str):
100-
detected_backend = device
101-
except ImportError:
102-
return ""
103-
104-
# TODO : some backend may not support priority 1, so keep priority 2 is necessary
105-
# Priority 2: Torch Global State
106-
if not detected_backend:
107-
check_priority = ["aipu", "npu", "cuda"]
108-
109-
# 3. parse according to benefit
110-
for candidate in check_priority:
111-
module = getattr(torch, candidate, None)
112-
if module and hasattr(module, "is_available") and module.is_available():
113-
detected_backend = candidate
114-
break
115-
116-
# (Normalization and Validation)
117-
canonical_backend = normalize_backend_name(detected_backend)
118-
119-
if not canonical_backend or canonical_backend not in SUPPORTED_BACKENDS:
120-
return ""
121-
122-
return canonical_backend
123-
124-
12570
# lazy load after first call hint trigger
12671
_global_hint_manager = None
12772

@@ -130,5 +75,7 @@ def hint_trigger(hook_name, *args, **kwargs):
13075
global _global_hint_manager
13176

13277
if _global_hint_manager is None:
133-
_global_hint_manager = HintManager(hint_get_flagtree_backend())
78+
# NVIDIA builds have no FlagTree backend marker.
79+
backend_name = FLAGTREE_BACKEND or "nvidia"
80+
_global_hint_manager = HintManager(backend_name)
13481
return _global_hint_manager.handler.trigger(hook_name, *args, **kwargs)

0 commit comments

Comments
 (0)