Kernels:
Trusted publisher
| # Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| # For a list of all contributors, visit: | |
| # https://github.com/fla-org/flash-linear-attention/graphs/contributors | |
| import logging | |
| import warnings | |
| import torch | |
| from ._config import FLA_CI_ENV | |
| logger = logging.getLogger(__name__) | |
| def get_abs_err(x, y): | |
| return (x.detach() - y.detach()).flatten().abs().max().item() | |
| def get_err_ratio(x, y): | |
| err = (x.detach() - y.detach()).flatten().square().mean().sqrt().item() | |
| base = (x.detach()).flatten().square().mean().sqrt().item() | |
| return err / (base + 1e-8) | |
| def assert_close(prefix, ref, tri, ratio, warning=False, err_atol=1e-6): | |
| abs_atol = get_abs_err(ref, tri) | |
| error_rate = get_err_ratio(ref, tri) | |
| msg = f"{prefix:>16} diff: {abs_atol:.6f} ratio: {error_rate:.6f}" | |
| logger.info(msg) | |
| if abs_atol <= err_atol: | |
| return | |
| assert not torch.isnan(ref).any(), f"{prefix}: NaN detected in ref" | |
| assert not torch.isnan(tri).any(), f"{prefix}: NaN detected in tri" | |
| if warning or (FLA_CI_ENV and (error_rate < 0.01 or abs_atol <= 0.3)): | |
| if error_rate > ratio: | |
| warnings.warn(msg) | |
| else: | |
| assert error_rate < ratio, msg | |