Coverage for src/flag_gems/runtime/backend/_ascend/__init__.py: 79%

1from backend_utils import VendorInfoBase # noqa: E402 

2 

3from .utils import CORE_NUM # noqa: F401 

4 

5 

6def get_triton_extra_name(): 

7 try: 

8 import triton 

9 from packaging import version 

10 

11 if version.parse(triton.__version__) <= version.parse("3.2.0"): 

12 return "ascend" 

13 else: 

14 return "cann" 

15 except Exception: 

16 return "ascend" 

17 

18 

19vendor_info = VendorInfoBase( 

20 vendor_name="ascend", 

21 device_name="npu", 

22 device_query_cmd="npu-smi info", 

23 dispatch_key="PrivateUse1", 

24 triton_extra_name=get_triton_extra_name(), 

25) 

26 

27CUSTOMIZED_UNUSED_OPS = ( 

28 "contiguous", 

29 "sort", 

30 "sort_stable", 

31 "copy_", 

32 "_to_copy", 

33) 

34 

35 

36__all__ = ["*"]