« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
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# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
17# pylint: disable=no-name-in-module,consider-using-from-import,pointless-string-statement
18# pylint: disable=redefined-outer-name,missing-public-type-hints,missing-public-docstring
19# pylint: disable=unused-argument,no-member,import-outside-toplevel
20# pylint: disable=invalid-name,missing-module-docstring,missing-function-docstring
21# pylint: disable=missing-class-docstring,broad-exception-caught,protected-access
22# pylint: disable=forbidden-backend-import
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
35import torch
36import triton
37import triton.language as tl
38import triton.language.extra.libdevice as tldevice
39import triton.runtime.driver as driver
41logger = logging.getLogger(__name__)
43FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1"
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.
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.
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).
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.
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")
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
78 def _make_wrapper(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:
79 cache: list = []
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
92 result = fn(*args, **kwargs)
93 cache.insert(0, (args, kwargs, result))
94 if len(cache) > maxsize:
95 cache.pop()
96 return result
98 return wrapper
100 if fn is not None:
101 return _make_wrapper(fn)
102 return _make_wrapper
105@tensor_cache
106def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor:
107 return cu_seqlens[1:] - cu_seqlens[:-1]
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)
116def get_abs_err(x, y):
117 return (x.detach() - y.detach()).flatten().abs().max().item()
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)
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 allow_warning = warning or (FLA_CI_ENV and (error_rate < 0.01 or abs_atol <= 0.3))
134 if allow_warning:
135 if error_rate > ratio:
136 warnings.warn(msg)
137 else:
138 if not error_rate < ratio:
139 raise AssertionError(msg)
142if hasattr(triton.language, '_experimental_make_tensor_descriptor'):
143 # For Triton 3.3.x
144 make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor
145elif hasattr(triton.language, 'make_tensor_descriptor'):
146 # For Triton 3.4.x and later
147 make_tensor_descriptor = triton.language.make_tensor_descriptor
148else:
149 """
150 Fallback implementation when TMA is not supported.
151 Returns None to indicate TMA descriptors are unavailable.
152 Just make triton compiler happy.
153 """
155 @triton.jit
156 def make_tensor_descriptor(
157 base,
158 shape,
159 strides,
160 block_shape,
161 _builder=None,
162 ):
163 return None
166@lru_cache(maxsize=None)
167def get_available_device() -> str:
168 try:
169 return triton.runtime.driver.active.get_current_target().backend
170 except Exception:
171 _cpu_device_warning()
172 return 'cpu'
175def map_triton_backend_to_torch_device() -> str:
176 backend = get_available_device() # 'cuda' | 'hip' | 'xpu' | 'cpu' | ...
177 return {'cuda': 'cuda', 'hip': 'cuda', 'xpu': 'xpu'}.get(backend, backend)
180device = get_available_device() if get_available_device() != 'hip' else 'cuda'
181device_torch_lib = getattr(torch, device)
182device_platform = get_available_device()
183is_amd = device_platform == 'hip'
184is_nvidia = device_platform == 'cuda'
185is_nvidia_hopper = is_nvidia and (
186 'NVIDIA H' in torch.cuda.get_device_name(0) or torch.cuda.get_device_capability()[0] >= 9
187)
189is_tf32_supported = is_nvidia and torch.cuda.get_device_capability(0)[0] >= 8
190is_tma_supported = (
191 (is_nvidia and torch.cuda.get_device_capability(0)[0] >= 9)
192 and os.environ.get('FLA_NO_USE_TMA', '0') != '1'
193 and (
194 hasattr(triton.language, '_experimental_make_tensor_descriptor')
195 or hasattr(triton.language, 'make_tensor_descriptor')
196 )
197)
199if is_nvidia and not is_tf32_supported:
200 # Make old card happy, since triton will use tf32 by default.
201 # This is a workaround for old nvidia card.
202 os.environ['TRITON_F32_DEFAULT'] = 'ieee'
205@lru_cache(maxsize=None)
206def check_pytorch_version(version_s: str = '2.4') -> bool:
207 return version.parse(torch.__version__) >= version.parse(version_s)
210if check_pytorch_version('2.4'):
211 device = 'cuda' if device == 'cpu' else device
212 autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device)
213 autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device)
215 def custom_device_ctx(index: int):
216 return device_torch_lib.device(index)
217else:
218 if device != 'cuda':
219 raise AssertionError('Only cuda device is supported for PyTorch version < 2.4.0.')
220 autocast_custom_fwd = device_torch_lib.amp.custom_fwd
221 autocast_custom_bwd = device_torch_lib.amp.custom_bwd
223 def custom_device_ctx(index: int):
224 return torch.cuda.device(index)
227def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:
228 """
229 A decorator to make sure all input tensors are contiguous and set the device based on input tensors.
230 """
232 @functools.wraps(fn)
233 def wrapper(*args, **kwargs):
234 contiguous_args = (i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args)
235 contiguous_kwargs = {k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()}
237 tensor = None
238 for arg in args:
239 if isinstance(arg, torch.Tensor):
240 tensor = arg
241 break
242 if tensor is None:
243 for value in kwargs.values():
244 if isinstance(value, torch.Tensor):
245 tensor = value
246 break
248 if tensor is not None:
249 ctx = custom_device_ctx(tensor.device.index)
250 else:
251 ctx = contextlib.nullcontext()
253 with ctx:
254 return fn(*contiguous_args, **contiguous_kwargs)
256 return wrapper
259def _cpu_device_warning():
260 warnings.warn(('Triton is not supported on current platform, roll back to CPU.'), stacklevel=1)
263@tensor_cache
264def prepare_chunk_offsets(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor:
265 return torch.cat([cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1)
268if os.environ.get('FLA_USE_FAST_OPS', '0') == '1':
269 exp = tldevice.fast_expf
270 exp2 = tldevice.exp2
271 log = tldevice.fast_logf
272 log2 = tldevice.fast_log2f
273else:
274 exp = tl.exp
275 exp2 = tl.math.exp2
276 log = tl.log
277 log2 = tl.log2
280def get_all_max_shared_mem():
281 try:
282 return [
283 triton.runtime.driver.active.utils.get_device_properties(i)['max_shared_mem']
284 for i in range(device_torch_lib.device_count())
285 ]
286 except Exception:
287 _cpu_device_warning()
288 return [-1]
291class Backend(Enum):
292 ADA = 101376 # RTX 4090
293 AMPERE = 166912 # A100
294 HOPPER = 232448 # H100
295 DEFAULT = 102400 # Default
297 @classmethod
298 def get_shared_memory(cls, arch: str) -> int:
299 try:
300 return cls[arch.upper()].value
301 except KeyError:
302 return cls.DEFAULT.value
305@lru_cache(maxsize=None)
306def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool:
307 try:
308 device_shared_mem_list = get_all_max_shared_mem()
309 max_shared_memory = device_shared_mem_list[tensor_idx]
310 return max_shared_memory >= Backend.get_shared_memory(arch)
311 except Exception:
312 return False
315def get_autotune_config(
316 multibuffer_list: tuple = (False,),
317 unit_flag_list: tuple = (False,),
318 limit_auto_multi_buffer_only_for_local_buffer_list: tuple = (False,),
319 limit_auto_multi_buffer_of_local_buffer_list: tuple = ("no-l0c",),
320 set_workspace_multibuffer_list: tuple = (2, 4),
321 enable_hivm_auto_cv_balance_list: tuple = (True,),
322 tile_mix_vector_loop_num_list: tuple = (2, 4),
323 tile_mix_cube_loop_num_list: tuple = (2, 4),
324):
325 configs = []
326 for (
327 multibuffer,
328 unit_flag,
329 limit_auto_multi_buffer_only_for_local_buffer,
330 limit_auto_multi_buffer_of_local_buffer,
331 ) in itertools.product(
332 list(multibuffer_list),
333 list(unit_flag_list),
334 list(limit_auto_multi_buffer_only_for_local_buffer_list),
335 list(limit_auto_multi_buffer_of_local_buffer_list),
336 ):
337 base_config_dict = {
338 'multibuffer': multibuffer,
339 'unit_flag': unit_flag,
340 'limit_auto_multi_buffer_only_for_local_buffer': limit_auto_multi_buffer_only_for_local_buffer,
341 'limit_auto_multi_buffer_of_local_buffer': limit_auto_multi_buffer_of_local_buffer,
342 }
344 if limit_auto_multi_buffer_only_for_local_buffer:
345 configs.append(triton.Config(base_config_dict))
346 else:
347 for (
348 set_workspace_multibuffer,
349 enable_hivm_auto_cv_balance,
350 tile_mix_vector_loop,
351 tile_mix_cube_loop,
352 ) in itertools.product(
353 list(set_workspace_multibuffer_list),
354 list(enable_hivm_auto_cv_balance_list),
355 list(tile_mix_vector_loop_num_list),
356 list(tile_mix_cube_loop_num_list),
357 ):
358 full_config_dict = base_config_dict.copy()
359 full_config_dict.update(
360 {
361 'set_workspace_multibuffer': set_workspace_multibuffer,
362 'enable_hivm_auto_cv_balance': enable_hivm_auto_cv_balance,
363 'tile_mix_vector_loop': tile_mix_vector_loop,
364 'tile_mix_cube_loop': tile_mix_cube_loop,
365 }
366 )
367 configs.append(triton.Config(full_config_dict))
368 return configs
371def get_npu_properties():
372 return driver.active.utils.get_device_properties(torch.npu.current_device())
375@functools.cache
376def get_vector_num() -> int:
377 import torch_npu
379 current_device = torch_npu.npu.current_device()
380 properties = driver.active.utils.get_device_properties(current_device)
381 return properties["num_vectorcore"]
384@lru_cache
385def is_arch35():
386 try:
387 import torch_npu
389 return "Ascend910_95" in torch_npu.npu.get_device_name() or "Ascend950" in torch_npu.npu.get_device_name()
390 except Exception:
391 return False