Kernels:
Trusted publisher
File size: 1,328 Bytes
e19323e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | # 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
|