Kernels:
Trusted publisher
File size: 5,136 Bytes
5855430 e0d830e 5855430 e0d830e 5855430 e0d830e 5855430 e0d830e 5855430 | 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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | # 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 sys
from importlib import import_module
from importlib.machinery import ModuleSpec
from types import ModuleType
from ..modules.causal_conv1d import ShortConvolution
from ..modules.fused_bitlinear import BitLinear, FusedBitLinear
from ..modules.fused_cross_entropy import FusedCrossEntropyLoss
from ..modules.fused_kl_div import FusedKLDivLoss
from ..modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss
from ..modules.long_conv import ImplicitLongConvolution, LongConvolution
from ..modules.mlp import GatedMLP
from ..modules.norm import (
FusedLayerNormGated,
FusedLayerNormSwishGate,
FusedLayerNormSwishGateLinear,
FusedRMSNormGated,
FusedRMSNormSwishGate,
FusedRMSNormSwishGateLinear,
GroupNorm,
GroupNormLinear,
L2Norm,
LayerNorm,
LayerNormLinear,
RMSNorm,
RMSNormLinear,
)
from ..modules.rotary import RotaryEmbedding
from ..modules.token_shift import TokenShift
_MODULE_ALIASES = {
'fused_norm_gate': ('norm.fused_norm_gate',),
'l2norm': ('norm.l2norm',),
'layernorm': ('norm.layernorm',),
'layernorm_gated': ('norm.layernorm_gated',),
'convolution': ('causal_conv1d', 'long_conv', 'causal_conv1d.cp', 'causal_conv1d.backends.cuda'),
'conv': ('causal_conv1d', 'long_conv'),
'conv.causal_conv1d': ('causal_conv1d.ops',),
'conv.short_conv': ('causal_conv1d.ops',),
'conv.long_conv': ('long_conv',),
'conv.cp': ('causal_conv1d.cp',),
'conv.cp.ops': ('causal_conv1d.cp',),
'conv.cuda': ('causal_conv1d.backends.cuda',),
'conv.cuda.ops': ('causal_conv1d.backends.cuda.ops',),
'conv.triton': ('causal_conv1d.ops',),
'conv.triton.ops': ('causal_conv1d.ops',),
'conv.triton.kernels': ('causal_conv1d.ops',),
}
for _old, _targets in _MODULE_ALIASES.items():
_fullname = f'{__name__}.{_old}'
if len(_targets) == 1:
_module = import_module(f'{__name__}.{_targets[0]}')
else:
_module = ModuleType(_fullname)
_module.__spec__ = ModuleSpec(_fullname, loader=None, is_package=True)
_module.__path__ = []
for _target in _targets:
_source = import_module(f'{__name__}.{_target}')
for _symbol in getattr(_source, '__all__', vars(_source)):
if not _symbol.startswith('_'):
vars(_module).setdefault(_symbol, getattr(_source, _symbol))
_module.__all__ = [name for name in vars(_module) if not name.startswith('_')]
sys.modules[_fullname] = _module
# conv.causal_conv1d is a function on the old package.
if _old != 'conv.causal_conv1d':
_parent, _, _child = _fullname.rpartition('.')
setattr(sys.modules[_parent], _child, _module)
# defer GRPO's compile policy until it is explicitly imported
for _name in ('activations', 'rotary', 'fused_cross_entropy', 'fused_kl_div', 'fused_linear_cross_entropy'):
_module = import_module(f'{__name__}.{_name}')
for _symbol, _value in vars(import_module(f'{__name__}.{_name}.ops')).items():
if not _symbol.startswith('_'):
vars(_module).setdefault(_symbol, _value)
_SYMBOL_ALIASES = {
'fused_bitlinear': {
'layer_norm_fwd_quant': 'norm.layernorm_quant.layer_norm_quant_fwd',
'layer_norm_bwd': 'norm.layernorm_quant.layer_norm_quant_bwd',
'LayerNormLinearQuantFn': 'LayerNormLinearQuantFunction',
'layer_norm_linear_quant_fn': 'layer_norm_linear_quant',
},
'fused_cross_entropy': {
'fused_cross_entropy_forward': 'cross_entropy_fwd',
'CrossEntropyLossFunction': 'FusedCrossEntropyFunction',
},
'fused_kl_div': {'fused_kl_div_forward': 'fused_kl_div_fwd', 'fused_kl_div_backward': 'fused_kl_div_bwd'},
'fused_linear_cross_entropy': {
'fused_linear_cross_entropy_forward': 'fused_linear_cross_entropy_fwd',
'fused_linear_cross_entropy_backward': 'fused_linear_cross_entropy_bwd',
},
}
for _name, _aliases in _SYMBOL_ALIASES.items():
_module = sys.modules[f'{__name__}.{_name}']
for _old, _new in _aliases.items():
_target, _, _symbol = _new.rpartition('.')
_source = import_module(f'{__name__}.{_target}') if _target else _module
setattr(_module, _old, getattr(_source, _symbol))
__all__ = [
'BitLinear',
'FusedBitLinear',
'FusedCrossEntropyLoss',
'FusedKLDivLoss',
'FusedLayerNormGated',
'FusedLayerNormSwishGate',
'FusedLayerNormSwishGateLinear',
'FusedLinearCrossEntropyLoss',
'FusedRMSNormGated',
'FusedRMSNormSwishGate',
'FusedRMSNormSwishGateLinear',
'GatedMLP',
'GroupNorm',
'GroupNormLinear',
'ImplicitLongConvolution',
'L2Norm',
'LayerNorm',
'LayerNormLinear',
'LongConvolution',
'RMSNorm',
'RMSNormLinear',
'RotaryEmbedding',
'ShortConvolution',
'TokenShift',
]
|