Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / custom_ops / gdn / triton / utils.py: 0%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15# -*- coding: utf-8 -*- 

16# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang 

17 

18# pylint: disable=no-name-in-module,consider-using-from-import,pointless-string-statement 

19# pylint: disable=redefined-outer-name,missing-public-type-hints,missing-public-docstring 

20# pylint: disable=unused-argument,no-member,import-outside-toplevel 

21# pylint: disable=invalid-name,missing-module-docstring,missing-function-docstring 

22# pylint: disable=missing-class-docstring,broad-exception-caught,protected-access 

23 

24import itertools 

25import contextlib 

26import os 

27import functools 

28import warnings 

29import logging 

30from enum import Enum 

31from functools import lru_cache 

32from typing import Any, Callable, Optional 

33from packaging import version 

34 

35import torch 

36import triton 

37import triton.language as tl 

38import triton.language.extra.libdevice as tldevice 

39import triton.runtime.driver as driver 

40 

41logger = logging.getLogger(__name__) 

42 

43FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1" 

44 

45 

46def tensor_cache(fn: Optional[Callable[..., torch.Tensor]] = None, *, maxsize: int = 1) -> Any: 

47 """ 

48 A decorator that caches the most recent results of a function with tensor inputs. 

49 

50 This decorator will store the outputs of the decorated function for the most recent 

51 set of input tensors, up to `maxsize` entries. If the function is called again with 

52 the same input tensors, it will return the cached result. 

53 

54 When maxsize=1 (default), the behavior is identical to caching only the most recent result. 

55 Can be used as @tensor_cache or @tensor_cache(maxsize=n). 

56 

57 Args: 

58 fn (Callable[..., torch.Tensor], optional): 

59 The function to be decorated when used without parentheses. 

60 maxsize (int): 

61 Maximum number of input combinations to cache. Default is 1. 

62 

63 Returns: 

64 Callable[..., torch.Tensor]: 

65 A wrapped version of the input function with caching. 

66 """ 

67 if maxsize < 1: 

68 raise ValueError("maxsize must be at least 1") 

69 

70 def _is_match(a: Any, b: Any) -> bool: 

71 if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor): 

72 return a is b 

73 try: 

74 return a == b 

75 except Exception: 

76 return a is b 

77 

78 def _make_wrapper(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: 

79 cache: list = [] 

80 

81 @functools.wraps(fn) 

82 def wrapper(*args: Any, **kwargs: Any) -> Any: 

83 for i, (cached_args, cached_kwargs, cached_result) in enumerate(cache): 

84 if len(args) == len(cached_args) and len(kwargs) == len(cached_kwargs): 

85 if all(_is_match(a, b) for a, b in zip(args, cached_args)) and all( 

86 k in cached_kwargs and _is_match(v, cached_kwargs[k]) for k, v in kwargs.items() 

87 ): 

88 if i != 0: 

89 cache.insert(0, cache.pop(i)) 

90 return cached_result 

91 

92 result = fn(*args, **kwargs) 

93 cache.insert(0, (args, kwargs, result)) 

94 if len(cache) > maxsize: 

95 cache.pop() 

96 return result 

97 

98 return wrapper 

99 

100 if fn is not None: 

101 return _make_wrapper(fn) 

102 return _make_wrapper 

103 

104 

105@tensor_cache 

106def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: 

107 return cu_seqlens[1:] - cu_seqlens[:-1] 

108 

109 

110@tensor_cache(maxsize=3) 

111def prepare_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor: 

112 indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()]) 

113 return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) 

114 

115 

116def get_abs_err(x, y): 

117 return (x.detach() - y.detach()).flatten().abs().max().item() 

118 

119 

120def get_err_ratio(x, y): 

121 err = (x.detach() - y.detach()).flatten().square().mean().sqrt().item() 

122 base = (x.detach()).flatten().square().mean().sqrt().item() 

123 return err / (base + 1e-8) 

124 

125 

126def assert_close(prefix, ref, tri, ratio, warning=False, err_atol=1e-6): 

127 abs_atol = get_abs_err(ref, tri) 

128 msg = f"{prefix:>16} diff: {abs_atol:.6f} ratio: {get_err_ratio(ref, tri):.6f}" 

129 logger.info(msg) 

130 error_rate = get_err_ratio(ref, tri) 

131 if abs_atol <= err_atol: 

132 return 

133 if warning or (FLA_CI_ENV and (error_rate < 0.01 or abs_atol <= 0.3)): 

134 if error_rate > ratio: 

135 warnings.warn(msg) 

136 else: 

137 assert error_rate < ratio, msg 

138 

139 

140if hasattr(triton.language, '_experimental_make_tensor_descriptor'): 

141 # For Triton 3.3.x 

142 make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor 

143elif hasattr(triton.language, 'make_tensor_descriptor'): 

144 # For Triton 3.4.x and later 

145 make_tensor_descriptor = triton.language.make_tensor_descriptor 

146else: 

147 """ 

148 Fallback implementation when TMA is not supported. 

149 Returns None to indicate TMA descriptors are unavailable. 

150 Just make triton compiler happy. 

151 """ 

152 

153 @triton.jit 

154 def make_tensor_descriptor( 

155 base, 

156 shape, 

157 strides, 

158 block_shape, 

159 _builder=None, 

160 ): 

161 return None 

162 

163 

164@lru_cache(maxsize=None) 

165def get_available_device() -> str: 

166 try: 

167 return triton.runtime.driver.active.get_current_target().backend 

168 except Exception: 

169 _cpu_device_warning() 

170 return 'cpu' 

171 

172 

173def map_triton_backend_to_torch_device() -> str: 

174 backend = get_available_device() # 'cuda' | 'hip' | 'xpu' | 'cpu' | ... 

175 return {'cuda': 'cuda', 'hip': 'cuda', 'xpu': 'xpu'}.get(backend, backend) 

176 

177 

178device = get_available_device() if get_available_device() != 'hip' else 'cuda' 

179device_torch_lib = getattr(torch, device) 

180device_platform = get_available_device() 

181is_amd = device_platform == 'hip' 

182is_nvidia = device_platform == 'cuda' 

183is_nvidia_hopper = is_nvidia and ( 

184 'NVIDIA H' in torch.cuda.get_device_name(0) or torch.cuda.get_device_capability()[0] >= 9 

185) 

186 

187is_tf32_supported = is_nvidia and torch.cuda.get_device_capability(0)[0] >= 8 

188is_tma_supported = ( 

189 (is_nvidia and torch.cuda.get_device_capability(0)[0] >= 9) 

190 and os.environ.get('FLA_NO_USE_TMA', '0') != '1' 

191 and ( 

192 hasattr(triton.language, '_experimental_make_tensor_descriptor') 

193 or hasattr(triton.language, 'make_tensor_descriptor') 

194 ) 

195) 

196 

197if is_nvidia and not is_tf32_supported: 

198 # Make old card happy, since triton will use tf32 by default. 

199 # This is a workaround for old nvidia card. 

200 os.environ['TRITON_F32_DEFAULT'] = 'ieee' 

201 

202 

203@lru_cache(maxsize=None) 

204def check_pytorch_version(version_s: str = '2.4') -> bool: 

205 return version.parse(torch.__version__) >= version.parse(version_s) 

206 

207 

208if check_pytorch_version('2.4'): 

209 device = 'cuda' if device == 'cpu' else device 

210 autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device) 

211 autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device) 

212 

213 def custom_device_ctx(index: int): 

214 return device_torch_lib.device(index) 

215else: 

216 assert device == 'cuda', 'Only cuda device is supported for PyTorch version < 2.4.0.' 

217 autocast_custom_fwd = device_torch_lib.amp.custom_fwd 

218 autocast_custom_bwd = device_torch_lib.amp.custom_bwd 

219 

220 def custom_device_ctx(index: int): 

221 return torch.cuda.device(index) 

222 

223 

224def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: 

225 """ 

226 A decorator to make sure all input tensors are contiguous and set the device based on input tensors. 

227 """ 

228 

229 @functools.wraps(fn) 

230 def wrapper(*args, **kwargs): 

231 contiguous_args = (i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args) 

232 contiguous_kwargs = {k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()} 

233 

234 tensor = None 

235 for arg in args: 

236 if isinstance(arg, torch.Tensor): 

237 tensor = arg 

238 break 

239 if tensor is None: 

240 for value in kwargs.values(): 

241 if isinstance(value, torch.Tensor): 

242 tensor = value 

243 break 

244 

245 if tensor is not None: 

246 ctx = custom_device_ctx(tensor.device.index) 

247 else: 

248 ctx = contextlib.nullcontext() 

249 

250 with ctx: 

251 return fn(*contiguous_args, **contiguous_kwargs) 

252 

253 return wrapper 

254 

255 

256def _cpu_device_warning(): 

257 warnings.warn(('Triton is not supported on current platform, roll back to CPU.'), stacklevel=1) 

258 

259 

260@tensor_cache 

261def prepare_chunk_offsets(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor: 

262 return torch.cat([cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1) 

263 

264 

265if os.environ.get('FLA_USE_FAST_OPS', '0') == '1': 

266 exp = tldevice.fast_expf 

267 exp2 = tldevice.exp2 

268 log = tldevice.fast_logf 

269 log2 = tldevice.fast_log2f 

270else: 

271 exp = tl.exp 

272 exp2 = tl.math.exp2 

273 log = tl.log 

274 log2 = tl.log2 

275 

276 

277def get_all_max_shared_mem(): 

278 try: 

279 return [ 

280 triton.runtime.driver.active.utils.get_device_properties(i)['max_shared_mem'] 

281 for i in range(device_torch_lib.device_count()) 

282 ] 

283 except Exception: 

284 _cpu_device_warning() 

285 return [-1] 

286 

287 

288class Backend(Enum): 

289 ADA = 101376 # RTX 4090 

290 AMPERE = 166912 # A100 

291 HOPPER = 232448 # H100 

292 DEFAULT = 102400 # Default 

293 

294 @classmethod 

295 def get_shared_memory(cls, arch: str) -> int: 

296 try: 

297 return cls[arch.upper()].value 

298 except KeyError: 

299 return cls.DEFAULT.value 

300 

301 

302@lru_cache(maxsize=None) 

303def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: 

304 try: 

305 device_shared_mem_list = get_all_max_shared_mem() 

306 max_shared_memory = device_shared_mem_list[tensor_idx] 

307 return max_shared_memory >= Backend.get_shared_memory(arch) 

308 except Exception: 

309 return False 

310 

311 

312def get_autotune_config( 

313 multibuffer_list: tuple = (False,), 

314 unit_flag_list: tuple = (False,), 

315 limit_auto_multi_buffer_only_for_local_buffer_list: tuple = (False,), 

316 limit_auto_multi_buffer_of_local_buffer_list: tuple = ("no-l0c",), 

317 set_workspace_multibuffer_list: tuple = (2, 4), 

318 enable_hivm_auto_cv_balance_list: tuple = (True,), 

319 tile_mix_vector_loop_num_list: tuple = (2, 4), 

320 tile_mix_cube_loop_num_list: tuple = (2, 4), 

321): 

322 configs = [] 

323 for ( 

324 multibuffer, 

325 unit_flag, 

326 limit_auto_multi_buffer_only_for_local_buffer, 

327 limit_auto_multi_buffer_of_local_buffer, 

328 ) in itertools.product( 

329 list(multibuffer_list), 

330 list(unit_flag_list), 

331 list(limit_auto_multi_buffer_only_for_local_buffer_list), 

332 list(limit_auto_multi_buffer_of_local_buffer_list), 

333 ): 

334 base_config_dict = { 

335 'multibuffer': multibuffer, 

336 'unit_flag': unit_flag, 

337 'limit_auto_multi_buffer_only_for_local_buffer': limit_auto_multi_buffer_only_for_local_buffer, 

338 'limit_auto_multi_buffer_of_local_buffer': limit_auto_multi_buffer_of_local_buffer, 

339 } 

340 

341 if limit_auto_multi_buffer_only_for_local_buffer: 

342 configs.append(triton.Config(base_config_dict)) 

343 else: 

344 for ( 

345 set_workspace_multibuffer, 

346 enable_hivm_auto_cv_balance, 

347 tile_mix_vector_loop, 

348 tile_mix_cube_loop, 

349 ) in itertools.product( 

350 list(set_workspace_multibuffer_list), 

351 list(enable_hivm_auto_cv_balance_list), 

352 list(tile_mix_vector_loop_num_list), 

353 list(tile_mix_cube_loop_num_list), 

354 ): 

355 full_config_dict = base_config_dict.copy() 

356 full_config_dict.update( 

357 { 

358 'set_workspace_multibuffer': set_workspace_multibuffer, 

359 'enable_hivm_auto_cv_balance': enable_hivm_auto_cv_balance, 

360 'tile_mix_vector_loop': tile_mix_vector_loop, 

361 'tile_mix_cube_loop': tile_mix_cube_loop, 

362 } 

363 ) 

364 configs.append(triton.Config(full_config_dict)) 

365 return configs 

366 

367 

368def get_npu_properties(): 

369 return driver.active.utils.get_device_properties(torch.npu.current_device()) 

370 

371 

372@functools.cache 

373def get_vector_num() -> int: 

374 import torch_npu 

375 

376 current_device = torch_npu.npu.current_device() 

377 properties = driver.active.utils.get_device_properties(current_device) 

378 return properties["num_vectorcore"] 

379 

380 

381@lru_cache 

382def is_arch35(): 

383 try: 

384 import torch_npu 

385 

386 return "Ascend910_95" in torch_npu.npu.get_device_name() or "Ascend950" in torch_npu.npu.get_device_name() 

387 except Exception: 

388 return False