ref(upstream): 搬运 3 大 GDN 上游仓库 — FLA naive ops + vllm GDN 子树 + xllm C++ 参考
来源:
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 有正确的实现可直接对接
This commit is contained in:
Claude
0
upstream_ref/vllm_gdn/gdn/__init__.py
Normal file
0
upstream_ref/vllm_gdn/gdn/__init__.py
Normal file
58
upstream_ref/vllm_gdn/gdn/base.py
Normal file
58
upstream_ref/vllm_gdn/gdn/base.py
Normal file
@@ -0,0 +1,58 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import torch
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from vllm.config import (
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.distributed import (
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from vllm.model_executor.custom_op import PluggableLayer
|
||||
from vllm.model_executor.layers.mamba.abstract import MambaBase
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import (
|
||||
MambaStateDtypeCalculator,
|
||||
)
|
||||
from vllm.model_executor.models.utils import extract_layer_index
|
||||
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
|
||||
|
||||
|
||||
class GatedDeltaNetAttention(PluggableLayer, MambaBase):
|
||||
"""Base class for GatedDeltaNet attention layer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: PretrainedConfig,
|
||||
vllm_config: VllmConfig,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.prefix = prefix
|
||||
self.tp_size = get_tensor_model_parallel_world_size()
|
||||
self.tp_rank = get_tensor_model_parallel_rank()
|
||||
self.layer_idx = extract_layer_index(prefix)
|
||||
self.hidden_size = config.hidden_size
|
||||
self.activation = config.hidden_act
|
||||
self.layer_norm_epsilon = config.rms_norm_eps
|
||||
self.model_config = vllm_config.model_config
|
||||
self.cache_config = vllm_config.cache_config
|
||||
self.quant_config = vllm_config.quant_config
|
||||
self.speculative_config = vllm_config.speculative_config
|
||||
self.num_spec = (
|
||||
self.speculative_config.num_speculative_tokens
|
||||
if self.speculative_config
|
||||
else 0
|
||||
)
|
||||
|
||||
@property
|
||||
def mamba_type(self) -> MambaAttentionBackendEnum:
|
||||
return MambaAttentionBackendEnum.GDN_ATTN
|
||||
|
||||
def get_state_dtype(self) -> tuple[torch.dtype, ...]:
|
||||
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
|
||||
self.model_config.dtype,
|
||||
self.cache_config.mamba_cache_dtype,
|
||||
self.cache_config.mamba_ssm_cache_dtype,
|
||||
)
|
||||
635
upstream_ref/vllm_gdn/gdn/kimi_gdn_linear_attn.py
Normal file
635
upstream_ref/vllm_gdn/gdn/kimi_gdn_linear_attn.py
Normal file
@@ -0,0 +1,635 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.distributed import divide, get_tensor_model_parallel_rank
|
||||
from vllm.forward_context import get_forward_context
|
||||
from vllm.model_executor.custom_op import PluggableLayer
|
||||
from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
|
||||
from vllm.model_executor.model_loader.weight_utils import (
|
||||
default_weight_loader,
|
||||
sharded_weight_loader,
|
||||
)
|
||||
from vllm.model_executor.parameter import BasevLLMParameter
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
|
||||
from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
|
||||
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
|
||||
|
||||
from ...linear import (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
ReplicatedLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from ..mamba_utils import (
|
||||
MambaStateDtypeCalculator,
|
||||
MambaStateShapeCalculator,
|
||||
is_conv_state_dim_first,
|
||||
)
|
||||
from ..ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update
|
||||
from ..ops.gather_initial_states import gather_initial_states
|
||||
|
||||
# Empirical lower bound for the KDA gate to avoid numerical underflow.
|
||||
_KDA_GATE_LOGBOUND_MIN = -5.0
|
||||
|
||||
|
||||
def a_log_weight_loader(
|
||||
shard_axis: int,
|
||||
) -> Callable[[torch.Tensor, torch.Tensor], None]:
|
||||
"""Load KDA A_log stored as either old 4D or current 1D weights."""
|
||||
|
||||
def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
shard_size = param.data.shape[shard_axis]
|
||||
start_idx = tp_rank * shard_size
|
||||
|
||||
if loaded_weight.dim() == 4:
|
||||
assert loaded_weight.shape[:2] == (1, 1), (
|
||||
f"Expected old A_log shape (1, 1, H, 1), got {loaded_weight.shape}"
|
||||
)
|
||||
assert loaded_weight.shape[-1] == 1, (
|
||||
f"Expected old A_log last dim to be 1, got {loaded_weight.shape}"
|
||||
)
|
||||
loaded_weight = loaded_weight.view(loaded_weight.shape[2])
|
||||
|
||||
loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size)
|
||||
return default_weight_loader(param, loaded_weight)
|
||||
|
||||
return loader
|
||||
|
||||
|
||||
def _make_fused_conv1d_weight_loader(
|
||||
dims: list[int],
|
||||
tp_size: int,
|
||||
tp_rank: int,
|
||||
) -> Callable[..., None]:
|
||||
sharded_dims = [dim // tp_size for dim in dims]
|
||||
|
||||
def weight_loader(
|
||||
param: torch.Tensor,
|
||||
loaded_weight: torch.Tensor,
|
||||
loaded_shard_id: int,
|
||||
) -> None:
|
||||
if loaded_weight.dim() == 2:
|
||||
loaded_weight = loaded_weight.unsqueeze(1)
|
||||
shard_size = sharded_dims[loaded_shard_id]
|
||||
source_start = tp_rank * shard_size
|
||||
target_start = sum(sharded_dims[:loaded_shard_id])
|
||||
loaded_shard = loaded_weight[source_start : source_start + shard_size]
|
||||
param.data[target_start : target_start + shard_size].copy_(loaded_shard)
|
||||
|
||||
return weight_loader
|
||||
|
||||
|
||||
class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear):
|
||||
"""Merged projection with one output replicated across TP ranks.
|
||||
|
||||
The replicated shard is represented as ``size * tp_size`` so the merged
|
||||
parameter reserves ``size`` local rows on every rank. Loading that shard
|
||||
from rank zero then gives every rank the complete checkpoint weight.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_size: int,
|
||||
output_sizes: list[int],
|
||||
replicated_shard_id: int,
|
||||
tp_size: int,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
self.replicated_shard_id = replicated_shard_id
|
||||
output_sizes = output_sizes.copy()
|
||||
output_sizes[replicated_shard_id] *= tp_size
|
||||
super().__init__(input_size, output_sizes, **kwargs)
|
||||
|
||||
def weight_loader(
|
||||
self,
|
||||
param: Parameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
loaded_shard_id: tuple[int, ...] | int | None = None,
|
||||
) -> None:
|
||||
tp_rank = self.tp_rank
|
||||
param_tp_rank = getattr(param, "tp_rank", None)
|
||||
if loaded_shard_id == self.replicated_shard_id:
|
||||
self.tp_rank = 0
|
||||
if param_tp_rank is not None:
|
||||
param.tp_rank = 0
|
||||
try:
|
||||
super().weight_loader(param, loaded_weight, loaded_shard_id)
|
||||
finally:
|
||||
self.tp_rank = tp_rank
|
||||
if param_tp_rank is not None:
|
||||
param.tp_rank = param_tp_rank
|
||||
|
||||
def weight_loader_v2(
|
||||
self,
|
||||
param: BasevLLMParameter,
|
||||
loaded_weight: torch.Tensor,
|
||||
loaded_shard_id: tuple[int, ...] | int | None = None,
|
||||
) -> None:
|
||||
tp_rank = self.tp_rank
|
||||
param_tp_rank = getattr(param, "tp_rank", None)
|
||||
if loaded_shard_id == self.replicated_shard_id:
|
||||
self.tp_rank = 0
|
||||
if param_tp_rank is not None:
|
||||
param.tp_rank = 0
|
||||
try:
|
||||
super().weight_loader_v2(param, loaded_weight, loaded_shard_id)
|
||||
finally:
|
||||
self.tp_rank = tp_rank
|
||||
if param_tp_rank is not None:
|
||||
param.tp_rank = param_tp_rank
|
||||
|
||||
|
||||
@PluggableLayer.register("kimi_gated_delta_net_attention")
|
||||
class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
def get_state_dtype(
|
||||
self,
|
||||
) -> tuple[torch.dtype, torch.dtype]:
|
||||
if self.model_config is None or self.cache_config is None:
|
||||
raise ValueError("model_config and cache_config must be set")
|
||||
return MambaStateDtypeCalculator.kda_state_dtype(
|
||||
self.model_config.dtype, self.cache_config.mamba_cache_dtype
|
||||
)
|
||||
|
||||
def get_state_shape(
|
||||
self,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...]]:
|
||||
return MambaStateShapeCalculator.kda_state_shape(
|
||||
self.tp_size,
|
||||
self.num_heads,
|
||||
self.head_dim,
|
||||
conv_kernel_size=self.conv_size,
|
||||
num_spec=self.num_spec,
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: KimiLinearConfig,
|
||||
vllm_config: VllmConfig,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__(config, vllm_config, prefix)
|
||||
|
||||
kda_config = config.linear_attn_config # type: ignore[attr-defined]
|
||||
assert kda_config is not None, "linear_attn_config must be set"
|
||||
self.head_dim = kda_config["head_dim"]
|
||||
self.num_heads = kda_config["num_heads"]
|
||||
assert self.num_heads % self.tp_size == 0
|
||||
self.local_num_heads = divide(self.num_heads, self.tp_size)
|
||||
|
||||
self.projection_size = self.head_dim * self.num_heads
|
||||
self.local_projection_size = divide(self.projection_size, self.tp_size)
|
||||
self.conv_size = kda_config["short_conv_kernel_size"]
|
||||
self.use_full_rank_gate = kda_config.get("use_full_rank_gate", False)
|
||||
|
||||
if self.use_full_rank_gate:
|
||||
# Keep f_a before the narrow beta shard, then pad each TP-local row
|
||||
# to select the aligned BF16 GEMM path. The padding also avoids an
|
||||
# Inductor correctness issue seen with the row-strided G view.
|
||||
qkvg_output_sizes = [self.projection_size] * 4
|
||||
in_proj_output_sizes = qkvg_output_sizes + [
|
||||
self.head_dim,
|
||||
self.num_heads,
|
||||
]
|
||||
local_output_size = (
|
||||
4 * self.local_projection_size + self.head_dim + self.local_num_heads
|
||||
)
|
||||
self.in_proj_padding = -local_output_size % 16
|
||||
if self.in_proj_padding:
|
||||
in_proj_output_sizes.append(self.in_proj_padding * self.tp_size)
|
||||
else:
|
||||
in_proj_output_sizes = [self.projection_size] * 3 + [
|
||||
self.num_heads,
|
||||
self.head_dim,
|
||||
]
|
||||
self.in_proj_padding = 0
|
||||
self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear(
|
||||
self.hidden_size,
|
||||
in_proj_output_sizes,
|
||||
replicated_shard_id=4,
|
||||
tp_size=self.tp_size,
|
||||
bias=False,
|
||||
quant_config=self.quant_config,
|
||||
prefix=f"{prefix}.in_proj_qkvgfab",
|
||||
)
|
||||
if self.in_proj_padding:
|
||||
self.in_proj_qkvgfab.weight.data[-self.in_proj_padding :].zero_()
|
||||
|
||||
self.f_b_proj = ColumnParallelLinear(
|
||||
self.head_dim,
|
||||
self.projection_size,
|
||||
bias=False,
|
||||
quant_config=self.quant_config,
|
||||
prefix=f"{prefix}.f_b_proj",
|
||||
)
|
||||
self.dt_bias = nn.Parameter(
|
||||
torch.empty(self.local_projection_size, dtype=torch.float32)
|
||||
)
|
||||
|
||||
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
|
||||
|
||||
# One packed parameter and cache let decode run a single conv update.
|
||||
# Prefill slices them back into Q/K/V to obtain dense outputs cheaply.
|
||||
self.conv1d = ColumnParallelLinear(
|
||||
input_size=self.conv_size,
|
||||
output_size=3 * self.projection_size,
|
||||
bias=False,
|
||||
params_dtype=torch.float32,
|
||||
prefix=f"{prefix}.conv1d",
|
||||
)
|
||||
self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
|
||||
delattr(self.conv1d.weight, "weight_loader")
|
||||
set_weight_attrs(
|
||||
self.conv1d.weight,
|
||||
{
|
||||
"weight_loader": _make_fused_conv1d_weight_loader(
|
||||
[self.projection_size] * 3,
|
||||
self.tp_size,
|
||||
self.tp_rank,
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
self.A_log = nn.Parameter(
|
||||
torch.empty(self.local_num_heads, dtype=torch.float32)
|
||||
)
|
||||
set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)})
|
||||
|
||||
self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None)
|
||||
if self.gate_lower_bound is not None:
|
||||
assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, (
|
||||
"KDA gate lower bound must be in "
|
||||
f"[{_KDA_GATE_LOGBOUND_MIN}, 0). "
|
||||
f"Got {self.gate_lower_bound}."
|
||||
)
|
||||
self.use_safe_gate = self.gate_lower_bound is not None
|
||||
additional_config = vllm_config.additional_config
|
||||
backend = (
|
||||
additional_config.get("kda_prefill_backend", "auto")
|
||||
if isinstance(additional_config, dict)
|
||||
else "auto"
|
||||
)
|
||||
backend = "triton" if backend == "auto" else backend
|
||||
assert backend == "triton", (
|
||||
"The shared Kimi GDN layer only supports the Triton KDA "
|
||||
f"prefill backend, got {backend!r}."
|
||||
)
|
||||
if not self.use_full_rank_gate:
|
||||
self.g_a_proj = ReplicatedLinear(
|
||||
self.hidden_size,
|
||||
self.head_dim,
|
||||
bias=False,
|
||||
quant_config=self.quant_config,
|
||||
prefix=f"{prefix}.g_a_proj",
|
||||
)
|
||||
self.g_b_proj = ColumnParallelLinear(
|
||||
self.head_dim,
|
||||
self.projection_size,
|
||||
bias=False,
|
||||
quant_config=self.quant_config,
|
||||
prefix=f"{prefix}.g_b_proj",
|
||||
)
|
||||
self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
|
||||
self.o_proj = RowParallelLinear(
|
||||
self.projection_size,
|
||||
self.hidden_size,
|
||||
bias=False,
|
||||
quant_config=self.quant_config,
|
||||
prefix=f"{prefix}.o_proj",
|
||||
)
|
||||
|
||||
compilation_config = vllm_config.compilation_config
|
||||
if prefix in compilation_config.static_forward_context:
|
||||
raise ValueError(f"Duplicate layer name: {prefix}")
|
||||
compilation_config.static_forward_context[prefix] = self
|
||||
|
||||
def rearrange_mixed_qkv(
|
||||
self, mixed_qkv: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
seq_len = mixed_qkv.shape[0]
|
||||
qkv = mixed_qkv.view(seq_len, 3, self.local_num_heads, self.head_dim)
|
||||
# Materialize all three row-strided inputs with one token-major to
|
||||
# QKV-major permutation. Each unbound tensor is then contiguous.
|
||||
qkv = qkv.permute(1, 0, 2, 3).contiguous().unsqueeze(1)
|
||||
return qkv.unbind(0)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
) -> None:
|
||||
num_tokens = hidden_states.size(0)
|
||||
projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
|
||||
if self.use_full_rank_gate:
|
||||
split_sizes = [
|
||||
3 * self.local_projection_size,
|
||||
self.local_projection_size,
|
||||
self.head_dim,
|
||||
self.local_num_heads,
|
||||
]
|
||||
if self.in_proj_padding:
|
||||
split_sizes.append(self.in_proj_padding)
|
||||
projected = projected_qkvgfab.split(split_sizes, dim=-1)
|
||||
mixed_qkv, g_proj_states, f_a, beta = projected[:4]
|
||||
else:
|
||||
mixed_qkv, beta, f_a = projected_qkvgfab.split(
|
||||
[
|
||||
3 * self.local_projection_size,
|
||||
self.local_num_heads,
|
||||
self.head_dim,
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
|
||||
|
||||
g1 = self.f_b_proj(f_a)[0]
|
||||
beta = beta.unsqueeze(0)
|
||||
g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)
|
||||
|
||||
g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
|
||||
|
||||
core_attn_out = torch.empty(
|
||||
(1, num_tokens, self.local_num_heads, self.head_dim),
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
self._forward(
|
||||
mixed_qkv=mixed_qkv,
|
||||
g1=g1,
|
||||
g2=g2,
|
||||
beta=beta,
|
||||
core_attn_out=core_attn_out,
|
||||
)
|
||||
core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
|
||||
output[:] = self.o_proj(core_attn_out)[0]
|
||||
|
||||
@eager_break_during_capture
|
||||
def _forward(
|
||||
self,
|
||||
mixed_qkv: torch.Tensor,
|
||||
g1: torch.Tensor,
|
||||
g2: torch.Tensor,
|
||||
beta: torch.Tensor,
|
||||
core_attn_out: torch.Tensor,
|
||||
) -> None:
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata_raw = forward_context.attn_metadata
|
||||
|
||||
if attn_metadata_raw is None:
|
||||
return
|
||||
|
||||
# Vendor-specific KDA kernels: AMD/ROCm and NVIDIA keep their own copies
|
||||
# under kimi_k3/{amd,nvidia}/ops so each can diverge independently.
|
||||
if current_platform.is_rocm():
|
||||
from vllm.models.kimi_k3.amd.ops.third_party.kda import (
|
||||
chunk_kda_with_fused_gate,
|
||||
fused_recurrent_kda,
|
||||
fused_recurrent_kda_packed_decode,
|
||||
)
|
||||
else:
|
||||
from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
|
||||
chunk_kda_with_fused_gate,
|
||||
fused_recurrent_kda,
|
||||
fused_recurrent_kda_packed_decode,
|
||||
)
|
||||
|
||||
assert isinstance(attn_metadata_raw, dict)
|
||||
attn_metadata_narrowed = attn_metadata_raw[self.prefix]
|
||||
assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)
|
||||
m = attn_metadata_narrowed
|
||||
has_initial_state = m.has_initial_state
|
||||
non_spec_query_start_loc = m.non_spec_query_start_loc
|
||||
non_spec_state_indices_tensor = m.non_spec_state_indices_tensor
|
||||
spec_sequence_masks = m.spec_sequence_masks
|
||||
spec_token_indx = m.spec_token_indx
|
||||
non_spec_token_indx = m.non_spec_token_indx
|
||||
spec_state_indices_tensor = m.spec_state_indices_tensor
|
||||
spec_query_start_loc = m.spec_query_start_loc
|
||||
num_accepted_tokens = m.num_accepted_tokens
|
||||
num_actual_tokens = m.num_actual_tokens
|
||||
mixed_qkv = mixed_qkv[:num_actual_tokens]
|
||||
g1 = g1[:, :num_actual_tokens]
|
||||
beta = beta[:, :num_actual_tokens]
|
||||
|
||||
constant_caches = self.kv_cache
|
||||
|
||||
conv_state, recurrent_state = constant_caches
|
||||
# conv_state must be (..., dim, width-1) for the conv kernels.
|
||||
# DS layout stores it that way directly; SD layout needs a transpose.
|
||||
if not is_conv_state_dim_first():
|
||||
conv_state = conv_state.transpose(-1, -2)
|
||||
|
||||
conv_weights = self.conv1d.weight.view(
|
||||
self.conv1d.weight.size(0), self.conv1d.weight.size(2)
|
||||
)
|
||||
q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split(
|
||||
self.local_projection_size, dim=0
|
||||
)
|
||||
q_conv_state, k_conv_state, v_conv_state = conv_state.split(
|
||||
self.local_projection_size, dim=-2
|
||||
)
|
||||
|
||||
# Split tokens into the multi-query spec-decode part and the remaining
|
||||
# (prefill / plain decode) part.
|
||||
if spec_sequence_masks is not None:
|
||||
if m.num_prefills == 0 and m.num_decodes == 0:
|
||||
mixed_qkv_spec = mixed_qkv
|
||||
g1_spec, beta_spec = g1, beta
|
||||
mixed_qkv_ns = g1_ns = beta_ns = None
|
||||
else:
|
||||
mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
|
||||
g1_spec = g1.index_select(1, spec_token_indx)
|
||||
beta_spec = beta.index_select(1, spec_token_indx)
|
||||
mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx)
|
||||
g1_ns = g1.index_select(1, non_spec_token_indx)
|
||||
beta_ns = beta.index_select(1, non_spec_token_indx)
|
||||
else:
|
||||
mixed_qkv_spec = g1_spec = beta_spec = None
|
||||
mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta
|
||||
|
||||
# ---------- spec-decode multi-query path ----------
|
||||
core_attn_out_spec = None
|
||||
if spec_sequence_masks is not None:
|
||||
assert spec_state_indices_tensor is not None
|
||||
assert spec_query_start_loc is not None
|
||||
spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes]
|
||||
spec_max_query_len = spec_state_indices_tensor.size(-1)
|
||||
|
||||
# Sibling beta and, for full-rank gates, output-gate views remain
|
||||
# live, so write the convolution output separately.
|
||||
spec_conv_out = torch.empty(
|
||||
mixed_qkv_spec.shape,
|
||||
dtype=mixed_qkv_spec.dtype,
|
||||
device=mixed_qkv_spec.device,
|
||||
)
|
||||
mixed_qkv_spec = causal_conv1d_update(
|
||||
mixed_qkv_spec,
|
||||
conv_state,
|
||||
conv_weights,
|
||||
self.conv1d.bias,
|
||||
activation="silu",
|
||||
conv_state_indices=spec_conv_indices,
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
query_start_loc=spec_query_start_loc,
|
||||
max_query_len=spec_max_query_len,
|
||||
validate_data=False,
|
||||
out=spec_conv_out,
|
||||
)
|
||||
q_spec, k_spec, v_spec = (
|
||||
rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
|
||||
for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1)
|
||||
)
|
||||
spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1]
|
||||
# Spec-only batches write directly into core_attn_out.
|
||||
spec_out = (
|
||||
core_attn_out[:, : q_spec.shape[1]]
|
||||
if m.num_prefills == 0 and m.num_decodes == 0
|
||||
else None
|
||||
)
|
||||
core_attn_out_spec, _ = fused_recurrent_kda(
|
||||
q=q_spec,
|
||||
k=k_spec,
|
||||
v=v_spec,
|
||||
raw_g=g1_spec,
|
||||
raw_beta=beta_spec,
|
||||
A_log=self.A_log,
|
||||
dt_bias=self.dt_bias,
|
||||
lower_bound=self.gate_lower_bound,
|
||||
initial_state=recurrent_state,
|
||||
cu_seqlens=spec_cu_seqlens,
|
||||
ssm_state_indices=spec_state_indices_tensor,
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
out=spec_out,
|
||||
)
|
||||
|
||||
# ---------- non-spec path (prefill or plain decode) ----------
|
||||
core_attn_out_non_spec = None
|
||||
if mixed_qkv_ns is not None:
|
||||
assert g1_ns is not None and beta_ns is not None
|
||||
if m.num_prefills > 0:
|
||||
q_ns, k_ns, v_ns = mixed_qkv_ns.split(
|
||||
self.local_projection_size, dim=-1
|
||||
)
|
||||
|
||||
# Packed prefill conv would require copying V solely to make
|
||||
# it dense for KDA. Separate calls accept the strided inputs
|
||||
# and produce dense Q/K/V without that extra traffic.
|
||||
# TODO: Use packed conv once every KDA prefill backend accepts
|
||||
# row-strided Q/K/V directly.
|
||||
def _prefill_conv(
|
||||
x: torch.Tensor,
|
||||
state: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return causal_conv1d_fn(
|
||||
x.transpose(0, 1),
|
||||
weight,
|
||||
None,
|
||||
activation="silu",
|
||||
conv_states=state,
|
||||
has_initial_state=has_initial_state,
|
||||
cache_indices=non_spec_state_indices_tensor,
|
||||
query_start_loc=non_spec_query_start_loc,
|
||||
metadata=m,
|
||||
).transpose(0, 1)
|
||||
|
||||
q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight)
|
||||
k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight)
|
||||
v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight)
|
||||
q_ns, k_ns, v_ns = (
|
||||
rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
|
||||
for x in (q_ns, k_ns, v_ns)
|
||||
)
|
||||
|
||||
assert non_spec_state_indices_tensor is not None
|
||||
assert has_initial_state is not None
|
||||
initial_state = gather_initial_states(
|
||||
recurrent_state,
|
||||
non_spec_state_indices_tensor,
|
||||
has_initial_state,
|
||||
)
|
||||
(
|
||||
core_attn_out_non_spec,
|
||||
last_recurrent_state,
|
||||
) = chunk_kda_with_fused_gate(
|
||||
q=q_ns,
|
||||
k=k_ns,
|
||||
v=v_ns,
|
||||
raw_g=g1_ns,
|
||||
raw_beta=beta_ns,
|
||||
A_log=self.A_log,
|
||||
g_bias=self.dt_bias,
|
||||
lower_bound=self.gate_lower_bound,
|
||||
initial_state=initial_state,
|
||||
output_final_state=True,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
cu_seqlens=non_spec_query_start_loc,
|
||||
)
|
||||
# Init cache
|
||||
recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
|
||||
|
||||
else:
|
||||
# pure-decode non-spec batch
|
||||
assert non_spec_state_indices_tensor is not None
|
||||
decode_conv_indices = non_spec_state_indices_tensor[
|
||||
: mixed_qkv_ns.size(0)
|
||||
]
|
||||
# Sibling beta and, for full-rank gates, output-gate views
|
||||
# remain live, so write the conv output separately.
|
||||
packed_conv_out = torch.empty(
|
||||
mixed_qkv_ns.shape,
|
||||
dtype=mixed_qkv_ns.dtype,
|
||||
device=mixed_qkv_ns.device,
|
||||
)
|
||||
mixed_qkv_ns = causal_conv1d_update(
|
||||
mixed_qkv_ns,
|
||||
conv_state,
|
||||
conv_weights,
|
||||
self.conv1d.bias,
|
||||
activation="silu",
|
||||
conv_state_indices=decode_conv_indices,
|
||||
validate_data=True,
|
||||
out=packed_conv_out,
|
||||
)
|
||||
core_attn_out_non_spec, _ = fused_recurrent_kda_packed_decode(
|
||||
mixed_qkv=mixed_qkv_ns,
|
||||
raw_g=g1_ns,
|
||||
raw_beta=beta_ns,
|
||||
A_log=self.A_log,
|
||||
dt_bias=self.dt_bias,
|
||||
lower_bound=self.gate_lower_bound,
|
||||
initial_state=recurrent_state,
|
||||
state_indices=decode_conv_indices,
|
||||
)
|
||||
|
||||
# ---------- merge spec and non-spec outputs ----------
|
||||
if core_attn_out_spec is not None and core_attn_out_non_spec is not None:
|
||||
# Mixed batches require indexed placement in the original order.
|
||||
merged = torch.empty(
|
||||
(1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
|
||||
dtype=core_attn_out_spec.dtype,
|
||||
device=core_attn_out_spec.device,
|
||||
)
|
||||
merged.index_copy_(1, spec_token_indx, core_attn_out_spec)
|
||||
merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
|
||||
core_attn_out[0, :num_actual_tokens] = merged[0, :num_actual_tokens]
|
||||
elif core_attn_out_non_spec is not None:
|
||||
core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
|
||||
0, :num_actual_tokens
|
||||
]
|
||||
else:
|
||||
assert core_attn_out_spec is not None
|
||||
core_attn_out.copy_(self.o_norm(core_attn_out, g2))
|
||||
634
upstream_ref/vllm_gdn/gdn/olmo_gdn_linear_attn.py
Normal file
634
upstream_ref/vllm_gdn/gdn/olmo_gdn_linear_attn.py
Normal file
@@ -0,0 +1,634 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
|
||||
from vllm.config import (
|
||||
VllmConfig,
|
||||
get_current_vllm_config,
|
||||
)
|
||||
from vllm.distributed import (
|
||||
divide,
|
||||
)
|
||||
from vllm.forward_context import ForwardContext, get_forward_context
|
||||
from vllm.model_executor.custom_op import PluggableLayer
|
||||
from vllm.model_executor.layers.layernorm import RMSNormGated
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
MergedColumnParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import (
|
||||
MambaStateShapeCalculator,
|
||||
is_conv_state_dim_first,
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
|
||||
causal_conv1d_fn,
|
||||
causal_conv1d_update,
|
||||
)
|
||||
from vllm.model_executor.model_loader.weight_utils import (
|
||||
sharded_weight_loader,
|
||||
)
|
||||
from vllm.model_executor.utils import set_weight_attrs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.third_party.flash_linear_attention.ops import (
|
||||
chunk_gated_delta_rule,
|
||||
fused_recurrent_gated_delta_rule,
|
||||
)
|
||||
from vllm.triton_utils import tl, triton
|
||||
from vllm.triton_utils.allocation import set_triton_allocator
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
|
||||
|
||||
|
||||
@PluggableLayer.register("olmo_hybrid_gated_delta_net_attention")
|
||||
class OlmoHybridGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
"""
|
||||
Gated DeltaNet linear attention layer for OLMo Hybrid.
|
||||
|
||||
This implements the linear attention mechanism that replaces sliding window
|
||||
attention in the hybrid architecture.
|
||||
"""
|
||||
|
||||
def get_state_shape(
|
||||
self,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
|
||||
return MambaStateShapeCalculator.gated_delta_net_state_shape(
|
||||
self.tp_size,
|
||||
self.num_k_heads,
|
||||
self.num_v_heads,
|
||||
self.head_k_dim,
|
||||
self.head_v_dim,
|
||||
self.conv_kernel_size,
|
||||
self.num_spec,
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config,
|
||||
vllm_config: VllmConfig,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__(config, vllm_config, prefix=prefix)
|
||||
|
||||
assert getattr(config, "linear_use_gate", True), (
|
||||
"OlmoHybridGatedDeltaNet requires linear_use_gate=True"
|
||||
)
|
||||
self.num_k_heads = config.linear_num_key_heads
|
||||
self.num_v_heads = config.linear_num_value_heads
|
||||
self.head_k_dim = config.linear_key_head_dim
|
||||
self.head_v_dim = config.linear_value_head_dim
|
||||
self.conv_kernel_size = config.linear_conv_kernel_dim
|
||||
self.key_dim = self.head_k_dim * self.num_k_heads
|
||||
self.value_dim = self.head_v_dim * self.num_v_heads
|
||||
self.allow_neg_eigval = getattr(config, "linear_allow_neg_eigval", False)
|
||||
|
||||
# Fused QKVG projection: 1 matmul instead of 4
|
||||
self.in_proj_qkvg = MergedColumnParallelLinear(
|
||||
input_size=self.hidden_size,
|
||||
output_sizes=[self.key_dim, self.key_dim, self.value_dim, self.value_dim],
|
||||
bias=False,
|
||||
quant_config=self.quant_config,
|
||||
prefix=f"{prefix}.in_proj_qkvg",
|
||||
)
|
||||
|
||||
# Separate B and A projections to preserve numerical precision.
|
||||
# Fusing these into one matmul changes FP accumulation order for the
|
||||
# gating scalars, which compounds through the GDN recurrent state.
|
||||
self.b_proj = ColumnParallelLinear(
|
||||
input_size=self.hidden_size,
|
||||
output_size=self.num_v_heads,
|
||||
bias=False,
|
||||
quant_config=self.quant_config,
|
||||
prefix=f"{prefix}.b_proj",
|
||||
)
|
||||
self.a_proj = ColumnParallelLinear(
|
||||
input_size=self.hidden_size,
|
||||
output_size=self.num_v_heads,
|
||||
bias=False,
|
||||
quant_config=self.quant_config,
|
||||
prefix=f"{prefix}.a_proj",
|
||||
)
|
||||
|
||||
# Fused conv1d: single parameter instead of 3
|
||||
self.conv_dim = self.key_dim * 2 + self.value_dim
|
||||
self.conv1d = ColumnParallelLinear(
|
||||
input_size=self.conv_kernel_size,
|
||||
output_size=self.conv_dim,
|
||||
bias=False,
|
||||
prefix=f"{prefix}.conv1d",
|
||||
)
|
||||
self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
|
||||
delattr(self.conv1d.weight, "weight_loader")
|
||||
set_weight_attrs(
|
||||
self.conv1d.weight,
|
||||
{
|
||||
"weight_loader": _make_fused_conv1d_weight_loader(
|
||||
[self.key_dim, self.key_dim, self.value_dim],
|
||||
self.tp_size,
|
||||
self.tp_rank,
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
self.dt_bias = nn.Parameter(
|
||||
torch.ones(self.num_v_heads // self.tp_size),
|
||||
)
|
||||
self.A_log = nn.Parameter(
|
||||
torch.empty(
|
||||
divide(self.num_v_heads, self.tp_size),
|
||||
)
|
||||
)
|
||||
|
||||
set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
|
||||
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
|
||||
|
||||
# use eps=1e-5 to match FLA's FusedRMSNormGated
|
||||
self.o_norm = RMSNormGated(
|
||||
self.head_v_dim,
|
||||
eps=1e-5,
|
||||
group_size=None,
|
||||
norm_before_gate=True,
|
||||
device=current_platform.current_device(),
|
||||
dtype=config.torch_dtype if hasattr(config, "torch_dtype") else None,
|
||||
)
|
||||
|
||||
self.o_proj = RowParallelLinear(
|
||||
self.value_dim,
|
||||
self.hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
quant_config=self.quant_config,
|
||||
prefix=f"{prefix}.o_proj",
|
||||
)
|
||||
|
||||
# FLA triton kernels need a PyTorch-backed allocator for scratch
|
||||
# memory (required by triton >= 3.x autotuner). Set once at init.
|
||||
set_triton_allocator(current_platform.current_device())
|
||||
|
||||
compilation_config = get_current_vllm_config().compilation_config
|
||||
if prefix in compilation_config.static_forward_context:
|
||||
raise ValueError(f"Duplicate layer name: {prefix}")
|
||||
compilation_config.static_forward_context[prefix] = self
|
||||
|
||||
def rearrange_mixed_qkv(self, mixed_qkv):
|
||||
if mixed_qkv is None:
|
||||
return None, None, None
|
||||
query, key, value = torch.split(
|
||||
mixed_qkv,
|
||||
[
|
||||
self.key_dim // self.tp_size,
|
||||
self.key_dim // self.tp_size,
|
||||
self.value_dim // self.tp_size,
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
num_k_heads = self.num_k_heads // self.tp_size
|
||||
num_v_heads = self.num_v_heads // self.tp_size
|
||||
|
||||
query = rearrange(query, "l (h d) -> 1 l h d", h=num_k_heads, d=self.head_k_dim)
|
||||
key = rearrange(key, "l (h d) -> 1 l h d", h=num_k_heads, d=self.head_k_dim)
|
||||
value = rearrange(value, "l (h d) -> 1 l h d", h=num_v_heads, d=self.head_v_dim)
|
||||
|
||||
# GQA expansion if needed
|
||||
if num_v_heads > num_k_heads:
|
||||
expand_ratio = num_v_heads // num_k_heads
|
||||
query = query.unsqueeze(3).expand(-1, -1, -1, expand_ratio, -1)
|
||||
query = query.reshape(1, query.shape[1], num_v_heads, self.head_k_dim)
|
||||
key = key.unsqueeze(3).expand(-1, -1, -1, expand_ratio, -1)
|
||||
key = key.reshape(1, key.shape[1], num_v_heads, self.head_k_dim)
|
||||
|
||||
return query.contiguous(), key.contiguous(), value.contiguous()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
):
|
||||
# NOTE: We wrap the ENTIRE linear attention forward (projections +
|
||||
# core recurrence + output norm + output projection) in a single
|
||||
# custom op, rather than just wrapping the recurrent core like
|
||||
# other GDN models (e.g. Qwen3Next) do.
|
||||
#
|
||||
# Why: torch.compile with inductor generates fused kernels for
|
||||
# matmuls and pointwise ops. These fused kernels can differ in
|
||||
# floating-point accumulation order from eager-mode cuBLAS,
|
||||
# introducing small numerical differences (~1e-7 per op). For
|
||||
# standard transformer attention this is harmless because each
|
||||
# position is computed independently. But for the GDN recurrent
|
||||
# state, these tiny input differences compound at every timestep
|
||||
# across the full sequence length, causing severe logprob
|
||||
# divergence (e.g. ~15% top-1 agreement with eager baseline).
|
||||
#
|
||||
# By making the full forward opaque to inductor, the projections
|
||||
# and output norm run with eager-mode kernels (cuBLAS, triton),
|
||||
# preserving numerical consistency. The tradeoff is reduced
|
||||
# compilation speedup (~1.5x vs ~3x), but logprob agreement
|
||||
# improves from ~15% to ~83% top-1 vs eager.
|
||||
#
|
||||
# The remaining ~17% divergence comes from inductor compiling
|
||||
# the MLP and transformer attention layers that are NOT wrapped
|
||||
# in custom ops -- their small precision differences propagate
|
||||
# as inputs to the GDN layers from outside.
|
||||
torch.ops.vllm.olmo_hybrid_gdn_full_forward(
|
||||
hidden_states,
|
||||
output,
|
||||
self.prefix,
|
||||
)
|
||||
|
||||
def _full_forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
):
|
||||
num_tokens = hidden_states.size(0)
|
||||
|
||||
# ============================================================
|
||||
# Part 1: Input Projection (2 fused matmuls instead of 6)
|
||||
# ============================================================
|
||||
projected_qkvg, _ = self.in_proj_qkvg(hidden_states)
|
||||
conv_dim_sharded = (self.key_dim * 2 + self.value_dim) // self.tp_size
|
||||
mixed_qkv = projected_qkvg[..., :conv_dim_sharded]
|
||||
gate = projected_qkvg[..., conv_dim_sharded:]
|
||||
|
||||
b, _ = self.b_proj(hidden_states)
|
||||
a, _ = self.a_proj(hidden_states)
|
||||
|
||||
# ============================================================
|
||||
# Part 2: Core Attention
|
||||
# ============================================================
|
||||
core_attn_out = torch.zeros(
|
||||
(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
self._forward_core(
|
||||
mixed_qkv=mixed_qkv,
|
||||
b=b,
|
||||
a=a,
|
||||
core_attn_out=core_attn_out,
|
||||
)
|
||||
|
||||
# ============================================================
|
||||
# Part 3: Output Projection
|
||||
# ============================================================
|
||||
gate = gate.view(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim)
|
||||
core_attn_out_flat = core_attn_out.reshape(-1, core_attn_out.shape[-1])
|
||||
gate_flat = gate.reshape(-1, gate.shape[-1])
|
||||
core_attn_out_normed = self.o_norm(core_attn_out_flat, gate_flat)
|
||||
core_attn_out = core_attn_out_normed.view(
|
||||
num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim
|
||||
)
|
||||
|
||||
core_attn_out = rearrange(core_attn_out, "l h d -> l (h d)")
|
||||
output[:num_tokens], _ = self.o_proj(core_attn_out)
|
||||
|
||||
def _forward_core(
|
||||
self,
|
||||
mixed_qkv: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
a: torch.Tensor,
|
||||
core_attn_out: torch.Tensor,
|
||||
):
|
||||
"""
|
||||
Core attention computation (called by custom op).
|
||||
"""
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = forward_context.attn_metadata
|
||||
|
||||
if attn_metadata is None:
|
||||
# V1 profile run
|
||||
return
|
||||
|
||||
assert isinstance(attn_metadata, dict)
|
||||
attn_metadata = attn_metadata[self.prefix] # type: ignore[assignment]
|
||||
assert isinstance(attn_metadata, GDNAttentionMetadata)
|
||||
has_initial_state = attn_metadata.has_initial_state
|
||||
spec_query_start_loc = attn_metadata.spec_query_start_loc
|
||||
non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
|
||||
spec_sequence_masks = attn_metadata.spec_sequence_masks
|
||||
spec_token_indx = attn_metadata.spec_token_indx
|
||||
non_spec_token_indx = attn_metadata.non_spec_token_indx
|
||||
spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor
|
||||
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor
|
||||
self_kv_cache = self.kv_cache
|
||||
# conv_state must be (..., dim, width-1) for the conv kernels.
|
||||
# DS layout stores it that way directly; SD layout needs a transpose.
|
||||
conv_state = (
|
||||
self_kv_cache[0]
|
||||
if is_conv_state_dim_first()
|
||||
else self_kv_cache[0].transpose(-1, -2)
|
||||
)
|
||||
ssm_state = self_kv_cache[1]
|
||||
num_actual_tokens = attn_metadata.num_actual_tokens
|
||||
num_accepted_tokens = attn_metadata.num_accepted_tokens
|
||||
|
||||
mixed_qkv = mixed_qkv[:num_actual_tokens]
|
||||
b = b[:num_actual_tokens]
|
||||
a = a[:num_actual_tokens]
|
||||
|
||||
conv_weights = self.conv1d.weight.view(
|
||||
self.conv1d.weight.size(0), self.conv1d.weight.size(2)
|
||||
)
|
||||
|
||||
if spec_sequence_masks is not None:
|
||||
if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
|
||||
mixed_qkv_spec = mixed_qkv
|
||||
mixed_qkv_non_spec = None
|
||||
else:
|
||||
mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
|
||||
mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx)
|
||||
else:
|
||||
mixed_qkv_spec = None
|
||||
mixed_qkv_non_spec = mixed_qkv
|
||||
|
||||
if spec_sequence_masks is not None:
|
||||
assert spec_query_start_loc is not None
|
||||
assert spec_state_indices_tensor is not None
|
||||
assert num_accepted_tokens is not None
|
||||
mixed_qkv_spec = causal_conv1d_update(
|
||||
mixed_qkv_spec,
|
||||
conv_state,
|
||||
conv_weights,
|
||||
None, # no bias
|
||||
self.activation,
|
||||
conv_state_indices=spec_state_indices_tensor[:, 0][
|
||||
: attn_metadata.num_spec_decodes
|
||||
],
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
query_start_loc=spec_query_start_loc,
|
||||
max_query_len=spec_state_indices_tensor.size(-1),
|
||||
validate_data=False,
|
||||
)
|
||||
|
||||
if attn_metadata.num_prefills > 0:
|
||||
assert mixed_qkv_non_spec is not None
|
||||
mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1)
|
||||
mixed_qkv_non_spec = causal_conv1d_fn(
|
||||
mixed_qkv_non_spec_T,
|
||||
conv_weights,
|
||||
None,
|
||||
activation=self.activation,
|
||||
conv_states=conv_state,
|
||||
has_initial_state=has_initial_state,
|
||||
cache_indices=non_spec_state_indices_tensor,
|
||||
query_start_loc=non_spec_query_start_loc,
|
||||
metadata=attn_metadata,
|
||||
).transpose(0, 1)
|
||||
elif attn_metadata.num_decodes > 0:
|
||||
assert non_spec_state_indices_tensor is not None
|
||||
mixed_qkv_non_spec = causal_conv1d_update(
|
||||
mixed_qkv_non_spec,
|
||||
conv_state,
|
||||
conv_weights,
|
||||
None,
|
||||
self.activation,
|
||||
conv_state_indices=non_spec_state_indices_tensor[
|
||||
: attn_metadata.num_decodes
|
||||
],
|
||||
validate_data=True,
|
||||
)
|
||||
else:
|
||||
mixed_qkv_non_spec = None
|
||||
|
||||
query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
|
||||
query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(
|
||||
mixed_qkv_non_spec
|
||||
)
|
||||
|
||||
g, beta = fused_olmo_hybrid_gdn_gating(
|
||||
self.A_log, a, b, self.dt_bias, self.allow_neg_eigval
|
||||
)
|
||||
|
||||
if spec_sequence_masks is not None:
|
||||
assert spec_token_indx is not None
|
||||
assert non_spec_token_indx is not None
|
||||
if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
|
||||
g_spec = g
|
||||
beta_spec = beta
|
||||
g_non_spec = None
|
||||
beta_non_spec = None
|
||||
else:
|
||||
g_spec = g.index_select(1, spec_token_indx)
|
||||
beta_spec = beta.index_select(1, spec_token_indx)
|
||||
g_non_spec = g.index_select(1, non_spec_token_indx)
|
||||
beta_non_spec = beta.index_select(1, non_spec_token_indx)
|
||||
else:
|
||||
g_spec = None
|
||||
beta_spec = None
|
||||
g_non_spec = g
|
||||
beta_non_spec = beta
|
||||
|
||||
if spec_sequence_masks is not None:
|
||||
assert spec_query_start_loc is not None
|
||||
assert spec_state_indices_tensor is not None
|
||||
assert num_accepted_tokens is not None
|
||||
core_attn_out_spec, last_recurrent_state = fused_recurrent_gated_delta_rule(
|
||||
q=query_spec,
|
||||
k=key_spec,
|
||||
v=value_spec,
|
||||
g=g_spec,
|
||||
beta=beta_spec,
|
||||
initial_state=ssm_state,
|
||||
inplace_final_state=True,
|
||||
cu_seqlens=spec_query_start_loc[: attn_metadata.num_spec_decodes + 1],
|
||||
ssm_state_indices=spec_state_indices_tensor,
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
else:
|
||||
core_attn_out_spec, last_recurrent_state = None, None
|
||||
|
||||
if attn_metadata.num_prefills > 0:
|
||||
assert non_spec_state_indices_tensor is not None
|
||||
assert has_initial_state is not None
|
||||
assert non_spec_query_start_loc is not None
|
||||
initial_state = ssm_state[non_spec_state_indices_tensor].contiguous()
|
||||
initial_state[~has_initial_state, ...] = 0
|
||||
(
|
||||
core_attn_out_non_spec,
|
||||
last_recurrent_state,
|
||||
) = chunk_gated_delta_rule(
|
||||
q=query_non_spec,
|
||||
k=key_non_spec,
|
||||
v=value_non_spec,
|
||||
g=g_non_spec,
|
||||
beta=beta_non_spec,
|
||||
initial_state=initial_state,
|
||||
output_final_state=True,
|
||||
cu_seqlens=non_spec_query_start_loc,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(
|
||||
ssm_state.dtype
|
||||
)
|
||||
elif attn_metadata.num_decodes > 0:
|
||||
assert non_spec_query_start_loc is not None
|
||||
assert non_spec_state_indices_tensor is not None
|
||||
core_attn_out_non_spec, last_recurrent_state = (
|
||||
fused_recurrent_gated_delta_rule(
|
||||
q=query_non_spec,
|
||||
k=key_non_spec,
|
||||
v=value_non_spec,
|
||||
g=g_non_spec,
|
||||
beta=beta_non_spec,
|
||||
initial_state=ssm_state,
|
||||
inplace_final_state=True,
|
||||
cu_seqlens=non_spec_query_start_loc[
|
||||
: attn_metadata.num_decodes + 1
|
||||
],
|
||||
ssm_state_indices=non_spec_state_indices_tensor,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
)
|
||||
else:
|
||||
core_attn_out_non_spec, last_recurrent_state = None, None
|
||||
|
||||
if spec_sequence_masks is not None and core_attn_out_non_spec is not None:
|
||||
merged_out = torch.empty(
|
||||
(1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
|
||||
dtype=core_attn_out_non_spec.dtype,
|
||||
device=core_attn_out_non_spec.device,
|
||||
)
|
||||
merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
|
||||
merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
|
||||
core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)
|
||||
elif spec_sequence_masks is not None:
|
||||
core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
|
||||
else:
|
||||
core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
|
||||
|
||||
|
||||
def _make_fused_conv1d_weight_loader(dims, tp_size, tp_rank):
|
||||
"""Weight loader for loading separate HF conv weights into a fused conv1d.
|
||||
|
||||
dims: list of original (un-sharded) dims per section,
|
||||
e.g. [key_dim, key_dim, value_dim]
|
||||
"""
|
||||
sharded_dims = [d // tp_size for d in dims]
|
||||
|
||||
def weight_loader(param, loaded_weight, loaded_shard_id=None):
|
||||
if loaded_weight.dim() == 2:
|
||||
loaded_weight = loaded_weight.unsqueeze(1)
|
||||
dim = dims[loaded_shard_id]
|
||||
shard_size = dim // tp_size
|
||||
tp_start = tp_rank * shard_size
|
||||
sharded_weight = loaded_weight[tp_start : tp_start + shard_size]
|
||||
offset = sum(sharded_dims[:loaded_shard_id])
|
||||
param.data[offset : offset + shard_size].copy_(sharded_weight)
|
||||
|
||||
return weight_loader
|
||||
|
||||
|
||||
def olmo_hybrid_gdn_full_forward(
|
||||
hidden_states: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> None:
|
||||
"""Full linear attention forward wrapped as a custom op.
|
||||
|
||||
Prevents inductor from compiling the projections around the GDN core,
|
||||
which would introduce numerical divergence that compounds through
|
||||
the recurrent state.
|
||||
"""
|
||||
forward_context: ForwardContext = get_forward_context()
|
||||
self = forward_context.no_compile_layers[layer_name]
|
||||
self._full_forward(
|
||||
hidden_states=hidden_states,
|
||||
output=output,
|
||||
)
|
||||
|
||||
|
||||
def olmo_hybrid_gdn_full_forward_fake(
|
||||
hidden_states: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> None:
|
||||
"""Fake implementation for torch.compile."""
|
||||
return
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="olmo_hybrid_gdn_full_forward",
|
||||
op_func=olmo_hybrid_gdn_full_forward,
|
||||
mutates_args=["output"],
|
||||
fake_impl=olmo_hybrid_gdn_full_forward_fake,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def fused_olmo_hybrid_gdn_gating_kernel(
|
||||
g,
|
||||
beta_output,
|
||||
A_log,
|
||||
a,
|
||||
b,
|
||||
dt_bias,
|
||||
seq_len,
|
||||
allow_neg_eigval: tl.constexpr,
|
||||
NUM_HEADS: tl.constexpr,
|
||||
beta: tl.constexpr,
|
||||
threshold: tl.constexpr,
|
||||
BLK_HEADS: tl.constexpr,
|
||||
):
|
||||
i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2)
|
||||
head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS)
|
||||
off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off
|
||||
mask = head_off < NUM_HEADS
|
||||
blk_A_log = tl.load(A_log + head_off, mask=mask)
|
||||
blk_a = tl.load(a + off, mask=mask)
|
||||
blk_b = tl.load(b + off, mask=mask)
|
||||
blk_bias = tl.load(dt_bias + head_off, mask=mask)
|
||||
|
||||
# g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
|
||||
x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)
|
||||
softplus_x = tl.where(
|
||||
beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x
|
||||
)
|
||||
blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x
|
||||
tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask)
|
||||
|
||||
# beta = self.b_proj(hidden_states).sigmoid()
|
||||
# if self.allow_neg_eigval: beta = beta * 2.0
|
||||
blk_beta_output = tl.sigmoid(blk_b.to(tl.float32))
|
||||
if allow_neg_eigval:
|
||||
blk_beta_output = blk_beta_output * 2.0
|
||||
tl.store(
|
||||
beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask
|
||||
)
|
||||
|
||||
|
||||
def fused_olmo_hybrid_gdn_gating(
|
||||
A_log: torch.Tensor,
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
dt_bias: torch.Tensor,
|
||||
allow_neg_eigval: bool = False,
|
||||
beta: float = 1.0,
|
||||
threshold: float = 20.0,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
batch, num_heads = a.shape
|
||||
seq_len = 1
|
||||
grid = (batch, seq_len, triton.cdiv(num_heads, 8))
|
||||
g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device)
|
||||
beta_output = torch.empty(1, batch, num_heads, dtype=torch.float32, device=b.device)
|
||||
fused_olmo_hybrid_gdn_gating_kernel[grid](
|
||||
g,
|
||||
beta_output,
|
||||
A_log,
|
||||
a,
|
||||
b,
|
||||
dt_bias,
|
||||
seq_len,
|
||||
allow_neg_eigval,
|
||||
num_heads,
|
||||
beta,
|
||||
threshold,
|
||||
8,
|
||||
num_warps=1,
|
||||
)
|
||||
return g, beta_output
|
||||
1751
upstream_ref/vllm_gdn/gdn/qwen_gdn_linear_attn.py
Normal file
1751
upstream_ref/vllm_gdn/gdn/qwen_gdn_linear_attn.py
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user