来源:
1. fla-org/flash-linear-attention (5538 stars)
→ upstream_ref/fla/ops/gated_delta_rule/naive.py (正确的纯 PyTorch GDN)
→ upstream_ref/fla/ops/gated_delta_rule/chunk.py (Triton chunk kernel)
→ upstream_ref/fla/layers/gated_deltanet.py (层集成)
2. vllm-project/vllm main (88717 stars)
→ upstream_ref/vllm_gdn/gdn/qwen_gdn_linear_attn.py (1751行, Qwen3.5 原生 GDN)
→ upstream_ref/vllm_gdn/ops/causal_conv1d.py (1289行, 正确的 Conv1d)
→ upstream_ref/vllm_gdn/third_party/ops/ (FLA Triton ops vendored)
→ upstream_ref/vllm_gdn/models/qwen3_5.py (vllm 最新 Qwen3.5 模型)
3. Deep-Spark/xllm (BI-V100 硬件厂商)
→ upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp (576行)
→ upstream_ref/xllm_latest/core/kernels/npu/npu_causal_conv1d.cpp
→ upstream_ref/xllm_latest/core/kernels/npu/npu_recurrent_gated_delta_rule.cpp
目的: 修复 corex_gdn.py Conv1d groups 接口不匹配问题
错误: conv1d_weight shape (2560,1,4) 被当成 (num_k_heads,1,4) 索引
conv_dim = key_dim*2 + value_dim = 10240, TP=4 后 2560
FLA naive.py 和 vllm qwen_gdn_linear_attn.py 有正确的实现可直接对接
53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
# 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
|
|
|
|
from fla.modules.convolution import ImplicitLongConvolution, LongConvolution, ShortConvolution
|
|
from fla.modules.fused_bitlinear import BitLinear, FusedBitLinear
|
|
from fla.modules.fused_cross_entropy import FusedCrossEntropyLoss
|
|
from fla.modules.fused_kl_div import FusedKLDivLoss
|
|
from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss
|
|
from fla.modules.fused_norm_gate import (
|
|
FusedLayerNormGated,
|
|
FusedLayerNormSwishGate,
|
|
FusedLayerNormSwishGateLinear,
|
|
FusedRMSNormGated,
|
|
FusedRMSNormSwishGate,
|
|
FusedRMSNormSwishGateLinear,
|
|
)
|
|
from fla.modules.l2norm import L2Norm
|
|
from fla.modules.layernorm import GroupNorm, GroupNormLinear, LayerNorm, LayerNormLinear, RMSNorm, RMSNormLinear
|
|
from fla.modules.mlp import GatedMLP
|
|
from fla.modules.rotary import RotaryEmbedding
|
|
from fla.modules.token_shift import TokenShift
|
|
|
|
__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',
|
|
]
|