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
733
upstream_ref/vllm_gdn/models/qwen3_5.py
Normal file
733
upstream_ref/vllm_gdn/models/qwen3_5.py
Normal file
@@ -0,0 +1,733 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Copyright 2025 The vLLM team.
|
||||
# Copyright 2025 The Qwen Team.
|
||||
# Copyright 2025 The HuggingFace Inc. team.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
||||
# and OPT implementations in this library. It has been modified from its
|
||||
# original forms to accommodate minor architectural differences compared
|
||||
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Inference-only Qwen3.5 Series compatible with HuggingFace weights."""
|
||||
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.compilation.decorators import support_torch_compile
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.distributed import (
|
||||
get_pp_group,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.layernorm import GemmaRMSNorm as Qwen3_5RMSNorm
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
|
||||
QwenGatedDeltaNetAttention,
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import (
|
||||
MambaStateCopyFunc,
|
||||
MambaStateCopyFuncCalculator,
|
||||
MambaStateDtypeCalculator,
|
||||
MambaStateShapeCalculator,
|
||||
)
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead,
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.tokenizers.registry import cached_tokenizer_from_config
|
||||
from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig
|
||||
from vllm.transformers_utils.configs.qwen3_5_moe import (
|
||||
Qwen3_5MoeConfig,
|
||||
Qwen3_5MoeTextConfig,
|
||||
)
|
||||
|
||||
from .interfaces import (
|
||||
HasInnerState,
|
||||
IsHybrid,
|
||||
MixtureOfExperts,
|
||||
MultiModalEmbeddings,
|
||||
SupportsEagle3,
|
||||
SupportsLoRA,
|
||||
SupportsMRoPE,
|
||||
SupportsPP,
|
||||
_require_is_multimodal,
|
||||
)
|
||||
from .qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
|
||||
from .qwen3_next import (
|
||||
Qwen3NextAttention,
|
||||
Qwen3NextDecoderLayer,
|
||||
Qwen3NextModel,
|
||||
Qwen3NextSparseMoeBlock,
|
||||
QwenNextMixtureOfExperts,
|
||||
_is_shared_expert_fse_compatible,
|
||||
)
|
||||
from .qwen3_vl import (
|
||||
Qwen3_VisionTransformer,
|
||||
Qwen3VLDummyInputsBuilder,
|
||||
Qwen3VLForConditionalGeneration,
|
||||
Qwen3VLMultiModalProcessor,
|
||||
Qwen3VLProcessingInfo,
|
||||
)
|
||||
from .utils import (
|
||||
AutoWeightsLoader,
|
||||
PPMissingLayer,
|
||||
WeightsMapper,
|
||||
_merge_multimodal_embeddings,
|
||||
extract_layer_index,
|
||||
make_empty_intermediate_tensors_factory,
|
||||
make_layers,
|
||||
maybe_fuse_shared_experts,
|
||||
maybe_prefix,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class Qwen3_5ProcessingInfo(Qwen3VLProcessingInfo):
|
||||
def get_hf_config(self):
|
||||
return self.ctx.get_hf_config(Qwen3_5Config)
|
||||
|
||||
|
||||
class Qwen3_5MoeProcessingInfo(Qwen3VLProcessingInfo):
|
||||
def get_hf_config(self):
|
||||
# transformers 5.x renames the top-level Qwen3.5-MoE config class to
|
||||
# Qwen3_5MoeTextConfig for text-only models, while transformers ≤4.x
|
||||
# returns Qwen3_5MoeConfig (the multimodal wrapper). Accept both so
|
||||
# that vLLM works regardless of which transformers version is installed.
|
||||
return self.ctx.get_hf_config((Qwen3_5MoeConfig, Qwen3_5MoeTextConfig))
|
||||
|
||||
|
||||
class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
layer_type: str,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super(Qwen3NextDecoderLayer, self).__init__()
|
||||
|
||||
config = vllm_config.model_config.hf_text_config
|
||||
model_config = vllm_config.model_config
|
||||
cache_config = vllm_config.cache_config
|
||||
parallel_config = vllm_config.parallel_config
|
||||
quant_config = vllm_config.quant_config
|
||||
|
||||
self.layer_type = layer_type
|
||||
self.layer_idx = extract_layer_index(prefix)
|
||||
is_moe_layer = config.model_type == "qwen3_5_moe_text"
|
||||
self.use_attn_reduce_scatter_for_moe = (
|
||||
parallel_config.use_sequence_parallel_moe
|
||||
and parallel_config.pipeline_parallel_size == 1
|
||||
and is_moe_layer
|
||||
)
|
||||
|
||||
if self.layer_type == "linear_attention":
|
||||
self.linear_attn = QwenGatedDeltaNetAttention(
|
||||
config=config,
|
||||
vllm_config=vllm_config,
|
||||
prefix=f"{prefix}.linear_attn",
|
||||
gqa_interleaved_layout=False,
|
||||
reduce_results=not self.use_attn_reduce_scatter_for_moe,
|
||||
)
|
||||
elif self.layer_type == "full_attention":
|
||||
self.self_attn = Qwen3NextAttention(
|
||||
config,
|
||||
model_config=model_config,
|
||||
cache_config=cache_config,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.self_attn",
|
||||
reduce_results=not self.use_attn_reduce_scatter_for_moe,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid layer_type {self.layer_type}")
|
||||
|
||||
# NOTE: Determine the MLP type based on the model type
|
||||
# Qwen3.5 use all layers for MLP / Qwen3.5-MoE use sparse MoE blocks
|
||||
if config.model_type == "qwen3_5_moe_text":
|
||||
self.mlp = Qwen3NextSparseMoeBlock(
|
||||
vllm_config=vllm_config,
|
||||
prefix=f"{prefix}.mlp",
|
||||
)
|
||||
elif config.model_type == "qwen3_5_text":
|
||||
self.mlp = Qwen3NextMLP(
|
||||
hidden_size=config.hidden_size,
|
||||
intermediate_size=config.intermediate_size,
|
||||
hidden_act=config.hidden_act,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.mlp",
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid model_type {config.model_type}")
|
||||
|
||||
self.input_layernorm = Qwen3_5RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
self.post_attention_layernorm = Qwen3_5RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
|
||||
self.layer_scale = getattr(config, "layer_scale", False)
|
||||
if self.layer_scale:
|
||||
self.attn_layer_scale = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
1,
|
||||
1,
|
||||
config.hidden_size,
|
||||
),
|
||||
)
|
||||
self.ffn_layer_scale = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
1,
|
||||
1,
|
||||
config.hidden_size,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@support_torch_compile(
|
||||
dynamic_arg_dims={
|
||||
"input_ids": 0,
|
||||
# positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
|
||||
# otherwise (seq_len, ).
|
||||
"positions": -1,
|
||||
"intermediate_tensors": 0,
|
||||
"inputs_embeds": 0,
|
||||
}
|
||||
)
|
||||
class Qwen3_5Model(Qwen3NextModel):
|
||||
# Qwen3.5 ships the GDN in_proj checkpoints separately (qwen3-next
|
||||
# pre-fuses them); fuse them on top of the qwen3-next QKV/gate_up mapping.
|
||||
hf_to_vllm_mapper = Qwen3NextModel.hf_to_vllm_mapper | WeightsMapper(
|
||||
orig_to_new_stacked={
|
||||
".in_proj_qkv": (".in_proj_qkvz", (0, 1, 2)),
|
||||
".in_proj_z": (".in_proj_qkvz", 3),
|
||||
".in_proj_b": (".in_proj_ba", 0),
|
||||
".in_proj_a": (".in_proj_ba", 1),
|
||||
}
|
||||
)
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super(Qwen3NextModel, self).__init__()
|
||||
|
||||
config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig = (
|
||||
vllm_config.model_config.hf_text_config
|
||||
)
|
||||
parallel_config = vllm_config.parallel_config
|
||||
|
||||
eplb_config = parallel_config.eplb_config
|
||||
self.num_redundant_experts = eplb_config.num_redundant_experts
|
||||
|
||||
self.config = config
|
||||
self.quant_config = vllm_config.quant_config
|
||||
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
self.vocab_size,
|
||||
config.hidden_size,
|
||||
)
|
||||
|
||||
def get_layer(prefix: str):
|
||||
return Qwen3_5DecoderLayer(
|
||||
vllm_config,
|
||||
layer_type=config.layer_types[extract_layer_index(prefix)],
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
self.start_layer, self.end_layer, self.layers = make_layers(
|
||||
config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
|
||||
)
|
||||
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
|
||||
["hidden_states", "residual"], config.hidden_size
|
||||
)
|
||||
|
||||
if get_pp_group().is_last_rank:
|
||||
self.norm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
else:
|
||||
self.norm = PPMissingLayer()
|
||||
|
||||
self.aux_hidden_state_layers: tuple[int, ...] = ()
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
# FSE must match construction (Qwen3NextSparseMoeBlock): reroute the
|
||||
# shared expert into the extra fused slot only when AITER FSE is both
|
||||
# requested and compatible with the quant spec.
|
||||
if "moe" in self.config.model_type:
|
||||
weights = maybe_fuse_shared_experts(
|
||||
weights,
|
||||
enabled=rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
|
||||
and _is_shared_expert_fse_compatible(self.quant_config),
|
||||
n_routed_experts=self.config.num_experts,
|
||||
n_shared_experts=1,
|
||||
ckpt_prefix="mlp.shared_expert",
|
||||
)
|
||||
loader = AutoWeightsLoader(self)
|
||||
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
|
||||
|
||||
|
||||
class Qwen3_5ForCausalLMBase(
|
||||
nn.Module,
|
||||
HasInnerState,
|
||||
IsHybrid,
|
||||
SupportsEagle3,
|
||||
SupportsLoRA,
|
||||
SupportsMRoPE,
|
||||
SupportsPP,
|
||||
):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
# GDN fused projections.
|
||||
"in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
|
||||
"in_proj_ba": ["in_proj_b", "in_proj_a"],
|
||||
}
|
||||
|
||||
# Some community text-only checkpoints keep the extraneous
|
||||
# `model.language_model.` prefix inherited from the VL training stack.
|
||||
# Strip it so both prefixed and clean checkpoints load correctly.
|
||||
hf_to_vllm_mapper = WeightsMapper(
|
||||
orig_to_new_prefix={"model.language_model.": "model."},
|
||||
)
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
config = vllm_config.model_config.hf_text_config
|
||||
self.vllm_config = vllm_config
|
||||
self.model_config = vllm_config.model_config
|
||||
cache_config = vllm_config.cache_config
|
||||
|
||||
scheduler_config = vllm_config.scheduler_config
|
||||
if cache_config.mamba_cache_mode == "all":
|
||||
raise NotImplementedError(
|
||||
"Qwen3.5 currently does not support 'all' prefix caching, "
|
||||
"please use '--mamba-cache-mode=align' instead"
|
||||
)
|
||||
self.quant_config = vllm_config.quant_config
|
||||
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.scheduler_config = scheduler_config
|
||||
self.model = Qwen3_5Model(
|
||||
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
|
||||
)
|
||||
|
||||
if get_pp_group().is_last_rank:
|
||||
if config.tie_word_embeddings:
|
||||
self.lm_head = self.model.embed_tokens
|
||||
else:
|
||||
self.lm_head = ParallelLMHead(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
quant_config=self.quant_config,
|
||||
prefix=maybe_prefix(prefix, "lm_head"),
|
||||
)
|
||||
else:
|
||||
self.lm_head = PPMissingLayer()
|
||||
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size)
|
||||
self.make_empty_intermediate_tensors = (
|
||||
self.model.make_empty_intermediate_tensors
|
||||
)
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.model.embed_input_ids(input_ids)
|
||||
|
||||
def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
|
||||
self.model.aux_hidden_state_layers = layers
|
||||
|
||||
def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
|
||||
num_layers = len(self.model.layers)
|
||||
return (2, num_layers // 2, num_layers - 3)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: object,
|
||||
):
|
||||
hidden_states = self.model(
|
||||
input_ids, positions, intermediate_tensors, inputs_embeds
|
||||
)
|
||||
|
||||
return hidden_states
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_dtype_from_config(
|
||||
cls,
|
||||
vllm_config: "VllmConfig",
|
||||
) -> tuple[torch.dtype, torch.dtype]:
|
||||
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
|
||||
vllm_config.model_config.dtype,
|
||||
vllm_config.cache_config.mamba_cache_dtype,
|
||||
vllm_config.cache_config.mamba_ssm_cache_dtype,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_shape_from_config(
|
||||
cls, vllm_config: "VllmConfig"
|
||||
) -> tuple[tuple[int, int], tuple[int, int]]:
|
||||
parallel_config = vllm_config.parallel_config
|
||||
hf_config = vllm_config.model_config.hf_text_config
|
||||
tp_size = parallel_config.tensor_parallel_size
|
||||
num_spec = (
|
||||
vllm_config.speculative_config.num_speculative_tokens
|
||||
if vllm_config.speculative_config
|
||||
else 0
|
||||
)
|
||||
return MambaStateShapeCalculator.gated_delta_net_state_shape(
|
||||
tp_size,
|
||||
hf_config.linear_num_key_heads,
|
||||
hf_config.linear_num_value_heads,
|
||||
hf_config.linear_key_head_dim,
|
||||
hf_config.linear_value_head_dim,
|
||||
hf_config.linear_conv_kernel_dim,
|
||||
num_spec,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_copy_func(
|
||||
cls,
|
||||
) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
|
||||
return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor | None:
|
||||
return self.logits_processor(self.lm_head, hidden_states)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(
|
||||
self,
|
||||
skip_prefixes=["mtp."],
|
||||
)
|
||||
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
|
||||
|
||||
def get_mrope_input_positions(
|
||||
self,
|
||||
input_tokens: list[int],
|
||||
mm_features: list[object],
|
||||
) -> tuple[torch.Tensor, int]:
|
||||
positions = torch.arange(len(input_tokens), dtype=torch.long)
|
||||
return positions.unsqueeze(0).expand(3, -1), 0
|
||||
|
||||
|
||||
class Qwen3_5ForCausalLM(Qwen3_5ForCausalLMBase):
|
||||
pass
|
||||
|
||||
|
||||
class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts):
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__(vllm_config=vllm_config, prefix=prefix)
|
||||
|
||||
# set MoE hyperparameters
|
||||
self.set_moe_parameters()
|
||||
|
||||
|
||||
########################################################
|
||||
# Qwen3_5-Dense
|
||||
########################################################
|
||||
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_processor(
|
||||
Qwen3VLMultiModalProcessor,
|
||||
info=Qwen3_5ProcessingInfo,
|
||||
dummy_inputs=Qwen3VLDummyInputsBuilder,
|
||||
)
|
||||
class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid):
|
||||
supports_multimodal_pruning = True
|
||||
|
||||
packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | {
|
||||
"in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
|
||||
"in_proj_ba": ["in_proj_b", "in_proj_a"],
|
||||
}
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
|
||||
# protocols have not __init__ method, so we need to use nn.Module.__init__
|
||||
nn.Module.__init__(self)
|
||||
config: Qwen3_5Config = vllm_config.model_config.hf_config
|
||||
quant_config = vllm_config.quant_config
|
||||
multimodal_config = vllm_config.model_config.multimodal_config
|
||||
|
||||
self.config = config
|
||||
self.model_config = vllm_config.model_config
|
||||
self.multimodal_config = multimodal_config
|
||||
self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
|
||||
self.is_multimodal_pruning_enabled = (
|
||||
multimodal_config.is_multimodal_pruning_enabled()
|
||||
)
|
||||
self.video_pruning_rate = self.multimodal_config.video_pruning_rate
|
||||
self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
|
||||
|
||||
# attributes needed by EVS-related functions inherited from Qwen3-VL
|
||||
self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes")
|
||||
self.deepstack_num_level = (
|
||||
len(config.vision_config.deepstack_visual_indexes)
|
||||
if self.use_deepstack
|
||||
else 0
|
||||
)
|
||||
self.visual_dim = config.vision_config.out_hidden_size
|
||||
self.multiscale_dim = self.visual_dim * self.deepstack_num_level
|
||||
|
||||
with self._mark_tower_model(vllm_config, {"image", "video"}):
|
||||
self.visual = Qwen3_VisionTransformer(
|
||||
config.vision_config,
|
||||
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, "visual"),
|
||||
)
|
||||
|
||||
with self._mark_language_model(vllm_config):
|
||||
self.language_model = Qwen3_5ForCausalLM(
|
||||
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model")
|
||||
)
|
||||
|
||||
self.make_empty_intermediate_tensors = (
|
||||
self.language_model.make_empty_intermediate_tensors
|
||||
)
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: MultiModalEmbeddings | None = None,
|
||||
*,
|
||||
is_multimodal: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
inputs_embeds = self._embed_text_input_ids(
|
||||
input_ids,
|
||||
self.language_model.embed_input_ids,
|
||||
is_multimodal=is_multimodal,
|
||||
)
|
||||
|
||||
if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
|
||||
return inputs_embeds
|
||||
|
||||
is_multimodal = _require_is_multimodal(is_multimodal)
|
||||
|
||||
inputs_embeds = _merge_multimodal_embeddings(
|
||||
inputs_embeds=inputs_embeds,
|
||||
multimodal_embeddings=multimodal_embeddings,
|
||||
is_multimodal=is_multimodal,
|
||||
)
|
||||
|
||||
return inputs_embeds
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: object,
|
||||
) -> torch.Tensor | IntermediateTensors:
|
||||
"""Run forward pass for Qwen3.5.
|
||||
|
||||
Args:
|
||||
input_ids: Flattened (concatenated) input_ids corresponding to a
|
||||
batch.
|
||||
positions: Flattened (concatenated) position ids corresponding to a
|
||||
batch.
|
||||
**NOTE**: If mrope is enabled (default setting for Qwen3VL
|
||||
opensource models), the shape will be `(3, seq_len)`,
|
||||
otherwise it will be `(seq_len,).
|
||||
intermediate_tensors: Intermediate tensors from previous pipeline
|
||||
stages.
|
||||
inputs_embeds: Pre-computed input embeddings.
|
||||
**kwargs: Additional keyword arguments including:
|
||||
- pixel_values: Pixel values to be fed to a model.
|
||||
`None` if no images are passed.
|
||||
- image_grid_thw: Tensor `(n_images, 3)` of image 3D grid in
|
||||
LLM. `None` if no images are passed.
|
||||
- pixel_values_videos: Pixel values of videos to be fed to a
|
||||
model. `None` if no videos are passed.
|
||||
- video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in
|
||||
LLM. `None` if no videos are passed.
|
||||
"""
|
||||
|
||||
if intermediate_tensors is not None:
|
||||
inputs_embeds = None
|
||||
|
||||
hidden_states = self.language_model.model(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
)
|
||||
|
||||
return hidden_states
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(
|
||||
self,
|
||||
skip_prefixes=["mtp."],
|
||||
)
|
||||
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_dtype_from_config(
|
||||
cls,
|
||||
vllm_config: "VllmConfig",
|
||||
) -> tuple[torch.dtype, torch.dtype]:
|
||||
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
|
||||
vllm_config.model_config.dtype,
|
||||
vllm_config.cache_config.mamba_cache_dtype,
|
||||
vllm_config.cache_config.mamba_ssm_cache_dtype,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_shape_from_config(
|
||||
cls, vllm_config: "VllmConfig"
|
||||
) -> tuple[tuple[int, int], tuple[int, int]]:
|
||||
parallel_config = vllm_config.parallel_config
|
||||
hf_config = vllm_config.model_config.hf_text_config
|
||||
tp_size = parallel_config.tensor_parallel_size
|
||||
num_spec = (
|
||||
vllm_config.speculative_config.num_speculative_tokens
|
||||
if vllm_config.speculative_config
|
||||
else 0
|
||||
)
|
||||
return MambaStateShapeCalculator.gated_delta_net_state_shape(
|
||||
tp_size,
|
||||
hf_config.linear_num_key_heads,
|
||||
hf_config.linear_num_value_heads,
|
||||
hf_config.linear_key_head_dim,
|
||||
hf_config.linear_value_head_dim,
|
||||
hf_config.linear_conv_kernel_dim,
|
||||
num_spec,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
|
||||
return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
|
||||
|
||||
|
||||
########################################################
|
||||
# Qwen3_5-MoE
|
||||
########################################################
|
||||
|
||||
|
||||
class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts):
|
||||
def update_physical_experts_metadata(
|
||||
self,
|
||||
num_physical_experts: int,
|
||||
num_local_physical_experts: int,
|
||||
) -> None:
|
||||
assert self.num_local_physical_experts == num_local_physical_experts
|
||||
self.num_physical_experts = num_physical_experts
|
||||
self.num_local_physical_experts = num_local_physical_experts
|
||||
self.num_redundant_experts = num_physical_experts - self.num_logical_experts
|
||||
for layer in self.language_model.model.layers:
|
||||
if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
|
||||
moe = layer.mlp
|
||||
moe.n_local_physical_experts = num_local_physical_experts
|
||||
moe.n_physical_experts = num_physical_experts
|
||||
moe.n_redundant_experts = self.num_redundant_experts
|
||||
moe.experts.update_expert_map()
|
||||
|
||||
def set_moe_parameters(self):
|
||||
self.moe_layers = []
|
||||
example_moe = None
|
||||
for layer in self.language_model.model.layers:
|
||||
if isinstance(layer, Qwen3_5DecoderLayer) and isinstance(
|
||||
layer.mlp, Qwen3NextSparseMoeBlock
|
||||
):
|
||||
example_moe = layer.mlp
|
||||
self.moe_layers.append(layer.mlp.experts)
|
||||
|
||||
if example_moe is None:
|
||||
raise RuntimeError(
|
||||
"No Qwen3_5 layer found in the language_model.model.layers."
|
||||
)
|
||||
|
||||
# Set MoE hyperparameters
|
||||
self.num_moe_layers = len(self.moe_layers)
|
||||
self.num_expert_groups = 1
|
||||
self.num_shared_experts = 0
|
||||
self.num_logical_experts = example_moe.n_logical_experts
|
||||
self.num_physical_experts = example_moe.n_physical_experts
|
||||
self.num_local_physical_experts = example_moe.n_local_physical_experts
|
||||
self.num_routed_experts = example_moe.n_routed_experts
|
||||
self.num_redundant_experts = example_moe.n_redundant_experts
|
||||
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_processor(
|
||||
Qwen3VLMultiModalProcessor,
|
||||
info=Qwen3_5MoeProcessingInfo,
|
||||
dummy_inputs=Qwen3VLDummyInputsBuilder,
|
||||
)
|
||||
class Qwen3_5MoeForConditionalGeneration(
|
||||
Qwen3_5ForConditionalGeneration, Qwen3_5_MoeMixtureOfExperts
|
||||
):
|
||||
# For MoE LoRA weights loading
|
||||
is_3d_moe_weight: bool = True
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
|
||||
# protocols have not __init__ method, so we need to use nn.Module.__init__
|
||||
nn.Module.__init__(self)
|
||||
config: Qwen3_5MoeConfig = vllm_config.model_config.hf_config
|
||||
quant_config = vllm_config.quant_config
|
||||
multimodal_config = vllm_config.model_config.multimodal_config
|
||||
|
||||
self.config = config
|
||||
self.model_config = vllm_config.model_config
|
||||
self.multimodal_config = multimodal_config
|
||||
self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
|
||||
self.is_multimodal_pruning_enabled = (
|
||||
multimodal_config.is_multimodal_pruning_enabled()
|
||||
)
|
||||
self.video_pruning_rate = self.multimodal_config.video_pruning_rate
|
||||
self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
|
||||
|
||||
# attributes needed by EVS-related functions inherited from Qwen3-VL
|
||||
self.use_deepstack = hasattr(config.vision_config, "deepstack_visual_indexes")
|
||||
self.deepstack_num_level = (
|
||||
len(config.vision_config.deepstack_visual_indexes)
|
||||
if self.use_deepstack
|
||||
else 0
|
||||
)
|
||||
self.visual_dim = config.vision_config.out_hidden_size
|
||||
self.multiscale_dim = self.visual_dim * self.deepstack_num_level
|
||||
|
||||
with self._mark_tower_model(vllm_config, {"image", "video"}):
|
||||
self.visual = Qwen3_VisionTransformer(
|
||||
config.vision_config,
|
||||
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, "visual"),
|
||||
)
|
||||
|
||||
with self._mark_language_model(vllm_config):
|
||||
self.language_model = Qwen3_5MoeForCausalLM(
|
||||
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model")
|
||||
)
|
||||
|
||||
self.make_empty_intermediate_tensors = (
|
||||
self.language_model.make_empty_intermediate_tensors
|
||||
)
|
||||
|
||||
# set MoE hyperparameters
|
||||
self.set_moe_parameters()
|
||||
883
upstream_ref/vllm_gdn/models/qwen3_next.py
Normal file
883
upstream_ref/vllm_gdn/models/qwen3_next.py
Normal file
@@ -0,0 +1,883 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Inference-only Qwen3Next model."""
|
||||
|
||||
from collections.abc import Iterable
|
||||
from itertools import islice
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.compilation.decorators import support_torch_compile
|
||||
from vllm.config import CacheConfig, ModelConfig, VllmConfig
|
||||
from vllm.distributed import (
|
||||
get_ep_group,
|
||||
get_pp_group,
|
||||
get_tensor_model_parallel_world_size,
|
||||
tensor_model_parallel_all_gather,
|
||||
tensor_model_parallel_reduce_scatter,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.attention import Attention
|
||||
from vllm.model_executor.layers.fused_moe import FusedMoEFactory
|
||||
from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate
|
||||
from vllm.model_executor.layers.layernorm import (
|
||||
GemmaRMSNorm as Qwen3NextRMSNorm,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import (
|
||||
QKVParallelLinear,
|
||||
ReplicatedLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
|
||||
QwenGatedDeltaNetAttention,
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import (
|
||||
MambaStateCopyFunc,
|
||||
MambaStateCopyFuncCalculator,
|
||||
MambaStateDtypeCalculator,
|
||||
MambaStateShapeCalculator,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead,
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
|
||||
from vllm.model_executor.models.utils import sequence_parallel_chunk
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
|
||||
from vllm.v1.attention.backend import AttentionType
|
||||
|
||||
from .interfaces import (
|
||||
EagleModelMixin,
|
||||
HasInnerState,
|
||||
IsHybrid,
|
||||
MixtureOfExperts,
|
||||
SupportsEagle3,
|
||||
SupportsLoRA,
|
||||
SupportsPP,
|
||||
)
|
||||
from .utils import (
|
||||
AutoWeightsLoader,
|
||||
PPMissingLayer,
|
||||
WeightsMapper,
|
||||
extract_layer_index,
|
||||
make_empty_intermediate_tensors_factory,
|
||||
make_layers,
|
||||
maybe_fuse_shared_experts,
|
||||
maybe_prefix,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
KVCache = tuple[torch.Tensor, torch.Tensor]
|
||||
|
||||
|
||||
def _is_shared_expert_fse_compatible(quant_config) -> bool:
|
||||
"""Check if shared expert can be fused with routed experts.
|
||||
|
||||
FSE requires that shared and routed expert weights use the same
|
||||
quantization format. Returns False when the shared expert is
|
||||
excluded from quantization (e.g. float32 shared in an MXFP4 model)
|
||||
or has a different quant spec than routed experts.
|
||||
"""
|
||||
if quant_config is None:
|
||||
return True
|
||||
# Quark stores its full config dict in quant_config.quant_config
|
||||
raw_config = getattr(quant_config, "quant_config", None)
|
||||
if not isinstance(raw_config, dict):
|
||||
return True
|
||||
exclude = raw_config.get("exclude", [])
|
||||
if not exclude:
|
||||
return True
|
||||
return not any("shared_expert." in str(e) for e in exclude)
|
||||
|
||||
|
||||
class Qwen3NextSparseMoeBlock(nn.Module):
|
||||
def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
|
||||
config = vllm_config.model_config.hf_text_config
|
||||
parallel_config = vllm_config.parallel_config
|
||||
quant_config = vllm_config.quant_config
|
||||
|
||||
self.tp_size = get_tensor_model_parallel_world_size()
|
||||
|
||||
self.ep_group = get_ep_group().device_group
|
||||
self.ep_rank = get_ep_group().rank_in_group
|
||||
self.ep_size = self.ep_group.size()
|
||||
self.n_routed_experts = config.num_experts
|
||||
|
||||
self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
|
||||
|
||||
if self.tp_size > config.num_experts:
|
||||
raise ValueError(
|
||||
f"Tensor parallel size {self.tp_size} is greater than "
|
||||
f"the number of experts {config.num_experts}."
|
||||
)
|
||||
|
||||
# Load balancing settings.
|
||||
eplb_config = vllm_config.parallel_config.eplb_config
|
||||
self.enable_eplb = parallel_config.enable_eplb
|
||||
|
||||
self.n_logical_experts = self.n_routed_experts
|
||||
self.n_redundant_experts = eplb_config.num_redundant_experts
|
||||
self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
|
||||
self.n_local_physical_experts = self.n_physical_experts // self.ep_size
|
||||
|
||||
self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
|
||||
self.physical_expert_end = (
|
||||
self.physical_expert_start + self.n_local_physical_experts
|
||||
)
|
||||
|
||||
self.gate = ReplicatedLinear(
|
||||
config.hidden_size,
|
||||
config.num_experts,
|
||||
bias=False,
|
||||
quant_config=None,
|
||||
prefix=f"{prefix}.gate",
|
||||
)
|
||||
|
||||
self.shared_expert_gate = ReplicatedLinear(
|
||||
config.hidden_size,
|
||||
1,
|
||||
bias=False,
|
||||
quant_config=None,
|
||||
prefix=f"{prefix}.shared_expert_gate",
|
||||
)
|
||||
|
||||
_fse_requested = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
|
||||
_fse_enabled = _fse_requested and _is_shared_expert_fse_compatible(quant_config)
|
||||
if _fse_requested and not _fse_enabled:
|
||||
logger.warning(
|
||||
"VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
|
||||
"shared expert has a different quantization spec than routed "
|
||||
"experts. Falling back to non-fused shared expert path."
|
||||
)
|
||||
if _fse_enabled or config.shared_expert_intermediate_size <= 0:
|
||||
self.shared_expert = None
|
||||
else:
|
||||
self.shared_expert = Qwen3NextMLP(
|
||||
hidden_size=config.hidden_size,
|
||||
intermediate_size=config.shared_expert_intermediate_size,
|
||||
hidden_act=config.hidden_act,
|
||||
quant_config=quant_config,
|
||||
reduce_results=False,
|
||||
expert_gate=self.shared_expert_gate,
|
||||
is_sequence_parallel=self.is_sequence_parallel,
|
||||
prefix=f"{prefix}.shared_expert",
|
||||
)
|
||||
|
||||
self.experts = FusedMoEFactory(
|
||||
shared_experts=self.shared_expert,
|
||||
gate=self.gate,
|
||||
num_experts=self.n_routed_experts,
|
||||
top_k=config.num_experts_per_tok,
|
||||
hidden_size=config.hidden_size,
|
||||
intermediate_size=config.moe_intermediate_size,
|
||||
renormalize=getattr(config, "norm_topk_prob", True),
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.experts",
|
||||
enable_eplb=self.enable_eplb,
|
||||
num_redundant_experts=self.n_redundant_experts,
|
||||
is_sequence_parallel=self.is_sequence_parallel,
|
||||
n_shared_experts=1 if self.shared_expert is None else None,
|
||||
shared_expert_gate=self.shared_expert_gate
|
||||
if self.shared_expert is None
|
||||
else None,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
already_sequence_parallel: bool = False,
|
||||
) -> torch.Tensor:
|
||||
# NOTE: hidden_states can have either 1D or 2D shape.
|
||||
orig_shape = hidden_states.shape
|
||||
num_tokens, hidden_dim = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, hidden_dim)
|
||||
|
||||
if self.is_sequence_parallel and not already_sequence_parallel:
|
||||
hidden_states = sequence_parallel_chunk(hidden_states)
|
||||
|
||||
if self.experts.is_internal_router:
|
||||
# In this case, the gate/router runs inside the MoERunner class
|
||||
final_hidden_states = self.experts(
|
||||
hidden_states=hidden_states, router_logits=hidden_states
|
||||
)
|
||||
else:
|
||||
# router_logits: (num_tokens, n_experts)
|
||||
router_logits, _ = self.gate(hidden_states)
|
||||
final_hidden_states = self.experts(
|
||||
hidden_states=hidden_states, router_logits=router_logits
|
||||
)
|
||||
|
||||
if self.is_sequence_parallel and not already_sequence_parallel:
|
||||
final_hidden_states = tensor_model_parallel_all_gather(
|
||||
final_hidden_states, 0
|
||||
)
|
||||
final_hidden_states = final_hidden_states[:num_tokens]
|
||||
|
||||
return final_hidden_states.view(orig_shape)
|
||||
|
||||
|
||||
class Qwen3NextAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: Qwen3NextConfig,
|
||||
model_config: ModelConfig | None = None,
|
||||
cache_config: CacheConfig | None = None,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
reduce_results: bool = True,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
self.total_num_heads = config.num_attention_heads
|
||||
assert self.total_num_heads % tp_size == 0
|
||||
self.num_heads = self.total_num_heads // tp_size
|
||||
self.total_num_kv_heads = config.num_key_value_heads
|
||||
if self.total_num_kv_heads >= tp_size:
|
||||
# Number of KV heads is greater than TP size, so we partition
|
||||
# the KV heads across multiple tensor parallel GPUs.
|
||||
assert self.total_num_kv_heads % tp_size == 0
|
||||
else:
|
||||
# Number of KV heads is less than TP size, so we replicate
|
||||
# the KV heads across multiple tensor parallel GPUs.
|
||||
assert tp_size % self.total_num_kv_heads == 0
|
||||
self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
|
||||
self.head_dim = config.head_dim or (self.hidden_size // self.num_heads)
|
||||
self.q_size = self.num_heads * self.head_dim
|
||||
self.kv_size = self.num_kv_heads * self.head_dim
|
||||
self.scaling = self.head_dim**-0.5
|
||||
self.dual_chunk_attention_config = getattr(
|
||||
config, "dual_chunk_attention_config", None
|
||||
)
|
||||
self.attn_output_gate = getattr(config, "attn_output_gate", True)
|
||||
|
||||
self.qkv_proj = QKVParallelLinear(
|
||||
config.hidden_size,
|
||||
self.head_dim,
|
||||
self.total_num_heads * (1 + self.attn_output_gate),
|
||||
self.total_num_kv_heads,
|
||||
bias=getattr(config, "qkv_bias", False),
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.qkv_proj",
|
||||
)
|
||||
|
||||
self.o_proj = RowParallelLinear(
|
||||
self.total_num_heads * self.head_dim,
|
||||
config.hidden_size,
|
||||
bias=False,
|
||||
reduce_results=reduce_results,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.o_proj",
|
||||
)
|
||||
|
||||
self.rotary_emb = get_rope(
|
||||
head_size=self.head_dim,
|
||||
max_position=config.max_position_embeddings,
|
||||
rope_parameters=config.rope_parameters,
|
||||
dual_chunk_attention_config=self.dual_chunk_attention_config,
|
||||
)
|
||||
|
||||
# Late-interaction retrieval models (e.g. ColQwen3.5) run BIDIRECTIONAL
|
||||
# attention on the full_attention layers; they set config.is_causal=False
|
||||
# via a VerifyAndUpdateConfig handler. Generation models leave is_causal
|
||||
# unset (-> causal/DECODER), so this is a no-op for them. Mirrors qwen3.py.
|
||||
attn_type = (
|
||||
AttentionType.DECODER
|
||||
if getattr(config, "is_causal", True)
|
||||
else AttentionType.ENCODER_ONLY
|
||||
)
|
||||
self.attn = Attention(
|
||||
self.num_heads,
|
||||
self.head_dim,
|
||||
self.scaling,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
cache_config=cache_config,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.attn",
|
||||
attn_type=attn_type,
|
||||
**{
|
||||
"layer_idx": extract_layer_index(prefix),
|
||||
"dual_chunk_attention_config": self.dual_chunk_attention_config,
|
||||
}
|
||||
if self.dual_chunk_attention_config
|
||||
else {},
|
||||
)
|
||||
|
||||
self.q_norm = Qwen3NextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
||||
self.k_norm = Qwen3NextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
||||
|
||||
# Fuse the gated split + QK-RMSNorm + (partial) NeoX RoPE + gate copy.
|
||||
# TODO: support MRoPE
|
||||
mm_config = model_config.multimodal_config if model_config else None
|
||||
text_only = mm_config is None or mm_config.language_model_only
|
||||
self.use_fused_qk_norm_rope_gate = (
|
||||
self.attn_output_gate
|
||||
and getattr(self.rotary_emb, "is_neox_style", False)
|
||||
and current_platform.is_cuda()
|
||||
and text_only
|
||||
)
|
||||
|
||||
def _project_qkv_gate(
|
||||
self,
|
||||
qkv: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
"""Return post-norm, post-RoPE (q, k, v) and the pre-sigmoid gate.
|
||||
|
||||
Dispatches between the fused Triton kernel and the eager
|
||||
split + QK-RMSNorm + RoPE path. ``gate`` is ``None`` when output
|
||||
gating is disabled.
|
||||
"""
|
||||
if self.use_fused_qk_norm_rope_gate:
|
||||
q_gate, k, v = qkv.split(
|
||||
[self.q_size * 2, self.kv_size, self.kv_size], dim=-1
|
||||
)
|
||||
# mRoPE passes positions as (3, n_tokens) for T/H/W. Fusion is only
|
||||
# enabled text-only, where the three rows are identical, so taking
|
||||
# the T row is exact. (1D positions pass through.)
|
||||
pos = positions[0] if positions.ndim == 2 else positions
|
||||
q, k, gate = fused_qk_rmsnorm_rope_gate(
|
||||
q_gate,
|
||||
k,
|
||||
self.q_norm.weight.float() + 1.0,
|
||||
self.k_norm.weight.float() + 1.0,
|
||||
self.rotary_emb.cos_sin_cache,
|
||||
pos,
|
||||
self.q_norm.variance_epsilon,
|
||||
self.num_heads,
|
||||
self.num_kv_heads,
|
||||
self.head_dim,
|
||||
self.rotary_emb.rotary_dim,
|
||||
)
|
||||
return q, k, v, gate
|
||||
|
||||
if self.attn_output_gate:
|
||||
q_gate, k, v = qkv.split(
|
||||
[self.q_size * 2, self.kv_size, self.kv_size], dim=-1
|
||||
)
|
||||
orig_shape = q_gate.shape[:-1]
|
||||
q_gate = q_gate.view(*orig_shape, self.num_heads, -1)
|
||||
q, gate = torch.chunk(q_gate, 2, dim=-1)
|
||||
q = q.reshape(*orig_shape, -1)
|
||||
gate = gate.reshape(*orig_shape, -1)
|
||||
else:
|
||||
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||||
gate = None
|
||||
|
||||
q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(
|
||||
-1, self.num_heads * self.head_dim
|
||||
)
|
||||
k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(
|
||||
-1, self.num_kv_heads * self.head_dim
|
||||
)
|
||||
q, k = self.rotary_emb(positions, q, k)
|
||||
return q, k, v, gate
|
||||
|
||||
def forward(
|
||||
self,
|
||||
positions: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
qkv, _ = self.qkv_proj(hidden_states)
|
||||
q, k, v, gate = self._project_qkv_gate(qkv, positions)
|
||||
attn_output = self.attn(q, k, v)
|
||||
if gate is not None:
|
||||
attn_output = attn_output * torch.sigmoid(gate)
|
||||
output, _ = self.o_proj(attn_output)
|
||||
return output
|
||||
|
||||
|
||||
class Qwen3NextDecoderLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
layer_type: str,
|
||||
prefix: str = "",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
config = vllm_config.model_config.hf_config
|
||||
model_config = vllm_config.model_config
|
||||
cache_config = vllm_config.cache_config
|
||||
quant_config = vllm_config.quant_config
|
||||
parallel_config = vllm_config.parallel_config
|
||||
|
||||
self.layer_type = layer_type
|
||||
self.layer_idx = extract_layer_index(prefix)
|
||||
|
||||
mlp_only_layers = (
|
||||
[] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers
|
||||
)
|
||||
is_moe_layer = (self.layer_idx not in mlp_only_layers) and (
|
||||
config.num_experts > 0
|
||||
and (self.layer_idx + 1) % config.decoder_sparse_step == 0
|
||||
)
|
||||
self.use_attn_reduce_scatter_for_moe = (
|
||||
parallel_config.use_sequence_parallel_moe
|
||||
and parallel_config.pipeline_parallel_size == 1
|
||||
and is_moe_layer
|
||||
)
|
||||
|
||||
if self.layer_type == "linear_attention":
|
||||
self.linear_attn = QwenGatedDeltaNetAttention(
|
||||
config,
|
||||
vllm_config=vllm_config,
|
||||
prefix=f"{prefix}.linear_attn",
|
||||
gqa_interleaved_layout=True,
|
||||
reduce_results=not self.use_attn_reduce_scatter_for_moe,
|
||||
)
|
||||
elif self.layer_type == "full_attention":
|
||||
self.self_attn = Qwen3NextAttention(
|
||||
config,
|
||||
model_config=model_config,
|
||||
cache_config=cache_config,
|
||||
quant_config=quant_config,
|
||||
reduce_results=not self.use_attn_reduce_scatter_for_moe,
|
||||
prefix=f"{prefix}.self_attn",
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid layer_type {self.layer_type}")
|
||||
|
||||
if is_moe_layer:
|
||||
self.mlp = Qwen3NextSparseMoeBlock(
|
||||
vllm_config=vllm_config,
|
||||
prefix=f"{prefix}.mlp",
|
||||
)
|
||||
else:
|
||||
self.mlp = Qwen3NextMLP(
|
||||
hidden_size=config.hidden_size,
|
||||
intermediate_size=config.intermediate_size,
|
||||
hidden_act=config.hidden_act,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.mlp",
|
||||
)
|
||||
|
||||
self.input_layernorm = Qwen3NextRMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
self.post_attention_layernorm = Qwen3NextRMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
|
||||
self.layer_scale = getattr(config, "layer_scale", False)
|
||||
if self.layer_scale:
|
||||
self.attn_layer_scale = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
1,
|
||||
1,
|
||||
config.hidden_size,
|
||||
),
|
||||
)
|
||||
self.ffn_layer_scale = torch.nn.Parameter(
|
||||
torch.zeros(
|
||||
1,
|
||||
1,
|
||||
config.hidden_size,
|
||||
),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
residual: torch.Tensor | None,
|
||||
positions: torch.Tensor = None,
|
||||
**kwargs: object,
|
||||
):
|
||||
full_num_tokens = positions.shape[-1]
|
||||
input_is_sequence_parallel = (
|
||||
self.use_attn_reduce_scatter_for_moe
|
||||
and residual is not None
|
||||
and hidden_states.shape[0] != full_num_tokens
|
||||
)
|
||||
|
||||
if residual is None:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
else:
|
||||
hidden_states, residual = self.input_layernorm(hidden_states, residual)
|
||||
|
||||
if input_is_sequence_parallel:
|
||||
hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
|
||||
hidden_states = hidden_states[:full_num_tokens]
|
||||
|
||||
if self.layer_type == "linear_attention":
|
||||
hidden_states = self.linear_attn(hidden_states=hidden_states)
|
||||
elif self.layer_type == "full_attention":
|
||||
hidden_states = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
positions=positions,
|
||||
)
|
||||
else:
|
||||
raise ValueError("Invalid layer_type")
|
||||
|
||||
if self.layer_scale:
|
||||
if len(hidden_states.shape) == 2:
|
||||
hidden_states = hidden_states * (
|
||||
self.attn_layer_scale.to(hidden_states.dtype)[0] + 1
|
||||
)
|
||||
else:
|
||||
hidden_states = hidden_states * (
|
||||
self.attn_layer_scale.to(hidden_states.dtype) + 1
|
||||
)
|
||||
|
||||
if self.use_attn_reduce_scatter_for_moe:
|
||||
tp_world_size = get_tensor_model_parallel_world_size()
|
||||
# small trick using minus, eg. -17 % 8 = 7
|
||||
sp_pad = (-hidden_states.shape[0]) % tp_world_size
|
||||
# pad if not divisible by world size
|
||||
hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad))
|
||||
hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)
|
||||
if not input_is_sequence_parallel:
|
||||
residual = sequence_parallel_chunk(residual)
|
||||
|
||||
# Fully Connected
|
||||
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
|
||||
if self.use_attn_reduce_scatter_for_moe:
|
||||
hidden_states = self.mlp(
|
||||
hidden_states,
|
||||
already_sequence_parallel=True,
|
||||
)
|
||||
else:
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
|
||||
if self.layer_scale:
|
||||
if len(hidden_states.shape) == 2:
|
||||
hidden_states = hidden_states * (
|
||||
self.ffn_layer_scale.to(hidden_states.dtype)[0] + 1
|
||||
)
|
||||
else:
|
||||
assert len(hidden_states.shape) == len(self.ffn_layer_scale.shape), (
|
||||
f"shape must be the same {len(hidden_states.shape)}, "
|
||||
f"{len(self.ffn_layer_scale.shape)}"
|
||||
)
|
||||
hidden_states = hidden_states * (
|
||||
self.ffn_layer_scale.to(hidden_states.dtype) + 1
|
||||
)
|
||||
|
||||
return hidden_states, residual
|
||||
|
||||
|
||||
def _all_gather_hidden_and_residual(
|
||||
hidden_states: torch.Tensor,
|
||||
residual: torch.Tensor | None,
|
||||
full_num_tokens: int,
|
||||
hidden_size: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
if residual is None:
|
||||
hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
|
||||
hidden_states = hidden_states[:full_num_tokens]
|
||||
return hidden_states, None
|
||||
|
||||
combined_states = torch.cat([hidden_states, residual], dim=-1)
|
||||
combined_states = tensor_model_parallel_all_gather(combined_states, 0)
|
||||
combined_states = combined_states[:full_num_tokens]
|
||||
hidden_states, residual = combined_states.split([hidden_size, hidden_size], dim=-1)
|
||||
return hidden_states, residual
|
||||
|
||||
|
||||
@support_torch_compile
|
||||
class Qwen3NextModel(nn.Module, EagleModelMixin):
|
||||
hf_to_vllm_mapper = WeightsMapper(
|
||||
orig_to_new_stacked={
|
||||
# weight_name: (param_name, shard_id)
|
||||
".q_proj": (".qkv_proj", "q"),
|
||||
".k_proj": (".qkv_proj", "k"),
|
||||
".v_proj": (".qkv_proj", "v"),
|
||||
".mlp.gate_proj": (".mlp.gate_up_proj", 0),
|
||||
".mlp.up_proj": (".mlp.gate_up_proj", 1),
|
||||
".shared_expert.gate_proj": (".shared_expert.gate_up_proj", 0),
|
||||
".shared_expert.up_proj": (".shared_expert.gate_up_proj", 1),
|
||||
}
|
||||
)
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
|
||||
config: Qwen3NextConfig = vllm_config.model_config.hf_text_config
|
||||
parallel_config = vllm_config.parallel_config
|
||||
|
||||
eplb_config = parallel_config.eplb_config
|
||||
self.num_redundant_experts = eplb_config.num_redundant_experts
|
||||
|
||||
self.config = config
|
||||
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = VocabParallelEmbedding(
|
||||
self.vocab_size,
|
||||
config.hidden_size,
|
||||
)
|
||||
|
||||
def get_layer(prefix: str):
|
||||
return Qwen3NextDecoderLayer(
|
||||
vllm_config,
|
||||
layer_type=config.layer_types[extract_layer_index(prefix)],
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
self.start_layer, self.end_layer, self.layers = make_layers(
|
||||
config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
|
||||
)
|
||||
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
|
||||
["hidden_states", "residual"], config.hidden_size
|
||||
)
|
||||
|
||||
if get_pp_group().is_last_rank:
|
||||
self.norm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
else:
|
||||
self.norm = PPMissingLayer()
|
||||
|
||||
self.aux_hidden_state_layers: tuple[int, ...] = ()
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.embed_tokens(input_ids)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
|
||||
if get_pp_group().is_first_rank:
|
||||
if inputs_embeds is not None:
|
||||
hidden_states = inputs_embeds
|
||||
else:
|
||||
hidden_states = self.embed_input_ids(input_ids)
|
||||
residual = None
|
||||
else:
|
||||
assert intermediate_tensors is not None
|
||||
hidden_states = intermediate_tensors["hidden_states"]
|
||||
residual = intermediate_tensors["residual"]
|
||||
|
||||
full_num_tokens = positions.shape[-1]
|
||||
aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
|
||||
for layer_idx, layer in enumerate(
|
||||
islice(self.layers, self.start_layer, self.end_layer),
|
||||
start=self.start_layer,
|
||||
):
|
||||
if (
|
||||
hidden_states.shape[0] != full_num_tokens
|
||||
and not layer.use_attn_reduce_scatter_for_moe
|
||||
):
|
||||
hidden_states, residual = _all_gather_hidden_and_residual(
|
||||
hidden_states,
|
||||
residual,
|
||||
full_num_tokens,
|
||||
self.config.hidden_size,
|
||||
)
|
||||
hidden_states, residual = layer(
|
||||
positions=positions,
|
||||
hidden_states=hidden_states,
|
||||
residual=residual,
|
||||
)
|
||||
if (layer_idx + 1) in self.aux_hidden_state_layers and hidden_states.shape[
|
||||
0
|
||||
] != full_num_tokens:
|
||||
hidden_states, residual = _all_gather_hidden_and_residual(
|
||||
hidden_states,
|
||||
residual,
|
||||
full_num_tokens,
|
||||
self.config.hidden_size,
|
||||
)
|
||||
self._maybe_add_hidden_state(
|
||||
aux_hidden_states, layer_idx + 1, hidden_states, residual
|
||||
)
|
||||
|
||||
if not get_pp_group().is_last_rank:
|
||||
return IntermediateTensors(
|
||||
{"hidden_states": hidden_states, "residual": residual}
|
||||
)
|
||||
if hidden_states.shape[0] != full_num_tokens:
|
||||
hidden_states, residual = _all_gather_hidden_and_residual(
|
||||
hidden_states,
|
||||
residual,
|
||||
full_num_tokens,
|
||||
self.config.hidden_size,
|
||||
)
|
||||
hidden_states, _ = self.norm(hidden_states, residual)
|
||||
if aux_hidden_states:
|
||||
return hidden_states, aux_hidden_states
|
||||
return hidden_states
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
weights = maybe_fuse_shared_experts(
|
||||
weights,
|
||||
n_routed_experts=getattr(self.config, "num_experts", 0),
|
||||
n_shared_experts=1,
|
||||
ckpt_prefix="mlp.shared_expert",
|
||||
)
|
||||
loader = AutoWeightsLoader(self)
|
||||
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
|
||||
|
||||
|
||||
class QwenNextMixtureOfExperts(MixtureOfExperts):
|
||||
def update_physical_experts_metadata(
|
||||
self,
|
||||
num_physical_experts: int,
|
||||
num_local_physical_experts: int,
|
||||
) -> None:
|
||||
assert self.num_local_physical_experts == num_local_physical_experts
|
||||
self.num_physical_experts = num_physical_experts
|
||||
self.num_local_physical_experts = num_local_physical_experts
|
||||
self.num_redundant_experts = num_physical_experts - self.num_logical_experts
|
||||
for layer in self.model.layers:
|
||||
if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
|
||||
moe = layer.mlp
|
||||
moe.n_local_physical_experts = num_local_physical_experts
|
||||
moe.n_physical_experts = num_physical_experts
|
||||
moe.n_redundant_experts = self.num_redundant_experts
|
||||
moe.experts.update_expert_map()
|
||||
|
||||
def set_moe_parameters(self):
|
||||
self.moe_layers = []
|
||||
example_moe = None
|
||||
for layer in self.model.layers:
|
||||
if isinstance(layer, Qwen3NextDecoderLayer) and isinstance(
|
||||
layer.mlp, Qwen3NextSparseMoeBlock
|
||||
):
|
||||
example_moe = layer.mlp
|
||||
self.moe_layers.append(layer.mlp.experts)
|
||||
|
||||
if example_moe is None:
|
||||
raise RuntimeError("No Qwen3Next layer found in the model.layers.")
|
||||
|
||||
# Set MoE hyperparameters
|
||||
self.num_moe_layers = len(self.moe_layers)
|
||||
self.num_expert_groups = 1
|
||||
self.num_shared_experts = 0
|
||||
self.num_logical_experts = example_moe.n_logical_experts
|
||||
self.num_physical_experts = example_moe.n_physical_experts
|
||||
self.num_local_physical_experts = example_moe.n_local_physical_experts
|
||||
self.num_routed_experts = example_moe.n_routed_experts
|
||||
self.num_redundant_experts = example_moe.n_redundant_experts
|
||||
|
||||
|
||||
class Qwen3NextForCausalLM(
|
||||
nn.Module,
|
||||
HasInnerState,
|
||||
SupportsLoRA,
|
||||
SupportsPP,
|
||||
QwenNextMixtureOfExperts,
|
||||
IsHybrid,
|
||||
SupportsEagle3,
|
||||
):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
"k_proj",
|
||||
"v_proj",
|
||||
],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"in_proj_qkvz": ["in_proj_qkvz"],
|
||||
"in_proj_ba": ["in_proj_ba"],
|
||||
}
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
config = vllm_config.model_config.hf_text_config
|
||||
self.vllm_config = vllm_config
|
||||
self.model_config = vllm_config.model_config
|
||||
cache_config = vllm_config.cache_config
|
||||
|
||||
scheduler_config = vllm_config.scheduler_config
|
||||
if cache_config.mamba_cache_mode == "all":
|
||||
raise NotImplementedError(
|
||||
"Qwen3Next currently does not support 'all' prefix caching, "
|
||||
"please use '--mamba-cache-mode=align' instead"
|
||||
)
|
||||
self.quant_config = vllm_config.quant_config
|
||||
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.scheduler_config = scheduler_config
|
||||
self.model = Qwen3NextModel(
|
||||
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
|
||||
)
|
||||
|
||||
self.lm_head = ParallelLMHead(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
prefix=maybe_prefix(prefix, "lm_head"),
|
||||
)
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size)
|
||||
self.make_empty_intermediate_tensors = (
|
||||
self.model.make_empty_intermediate_tensors
|
||||
)
|
||||
|
||||
# Set MoE hyperparameters
|
||||
self.set_moe_parameters()
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return self.model.embed_input_ids(input_ids)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: object,
|
||||
):
|
||||
hidden_states = self.model(
|
||||
input_ids, positions, intermediate_tensors, inputs_embeds
|
||||
)
|
||||
|
||||
return hidden_states
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_dtype_from_config(
|
||||
cls,
|
||||
vllm_config: "VllmConfig",
|
||||
) -> tuple[torch.dtype, torch.dtype]:
|
||||
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
|
||||
vllm_config.model_config.dtype,
|
||||
vllm_config.cache_config.mamba_cache_dtype,
|
||||
vllm_config.cache_config.mamba_ssm_cache_dtype,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_shape_from_config(
|
||||
cls, vllm_config: "VllmConfig"
|
||||
) -> tuple[tuple[int, int], tuple[int, int]]:
|
||||
parallel_config = vllm_config.parallel_config
|
||||
hf_config = vllm_config.model_config.hf_text_config
|
||||
tp_size = parallel_config.tensor_parallel_size
|
||||
num_spec = (
|
||||
vllm_config.speculative_config.num_speculative_tokens
|
||||
if vllm_config.speculative_config
|
||||
else 0
|
||||
)
|
||||
return MambaStateShapeCalculator.gated_delta_net_state_shape(
|
||||
tp_size,
|
||||
hf_config.linear_num_key_heads,
|
||||
hf_config.linear_num_value_heads,
|
||||
hf_config.linear_key_head_dim,
|
||||
hf_config.linear_value_head_dim,
|
||||
hf_config.linear_conv_kernel_dim,
|
||||
num_spec,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
|
||||
return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
) -> torch.Tensor | None:
|
||||
return self.logits_processor(self.lm_head, hidden_states)
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(self, skip_prefixes=["mtp."])
|
||||
return loader.load_weights(weights)
|
||||
1289
upstream_ref/vllm_gdn/ops/causal_conv1d.py
Normal file
1289
upstream_ref/vllm_gdn/ops/causal_conv1d.py
Normal file
File diff suppressed because it is too large
Load Diff
8
upstream_ref/vllm_gdn/third_party/__init__.py
vendored
Normal file
8
upstream_ref/vllm_gdn/third_party/__init__.py
vendored
Normal file
@@ -0,0 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
||||
#
|
||||
# This file contains code copied from the flash-linear-attention project.
|
||||
# The original source code was licensed under the MIT license and included
|
||||
# the following copyright notice:
|
||||
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
||||
25
upstream_ref/vllm_gdn/third_party/ops/__init__.py
vendored
Normal file
25
upstream_ref/vllm_gdn/third_party/ops/__init__.py
vendored
Normal file
@@ -0,0 +1,25 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
||||
#
|
||||
# This file contains code copied from the flash-linear-attention project.
|
||||
# The original source code was licensed under the MIT license and included
|
||||
# the following copyright notice:
|
||||
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
||||
from .chunk import chunk_gated_delta_rule
|
||||
from .fused_gdn_prefill_post_conv import fused_post_conv_prep
|
||||
from .fused_recurrent import (
|
||||
fused_recurrent_gated_delta_rule,
|
||||
fused_recurrent_gated_delta_rule_packed_decode,
|
||||
)
|
||||
from .fused_sigmoid_gating import fused_sigmoid_gating_delta_rule_update
|
||||
from .layernorm_guard import RMSNormGated
|
||||
|
||||
__all__ = [
|
||||
"RMSNormGated",
|
||||
"chunk_gated_delta_rule",
|
||||
"fused_recurrent_gated_delta_rule",
|
||||
"fused_recurrent_gated_delta_rule_packed_decode",
|
||||
"fused_post_conv_prep",
|
||||
"fused_sigmoid_gating_delta_rule_update",
|
||||
]
|
||||
245
upstream_ref/vllm_gdn/third_party/ops/chunk.py
vendored
Normal file
245
upstream_ref/vllm_gdn/third_party/ops/chunk.py
vendored
Normal file
@@ -0,0 +1,245 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
||||
#
|
||||
# This file contains code copied from the flash-linear-attention project.
|
||||
# The original source code was licensed under the MIT license and included
|
||||
# the following copyright notice:
|
||||
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
||||
# ruff: noqa: E501
|
||||
|
||||
import torch
|
||||
|
||||
from .chunk_delta_h import chunk_gated_delta_rule_fwd_h
|
||||
from .chunk_o import chunk_fwd_o
|
||||
from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
|
||||
from .cumsum import chunk_local_cumsum
|
||||
from .l2norm import l2norm_fwd
|
||||
from .solve_tril import solve_tril
|
||||
from .utils import FLA_CHUNK_SIZE, SUPPRESS_LEVEL, input_guard
|
||||
from .wy_fast import recompute_w_u_fwd
|
||||
|
||||
|
||||
def chunk_gated_delta_rule_fwd(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
beta: torch.Tensor,
|
||||
scale: float,
|
||||
initial_state: torch.Tensor,
|
||||
output_final_state: bool,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
chunk_indices: torch.Tensor | None = None,
|
||||
chunk_offsets: torch.Tensor | None = None,
|
||||
core_attn_out: torch.Tensor | None = None,
|
||||
):
|
||||
g = chunk_local_cumsum(
|
||||
g, chunk_size=FLA_CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices
|
||||
)
|
||||
# obtain WY representation. u is actually the new v.
|
||||
A = chunk_scaled_dot_kkt_fwd(
|
||||
k=k,
|
||||
beta=beta,
|
||||
g=g,
|
||||
cu_seqlens=cu_seqlens,
|
||||
chunk_indices=chunk_indices,
|
||||
output_dtype=torch.float32,
|
||||
)
|
||||
A = solve_tril(
|
||||
A=A, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, output_dtype=k.dtype
|
||||
)
|
||||
w, u = recompute_w_u_fwd(
|
||||
k=k,
|
||||
v=v,
|
||||
beta=beta,
|
||||
A=A,
|
||||
g_cumsum=g,
|
||||
cu_seqlens=cu_seqlens,
|
||||
chunk_indices=chunk_indices,
|
||||
)
|
||||
h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
|
||||
k=k,
|
||||
w=w,
|
||||
u=u,
|
||||
g=g,
|
||||
initial_state=initial_state,
|
||||
output_final_state=output_final_state,
|
||||
cu_seqlens=cu_seqlens,
|
||||
chunk_indices=chunk_indices,
|
||||
chunk_offsets=chunk_offsets,
|
||||
)
|
||||
o = chunk_fwd_o(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v_new,
|
||||
h=h,
|
||||
g=g,
|
||||
scale=scale,
|
||||
cu_seqlens=cu_seqlens,
|
||||
chunk_indices=chunk_indices,
|
||||
core_attn_out=core_attn_out,
|
||||
)
|
||||
if SUPPRESS_LEVEL < 3:
|
||||
return g, o, A, final_state, None, None, None
|
||||
elif SUPPRESS_LEVEL >= 3:
|
||||
return g, o, A, final_state, w, h, v_new
|
||||
|
||||
|
||||
class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
|
||||
@staticmethod
|
||||
@input_guard
|
||||
@torch.amp.custom_fwd(device_type="cuda")
|
||||
def forward(
|
||||
ctx,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
beta: torch.Tensor,
|
||||
scale: float,
|
||||
initial_state: torch.Tensor,
|
||||
output_final_state: bool,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
chunk_indices: torch.Tensor | None = None,
|
||||
chunk_offsets: torch.Tensor | None = None,
|
||||
use_qk_l2norm_in_kernel: bool = False,
|
||||
core_attn_out: torch.Tensor | None = None,
|
||||
):
|
||||
if use_qk_l2norm_in_kernel:
|
||||
q = l2norm_fwd(q)
|
||||
k = l2norm_fwd(k)
|
||||
|
||||
g, o, A, final_state, w, h, v_new = chunk_gated_delta_rule_fwd(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=g,
|
||||
beta=beta,
|
||||
scale=scale,
|
||||
initial_state=initial_state,
|
||||
output_final_state=output_final_state,
|
||||
cu_seqlens=cu_seqlens,
|
||||
chunk_indices=chunk_indices,
|
||||
chunk_offsets=chunk_offsets,
|
||||
core_attn_out=core_attn_out,
|
||||
)
|
||||
ctx.scale = scale
|
||||
ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
|
||||
if core_attn_out is not None:
|
||||
assert not torch.is_grad_enabled(), (
|
||||
"core_attn_out buffer reuse is only supported for inference"
|
||||
)
|
||||
assert q.dtype == o.dtype, "Incompatible dtype for inplace computation"
|
||||
return o.to(q.dtype), final_state
|
||||
|
||||
|
||||
@torch.compiler.disable
|
||||
def chunk_gated_delta_rule(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
beta: torch.Tensor,
|
||||
scale: float = None,
|
||||
initial_state: torch.Tensor = None,
|
||||
output_final_state: bool = False,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
chunk_indices: torch.Tensor | None = None,
|
||||
chunk_offsets: torch.Tensor | None = None,
|
||||
use_qk_l2norm_in_kernel: bool = False,
|
||||
core_attn_out: torch.Tensor | None = None,
|
||||
):
|
||||
r"""
|
||||
Args:
|
||||
q (torch.Tensor):
|
||||
Queries of shape `[B, T, H, K]`.
|
||||
k (torch.Tensor):
|
||||
Keys of shape `[B, T, H, K]`.
|
||||
v (torch.Tensor):
|
||||
Values of shape `[B, T, H, V]`.
|
||||
g (torch.Tensor):
|
||||
(forget) Gating tensor (in log space!) of shape `[B, T, H]`.
|
||||
beta (torch.Tensor):
|
||||
Betas of shape `[B, T, H]`.
|
||||
scale (Optional[int]):
|
||||
Scale factor for the RetNet attention scores.
|
||||
If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
|
||||
initial_state (Optional[torch.Tensor]):
|
||||
Initial state of shape `[N, H, V, K]` for `N` input sequences.
|
||||
For equal-length input sequences, `N` equals the batch size `B`.
|
||||
Default: `None`.
|
||||
output_final_state (Optional[bool]):
|
||||
Whether to output the final state of shape `[N, H, V, K]`. Default: `False`.
|
||||
cu_seqlens (torch.Tensor):
|
||||
Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
|
||||
consistent with the FlashAttention API.
|
||||
Returns:
|
||||
o (torch.Tensor):
|
||||
Outputs of shape `[B, T, H, V]`.
|
||||
final_state (torch.Tensor):
|
||||
Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`.
|
||||
|
||||
Examples::
|
||||
>>> import torch
|
||||
>>> import torch.nn.functional as F
|
||||
>>> from einops import rearrange
|
||||
>>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule
|
||||
# inputs with equal lengths
|
||||
>>> B, T, H, K, V = 4, 2048, 4, 512, 512
|
||||
>>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
|
||||
>>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)
|
||||
>>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')
|
||||
>>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid()
|
||||
>>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda'))
|
||||
>>> h0 = torch.randn(B, H, V, K, dtype=torch.bfloat16, device='cuda')
|
||||
>>> o, ht = chunk_gated_delta_rule(
|
||||
q, k, v, g, beta,
|
||||
initial_state=h0,
|
||||
output_final_state=True
|
||||
)
|
||||
# for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
|
||||
>>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))
|
||||
# for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
|
||||
>>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.int32)
|
||||
>>> o_var, ht_var = chunk_gated_delta_rule(
|
||||
q, k, v, g, beta,
|
||||
initial_state=h0,
|
||||
output_final_state=True,
|
||||
cu_seqlens=cu_seqlens
|
||||
)
|
||||
"""
|
||||
assert q.dtype == k.dtype == v.dtype
|
||||
assert q.dtype != torch.float32, (
|
||||
"ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
|
||||
)
|
||||
assert len(beta.shape) == 3, "beta must be of shape [B, T, H]."
|
||||
if cu_seqlens is not None:
|
||||
if q.shape[0] != 1:
|
||||
raise ValueError(
|
||||
f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
|
||||
f"Please flatten variable-length inputs before processing."
|
||||
)
|
||||
if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
|
||||
raise ValueError(
|
||||
f"The number of initial states is expected to be equal to the number of input sequences, "
|
||||
f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}."
|
||||
)
|
||||
if scale is None:
|
||||
scale = k.shape[-1] ** -0.5
|
||||
o, final_state = ChunkGatedDeltaRuleFunction.apply(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
g,
|
||||
beta,
|
||||
scale,
|
||||
initial_state,
|
||||
output_final_state,
|
||||
cu_seqlens,
|
||||
chunk_indices,
|
||||
chunk_offsets,
|
||||
use_qk_l2norm_in_kernel,
|
||||
core_attn_out,
|
||||
)
|
||||
return o, final_state
|
||||
282
upstream_ref/vllm_gdn/third_party/ops/cumsum.py
vendored
Normal file
282
upstream_ref/vllm_gdn/third_party/ops/cumsum.py
vendored
Normal file
@@ -0,0 +1,282 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
||||
#
|
||||
# This file contains code copied from the flash-linear-attention project.
|
||||
# The original source code was licensed under the MIT license and included
|
||||
# the following copyright notice:
|
||||
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
||||
# ruff: noqa: E501
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
from .index import prepare_chunk_indices
|
||||
from .utils import check_shared_mem, input_guard
|
||||
|
||||
BS_LIST = [32, 64] if check_shared_mem() else [16, 32]
|
||||
|
||||
|
||||
@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
|
||||
@triton.autotune(
|
||||
configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]],
|
||||
key=["B", "H", "BT", "IS_VARLEN", "REVERSE"],
|
||||
)
|
||||
@triton.jit(do_not_specialize=["T"])
|
||||
def chunk_local_cumsum_scalar_kernel(
|
||||
s,
|
||||
o,
|
||||
cu_seqlens,
|
||||
chunk_indices,
|
||||
T,
|
||||
B: tl.constexpr,
|
||||
H: tl.constexpr,
|
||||
BT: tl.constexpr,
|
||||
REVERSE: tl.constexpr,
|
||||
IS_VARLEN: tl.constexpr,
|
||||
HEAD_FIRST: tl.constexpr,
|
||||
):
|
||||
i_t, i_bh = tl.program_id(0), tl.program_id(1)
|
||||
i_b, i_h = i_bh // H, i_bh % H
|
||||
if IS_VARLEN:
|
||||
i_n, i_t = (
|
||||
tl.load(chunk_indices + i_t * 2).to(tl.int32),
|
||||
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
|
||||
)
|
||||
bos, eos = (
|
||||
tl.load(cu_seqlens + i_n).to(tl.int32),
|
||||
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
|
||||
)
|
||||
T = eos - bos
|
||||
else:
|
||||
bos, eos = i_b * T, i_b * T + T
|
||||
|
||||
if HEAD_FIRST:
|
||||
p_s = tl.make_block_ptr(
|
||||
s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,)
|
||||
)
|
||||
p_o = tl.make_block_ptr(
|
||||
o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,)
|
||||
)
|
||||
else:
|
||||
p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
|
||||
p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
|
||||
# [BT]
|
||||
b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32)
|
||||
b_o = tl.cumsum(b_s, axis=0)
|
||||
if REVERSE:
|
||||
b_z = tl.sum(b_s, axis=0)
|
||||
b_o = -b_o + b_z[None] + b_s
|
||||
tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,))
|
||||
|
||||
|
||||
@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({"BS": BS}, num_warps=num_warps)
|
||||
for BS in BS_LIST
|
||||
for num_warps in [2, 4, 8]
|
||||
],
|
||||
key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE"],
|
||||
)
|
||||
@triton.jit(do_not_specialize=["T"])
|
||||
def chunk_local_cumsum_vector_kernel(
|
||||
s,
|
||||
o,
|
||||
cu_seqlens,
|
||||
chunk_indices,
|
||||
T,
|
||||
B: tl.constexpr,
|
||||
H: tl.constexpr,
|
||||
S: tl.constexpr,
|
||||
BT: tl.constexpr,
|
||||
BS: tl.constexpr,
|
||||
REVERSE: tl.constexpr,
|
||||
IS_VARLEN: tl.constexpr,
|
||||
HEAD_FIRST: tl.constexpr,
|
||||
):
|
||||
i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
|
||||
i_b, i_h = i_bh // H, i_bh % H
|
||||
if IS_VARLEN:
|
||||
i_n, i_t = (
|
||||
tl.load(chunk_indices + i_t * 2).to(tl.int32),
|
||||
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
|
||||
)
|
||||
bos, eos = (
|
||||
tl.load(cu_seqlens + i_n).to(tl.int32),
|
||||
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
|
||||
)
|
||||
T = eos - bos
|
||||
else:
|
||||
bos, eos = i_b * T, i_b * T + T
|
||||
|
||||
o_i = tl.arange(0, BT)
|
||||
if REVERSE:
|
||||
m_s = tl.where(o_i[:, None] <= o_i[None, :], 1.0, 0.0)
|
||||
else:
|
||||
m_s = tl.where(o_i[:, None] >= o_i[None, :], 1.0, 0.0)
|
||||
|
||||
if HEAD_FIRST:
|
||||
p_s = tl.make_block_ptr(
|
||||
s + (bos * H + i_h * T) * S,
|
||||
(T, S),
|
||||
(S, 1),
|
||||
(i_t * BT, i_s * BS),
|
||||
(BT, BS),
|
||||
(1, 0),
|
||||
)
|
||||
p_o = tl.make_block_ptr(
|
||||
o + (bos * H + i_h * T) * S,
|
||||
(T, S),
|
||||
(S, 1),
|
||||
(i_t * BT, i_s * BS),
|
||||
(BT, BS),
|
||||
(1, 0),
|
||||
)
|
||||
else:
|
||||
p_s = tl.make_block_ptr(
|
||||
s + (bos * H + i_h) * S,
|
||||
(T, S),
|
||||
(H * S, 1),
|
||||
(i_t * BT, i_s * BS),
|
||||
(BT, BS),
|
||||
(1, 0),
|
||||
)
|
||||
p_o = tl.make_block_ptr(
|
||||
o + (bos * H + i_h) * S,
|
||||
(T, S),
|
||||
(H * S, 1),
|
||||
(i_t * BT, i_s * BS),
|
||||
(BT, BS),
|
||||
(1, 0),
|
||||
)
|
||||
# [BT, BS]
|
||||
b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
|
||||
b_o = tl.dot(m_s, b_s, allow_tf32=False)
|
||||
tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
|
||||
|
||||
|
||||
def chunk_local_cumsum_scalar(
|
||||
g: torch.Tensor,
|
||||
chunk_size: int,
|
||||
reverse: bool = False,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
chunk_indices: torch.Tensor | None = None,
|
||||
head_first: bool = False,
|
||||
output_dtype: torch.dtype | None = torch.float,
|
||||
) -> torch.Tensor:
|
||||
if head_first:
|
||||
B, H, T = g.shape
|
||||
else:
|
||||
B, T, H = g.shape
|
||||
assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
|
||||
"chunk_size must be a power of 2"
|
||||
)
|
||||
if chunk_indices is None and cu_seqlens is not None:
|
||||
chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
|
||||
BT = chunk_size
|
||||
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
|
||||
g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
|
||||
grid = (NT, B * H)
|
||||
chunk_local_cumsum_scalar_kernel[grid](
|
||||
g_org,
|
||||
g,
|
||||
cu_seqlens,
|
||||
chunk_indices,
|
||||
T=T,
|
||||
B=B,
|
||||
H=H,
|
||||
BT=BT,
|
||||
HEAD_FIRST=head_first,
|
||||
REVERSE=reverse,
|
||||
)
|
||||
return g
|
||||
|
||||
|
||||
def chunk_local_cumsum_vector(
|
||||
g: torch.Tensor,
|
||||
chunk_size: int,
|
||||
reverse: bool = False,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
chunk_indices: torch.Tensor | None = None,
|
||||
head_first: bool = False,
|
||||
output_dtype: torch.dtype | None = torch.float,
|
||||
) -> torch.Tensor:
|
||||
if head_first:
|
||||
B, H, T, S = g.shape
|
||||
else:
|
||||
B, T, H, S = g.shape
|
||||
assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
|
||||
"chunk_size must be a power of 2"
|
||||
)
|
||||
if chunk_indices is None and cu_seqlens is not None:
|
||||
chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
|
||||
BT = chunk_size
|
||||
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
|
||||
|
||||
g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
|
||||
|
||||
def grid(meta):
|
||||
return (triton.cdiv(meta["S"], meta["BS"]), NT, B * H)
|
||||
|
||||
# keep cumulative normalizer in fp32
|
||||
# this kernel is equivalent to
|
||||
# g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1)
|
||||
chunk_local_cumsum_vector_kernel[grid](
|
||||
g_org,
|
||||
g,
|
||||
cu_seqlens,
|
||||
chunk_indices,
|
||||
T=T,
|
||||
B=B,
|
||||
H=H,
|
||||
S=S,
|
||||
BT=BT,
|
||||
HEAD_FIRST=head_first,
|
||||
REVERSE=reverse,
|
||||
)
|
||||
return g
|
||||
|
||||
|
||||
@input_guard
|
||||
def chunk_local_cumsum(
|
||||
g: torch.Tensor,
|
||||
chunk_size: int,
|
||||
reverse: bool = False,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
chunk_indices: torch.Tensor | None = None,
|
||||
head_first: bool = False,
|
||||
output_dtype: torch.dtype | None = torch.float,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if cu_seqlens is not None:
|
||||
assert g.shape[0] == 1, (
|
||||
"Only batch size 1 is supported when cu_seqlens are provided"
|
||||
)
|
||||
if len(g.shape) == 3:
|
||||
return chunk_local_cumsum_scalar(
|
||||
g,
|
||||
chunk_size,
|
||||
reverse,
|
||||
cu_seqlens,
|
||||
chunk_indices,
|
||||
head_first,
|
||||
output_dtype,
|
||||
)
|
||||
elif len(g.shape) == 4:
|
||||
return chunk_local_cumsum_vector(
|
||||
g,
|
||||
chunk_size,
|
||||
reverse,
|
||||
cu_seqlens,
|
||||
chunk_indices,
|
||||
head_first,
|
||||
output_dtype,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported input shape {g.shape}. "
|
||||
f"which should be (B, T, H, D) if `head_first=False` "
|
||||
f"or (B, H, T, D) otherwise"
|
||||
)
|
||||
248
upstream_ref/vllm_gdn/third_party/ops/fused_gdn_prefill_post_conv.py
vendored
Normal file
248
upstream_ref/vllm_gdn/third_party/ops/fused_gdn_prefill_post_conv.py
vendored
Normal file
@@ -0,0 +1,248 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Fused post-conv1d preparation for GDN prefill.
|
||||
|
||||
Replaces the chain:
|
||||
split → rearrange → contiguous * 3 → l2norm * 2 → gating
|
||||
with a **single Triton kernel** that reads the conv'd mixed_qkv output
|
||||
and writes directly to q/k/v/g/beta in the target contiguous layout.
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fused_post_conv_kernel(
|
||||
# ---- inputs ----
|
||||
mixed_qkv_ptr, # [L, qkv_dim] conv'd output (contiguous)
|
||||
a_ptr, # [L, HV]
|
||||
b_ptr, # [L, HV]
|
||||
# ---- params ----
|
||||
A_log_ptr, # [HV]
|
||||
dt_bias_ptr, # [HV]
|
||||
# ---- outputs ----
|
||||
q_ptr, # [L, H, K] contiguous
|
||||
k_ptr, # [L, H, K] contiguous
|
||||
v_ptr, # [L, HV, V] contiguous
|
||||
g_ptr, # [L, HV] float32
|
||||
beta_ptr, # [L, HV] float32
|
||||
# ---- strides ----
|
||||
stride_x_tok, # qkv_dim
|
||||
stride_a_tok, # HV
|
||||
stride_b_tok, # HV
|
||||
stride_q_tok, # H * K
|
||||
stride_k_tok, # H * K
|
||||
stride_v_tok, # HV * V
|
||||
# ---- dims ----
|
||||
L,
|
||||
H: tl.constexpr,
|
||||
HV: tl.constexpr,
|
||||
K: tl.constexpr,
|
||||
V: tl.constexpr,
|
||||
APPLY_L2NORM: tl.constexpr,
|
||||
L2NORM_EPS: tl.constexpr,
|
||||
OUTPUT_G_EXP: tl.constexpr,
|
||||
SOFTPLUS_THRESHOLD: tl.constexpr,
|
||||
BLOCK_T: tl.constexpr,
|
||||
BK: tl.constexpr,
|
||||
BV: tl.constexpr,
|
||||
):
|
||||
"""Single fused kernel for post-conv1d preparation.
|
||||
|
||||
Grid: (ceil(L, BLOCK_T), H + HV)
|
||||
- program_id(1) in [0, H): Q/K head processing + l2norm
|
||||
- program_id(1) in [H, H+HV): V head processing + gating
|
||||
"""
|
||||
i_tb = tl.program_id(0)
|
||||
i_head = tl.program_id(1)
|
||||
|
||||
HK: tl.constexpr = H * K
|
||||
|
||||
offs_t = i_tb * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T]
|
||||
mask_t = offs_t < L
|
||||
|
||||
if i_head < H:
|
||||
# ============ Q/K head processing ============
|
||||
i_h = i_head
|
||||
offs_k = tl.arange(0, BK) # [BK]
|
||||
mask_k = offs_k < K
|
||||
mask_2d = mask_t[:, None] & mask_k[None, :] # [BLOCK_T, BK]
|
||||
|
||||
# Load Q features: mixed_qkv[t, i_h*K + k]
|
||||
q_offsets = offs_t[:, None] * stride_x_tok + i_h * K + offs_k[None, :]
|
||||
q_f32 = tl.load(mixed_qkv_ptr + q_offsets, mask=mask_2d, other=0).to(tl.float32)
|
||||
|
||||
# Load K features: mixed_qkv[t, HK + i_h*K + k]
|
||||
k_offsets = offs_t[:, None] * stride_x_tok + HK + i_h * K + offs_k[None, :]
|
||||
k_f32 = tl.load(mixed_qkv_ptr + k_offsets, mask=mask_2d, other=0).to(tl.float32)
|
||||
|
||||
if APPLY_L2NORM:
|
||||
q_sq_sum = tl.sum(q_f32 * q_f32, axis=1) # [BLOCK_T]
|
||||
q_inv = 1.0 / tl.sqrt(q_sq_sum + L2NORM_EPS)
|
||||
q_f32 = q_f32 * q_inv[:, None]
|
||||
|
||||
k_sq_sum = tl.sum(k_f32 * k_f32, axis=1)
|
||||
k_inv = 1.0 / tl.sqrt(k_sq_sum + L2NORM_EPS)
|
||||
k_f32 = k_f32 * k_inv[:, None]
|
||||
|
||||
# Store Q
|
||||
q_out = offs_t[:, None] * stride_q_tok + i_h * K + offs_k[None, :]
|
||||
tl.store(
|
||||
q_ptr + q_out,
|
||||
q_f32.to(q_ptr.dtype.element_ty),
|
||||
mask=mask_2d,
|
||||
)
|
||||
|
||||
# Store K
|
||||
k_out = offs_t[:, None] * stride_k_tok + i_h * K + offs_k[None, :]
|
||||
tl.store(
|
||||
k_ptr + k_out,
|
||||
k_f32.to(k_ptr.dtype.element_ty),
|
||||
mask=mask_2d,
|
||||
)
|
||||
else:
|
||||
# ============ V head + gating processing ============
|
||||
i_hv = i_head - H
|
||||
offs_v = tl.arange(0, BV) # [BV]
|
||||
mask_v = offs_v < V
|
||||
mask_2d = mask_t[:, None] & mask_v[None, :] # [BLOCK_T, BV]
|
||||
|
||||
V_OFFSET: tl.constexpr = 2 * H * K
|
||||
|
||||
# Load V features: mixed_qkv[t, 2*H*K + i_hv*V + v]
|
||||
v_offsets = (
|
||||
offs_t[:, None] * stride_x_tok + V_OFFSET + i_hv * V + offs_v[None, :]
|
||||
)
|
||||
v_vals = tl.load(mixed_qkv_ptr + v_offsets, mask=mask_2d, other=0)
|
||||
|
||||
# Store V
|
||||
v_out = offs_t[:, None] * stride_v_tok + i_hv * V + offs_v[None, :]
|
||||
tl.store(v_ptr + v_out, v_vals, mask=mask_2d)
|
||||
|
||||
# Gating: one scalar per (token, v-head)
|
||||
A_log_val = tl.load(A_log_ptr + i_hv).to(tl.float32)
|
||||
dt_bias_val = tl.load(dt_bias_ptr + i_hv).to(tl.float32)
|
||||
|
||||
a_offsets = offs_t * stride_a_tok + i_hv
|
||||
b_offsets = offs_t * stride_b_tok + i_hv
|
||||
a_vals = tl.load(a_ptr + a_offsets, mask=mask_t, other=0).to(tl.float32)
|
||||
b_vals = tl.load(b_ptr + b_offsets, mask=mask_t, other=0).to(tl.float32)
|
||||
|
||||
# g = -exp(A_log) * softplus(a + dt_bias)
|
||||
x = a_vals + dt_bias_val
|
||||
sp = tl.where(x > 0, x + tl.log(1.0 + tl.exp(-x)), tl.log(1.0 + tl.exp(x)))
|
||||
sp = tl.where(x <= SOFTPLUS_THRESHOLD, sp, x)
|
||||
g_vals = -tl.exp(A_log_val) * sp
|
||||
|
||||
if OUTPUT_G_EXP:
|
||||
g_vals = tl.exp(g_vals)
|
||||
|
||||
beta_vals = tl.sigmoid(b_vals)
|
||||
|
||||
gb_offsets = offs_t * HV + i_hv
|
||||
tl.store(g_ptr + gb_offsets, g_vals, mask=mask_t)
|
||||
tl.store(beta_ptr + gb_offsets, beta_vals, mask=mask_t)
|
||||
|
||||
|
||||
def fused_post_conv_prep(
|
||||
conv_output: torch.Tensor, # [L, qkv_dim] conv'd mixed_qkv
|
||||
a: torch.Tensor, # [L, HV]
|
||||
b: torch.Tensor, # [L, HV]
|
||||
A_log: torch.Tensor, # [HV]
|
||||
dt_bias: torch.Tensor, # [HV]
|
||||
num_k_heads: int,
|
||||
head_k_dim: int,
|
||||
head_v_dim: int,
|
||||
apply_l2norm: bool = True,
|
||||
output_g_exp: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Fused post-conv1d prep: split + l2norm + gating in one kernel.
|
||||
|
||||
Args:
|
||||
conv_output: [L, qkv_dim] contiguous conv'd mixed_qkv
|
||||
a: [L, HV] gating input
|
||||
b: [L, HV] gating input
|
||||
A_log: [HV] log decay parameter
|
||||
dt_bias: [HV] dt bias parameter
|
||||
num_k_heads: number of K heads (H)
|
||||
head_k_dim: dimension per K head (K)
|
||||
head_v_dim: dimension per V head (V)
|
||||
apply_l2norm: whether to L2-normalize q and k
|
||||
output_g_exp: if True, output exp(g) instead of g (for FlashInfer)
|
||||
|
||||
Returns:
|
||||
q: [L, H, K] contiguous, optionally l2-normalized
|
||||
k: [L, H, K] contiguous, optionally l2-normalized
|
||||
v: [L, HV, V] contiguous
|
||||
g: [L, HV] float32
|
||||
beta: [L, HV] float32
|
||||
"""
|
||||
L = conv_output.shape[0]
|
||||
qkv_dim = conv_output.shape[1]
|
||||
H = num_k_heads
|
||||
K = head_k_dim
|
||||
V = head_v_dim
|
||||
HV = A_log.shape[0]
|
||||
dtype = conv_output.dtype
|
||||
device = conv_output.device
|
||||
|
||||
assert qkv_dim == 2 * H * K + HV * V, (
|
||||
f"qkv_dim={qkv_dim} != 2*H*K + HV*V = {2 * H * K + HV * V}"
|
||||
)
|
||||
|
||||
# Allocate outputs in target contiguous layout
|
||||
q = torch.empty(L, H, K, dtype=dtype, device=device)
|
||||
k = torch.empty(L, H, K, dtype=dtype, device=device)
|
||||
v = torch.empty(L, HV, V, dtype=dtype, device=device)
|
||||
g = torch.empty(L, HV, dtype=torch.float32, device=device)
|
||||
beta = torch.empty(L, HV, dtype=torch.float32, device=device)
|
||||
|
||||
if L == 0:
|
||||
return q, k, v, g, beta
|
||||
|
||||
# ---- Kernel config ----
|
||||
BK = triton.next_power_of_2(K)
|
||||
BV = triton.next_power_of_2(V)
|
||||
BLOCK_T = 16 # tokens per block
|
||||
|
||||
# Single kernel: blocks [0,H) do Q/K, blocks [H, H+HV) do V+gating
|
||||
grid = (triton.cdiv(L, BLOCK_T), H + HV)
|
||||
_fused_post_conv_kernel[grid](
|
||||
mixed_qkv_ptr=conv_output,
|
||||
a_ptr=a,
|
||||
b_ptr=b,
|
||||
A_log_ptr=A_log,
|
||||
dt_bias_ptr=dt_bias,
|
||||
q_ptr=q,
|
||||
k_ptr=k,
|
||||
v_ptr=v,
|
||||
g_ptr=g,
|
||||
beta_ptr=beta,
|
||||
stride_x_tok=conv_output.stride(0),
|
||||
stride_a_tok=a.stride(0),
|
||||
stride_b_tok=b.stride(0),
|
||||
stride_q_tok=q.stride(0),
|
||||
stride_k_tok=k.stride(0),
|
||||
stride_v_tok=v.stride(0),
|
||||
L=L,
|
||||
H=H,
|
||||
HV=HV,
|
||||
K=K,
|
||||
V=V,
|
||||
APPLY_L2NORM=apply_l2norm,
|
||||
L2NORM_EPS=1e-6,
|
||||
OUTPUT_G_EXP=output_g_exp,
|
||||
SOFTPLUS_THRESHOLD=20.0,
|
||||
BLOCK_T=BLOCK_T,
|
||||
BK=BK,
|
||||
BV=BV,
|
||||
num_warps=4,
|
||||
num_stages=2,
|
||||
)
|
||||
|
||||
return q, k, v, g, beta
|
||||
619
upstream_ref/vllm_gdn/third_party/ops/fused_recurrent.py
vendored
Normal file
619
upstream_ref/vllm_gdn/third_party/ops/fused_recurrent.py
vendored
Normal file
@@ -0,0 +1,619 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
||||
#
|
||||
# This file contains code copied from the flash-linear-attention project.
|
||||
# The original source code was licensed under the MIT license and included
|
||||
# the following copyright notice:
|
||||
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
||||
# ruff: noqa: E501
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
from .op import exp, log
|
||||
|
||||
|
||||
@triton.heuristics(
|
||||
{
|
||||
"USE_INITIAL_STATE": lambda args: args["h0"] is not None,
|
||||
"IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
|
||||
"IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None,
|
||||
"IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None,
|
||||
}
|
||||
)
|
||||
@triton.jit(do_not_specialize=["N", "T"])
|
||||
def fused_recurrent_gated_delta_rule_fwd_kernel(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
g,
|
||||
beta,
|
||||
o,
|
||||
h0,
|
||||
ht,
|
||||
cu_seqlens,
|
||||
ssm_state_indices,
|
||||
num_accepted_tokens,
|
||||
scale,
|
||||
N: tl.int64, # num of sequences
|
||||
T: tl.int64, # num of tokens
|
||||
B: tl.constexpr,
|
||||
H: tl.constexpr,
|
||||
HV: tl.constexpr,
|
||||
K: tl.constexpr,
|
||||
V: tl.constexpr,
|
||||
BK: tl.constexpr,
|
||||
BV: tl.constexpr,
|
||||
stride_init_state_token: tl.constexpr,
|
||||
stride_final_state_token: tl.constexpr,
|
||||
stride_indices_seq: tl.constexpr,
|
||||
stride_indices_tok: tl.constexpr,
|
||||
USE_INITIAL_STATE: tl.constexpr, # whether to use initial state
|
||||
INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace
|
||||
IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar,
|
||||
USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
|
||||
IS_VARLEN: tl.constexpr,
|
||||
IS_CONTINUOUS_BATCHING: tl.constexpr,
|
||||
IS_SPEC_DECODING: tl.constexpr,
|
||||
IS_KDA: tl.constexpr,
|
||||
):
|
||||
i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
|
||||
i_n, i_hv = i_nh // HV, i_nh % HV
|
||||
i_h = i_hv // (HV // H)
|
||||
if IS_VARLEN:
|
||||
bos, eos = (
|
||||
tl.load(cu_seqlens + i_n).to(tl.int64),
|
||||
tl.load(cu_seqlens + i_n + 1).to(tl.int64),
|
||||
)
|
||||
all = T
|
||||
T = eos - bos
|
||||
else:
|
||||
bos, eos = i_n * T, i_n * T + T
|
||||
all = B * T
|
||||
|
||||
if T == 0:
|
||||
# no tokens to process for this sequence
|
||||
return
|
||||
|
||||
o_k = i_k * BK + tl.arange(0, BK)
|
||||
o_v = i_v * BV + tl.arange(0, BV)
|
||||
|
||||
p_q = q + (bos * H + i_h) * K + o_k
|
||||
p_k = k + (bos * H + i_h) * K + o_k
|
||||
p_v = v + (bos * HV + i_hv) * V + o_v
|
||||
if IS_BETA_HEADWISE:
|
||||
p_beta = beta + (bos * HV + i_hv) * V + o_v
|
||||
else:
|
||||
p_beta = beta + bos * HV + i_hv
|
||||
|
||||
if not IS_KDA:
|
||||
p_g = g + bos * HV + i_hv
|
||||
else:
|
||||
p_gk = g + (bos * HV + i_hv) * K + o_k
|
||||
|
||||
p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v
|
||||
|
||||
mask_k = o_k < K
|
||||
mask_v = o_v < V
|
||||
mask_h = mask_v[:, None] & mask_k[None, :]
|
||||
|
||||
b_h = tl.zeros([BV, BK], dtype=tl.float32)
|
||||
if USE_INITIAL_STATE:
|
||||
if IS_CONTINUOUS_BATCHING:
|
||||
if IS_SPEC_DECODING:
|
||||
i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
|
||||
else:
|
||||
i_t = 0
|
||||
# Load state index and check for invalid entries
|
||||
state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(
|
||||
tl.int64
|
||||
)
|
||||
# Skip if state index is invalid (NULL_BLOCK_ID=0)
|
||||
if state_idx <= 0:
|
||||
return
|
||||
p_h0 = h0 + state_idx * stride_init_state_token
|
||||
else:
|
||||
p_h0 = h0 + bos * HV * V * K
|
||||
p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
|
||||
b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
|
||||
|
||||
for i_t in range(0, T):
|
||||
b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
|
||||
b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
|
||||
b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
|
||||
|
||||
if USE_QK_L2NORM_IN_KERNEL:
|
||||
b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
|
||||
b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
|
||||
b_q = b_q * scale
|
||||
# [BV, BK]
|
||||
if not IS_KDA:
|
||||
b_g = tl.load(p_g).to(tl.float32)
|
||||
b_h *= exp(b_g)
|
||||
else:
|
||||
b_gk = tl.load(p_gk).to(tl.float32)
|
||||
b_h *= exp(b_gk[None, :])
|
||||
# [BV]
|
||||
b_v -= tl.sum(b_h * b_k[None, :], 1)
|
||||
if IS_BETA_HEADWISE:
|
||||
b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
|
||||
else:
|
||||
b_beta = tl.load(p_beta).to(tl.float32)
|
||||
b_v *= b_beta
|
||||
# [BV, BK]
|
||||
b_h += b_v[:, None] * b_k[None, :]
|
||||
# [BV]
|
||||
b_o = tl.sum(b_h * b_q[None, :], 1)
|
||||
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
|
||||
|
||||
# keep the states for multi-query tokens
|
||||
if INPLACE_FINAL_STATE:
|
||||
# Load state index and check for invalid entries
|
||||
final_state_idx = tl.load(
|
||||
ssm_state_indices + i_n * stride_indices_seq + i_t
|
||||
).to(tl.int64)
|
||||
# Only store if state index is valid (not NULL_BLOCK_ID=0)
|
||||
if final_state_idx > 0:
|
||||
p_ht = ht + final_state_idx * stride_final_state_token
|
||||
p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
|
||||
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
|
||||
else:
|
||||
p_ht = ht + (bos + i_t) * stride_final_state_token
|
||||
p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
|
||||
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
|
||||
|
||||
p_q += H * K
|
||||
p_k += H * K
|
||||
p_o += HV * V
|
||||
p_v += HV * V
|
||||
if not IS_KDA:
|
||||
p_g += HV
|
||||
else:
|
||||
p_gk += HV * K
|
||||
p_beta += HV * (V if IS_BETA_HEADWISE else 1)
|
||||
|
||||
|
||||
def fused_recurrent_gated_delta_rule_fwd(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
beta: torch.Tensor,
|
||||
scale: float,
|
||||
initial_state: torch.Tensor,
|
||||
inplace_final_state: bool = True,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
ssm_state_indices: torch.Tensor | None = None,
|
||||
num_accepted_tokens: torch.Tensor | None = None,
|
||||
use_qk_l2norm_in_kernel: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
B, T, H, K, V = *k.shape, v.shape[-1]
|
||||
HV = v.shape[2]
|
||||
N = B if cu_seqlens is None else len(cu_seqlens) - 1
|
||||
BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32)
|
||||
NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
|
||||
assert NK == 1, "NK > 1 is not supported yet"
|
||||
num_stages = 3
|
||||
num_warps = 1
|
||||
|
||||
o = q.new_empty(NK, *v.shape)
|
||||
if inplace_final_state:
|
||||
final_state = initial_state
|
||||
else:
|
||||
final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype)
|
||||
|
||||
stride_init_state_token = initial_state.stride(0)
|
||||
stride_final_state_token = final_state.stride(0)
|
||||
|
||||
if ssm_state_indices is None:
|
||||
stride_indices_seq, stride_indices_tok = 1, 1
|
||||
elif ssm_state_indices.ndim == 1:
|
||||
stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1
|
||||
else:
|
||||
stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()
|
||||
|
||||
grid = (NK, NV, N * HV)
|
||||
fused_recurrent_gated_delta_rule_fwd_kernel[grid](
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
g=g,
|
||||
beta=beta,
|
||||
o=o,
|
||||
h0=initial_state,
|
||||
ht=final_state,
|
||||
cu_seqlens=cu_seqlens,
|
||||
ssm_state_indices=ssm_state_indices,
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
scale=scale,
|
||||
N=N,
|
||||
T=T,
|
||||
B=B,
|
||||
H=H,
|
||||
HV=HV,
|
||||
K=K,
|
||||
V=V,
|
||||
BK=BK,
|
||||
BV=BV,
|
||||
stride_init_state_token=stride_init_state_token,
|
||||
stride_final_state_token=stride_final_state_token,
|
||||
stride_indices_seq=stride_indices_seq,
|
||||
stride_indices_tok=stride_indices_tok,
|
||||
IS_BETA_HEADWISE=beta.ndim == v.ndim,
|
||||
USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
|
||||
INPLACE_FINAL_STATE=inplace_final_state,
|
||||
IS_KDA=False,
|
||||
num_warps=num_warps,
|
||||
num_stages=num_stages,
|
||||
)
|
||||
o = o.squeeze(0)
|
||||
return o, final_state
|
||||
|
||||
|
||||
@triton.jit
|
||||
def fused_recurrent_gated_delta_rule_packed_decode_kernel(
|
||||
mixed_qkv,
|
||||
a,
|
||||
b,
|
||||
A_log,
|
||||
dt_bias,
|
||||
o,
|
||||
h0,
|
||||
ht,
|
||||
ssm_state_indices,
|
||||
scale,
|
||||
stride_mixed_qkv_tok: tl.constexpr,
|
||||
stride_a_tok: tl.constexpr,
|
||||
stride_b_tok: tl.constexpr,
|
||||
stride_init_state_token: tl.constexpr,
|
||||
stride_final_state_token: tl.constexpr,
|
||||
stride_indices_seq: tl.constexpr,
|
||||
H: tl.constexpr,
|
||||
HV: tl.constexpr,
|
||||
K: tl.constexpr,
|
||||
V: tl.constexpr,
|
||||
BK: tl.constexpr,
|
||||
BV: tl.constexpr,
|
||||
SOFTPLUS_THRESHOLD: tl.constexpr,
|
||||
USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
|
||||
):
|
||||
i_v, i_nh = tl.program_id(0), tl.program_id(1)
|
||||
i_n, i_hv = i_nh // HV, i_nh % HV
|
||||
i_h = i_hv // (HV // H)
|
||||
|
||||
o_k = tl.arange(0, BK)
|
||||
o_v = i_v * BV + tl.arange(0, BV)
|
||||
mask_k = o_k < K
|
||||
mask_v = o_v < V
|
||||
mask_h = mask_v[:, None] & mask_k[None, :]
|
||||
|
||||
state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq).to(tl.int64)
|
||||
p_o = o + (i_n * HV + i_hv) * V + o_v
|
||||
|
||||
# Skip if state index is invalid (NULL_BLOCK_ID=0)
|
||||
if state_idx <= 0:
|
||||
zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty)
|
||||
tl.store(p_o, zero, mask=mask_v)
|
||||
return
|
||||
|
||||
p_h0 = h0 + state_idx * stride_init_state_token
|
||||
p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
|
||||
b_h = tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
|
||||
|
||||
p_mixed = mixed_qkv + i_n * stride_mixed_qkv_tok
|
||||
q_off = i_h * K + o_k
|
||||
k_off = (H * K) + i_h * K + o_k
|
||||
v_off = (2 * H * K) + i_hv * V + o_v
|
||||
b_q = tl.load(p_mixed + q_off, mask=mask_k, other=0).to(tl.float32)
|
||||
b_k = tl.load(p_mixed + k_off, mask=mask_k, other=0).to(tl.float32)
|
||||
b_v = tl.load(p_mixed + v_off, mask=mask_v, other=0).to(tl.float32)
|
||||
|
||||
if USE_QK_L2NORM_IN_KERNEL:
|
||||
b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
|
||||
b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
|
||||
b_q = b_q * scale
|
||||
|
||||
a_val = tl.load(a + i_n * stride_a_tok + i_hv).to(tl.float32)
|
||||
b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32)
|
||||
A_log_val = tl.load(A_log + i_hv).to(tl.float32)
|
||||
dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32)
|
||||
x = a_val + dt_bias_val
|
||||
softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x)
|
||||
g_val = -tl.exp(A_log_val) * softplus_x
|
||||
beta_val = tl.sigmoid(b_val).to(b.dtype.element_ty).to(tl.float32)
|
||||
|
||||
b_h *= exp(g_val)
|
||||
b_v -= tl.sum(b_h * b_k[None, :], 1)
|
||||
b_v *= beta_val
|
||||
b_h += b_v[:, None] * b_k[None, :]
|
||||
b_o = tl.sum(b_h * b_q[None, :], 1)
|
||||
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
|
||||
|
||||
p_ht = ht + state_idx * stride_final_state_token
|
||||
p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
|
||||
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
|
||||
|
||||
|
||||
def fused_recurrent_gated_delta_rule_packed_decode(
|
||||
mixed_qkv: torch.Tensor,
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
A_log: torch.Tensor,
|
||||
dt_bias: torch.Tensor,
|
||||
scale: float,
|
||||
initial_state: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
ssm_state_indices: torch.Tensor,
|
||||
use_qk_l2norm_in_kernel: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if mixed_qkv.ndim != 2:
|
||||
raise ValueError(
|
||||
f"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim})."
|
||||
)
|
||||
if mixed_qkv.stride(-1) != 1:
|
||||
raise ValueError("`mixed_qkv` must be contiguous in the last dim.")
|
||||
if a.ndim != 2 or b.ndim != 2:
|
||||
raise ValueError(
|
||||
f"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim})."
|
||||
)
|
||||
if a.stride(-1) != 1 or b.stride(-1) != 1:
|
||||
raise ValueError("`a`/`b` must be contiguous in the last dim.")
|
||||
if A_log.ndim != 1 or dt_bias.ndim != 1:
|
||||
raise ValueError("`A_log`/`dt_bias` must be 1D tensors.")
|
||||
if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:
|
||||
raise ValueError("`A_log`/`dt_bias` must be contiguous.")
|
||||
if ssm_state_indices.ndim != 1:
|
||||
raise ValueError(
|
||||
f"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim})."
|
||||
)
|
||||
if not out.is_contiguous():
|
||||
raise ValueError("`out` must be contiguous.")
|
||||
|
||||
dev = mixed_qkv.device
|
||||
if (
|
||||
a.device != dev
|
||||
or b.device != dev
|
||||
or A_log.device != dev
|
||||
or dt_bias.device != dev
|
||||
or initial_state.device != dev
|
||||
or out.device != dev
|
||||
or ssm_state_indices.device != dev
|
||||
):
|
||||
raise ValueError("All inputs must be on the same device.")
|
||||
|
||||
B = mixed_qkv.shape[0]
|
||||
if a.shape[0] != B or b.shape[0] != B:
|
||||
raise ValueError(
|
||||
"Mismatched batch sizes: "
|
||||
f"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, b.shape[0]={b.shape[0]}."
|
||||
)
|
||||
if ssm_state_indices.shape[0] != B:
|
||||
raise ValueError(
|
||||
f"`ssm_state_indices` must have shape [B] (got {tuple(ssm_state_indices.shape)}; expected ({B},))."
|
||||
)
|
||||
|
||||
if initial_state.ndim != 4:
|
||||
raise ValueError(
|
||||
f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})."
|
||||
)
|
||||
if initial_state.stride(-1) != 1:
|
||||
raise ValueError("`initial_state` must be contiguous in the last dim.")
|
||||
HV, V, K = initial_state.shape[-3:]
|
||||
if a.shape[1] != HV or b.shape[1] != HV:
|
||||
raise ValueError(
|
||||
f"`a`/`b` must have shape [B, HV] with HV={HV} (got a.shape={tuple(a.shape)}, b.shape={tuple(b.shape)})."
|
||||
)
|
||||
if A_log.numel() != HV or dt_bias.numel() != HV:
|
||||
raise ValueError(
|
||||
f"`A_log` and `dt_bias` must have {HV} elements (got A_log.numel()={A_log.numel()}, dt_bias.numel()={dt_bias.numel()})."
|
||||
)
|
||||
if out.shape != (B, 1, HV, V):
|
||||
raise ValueError(
|
||||
f"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)})."
|
||||
)
|
||||
|
||||
qkv_dim = mixed_qkv.shape[1]
|
||||
qk_dim = qkv_dim - HV * V
|
||||
if qk_dim <= 0 or qk_dim % 2 != 0:
|
||||
raise ValueError(
|
||||
f"Invalid packed `mixed_qkv` last dim={qkv_dim} for HV={HV}, V={V}."
|
||||
)
|
||||
q_dim = qk_dim // 2
|
||||
if q_dim % K != 0:
|
||||
raise ValueError(f"Invalid packed Q size {q_dim}: must be divisible by K={K}.")
|
||||
H = q_dim // K
|
||||
if H <= 0 or HV % H != 0:
|
||||
raise ValueError(
|
||||
f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}."
|
||||
)
|
||||
|
||||
BK = triton.next_power_of_2(K)
|
||||
if triton.cdiv(K, BK) != 1:
|
||||
raise ValueError(
|
||||
f"Packed decode kernel only supports NK=1 (got K={K}, BK={BK})."
|
||||
)
|
||||
BV = min(triton.next_power_of_2(V), 32)
|
||||
num_stages = 3
|
||||
num_warps = 1
|
||||
|
||||
stride_mixed_qkv_tok = mixed_qkv.stride(0)
|
||||
stride_a_tok = a.stride(0)
|
||||
stride_b_tok = b.stride(0)
|
||||
stride_init_state_token = initial_state.stride(0)
|
||||
stride_final_state_token = initial_state.stride(0)
|
||||
stride_indices_seq = ssm_state_indices.stride(0)
|
||||
|
||||
NV = triton.cdiv(V, BV)
|
||||
grid = (NV, B * HV)
|
||||
fused_recurrent_gated_delta_rule_packed_decode_kernel[grid](
|
||||
mixed_qkv=mixed_qkv,
|
||||
a=a,
|
||||
b=b,
|
||||
A_log=A_log,
|
||||
dt_bias=dt_bias,
|
||||
o=out,
|
||||
h0=initial_state,
|
||||
ht=initial_state,
|
||||
ssm_state_indices=ssm_state_indices,
|
||||
scale=scale,
|
||||
stride_mixed_qkv_tok=stride_mixed_qkv_tok,
|
||||
stride_a_tok=stride_a_tok,
|
||||
stride_b_tok=stride_b_tok,
|
||||
stride_init_state_token=stride_init_state_token,
|
||||
stride_final_state_token=stride_final_state_token,
|
||||
stride_indices_seq=stride_indices_seq,
|
||||
H=H,
|
||||
HV=HV,
|
||||
K=K,
|
||||
V=V,
|
||||
BK=BK,
|
||||
BV=BV,
|
||||
SOFTPLUS_THRESHOLD=20.0,
|
||||
USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
|
||||
num_warps=num_warps,
|
||||
num_stages=num_stages,
|
||||
)
|
||||
return out, initial_state
|
||||
|
||||
|
||||
class FusedRecurrentFunction(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
beta: torch.Tensor,
|
||||
scale: float,
|
||||
initial_state: torch.Tensor,
|
||||
inplace_final_state: bool = True,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
ssm_state_indices: torch.Tensor | None = None,
|
||||
num_accepted_tokens: torch.Tensor | None = None,
|
||||
use_qk_l2norm_in_kernel: bool = False,
|
||||
):
|
||||
o, final_state = fused_recurrent_gated_delta_rule_fwd(
|
||||
q=q.contiguous(),
|
||||
k=k.contiguous(),
|
||||
v=v.contiguous(),
|
||||
g=g.contiguous(),
|
||||
beta=beta.contiguous(),
|
||||
scale=scale,
|
||||
initial_state=initial_state,
|
||||
inplace_final_state=inplace_final_state,
|
||||
cu_seqlens=cu_seqlens,
|
||||
ssm_state_indices=ssm_state_indices,
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
|
||||
)
|
||||
|
||||
return o, final_state
|
||||
|
||||
|
||||
def fused_recurrent_gated_delta_rule(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
g: torch.Tensor,
|
||||
beta: torch.Tensor = None,
|
||||
scale: float = None,
|
||||
initial_state: torch.Tensor = None,
|
||||
inplace_final_state: bool = True,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
ssm_state_indices: torch.Tensor | None = None,
|
||||
num_accepted_tokens: torch.Tensor | None = None,
|
||||
use_qk_l2norm_in_kernel: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
r"""
|
||||
Args:
|
||||
q (torch.Tensor):
|
||||
queries of shape `[B, T, H, K]`.
|
||||
k (torch.Tensor):
|
||||
keys of shape `[B, T, H, K]`.
|
||||
v (torch.Tensor):
|
||||
values of shape `[B, T, HV, V]`.
|
||||
GVA is applied if `HV > H`.
|
||||
g (torch.Tensor):
|
||||
g (decays) of shape `[B, T, HV]`.
|
||||
beta (torch.Tensor):
|
||||
betas of shape `[B, T, HV]`.
|
||||
scale (Optional[int]):
|
||||
Scale factor for the RetNet attention scores.
|
||||
If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
|
||||
initial_state (Optional[torch.Tensor]):
|
||||
Initial state of shape `[N, HV, V, K]` for `N` input sequences.
|
||||
For equal-length input sequences, `N` equals the batch size `B`.
|
||||
Default: `None`.
|
||||
inplace_final_state: bool:
|
||||
Whether to store the final state in-place to save memory.
|
||||
Default: `True`.
|
||||
cu_seqlens (torch.Tensor):
|
||||
Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
|
||||
consistent with the FlashAttention API.
|
||||
ssm_state_indices (Optional[torch.Tensor]):
|
||||
Indices to map the input sequences to the initial/final states.
|
||||
num_accepted_tokens (Optional[torch.Tensor]):
|
||||
Number of accepted tokens for each sequence during decoding.
|
||||
|
||||
Returns:
|
||||
o (torch.Tensor):
|
||||
Outputs of shape `[B, T, HV, V]`.
|
||||
final_state (torch.Tensor):
|
||||
Final state of shape `[N, HV, V, K]`.
|
||||
|
||||
Examples::
|
||||
>>> import torch
|
||||
>>> import torch.nn.functional as F
|
||||
>>> from einops import rearrange
|
||||
>>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
|
||||
# inputs with equal lengths
|
||||
>>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512
|
||||
>>> q = torch.randn(B, T, H, K, device='cuda')
|
||||
>>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)
|
||||
>>> v = torch.randn(B, T, HV, V, device='cuda')
|
||||
>>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))
|
||||
>>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()
|
||||
>>> h0 = torch.randn(B, HV, V, K, device='cuda')
|
||||
>>> o, ht = fused_gated_recurrent_delta_rule(
|
||||
q, k, v, g, beta,
|
||||
initial_state=h0,
|
||||
)
|
||||
# for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
|
||||
>>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))
|
||||
# for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
|
||||
>>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.int32)
|
||||
>>> o_var, ht_var = fused_gated_recurrent_delta_rule(
|
||||
q, k, v, g, beta,
|
||||
initial_state=h0,
|
||||
cu_seqlens=cu_seqlens
|
||||
)
|
||||
"""
|
||||
if cu_seqlens is not None and q.shape[0] != 1:
|
||||
raise ValueError(
|
||||
f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
|
||||
f"Please flatten variable-length inputs before processing."
|
||||
)
|
||||
if scale is None:
|
||||
scale = k.shape[-1] ** -0.5
|
||||
else:
|
||||
assert scale > 0, "scale must be positive"
|
||||
if beta is None:
|
||||
beta = torch.ones_like(q[..., 0])
|
||||
o, final_state = FusedRecurrentFunction.apply(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
g,
|
||||
beta,
|
||||
scale,
|
||||
initial_state,
|
||||
inplace_final_state,
|
||||
cu_seqlens,
|
||||
ssm_state_indices,
|
||||
num_accepted_tokens,
|
||||
use_qk_l2norm_in_kernel,
|
||||
)
|
||||
return o, final_state
|
||||
151
upstream_ref/vllm_gdn/third_party/ops/l2norm.py
vendored
Normal file
151
upstream_ref/vllm_gdn/third_party/ops/l2norm.py
vendored
Normal file
@@ -0,0 +1,151 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
||||
#
|
||||
# This file contains code copied from the flash-linear-attention project.
|
||||
# The original source code was licensed under the MIT license and included
|
||||
# the following copyright notice:
|
||||
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
BT_LIST = [8, 16, 32, 64, 128]
|
||||
|
||||
USE_DEFAULT_FLA_NORM = int(os.getenv("USE_DEFAULT_FLA_NORM", "0"))
|
||||
|
||||
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16, 32]
|
||||
],
|
||||
key=["D"],
|
||||
)
|
||||
@triton.jit
|
||||
def l2norm_fwd_kernel1(
|
||||
x,
|
||||
y,
|
||||
D,
|
||||
BD: tl.constexpr,
|
||||
eps,
|
||||
):
|
||||
i_t = tl.program_id(0)
|
||||
x += i_t * D
|
||||
y += i_t * D
|
||||
# Compute mean and variance
|
||||
cols = tl.arange(0, BD)
|
||||
mask = cols < D
|
||||
b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32)
|
||||
b_var = tl.sum(b_x * b_x, axis=0)
|
||||
b_rstd = 1 / tl.sqrt(b_var + eps)
|
||||
# tl.store(Rstd + i_t, rstd)
|
||||
# Normalize and apply linear transformation
|
||||
b_y = b_x * b_rstd
|
||||
tl.store(y + cols, b_y, mask=mask)
|
||||
|
||||
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({"BT": BT}, num_warps=num_warps)
|
||||
for num_warps in [1, 2, 4, 8, 16]
|
||||
for BT in BT_LIST
|
||||
],
|
||||
key=["D"],
|
||||
)
|
||||
@triton.jit(do_not_specialize=["NB"])
|
||||
def l2norm_fwd_kernel(
|
||||
x,
|
||||
y,
|
||||
eps,
|
||||
NB,
|
||||
T,
|
||||
D: tl.constexpr,
|
||||
BT: tl.constexpr,
|
||||
BD: tl.constexpr,
|
||||
):
|
||||
i_t = tl.program_id(0)
|
||||
p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
|
||||
b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
|
||||
b_var = tl.sum(b_x * b_x, axis=1)
|
||||
b_y = b_x / tl.sqrt(b_var + eps)[:, None]
|
||||
p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
|
||||
tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
|
||||
|
||||
|
||||
@triton.jit
|
||||
def l2norm_fwd_kernel2(
|
||||
X, Y, eps, M, N: tl.constexpr, BD: tl.constexpr, MBLOCK: tl.constexpr
|
||||
):
|
||||
xoffset = tl.program_id(0) * MBLOCK
|
||||
row_idx = xoffset + tl.arange(0, MBLOCK)[:, None]
|
||||
xmask = row_idx < M
|
||||
rindex = tl.arange(0, BD)[None, :]
|
||||
cmask = rindex < N
|
||||
mask = xmask & cmask
|
||||
xs = tl.load(X + (rindex + N * row_idx), mask, other=0.0).to(tl.float32)
|
||||
square = tl.broadcast_to(xs * xs, [MBLOCK, BD])
|
||||
square_sum = tl.sum(tl.where(xmask, square, 0), 1)[:, None]
|
||||
rsqrt = tl.rsqrt(square_sum + eps)
|
||||
tl.store(Y + (rindex + N * row_idx), xs * rsqrt, mask)
|
||||
|
||||
|
||||
def l2norm_fwd(
|
||||
x: torch.Tensor, eps: float = 1e-6, output_dtype: torch.dtype | None = None
|
||||
):
|
||||
x_shape_og = x.shape
|
||||
x = x.view(-1, x.shape[-1])
|
||||
# allocate output
|
||||
if output_dtype is None:
|
||||
y = torch.empty_like(x)
|
||||
else:
|
||||
y = torch.empty_like(x, dtype=output_dtype)
|
||||
assert y.stride(-1) == 1
|
||||
T, D = x.shape[0], x.shape[-1]
|
||||
# rstd = torch.empty((T,), dtype=torch.float32, device=x.device)
|
||||
# Less than 64KB per feature: enqueue fused kernel
|
||||
MAX_FUSED_SIZE = 65536 // x.element_size()
|
||||
BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
|
||||
if D > BD:
|
||||
raise RuntimeError("This layer doesn't support feature dim >= 64KB.")
|
||||
|
||||
if not USE_DEFAULT_FLA_NORM:
|
||||
MBLOCK = 32
|
||||
# M, N = x.shape
|
||||
l2norm_fwd_kernel2[(triton.cdiv(T, MBLOCK),)](
|
||||
x,
|
||||
y,
|
||||
eps,
|
||||
T,
|
||||
D,
|
||||
BD,
|
||||
MBLOCK,
|
||||
)
|
||||
else:
|
||||
if D <= 512:
|
||||
NB = triton.cdiv(T, 2048)
|
||||
|
||||
def grid(meta):
|
||||
return (triton.cdiv(T, meta["BT"]),)
|
||||
|
||||
l2norm_fwd_kernel[grid](
|
||||
x,
|
||||
y,
|
||||
eps,
|
||||
NB=NB,
|
||||
T=T,
|
||||
D=D,
|
||||
BD=BD,
|
||||
)
|
||||
else:
|
||||
l2norm_fwd_kernel1[(T,)](
|
||||
x,
|
||||
y,
|
||||
eps=eps,
|
||||
D=D,
|
||||
BD=BD,
|
||||
)
|
||||
|
||||
return y.view(x_shape_og)
|
||||
200
upstream_ref/vllm_gdn/third_party/ops/utils.py
vendored
Normal file
200
upstream_ref/vllm_gdn/third_party/ops/utils.py
vendored
Normal file
@@ -0,0 +1,200 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
||||
#
|
||||
# This file contains code copied from the flash-linear-attention project.
|
||||
# The original source code was licensed under the MIT license and included
|
||||
# the following copyright notice:
|
||||
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
||||
# ruff: noqa: E501
|
||||
import contextlib
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from enum import Enum
|
||||
from typing import Any, Literal
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import triton
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
COMPILER_MODE = os.getenv("FLA_COMPILER_MODE") == "1"
|
||||
FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1"
|
||||
|
||||
SUPPRESS_LEVEL = int(os.getenv("GDN_RECOMPUTE_SUPPRESS_LEVEL", "0"))
|
||||
|
||||
# Default chunk size used across FLA triton kernels (kda, chunk, chunk_o, etc.)
|
||||
FLA_CHUNK_SIZE = 64
|
||||
|
||||
|
||||
def tensor_cache(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:
|
||||
"""
|
||||
A decorator that caches the most recent results of a function with tensor inputs.
|
||||
|
||||
This decorator will store the output of the decorated function for the most recent set of input tensors.
|
||||
The cache is limited to a fixed size (default is 4). When the cache is full, the oldest entry will be removed.
|
||||
|
||||
Args:
|
||||
fn (Callable[..., torch.Tensor]):
|
||||
The function to be decorated. It should take tensor inputs and return tensor outputs.
|
||||
|
||||
Returns:
|
||||
Callable[..., torch.Tensor]:
|
||||
A wrapped version of the input function with single-entry caching.
|
||||
"""
|
||||
|
||||
cache_entries: tuple[tuple | None, dict | None, Any] = []
|
||||
cache_size = 8
|
||||
|
||||
@functools.wraps(fn)
|
||||
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
nonlocal cache_entries, cache_size
|
||||
for i, entry in enumerate(cache_entries):
|
||||
last_args, last_kwargs, last_result = entry
|
||||
if (
|
||||
len(args) == len(last_args)
|
||||
and len(kwargs) == len(last_kwargs)
|
||||
and all(a is b for a, b in zip(args, last_args))
|
||||
and all(
|
||||
k in last_kwargs and v is last_kwargs[k] for k, v in kwargs.items()
|
||||
)
|
||||
):
|
||||
cache_entries = (
|
||||
cache_entries[:i]
|
||||
+ cache_entries[i + 1 :]
|
||||
+ [(args, kwargs, last_result)]
|
||||
)
|
||||
return last_result
|
||||
|
||||
result = fn(*args, **kwargs)
|
||||
|
||||
if len(cache_entries) >= cache_size:
|
||||
cache_entries = cache_entries[1:]
|
||||
cache_entries.append((args, kwargs, result))
|
||||
return result
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:
|
||||
"""
|
||||
A decorator to make sure all input tensors are contiguous and set the device based on input tensors.
|
||||
"""
|
||||
|
||||
@functools.wraps(fn)
|
||||
def wrapper(*args, **kwargs):
|
||||
contiguous_args = (
|
||||
i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args
|
||||
)
|
||||
contiguous_kwargs = {
|
||||
k: (v if not isinstance(v, torch.Tensor) else v.contiguous())
|
||||
for k, v in kwargs.items()
|
||||
}
|
||||
|
||||
tensor = None
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor):
|
||||
tensor = arg
|
||||
break
|
||||
if tensor is None:
|
||||
for value in kwargs.values():
|
||||
if isinstance(value, torch.Tensor):
|
||||
tensor = value
|
||||
break
|
||||
|
||||
if tensor is not None:
|
||||
ctx = torch.accelerator.device_index(tensor.device.index)
|
||||
else:
|
||||
ctx = contextlib.nullcontext()
|
||||
|
||||
with ctx:
|
||||
return fn(*contiguous_args, **contiguous_kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@functools.cache
|
||||
def get_available_device() -> str:
|
||||
try:
|
||||
return triton.runtime.driver.active.get_current_target().backend
|
||||
except (RuntimeError, AttributeError):
|
||||
return "cpu"
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _check_platform() -> Literal["nvidia", "amd", "intel", "musa"]:
|
||||
device = get_available_device()
|
||||
mapping = {
|
||||
"cuda": "nvidia",
|
||||
"hip": "amd",
|
||||
"xpu": "intel",
|
||||
}
|
||||
# return the mapped value, or the original if not found
|
||||
return mapping.get(device, device)
|
||||
|
||||
|
||||
# For AMD GPUs, the triton backend is 'hip', while for Nvidia GPUs, the triton backend is 'cuda'.
|
||||
# However, the torch backend is 'cuda' for both Nvidia and AMD GPUs.
|
||||
# Therefore, we need to check the triton backend to determine the actual GPU vendor.
|
||||
device = "cuda" if current_platform.is_cuda_alike() else get_available_device()
|
||||
device_torch_lib = getattr(torch, device, None)
|
||||
device_platform = _check_platform()
|
||||
|
||||
is_amd = device_platform == "amd"
|
||||
is_intel = device_platform == "intel"
|
||||
is_nvidia = device_platform == "nvidia"
|
||||
is_intel_alchemist = is_intel and "Intel(R) Arc(TM) A" in torch.xpu.get_device_name(0)
|
||||
is_nvidia_hopper = is_nvidia and (
|
||||
"NVIDIA H" in torch.cuda.get_device_name(0)
|
||||
or torch.cuda.get_device_capability()[0] >= 9
|
||||
)
|
||||
use_cuda_graph = is_nvidia and os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1"
|
||||
is_gather_supported = hasattr(triton.language, "gather")
|
||||
is_tma_supported = (
|
||||
is_nvidia_hopper
|
||||
and os.getenv("FLA_USE_TMA", "0") == "1"
|
||||
and (
|
||||
hasattr(triton.language, "_experimental_make_tensor_descriptor")
|
||||
or hasattr(triton.language, "make_tensor_descriptor")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def get_all_max_shared_mem():
|
||||
try:
|
||||
return [
|
||||
triton.runtime.driver.active.utils.get_device_properties(i)[
|
||||
"max_shared_mem"
|
||||
]
|
||||
for i in range(device_torch_lib.device_count())
|
||||
]
|
||||
except BaseException:
|
||||
return [-1]
|
||||
|
||||
|
||||
class Backend(Enum):
|
||||
ADA = 101376 # RTX 4090
|
||||
AMPERE = 166912 # A100
|
||||
HOPPER = 232448 # H100
|
||||
DEFAULT = 102400 # Default
|
||||
|
||||
@classmethod
|
||||
def get_shared_memory(cls, arch: str) -> int:
|
||||
try:
|
||||
return cls[arch.upper()].value
|
||||
except KeyError:
|
||||
return cls.DEFAULT.value
|
||||
|
||||
|
||||
@functools.cache
|
||||
def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool:
|
||||
try:
|
||||
device_shared_mem_list = get_all_max_shared_mem()
|
||||
max_shared_memory = device_shared_mem_list[tensor_idx]
|
||||
return max_shared_memory >= Backend.get_shared_memory(arch)
|
||||
except Exception:
|
||||
return False
|
||||
538
upstream_ref/vllm_gdn/v1/attention/backends/gdn_attn.py
Normal file
538
upstream_ref/vllm_gdn/v1/attention/backends/gdn_attn.py
Normal file
@@ -0,0 +1,538 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Backend for GatedDeltaNet attention."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.utils.torch_utils import async_tensor_h2d
|
||||
from vllm.v1.attention.backend import (
|
||||
AttentionBackend,
|
||||
AttentionCGSupport,
|
||||
AttentionMetadataBuilder,
|
||||
CommonAttentionMetadata,
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import (
|
||||
NULL_BLOCK_ID,
|
||||
compute_causal_conv1d_metadata,
|
||||
mamba_get_block_table_tensor,
|
||||
split_decodes_and_prefills,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import MambaSpec
|
||||
|
||||
|
||||
class GDNAttentionBackend(AttentionBackend):
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "GDN_ATTN"
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["GDNAttentionMetadataBuilder"]:
|
||||
return GDNAttentionMetadataBuilder
|
||||
|
||||
@classmethod
|
||||
def is_ssm(cls) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
@dataclass
|
||||
class GDNAttentionMetadata:
|
||||
num_prefills: int
|
||||
num_prefill_tokens: int
|
||||
num_decodes: int
|
||||
num_decode_tokens: int
|
||||
num_spec_decodes: int
|
||||
num_spec_decode_tokens: int
|
||||
num_actual_tokens: int
|
||||
|
||||
has_initial_state: torch.Tensor | None = None
|
||||
|
||||
spec_query_start_loc: torch.Tensor | None = None # shape: [num_spec_decodes + 1,]
|
||||
non_spec_query_start_loc: torch.Tensor | None = (
|
||||
None # shape: [batch - num_spec_decodes + 1,]
|
||||
)
|
||||
|
||||
spec_state_indices_tensor: torch.Tensor | None = None # shape: [batch, num_spec]
|
||||
non_spec_state_indices_tensor: torch.Tensor | None = (
|
||||
None # shape: [batch - num_spec_decodes,]
|
||||
)
|
||||
spec_sequence_masks: torch.Tensor | None = None # shape: [batch,]
|
||||
spec_token_indx: torch.Tensor | None = None
|
||||
non_spec_token_indx: torch.Tensor | None = None
|
||||
|
||||
num_accepted_tokens: torch.Tensor | None = None # shape: [batch,]
|
||||
|
||||
# Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices)
|
||||
chunk_indices: torch.Tensor | None = None
|
||||
chunk_offsets: torch.Tensor | None = None
|
||||
# Chunk-kernel inputs for prefill
|
||||
prefill_query_start_loc: torch.Tensor | None = None
|
||||
prefill_state_indices: torch.Tensor | None = None
|
||||
prefill_has_initial_state: torch.Tensor | None = None
|
||||
|
||||
# The following attributes are for triton implementation of causal_conv1d
|
||||
nums_dict: dict | None = None
|
||||
batch_ptr: torch.Tensor | None = None
|
||||
token_chunk_offset_ptr: torch.Tensor | None = None
|
||||
|
||||
|
||||
class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]):
|
||||
kv_cache_spec: MambaSpec
|
||||
_cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
|
||||
|
||||
reorder_batch_threshold: int = 1
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kv_cache_spec: MambaSpec,
|
||||
layer_names: list[str],
|
||||
vllm_config: VllmConfig,
|
||||
device: torch.device,
|
||||
):
|
||||
self.vllm_config = vllm_config
|
||||
self.compilation_config = vllm_config.compilation_config
|
||||
self.speculative_config = vllm_config.speculative_config
|
||||
self.kv_cache_spec = kv_cache_spec
|
||||
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
|
||||
_resolve_gdn_prefill_backend,
|
||||
)
|
||||
|
||||
self.gdn_prefill_backend: Literal["triton", "flashinfer", "cutedsl"]
|
||||
_, self.gdn_prefill_backend = _resolve_gdn_prefill_backend(vllm_config)
|
||||
|
||||
if self.speculative_config:
|
||||
assert self.speculative_config.num_speculative_tokens is not None
|
||||
self.num_spec: int = self.speculative_config.num_speculative_tokens
|
||||
else:
|
||||
self.num_spec = 0
|
||||
self.use_spec_decode: bool = self.num_spec > 0
|
||||
self._init_reorder_batch_threshold(1, self.use_spec_decode)
|
||||
|
||||
self.use_full_cuda_graph: bool = (
|
||||
self.compilation_config.cudagraph_mode.has_full_cudagraphs()
|
||||
)
|
||||
|
||||
self.decode_cudagraph_max_bs: int = (
|
||||
self.vllm_config.scheduler_config.max_num_seqs * (self.num_spec + 1)
|
||||
)
|
||||
if self.compilation_config.max_cudagraph_capture_size is not None:
|
||||
self.decode_cudagraph_max_bs = min(
|
||||
self.decode_cudagraph_max_bs,
|
||||
self.compilation_config.max_cudagraph_capture_size,
|
||||
)
|
||||
|
||||
self.spec_state_indices_tensor: torch.Tensor = torch.empty(
|
||||
(self.decode_cudagraph_max_bs, self.num_spec + 1),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
self.non_spec_state_indices_tensor: torch.Tensor = torch.empty(
|
||||
(self.decode_cudagraph_max_bs,),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
self.spec_sequence_masks: torch.Tensor = torch.empty(
|
||||
(self.decode_cudagraph_max_bs,),
|
||||
dtype=torch.bool,
|
||||
device=device,
|
||||
)
|
||||
self.spec_token_indx: torch.Tensor = torch.empty(
|
||||
(self.decode_cudagraph_max_bs * (self.num_spec + 1),),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
self.non_spec_token_indx: torch.Tensor = torch.empty(
|
||||
(self.decode_cudagraph_max_bs * (self.num_spec + 1),),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
self.spec_query_start_loc: torch.Tensor = torch.empty(
|
||||
(self.decode_cudagraph_max_bs + 1,),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
self.non_spec_query_start_loc: torch.Tensor = torch.empty(
|
||||
(self.decode_cudagraph_max_bs + 1,),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
self.num_accepted_tokens: torch.Tensor = torch.empty(
|
||||
(self.decode_cudagraph_max_bs,),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
def build( # type: ignore[override]
|
||||
self,
|
||||
common_prefix_len: int,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
num_accepted_tokens: torch.Tensor | None = None,
|
||||
num_decode_draft_tokens_cpu: torch.Tensor | None = None,
|
||||
fast_build: bool = False,
|
||||
) -> GDNAttentionMetadata:
|
||||
m = common_attn_metadata
|
||||
|
||||
query_start_loc = m.query_start_loc
|
||||
query_start_loc_cpu = m.query_start_loc_cpu
|
||||
context_lens_tensor = m.compute_num_computed_tokens()
|
||||
nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
|
||||
block_table_tensor = mamba_get_block_table_tensor(
|
||||
m.block_table_tensor,
|
||||
m.seq_lens,
|
||||
self.kv_cache_spec,
|
||||
self.vllm_config.cache_config.mamba_cache_mode,
|
||||
)
|
||||
|
||||
spec_sequence_masks_cpu: torch.Tensor | None = None
|
||||
if (
|
||||
not self.use_spec_decode
|
||||
or num_decode_draft_tokens_cpu is None
|
||||
or num_decode_draft_tokens_cpu[num_decode_draft_tokens_cpu >= 0]
|
||||
.sum()
|
||||
.item()
|
||||
== 0
|
||||
):
|
||||
spec_sequence_masks = None
|
||||
num_spec_decodes = 0
|
||||
else:
|
||||
spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
|
||||
num_spec_decodes = spec_sequence_masks_cpu.sum().item()
|
||||
if num_spec_decodes == 0:
|
||||
spec_sequence_masks = None
|
||||
spec_sequence_masks_cpu = None
|
||||
else:
|
||||
spec_sequence_masks = async_tensor_h2d(
|
||||
spec_sequence_masks_cpu, device=query_start_loc.device
|
||||
)
|
||||
|
||||
if spec_sequence_masks is None:
|
||||
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
||||
split_decodes_and_prefills(m, decode_threshold=1)
|
||||
)
|
||||
num_spec_decode_tokens = 0
|
||||
spec_token_indx = None
|
||||
non_spec_token_indx = None
|
||||
spec_state_indices_tensor = None
|
||||
non_spec_state_indices_tensor = block_table_tensor[:, 0]
|
||||
spec_query_start_loc = None
|
||||
non_spec_query_start_loc = query_start_loc
|
||||
non_spec_query_start_loc_cpu = query_start_loc_cpu
|
||||
num_accepted_tokens = None
|
||||
else:
|
||||
query_lens = query_start_loc[1:] - query_start_loc[:-1]
|
||||
assert spec_sequence_masks_cpu is not None
|
||||
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
|
||||
|
||||
# Use CPU tensors to avoid CPU-GPU sync
|
||||
non_spec_query_lens_cpu = query_lens_cpu[~spec_sequence_masks_cpu]
|
||||
num_decodes = (non_spec_query_lens_cpu == 1).sum().item()
|
||||
# Exclude zero-length padded sequences from prefill count.
|
||||
num_zero_len = (non_spec_query_lens_cpu == 0).sum().item()
|
||||
num_prefills = non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len
|
||||
num_decode_tokens = num_decodes
|
||||
num_prefill_tokens = (
|
||||
non_spec_query_lens_cpu.sum().item() - num_decode_tokens
|
||||
)
|
||||
num_spec_decode_tokens = (
|
||||
query_lens_cpu.sum().item() - num_prefill_tokens - num_decode_tokens
|
||||
)
|
||||
|
||||
# num_decodes and num_spec_decodes are mutually exclusive.
|
||||
# Reclassify non-spec decodes as prefills when spec decodes
|
||||
# exist — the prefill kernel handles 1-token sequences with
|
||||
# initial state correctly, producing identical results.
|
||||
if num_decodes > 0 and num_spec_decodes > 0:
|
||||
num_prefills += num_decodes
|
||||
num_prefill_tokens += num_decode_tokens
|
||||
num_decodes = 0
|
||||
num_decode_tokens = 0
|
||||
|
||||
if num_prefills == 0 and num_decodes == 0:
|
||||
spec_token_size = min(
|
||||
num_spec_decodes * (self.num_spec + 1),
|
||||
query_start_loc_cpu[-1].item(),
|
||||
)
|
||||
spec_token_indx = torch.arange(
|
||||
spec_token_size,
|
||||
dtype=torch.int32,
|
||||
device=query_start_loc.device,
|
||||
)
|
||||
non_spec_token_indx = torch.empty(
|
||||
0, dtype=torch.int32, device=query_start_loc.device
|
||||
)
|
||||
# Filter by spec_sequence_masks to exclude padded sequences
|
||||
spec_state_indices_tensor = block_table_tensor[
|
||||
spec_sequence_masks_cpu, : self.num_spec + 1
|
||||
]
|
||||
non_spec_state_indices_tensor = None
|
||||
# Padded sequences are always at the back, so the first
|
||||
# num_spec_decodes + 1 entries of query_start_loc already
|
||||
# contain the correct cumulative token counts.
|
||||
spec_query_start_loc = query_start_loc[: num_spec_decodes + 1]
|
||||
non_spec_query_start_loc = None
|
||||
non_spec_query_start_loc_cpu = None
|
||||
else:
|
||||
spec_token_masks = torch.repeat_interleave(
|
||||
spec_sequence_masks,
|
||||
query_lens,
|
||||
output_size=query_start_loc_cpu[-1].item(),
|
||||
)
|
||||
index = torch.argsort(spec_token_masks, stable=True)
|
||||
num_non_spec_tokens = num_prefill_tokens + num_decode_tokens
|
||||
non_spec_token_indx = index[:num_non_spec_tokens]
|
||||
spec_token_indx = index[num_non_spec_tokens:]
|
||||
|
||||
spec_state_indices_tensor = block_table_tensor[
|
||||
spec_sequence_masks_cpu, : self.num_spec + 1
|
||||
]
|
||||
non_spec_state_indices_tensor = block_table_tensor[
|
||||
~spec_sequence_masks_cpu, 0
|
||||
]
|
||||
|
||||
spec_query_start_loc = torch.zeros(
|
||||
num_spec_decodes + 1,
|
||||
dtype=torch.int32,
|
||||
device=query_start_loc.device,
|
||||
)
|
||||
torch.cumsum(
|
||||
query_lens[spec_sequence_masks_cpu],
|
||||
dim=0,
|
||||
out=spec_query_start_loc[1:],
|
||||
)
|
||||
non_spec_query_start_loc = torch.zeros(
|
||||
query_lens.size(0) - num_spec_decodes + 1,
|
||||
dtype=torch.int32,
|
||||
device=query_start_loc.device,
|
||||
)
|
||||
torch.cumsum(
|
||||
query_lens[~spec_sequence_masks_cpu],
|
||||
dim=0,
|
||||
out=non_spec_query_start_loc[1:],
|
||||
)
|
||||
non_spec_query_start_loc_cpu = torch.zeros(
|
||||
query_lens_cpu.size(0) - num_spec_decodes + 1,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
torch.cumsum(
|
||||
query_lens_cpu[~spec_sequence_masks_cpu],
|
||||
dim=0,
|
||||
out=non_spec_query_start_loc_cpu[1:],
|
||||
)
|
||||
|
||||
assert num_accepted_tokens is not None
|
||||
num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu]
|
||||
|
||||
chunk_indices: torch.Tensor | None = None
|
||||
chunk_offsets: torch.Tensor | None = None
|
||||
prefill_query_start_loc: torch.Tensor | None = None
|
||||
prefill_state_indices: torch.Tensor | None = None
|
||||
prefill_has_initial_state: torch.Tensor | None = None
|
||||
if num_prefills > 0:
|
||||
from vllm.third_party.flash_linear_attention.ops.utils import (
|
||||
FLA_CHUNK_SIZE,
|
||||
)
|
||||
|
||||
# In a mixed non-spec batch, decodes are peeled off to the recurrent
|
||||
# kernel (decode-first front slice), so build chunk metadata from the
|
||||
# rebased prefill-only cu_seqlens; otherwise use the full non-spec one.
|
||||
# _forward_core keys off the same condition, so they agree.
|
||||
if spec_sequence_masks is None and num_decodes > 0:
|
||||
assert non_spec_query_start_loc is not None
|
||||
assert non_spec_query_start_loc_cpu is not None
|
||||
assert non_spec_state_indices_tensor is not None
|
||||
prefill_query_start_loc = (
|
||||
non_spec_query_start_loc[num_decodes:] - num_decode_tokens
|
||||
)
|
||||
prefill_query_start_loc_cpu = (
|
||||
non_spec_query_start_loc_cpu[num_decodes:] - num_decode_tokens
|
||||
)
|
||||
prefill_state_indices = non_spec_state_indices_tensor[num_decodes:]
|
||||
else:
|
||||
prefill_query_start_loc = non_spec_query_start_loc
|
||||
prefill_query_start_loc_cpu = non_spec_query_start_loc_cpu
|
||||
prefill_state_indices = non_spec_state_indices_tensor
|
||||
|
||||
if self.gdn_prefill_backend == "cutedsl":
|
||||
from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
|
||||
prepare_metadata_cutedsl,
|
||||
)
|
||||
|
||||
assert prefill_query_start_loc is not None
|
||||
assert prefill_query_start_loc_cpu is not None
|
||||
total_tokens = int(prefill_query_start_loc_cpu[-1].item())
|
||||
chunk_indices, chunk_offsets = prepare_metadata_cutedsl(
|
||||
prefill_query_start_loc,
|
||||
total_tokens,
|
||||
FLA_CHUNK_SIZE,
|
||||
)
|
||||
else:
|
||||
gpu_device = query_start_loc.device
|
||||
# Only prefill batches use FLA chunk ops.
|
||||
# Pre-compute on CPU and async-copy to GPU to avoid
|
||||
# GPU→CPU sync (.tolist()) in prepare_chunk_indices.
|
||||
from vllm.third_party.flash_linear_attention.ops.index import (
|
||||
prepare_chunk_indices,
|
||||
prepare_chunk_offsets,
|
||||
)
|
||||
|
||||
assert prefill_query_start_loc_cpu is not None
|
||||
chunk_indices = async_tensor_h2d(
|
||||
prepare_chunk_indices(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
|
||||
device=gpu_device,
|
||||
)
|
||||
chunk_offsets = async_tensor_h2d(
|
||||
prepare_chunk_offsets(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
|
||||
device=gpu_device,
|
||||
)
|
||||
|
||||
if num_prefills > 0:
|
||||
has_initial_state = context_lens_tensor > 0
|
||||
if spec_sequence_masks_cpu is not None:
|
||||
has_initial_state = has_initial_state[~spec_sequence_masks_cpu]
|
||||
assert non_spec_query_start_loc_cpu is not None
|
||||
nums_dict, batch_ptr, token_chunk_offset_ptr = (
|
||||
compute_causal_conv1d_metadata(
|
||||
non_spec_query_start_loc_cpu,
|
||||
device=query_start_loc.device,
|
||||
)
|
||||
)
|
||||
if spec_sequence_masks is None and num_decodes > 0:
|
||||
prefill_has_initial_state = has_initial_state[num_decodes:]
|
||||
else:
|
||||
prefill_has_initial_state = has_initial_state
|
||||
else:
|
||||
has_initial_state = None
|
||||
|
||||
# Function code counted on either presency non-spec decode or spec decode,
|
||||
# but not both.
|
||||
assert not (num_decodes > 0 and num_spec_decodes > 0), (
|
||||
f"num_decodes: {num_decodes}, num_spec_decodes: {num_spec_decodes}"
|
||||
)
|
||||
|
||||
# Prepare per-request tensors for cudagraph. m.num_actual_tokens is
|
||||
# token-padded for FULL graph replay, but the GDN state/query/accepted
|
||||
# metadata below is indexed by request.
|
||||
batch_size = m.num_reqs
|
||||
|
||||
if (
|
||||
self.use_full_cuda_graph
|
||||
and num_prefills == 0
|
||||
and num_decodes == 0
|
||||
and num_spec_decodes <= self.decode_cudagraph_max_bs
|
||||
and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
|
||||
):
|
||||
assert spec_sequence_masks is not None
|
||||
self.spec_state_indices_tensor[:num_spec_decodes].copy_(
|
||||
spec_state_indices_tensor, non_blocking=True
|
||||
)
|
||||
spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size]
|
||||
spec_state_indices_tensor[num_spec_decodes:].fill_(NULL_BLOCK_ID)
|
||||
|
||||
self.spec_sequence_masks[:num_spec_decodes].copy_(
|
||||
spec_sequence_masks[:num_spec_decodes], non_blocking=True
|
||||
)
|
||||
spec_sequence_masks = self.spec_sequence_masks[:batch_size]
|
||||
spec_sequence_masks[num_spec_decodes:].fill_(False)
|
||||
|
||||
assert non_spec_token_indx is not None and spec_token_indx is not None
|
||||
self.non_spec_token_indx[: non_spec_token_indx.size(0)].copy_(
|
||||
non_spec_token_indx, non_blocking=True
|
||||
)
|
||||
non_spec_token_indx = self.non_spec_token_indx[
|
||||
: non_spec_token_indx.size(0)
|
||||
]
|
||||
|
||||
self.spec_token_indx[: spec_token_indx.size(0)].copy_(
|
||||
spec_token_indx, non_blocking=True
|
||||
)
|
||||
spec_token_indx = self.spec_token_indx[: spec_token_indx.size(0)]
|
||||
|
||||
self.spec_query_start_loc[: num_spec_decodes + 1].copy_(
|
||||
spec_query_start_loc, non_blocking=True
|
||||
)
|
||||
spec_num_query_tokens = spec_query_start_loc[-1] # type: ignore[index]
|
||||
spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1]
|
||||
spec_query_start_loc[num_spec_decodes + 1 :].fill_(spec_num_query_tokens)
|
||||
|
||||
self.num_accepted_tokens[:num_spec_decodes].copy_(
|
||||
num_accepted_tokens, non_blocking=True
|
||||
)
|
||||
num_accepted_tokens = self.num_accepted_tokens[:batch_size]
|
||||
num_accepted_tokens[num_spec_decodes:].fill_(1)
|
||||
|
||||
if (
|
||||
self.use_full_cuda_graph
|
||||
and num_prefills == 0
|
||||
and num_spec_decodes == 0
|
||||
and num_decodes <= self.decode_cudagraph_max_bs
|
||||
):
|
||||
self.non_spec_state_indices_tensor[:num_decodes].copy_(
|
||||
non_spec_state_indices_tensor, non_blocking=True
|
||||
)
|
||||
non_spec_state_indices_tensor = self.non_spec_state_indices_tensor[
|
||||
:batch_size
|
||||
]
|
||||
non_spec_state_indices_tensor[num_decodes:].fill_(NULL_BLOCK_ID)
|
||||
|
||||
self.non_spec_query_start_loc[: num_decodes + 1].copy_(
|
||||
non_spec_query_start_loc, non_blocking=True
|
||||
)
|
||||
non_spec_num_query_tokens = non_spec_query_start_loc[-1] # type: ignore[index]
|
||||
non_spec_query_start_loc = self.non_spec_query_start_loc[: batch_size + 1]
|
||||
non_spec_query_start_loc[num_decodes + 1 :].fill_(non_spec_num_query_tokens)
|
||||
|
||||
attn_metadata = GDNAttentionMetadata(
|
||||
num_prefills=num_prefills,
|
||||
num_prefill_tokens=num_prefill_tokens,
|
||||
num_decodes=num_decodes,
|
||||
num_decode_tokens=num_decode_tokens,
|
||||
num_spec_decodes=num_spec_decodes,
|
||||
num_spec_decode_tokens=num_spec_decode_tokens,
|
||||
num_actual_tokens=m.num_actual_tokens,
|
||||
has_initial_state=has_initial_state,
|
||||
chunk_indices=chunk_indices,
|
||||
chunk_offsets=chunk_offsets,
|
||||
prefill_query_start_loc=prefill_query_start_loc,
|
||||
prefill_state_indices=prefill_state_indices,
|
||||
prefill_has_initial_state=prefill_has_initial_state,
|
||||
spec_query_start_loc=spec_query_start_loc,
|
||||
non_spec_query_start_loc=non_spec_query_start_loc,
|
||||
spec_state_indices_tensor=spec_state_indices_tensor,
|
||||
non_spec_state_indices_tensor=non_spec_state_indices_tensor,
|
||||
spec_sequence_masks=spec_sequence_masks,
|
||||
spec_token_indx=spec_token_indx,
|
||||
non_spec_token_indx=non_spec_token_indx,
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
nums_dict=nums_dict,
|
||||
batch_ptr=batch_ptr,
|
||||
token_chunk_offset_ptr=token_chunk_offset_ptr,
|
||||
)
|
||||
return attn_metadata
|
||||
|
||||
def build_for_cudagraph_capture(
|
||||
self, common_attn_metadata: CommonAttentionMetadata
|
||||
):
|
||||
"""
|
||||
This method builds the metadata for full cudagraph capture.
|
||||
Currently, only decode is supported for full cudagraphs with Mamba.
|
||||
"""
|
||||
m = common_attn_metadata
|
||||
|
||||
assert (
|
||||
m.num_reqs <= self.decode_cudagraph_max_bs
|
||||
and m.num_actual_tokens <= self.decode_cudagraph_max_bs
|
||||
), (
|
||||
f"GDN only supports decode-only full CUDAGraph capture. "
|
||||
f"Make sure batch size ({m.num_reqs}) <= "
|
||||
f"cudagraph capture sizes ({self.decode_cudagraph_max_bs}), "
|
||||
f"and number of tokens ({m.num_actual_tokens}) <= "
|
||||
f"cudagraph capture sizes ({self.decode_cudagraph_max_bs})."
|
||||
)
|
||||
|
||||
num_accepted_tokens = torch.diff(m.query_start_loc)
|
||||
num_decode_draft_tokens_cpu = (num_accepted_tokens - 1).cpu()
|
||||
|
||||
return self.build(0, m, num_accepted_tokens, num_decode_draft_tokens_cpu)
|
||||
Reference in New Issue
Block a user