From 6cdf2ec87b3b5bddbf3b1d003ccd28f5c8de0952 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 11 Aug 2026 03:55:50 +0000 Subject: [PATCH] =?UTF-8?q?ref(upstream):=20=E6=90=AC=E8=BF=90=203=20?= =?UTF-8?q?=E5=A4=A7=20GDN=20=E4=B8=8A=E6=B8=B8=E4=BB=93=E5=BA=93=20?= =?UTF-8?q?=E2=80=94=20FLA=20naive=20ops=20+=20vllm=20GDN=20=E5=AD=90?= =?UTF-8?q?=E6=A0=91=20+=20xllm=20C++=20=E5=8F=82=E8=80=83?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 来源: 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 有正确的实现可直接对接 --- upstream_ref/fla/layers/gated_deltanet.py | 364 ++++ upstream_ref/fla/modules/__init__.py | 52 + upstream_ref/fla/modules/convolution.py | 42 + upstream_ref/fla/ops/__init__.py | 91 + .../fla/ops/gated_delta_rule/__init__.py | 17 + .../ops/gated_delta_rule/backends/__init__.py | 20 + .../gated_delta_rule/backends/flash_qla.py | 132 ++ .../backends/triton_ascend/__init__.py | 96 + .../backends/triton_ascend/chunk_fwd.py | 62 + .../backends/triton_ascend/gate.py | 375 ++++ .../backends/triton_ascend/wy_fast.py | 823 ++++++++ .../fla/ops/gated_delta_rule/chunk.py | 591 ++++++ .../fla/ops/gated_delta_rule/chunk_fwd.py | 425 ++++ .../ops/gated_delta_rule/fused_recurrent.py | 478 +++++ upstream_ref/fla/ops/gated_delta_rule/gate.py | 344 ++++ .../fla/ops/gated_delta_rule/naive.py | 161 ++ .../fla/ops/gated_delta_rule/wy_fast.py | 347 ++++ upstream_ref/vllm_gdn/gdn/__init__.py | 0 upstream_ref/vllm_gdn/gdn/base.py | 58 + .../vllm_gdn/gdn/kimi_gdn_linear_attn.py | 635 ++++++ .../vllm_gdn/gdn/olmo_gdn_linear_attn.py | 634 ++++++ .../vllm_gdn/gdn/qwen_gdn_linear_attn.py | 1751 +++++++++++++++++ upstream_ref/vllm_gdn/models/qwen3_5.py | 733 +++++++ upstream_ref/vllm_gdn/models/qwen3_next.py | 883 +++++++++ upstream_ref/vllm_gdn/ops/causal_conv1d.py | 1289 ++++++++++++ upstream_ref/vllm_gdn/third_party/__init__.py | 8 + .../vllm_gdn/third_party/ops/__init__.py | 25 + .../vllm_gdn/third_party/ops/chunk.py | 245 +++ .../vllm_gdn/third_party/ops/cumsum.py | 282 +++ .../ops/fused_gdn_prefill_post_conv.py | 248 +++ .../third_party/ops/fused_recurrent.py | 619 ++++++ .../vllm_gdn/third_party/ops/l2norm.py | 151 ++ .../vllm_gdn/third_party/ops/utils.py | 200 ++ .../v1/attention/backends/gdn_attn.py | 538 +++++ .../core/kernels/npu/npu_causal_conv1d.cpp | 59 + .../npu/npu_recurrent_gated_delta_rule.cpp | 83 + .../core/layers/mlu/qwen3_5_attention.cpp | 236 +++ .../core/layers/mlu/qwen3_5_attention.h | 79 + .../core/layers/mlu/qwen3_5_decoder_layer.cpp | 193 ++ .../core/layers/mlu/qwen3_5_decoder_layer.h | 73 + .../core/layers/mlu/qwen3_5_fused_moe.cpp | 209 ++ .../core/layers/mlu/qwen3_5_fused_moe.h | 47 + .../npu_torch/qwen3_5_gated_delta_net.cpp | 58 +- .../npu_torch/qwen3_5_gated_delta_net.h | 16 +- .../npu_torch/qwen3_gated_delta_net_base.cpp | 858 ++------ .../npu_torch/qwen3_gated_delta_net_base.h | 38 +- 46 files changed, 13857 insertions(+), 811 deletions(-) create mode 100644 upstream_ref/fla/layers/gated_deltanet.py create mode 100644 upstream_ref/fla/modules/__init__.py create mode 100644 upstream_ref/fla/modules/convolution.py create mode 100644 upstream_ref/fla/ops/__init__.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/__init__.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/backends/__init__.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/backends/flash_qla.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/__init__.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/chunk_fwd.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/gate.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/chunk.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/chunk_fwd.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/fused_recurrent.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/gate.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/naive.py create mode 100644 upstream_ref/fla/ops/gated_delta_rule/wy_fast.py create mode 100644 upstream_ref/vllm_gdn/gdn/__init__.py create mode 100644 upstream_ref/vllm_gdn/gdn/base.py create mode 100644 upstream_ref/vllm_gdn/gdn/kimi_gdn_linear_attn.py create mode 100644 upstream_ref/vllm_gdn/gdn/olmo_gdn_linear_attn.py create mode 100644 upstream_ref/vllm_gdn/gdn/qwen_gdn_linear_attn.py create mode 100644 upstream_ref/vllm_gdn/models/qwen3_5.py create mode 100644 upstream_ref/vllm_gdn/models/qwen3_next.py create mode 100644 upstream_ref/vllm_gdn/ops/causal_conv1d.py create mode 100644 upstream_ref/vllm_gdn/third_party/__init__.py create mode 100644 upstream_ref/vllm_gdn/third_party/ops/__init__.py create mode 100644 upstream_ref/vllm_gdn/third_party/ops/chunk.py create mode 100644 upstream_ref/vllm_gdn/third_party/ops/cumsum.py create mode 100644 upstream_ref/vllm_gdn/third_party/ops/fused_gdn_prefill_post_conv.py create mode 100644 upstream_ref/vllm_gdn/third_party/ops/fused_recurrent.py create mode 100644 upstream_ref/vllm_gdn/third_party/ops/l2norm.py create mode 100644 upstream_ref/vllm_gdn/third_party/ops/utils.py create mode 100644 upstream_ref/vllm_gdn/v1/attention/backends/gdn_attn.py create mode 100644 upstream_ref/xllm_latest/core/kernels/npu/npu_causal_conv1d.cpp create mode 100644 upstream_ref/xllm_latest/core/kernels/npu/npu_recurrent_gated_delta_rule.cpp create mode 100644 upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.cpp create mode 100644 upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.h create mode 100644 upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_decoder_layer.cpp create mode 100644 upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_decoder_layer.h create mode 100644 upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.cpp create mode 100644 upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.h diff --git a/upstream_ref/fla/layers/gated_deltanet.py b/upstream_ref/fla/layers/gated_deltanet.py new file mode 100644 index 00000000..c631f92b --- /dev/null +++ b/upstream_ref/fla/layers/gated_deltanet.py @@ -0,0 +1,364 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange +from torch.nn import functional as F + +from fla.layers.utils import get_layer_cache, repad_hidden_states, unpad_hidden_states, update_layer_cache +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.modules.convolution import causal_conv1d +from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +class GatedDeltaNet(nn.Module): + """ + Gated Delta Networks (GDN) layer implementation. + + Reference: `Gated Delta Networks: Improving Mamba2 with Delta Rule `_ + + Similar to Mamba2, each layer contains around 6*hidden_size*hidden_size parameters. + + Parameter allocation when use_gate=True: + - 0.75 * hidden_size * hidden_size for the q_proj and k_proj each + - 1.5 * hidden_size * hidden_size for the v_proj, g_proj and o_proj each + - Others are ignorably small. + - In total = 0.75 * 2 + 1.5 * 3 = 6 * hidden_size * hidden_size + NOTE: num_heads * head_dim = 0.75 * hidden_size, please make sure to set the correct num_heads and head_dim. + + Parameter allocation when use_gate=False: + - 1 * hidden_size * hidden_size for the q_proj and k_proj each + - 2 * hidden_size * hidden_size for the v_proj and o_proj each + - Others are ignorably small. + - In total = 1 * 2 + 2 * 2 = 6 * hidden_size * hidden_size + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + expand_v (float, Optional): + The expansion ratio for the value dimension. Default: 2.0. + head_dim (int, Optional): + The dimension of each head. Default: 256. + num_heads (int, Optional): + The number of heads. Default: 6. + num_v_heads (int, Optional): + The number of heads for the value projection, equal to `num_heads` if `None`. + GVA (Grouped Value Attention) is applied if `num_v_heads` > `num_heads`, + where `num_v_heads` must be divisible by `num_heads`. + The kernels natively support GVA by mapping multiple value heads to each query/key head. + Default: `None`. + mode (str, Optional): + Which Gated DeltaNet kernel to use. + Currently available: `chunk` and `fused_recurrent`. + Default: `chunk`. + use_gate (bool, Optional): + Whether to use output gate. Default: `True`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `True`. + allow_neg_eigval (bool, Optional): + Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2. + See reference: + `Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues `_ + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the normalization layer. Default: 1e-5. + """ + + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 2, + head_dim: int = 256, + num_heads: int = 6, + num_v_heads: int = None, + mode: str = 'chunk', + use_gate: bool = True, + use_short_conv: bool = True, + allow_neg_eigval: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + layer_idx: int = None, + norm_eps: float = 1e-5, + **kwargs, + ) -> GatedDeltaNet: + super().__init__() + + self.mode = mode + self.allow_neg_eigval = allow_neg_eigval + self.hidden_size = hidden_size + self.expand_v = expand_v + + self.use_gate = use_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads + + self.head_k_dim = head_dim + self.head_v_dim = int(self.head_dim * self.expand_v) + self.key_dim = int(self.num_heads * self.head_k_dim) + self.value_dim = int(self.num_v_heads * self.head_v_dim) + self.layer_idx = layer_idx + + # Consistency check: Ensure expand_v produces integer values + if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. " + f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.", + ) + if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0: + raise ValueError( + f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.", + ) + + if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. " + f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.", + ) + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.a_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) + self.b_proj = nn.Linear(hidden_size, self.num_v_heads, bias=False) + + A = torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log._no_weight_decay = True + # hard coded for now + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.num_v_heads) * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min), + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + # Just to be explicit. Without this we already don't put wd on dt_bias because of the check + # name.endswith("bias") in param_grouping.py + self.dt_bias._no_weight_decay = True + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.k_conv1d = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + self.v_conv1d = ShortConvolution( + hidden_size=self.value_dim, + kernel_size=conv_size, + bias=conv_bias, + activation='silu', + ) + else: + warnings.warn( + "ShortConvolution is crucial to the performance. " + "Do not turn it off, i.e., setting `use_short_conv=False` unless you know what you are doing.", + ) + if use_gate: + self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps, dtype=torch.float32) + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def _use_fused_qkv_conv( + self, + last_state: dict | None, + use_cache: bool | None, + cu_seqlens: torch.Tensor | None, + ) -> bool: + # The dense no-cache q/k/v short convolutions can collapse into a single + # causal_conv1d only when there is no cache/varlen state and the three convs + # share the same backend, activation, and kernel size. + if not (self.use_short_conv and last_state is None and not use_cache and cu_seqlens is None): + return False + return ( + self.q_conv1d.backend == self.k_conv1d.backend == self.v_conv1d.backend and + self.q_conv1d.activation == self.k_conv1d.activation == self.v_conv1d.activation and + self.q_conv1d.kernel_size == self.k_conv1d.kernel_size == self.v_conv1d.kernel_size + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + # change to inference mode. + mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + last_state = get_layer_cache(self, past_key_values) + + cu_seqlens = kwargs.get('cu_seqlens') + hidden_states, indices, cu_seqlens = unpad_hidden_states(hidden_states, cu_seqlens, attention_mask, q_len) + + conv_state_q, conv_state_k, conv_state_v = None, None, None + if self._use_fused_qkv_conv(last_state, use_cache, cu_seqlens): + qkv = torch.cat( + [ + self.q_proj(hidden_states), + self.k_proj(hidden_states), + self.v_proj(hidden_states), + ], + dim=-1, + ) + qkv_weight = torch.cat( + [ + self.q_conv1d.weight.squeeze(1), + self.k_conv1d.weight.squeeze(1), + self.v_conv1d.weight.squeeze(1), + ], + dim=0, + ) + if self.conv_bias: + qkv_bias = torch.cat([self.q_conv1d.bias, self.k_conv1d.bias, self.v_conv1d.bias], dim=0) + else: + qkv_bias = None + qkv, _ = causal_conv1d( + x=qkv, + weight=qkv_weight, + bias=qkv_bias, + activation=self.q_conv1d.activation, + backend=self.q_conv1d.backend, + ) + q, k, v = torch.split(qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1) + elif self.use_short_conv: + if last_state is not None: + conv_state_q, conv_state_k, conv_state_v = last_state['conv_state'] + q, conv_state_q = self.q_conv1d( + x=self.q_proj(hidden_states), + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k, conv_state_k = self.k_conv1d( + x=self.k_proj(hidden_states), + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + v, conv_state_v = self.v_conv1d( + x=self.v_proj(hidden_states), + cache=conv_state_v, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + q = F.silu(self.q_proj(hidden_states)) + k = F.silu(self.k_proj(hidden_states)) + v = F.silu(self.v_proj(hidden_states)) + + q, k = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + + beta = self.b_proj(hidden_states) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + if mode == 'chunk': + o, recurrent_state = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=self.a_proj(hidden_states), + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=use_cache, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + allow_neg_eigval=self.allow_neg_eigval, + state_v_first=True, + cu_seqlens=cu_seqlens, + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=self.a_proj(hidden_states), + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_state, + output_final_state=use_cache, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + allow_neg_eigval=self.allow_neg_eigval, + state_v_first=True, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + update_layer_cache( + self, + past_key_values, + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k, conv_state_v) if self.use_short_conv else None, + offset=q_len, + ) + + if self.use_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + o = repad_hidden_states(o, indices, batch_size, q_len) + + return o, None, past_key_values diff --git a/upstream_ref/fla/modules/__init__.py b/upstream_ref/fla/modules/__init__.py new file mode 100644 index 00000000..86129dcb --- /dev/null +++ b/upstream_ref/fla/modules/__init__.py @@ -0,0 +1,52 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from fla.modules.convolution import ImplicitLongConvolution, LongConvolution, ShortConvolution +from fla.modules.fused_bitlinear import BitLinear, FusedBitLinear +from fla.modules.fused_cross_entropy import FusedCrossEntropyLoss +from fla.modules.fused_kl_div import FusedKLDivLoss +from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss +from fla.modules.fused_norm_gate import ( + FusedLayerNormGated, + FusedLayerNormSwishGate, + FusedLayerNormSwishGateLinear, + FusedRMSNormGated, + FusedRMSNormSwishGate, + FusedRMSNormSwishGateLinear, +) +from fla.modules.l2norm import L2Norm +from fla.modules.layernorm import GroupNorm, GroupNormLinear, LayerNorm, LayerNormLinear, RMSNorm, RMSNormLinear +from fla.modules.mlp import GatedMLP +from fla.modules.rotary import RotaryEmbedding +from fla.modules.token_shift import TokenShift + +__all__ = [ + 'BitLinear', + 'FusedBitLinear', + 'FusedCrossEntropyLoss', + 'FusedKLDivLoss', + 'FusedLayerNormGated', + 'FusedLayerNormSwishGate', + 'FusedLayerNormSwishGateLinear', + 'FusedLinearCrossEntropyLoss', + 'FusedRMSNormGated', + 'FusedRMSNormSwishGate', + 'FusedRMSNormSwishGateLinear', + 'GatedMLP', + 'GroupNorm', + 'GroupNormLinear', + 'ImplicitLongConvolution', + 'L2Norm', + 'LayerNorm', + 'LayerNormLinear', + 'LongConvolution', + 'RMSNorm', + 'RMSNormLinear', + 'RotaryEmbedding', + 'ShortConvolution', + 'TokenShift', +] diff --git a/upstream_ref/fla/modules/convolution.py b/upstream_ref/fla/modules/convolution.py new file mode 100644 index 00000000..8f402f14 --- /dev/null +++ b/upstream_ref/fla/modules/convolution.py @@ -0,0 +1,42 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from fla.modules.conv import ( + ImplicitLongConvolution, + LongConvolution, + PositionalEmbedding, + ShortConvolution, + causal_conv1d, + fft_conv, +) +from fla.modules.conv.cp import CausalConv1dFunctionCP, causal_conv1d_cp +from fla.modules.conv.cuda import FastCausalConv1dFn, fast_causal_conv1d_fn +from fla.modules.conv.triton import ( + CausalConv1dFunction, + causal_conv1d_bwd, + causal_conv1d_fwd, + causal_conv1d_update, + causal_conv1d_update_states, +) + +__all__ = [ + 'CausalConv1dFunction', + 'CausalConv1dFunctionCP', + 'FastCausalConv1dFn', + 'ImplicitLongConvolution', + 'LongConvolution', + 'PositionalEmbedding', + 'ShortConvolution', + 'causal_conv1d', + 'causal_conv1d_bwd', + 'causal_conv1d_cp', + 'causal_conv1d_fwd', + 'causal_conv1d_update', + 'causal_conv1d_update_states', + 'fast_causal_conv1d_fn', + 'fft_conv', +] diff --git a/upstream_ref/fla/ops/__init__.py b/upstream_ref/fla/ops/__init__.py new file mode 100644 index 00000000..3f962e5d --- /dev/null +++ b/upstream_ref/fla/ops/__init__.py @@ -0,0 +1,91 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from .abc import chunk_abc +from .attn import parallel_attn +from .attnres import fused_attnres +from .based import fused_chunk_based, parallel_based +from .comba import chunk_comba, fused_recurrent_comba +from .delta_rule import chunk_delta_rule, fused_chunk_delta_rule, fused_recurrent_delta_rule +from .forgetting_attn import parallel_forgetting_attn +from .gated_delta_rule import chunk_gated_delta_rule, chunk_gdn, fused_recurrent_gated_delta_rule, fused_recurrent_gdn +from .generalized_delta_rule import ( + chunk_dplr_delta_rule, + chunk_iplr_delta_rule, + fused_recurrent_dplr_delta_rule, + fused_recurrent_iplr_delta_rule, +) +from .gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla +from .gsa import chunk_gsa, fused_recurrent_gsa +from .hgrn import fused_recurrent_hgrn +from .kda import chunk_kda, fused_recurrent_kda +from .lightning_attn import chunk_lightning_attn, fused_recurrent_lightning_attn +from .linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn +from .log_linear_attn import chunk_log_linear_attn +from .mesa_net import chunk_mesa_net +from .nsa import parallel_nsa +from .parallax import parallel_parallax +from .path_attn import parallel_path_attn +from .retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention +from .rwkv6 import chunk_rwkv6, fused_recurrent_rwkv6 +from .rwkv7 import chunk_rwkv7, fused_recurrent_rwkv7 +from .simple_gla import chunk_simple_gla, fused_chunk_simple_gla, fused_recurrent_simple_gla, parallel_simple_gla +from .wall_attn import parallel_wall_attn, parallel_wall_attn_decode + +__all__ = [ + 'chunk_abc', + 'chunk_comba', + 'chunk_delta_rule', + 'chunk_dplr_delta_rule', + 'chunk_gated_delta_rule', + 'chunk_gdn', + 'chunk_gla', + 'chunk_gsa', + 'chunk_iplr_delta_rule', + 'chunk_kda', + 'chunk_lightning_attn', + 'chunk_linear_attn', + 'chunk_log_linear_attn', + 'chunk_mesa_net', + 'chunk_retention', + 'chunk_rwkv6', + 'chunk_rwkv7', + 'chunk_simple_gla', + 'fused_attnres', + 'fused_chunk_based', + 'fused_chunk_delta_rule', + 'fused_chunk_gla', + 'fused_chunk_linear_attn', + 'fused_chunk_retention', + 'fused_chunk_simple_gla', + 'fused_recurrent_comba', + 'fused_recurrent_delta_rule', + 'fused_recurrent_dplr_delta_rule', + 'fused_recurrent_gated_delta_rule', + 'fused_recurrent_gdn', + 'fused_recurrent_gla', + 'fused_recurrent_gsa', + 'fused_recurrent_hgrn', + 'fused_recurrent_iplr_delta_rule', + 'fused_recurrent_kda', + 'fused_recurrent_lightning_attn', + 'fused_recurrent_linear_attn', + 'fused_recurrent_retention', + 'fused_recurrent_rwkv6', + 'fused_recurrent_rwkv7', + 'fused_recurrent_simple_gla', + 'parallel_attn', + 'parallel_based', + 'parallel_forgetting_attn', + 'parallel_nsa', + 'parallel_parallax', + 'parallel_path_attn', + 'parallel_retention', + 'parallel_simple_gla', + 'parallel_wall_attn', + 'parallel_wall_attn_decode', +] diff --git a/upstream_ref/fla/ops/gated_delta_rule/__init__.py b/upstream_ref/fla/ops/gated_delta_rule/__init__.py new file mode 100644 index 00000000..7e65713b --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from .chunk import chunk_gated_delta_rule, chunk_gdn +from .fused_recurrent import fused_recurrent_gated_delta_rule, fused_recurrent_gdn +from .naive import naive_chunk_gated_delta_rule, naive_recurrent_gated_delta_rule + +__all__ = [ + "chunk_gated_delta_rule", "chunk_gdn", + "fused_recurrent_gated_delta_rule", "fused_recurrent_gdn", + "naive_chunk_gated_delta_rule", + "naive_recurrent_gated_delta_rule", +] diff --git a/upstream_ref/fla/ops/gated_delta_rule/backends/__init__.py b/upstream_ref/fla/ops/gated_delta_rule/backends/__init__.py new file mode 100644 index 00000000..b3d7c1dd --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/backends/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""GDR backends.""" + +from fla.ops.backends import BackendRegistry, dispatch +from fla.ops.gated_delta_rule.backends.flash_qla import FlashQLABackend +from fla.ops.gated_delta_rule.backends.triton_ascend import TritonAscendGDNBackend + +gdr_registry = BackendRegistry("gated_delta_rule") + +gdr_registry.register(TritonAscendGDNBackend()) +gdr_registry.register(FlashQLABackend()) + + +__all__ = ['dispatch', 'gdr_registry'] diff --git a/upstream_ref/fla/ops/gated_delta_rule/backends/flash_qla.py b/upstream_ref/fla/ops/gated_delta_rule/backends/flash_qla.py new file mode 100644 index 00000000..cdf158c7 --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/backends/flash_qla.py @@ -0,0 +1,132 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors +# +# Copyright (c) 2026 Qwen Team, Alibaba Cloud + +"""FlashQLA backend for chunk_gated_delta_rule.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from fla.ops.backends import BaseBackend +from fla.utils import IS_NVIDIA_HOPPER, IS_NVIDIA_SM100, IS_NVIDIA_SM120 + +if TYPE_CHECKING: + from fla.ops.cp import FLACPContext + + +def _needs_backward(*tensors: torch.Tensor | None) -> bool: + """Whether autograd would later require a backward pass over these inputs.""" + if not torch.is_grad_enabled(): + return False + return any(isinstance(t, torch.Tensor) and t.requires_grad for t in tensors) + + +class FlashQLABackend(BaseBackend): + """Copyright (c) 2026 Qwen Team, Alibaba Cloud + + Fused TileLang forward and backward with intra-card CP (replaces the multi-kernel Triton path). + https://github.com/QwenLM/FlashQLA + + SM90/SM100/SM103 run both directions. SM120 (consumer/workstation Blackwell) ships a + bfloat16 forward kernel only, so it is dispatched exclusively for grad-free bf16 calls + (inference, frozen weights) and falls back to Triton otherwise. + + Disable with ``FLA_FLASH_QLA=0``. + """ + + backend_type = "flash_qla" + package_name = "flash_qla" + env_var = "FLA_FLASH_QLA" + default_enable = True + priority = 3 + + def chunk_gated_delta_rule_verifier( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, + ) -> tuple[bool, str | None]: + if not (IS_NVIDIA_HOPPER or IS_NVIDIA_SM100 or IS_NVIDIA_SM120): + return False, "FlashQLA requires NVIDIA SM90, SM100/SM103 or SM120" + if IS_NVIDIA_SM120 and _needs_backward(q, k, v, g, beta, initial_state): + return False, "FlashQLA on SM120 implements the forward pass only, but an input requires grad" + if q.dtype != torch.float16 and q.dtype != torch.bfloat16: + return False, f"FlashQLA requires dtype float16 or bfloat16, got {q.dtype}" + if not (q.dtype == k.dtype == v.dtype): + return False, f"FlashQLA requires q, k, v to have the same dtype, got {q.dtype}, {k.dtype}, {v.dtype}" + # NOTE: the masked tail-store in FlashQLA's blackwell_sm120 forward kernel emits + # tl::pack_float16x4 on cutlass::half_t, which fails to compile under nvcc. + if IS_NVIDIA_SM120 and q.dtype == torch.float16: + return False, "FlashQLA's SM120 forward kernel does not compile for float16" + if q.shape[-1] != 128: + return False, f"FlashQLA requires K=128, got {q.shape[-1]}" + if v.shape[-1] != 128: + return False, f"FlashQLA requires V=128, got {v.shape[-1]}" + if kwargs.get('use_gate_in_kernel'): + return False, "FlashQLA does not support use_gate_in_kernel" + if use_beta_sigmoid_in_kernel: + return False, "FlashQLA does not support use_beta_sigmoid_in_kernel" + if allow_neg_eigval: + return False, "FlashQLA does not support allow_neg_eigval" + if 'transpose_state_layout' in kwargs: + return False, "FlashQLA does not support the deprecated transpose_state_layout" + if cp_context is not None: + return False, "FlashQLA does not support inter-card context parallel" + return True, None + + def chunk_gated_delta_rule( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, + ): + import flash_qla + + return flash_qla.chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + auto_cp=True, + ) diff --git a/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/__init__.py b/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/__init__.py new file mode 100644 index 00000000..aab6adbb --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/__init__.py @@ -0,0 +1,96 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""Triton-Ascend Ascend NPU backend for GDN gated_delta_rule ops.""" + +from __future__ import annotations + +import torch + +from fla.ops.backends import BaseBackend + + +class TritonAscendGDNBackend(BaseBackend): + """Ascend NPU backend for GDN gate and WY-representation kernels.""" + + backend_type = "triton_ascend" + package_name = None + env_var = None + priority = 0 + + @classmethod + def is_available(cls) -> bool: + from fla.utils import IS_NPU + return IS_NPU + + def gdn_gate_fwd_verifier(self, *args, **kwargs): + return True, None + + def gdn_gate_fwd(self, *args, **kwargs): + from fla.ops.gated_delta_rule.backends.triton_ascend.gate import gdn_gate_fwd_npu + return gdn_gate_fwd_npu(*args, **kwargs) + + def gdn_gate_chunk_cumsum_verifier(self, *args, **kwargs): + return True, None + + def gdn_gate_chunk_cumsum(self, *args, **kwargs): + from fla.ops.gated_delta_rule.backends.triton_ascend.gate import gdn_gate_chunk_cumsum_npu + return gdn_gate_chunk_cumsum_npu(*args, **kwargs) + + def gdn_gate_bwd_verifier(self, *args, **kwargs): + return True, None + + def gdn_gate_bwd(self, *args, **kwargs): + from fla.ops.gated_delta_rule.backends.triton_ascend.gate import gdn_gate_bwd_npu + return gdn_gate_bwd_npu(*args, **kwargs) + + def recompute_w_u_fwd_verifier( + self, + k, + v, + beta, + A, + g=None, + cu_seqlens=None, + chunk_indices=None, + ) -> tuple[bool, str | None]: + from fla.utils import IS_NPU + if not IS_NPU: + return False, "not running on NPU" + if k.device.type != "npu": + return False, "input device is not NPU" + if all(t.dtype in (torch.float32, torch.float16, torch.bfloat16) + for t in (k, v, beta, A)): + return True, None + return False, "unsupported dtype for NPU recompute_w_u_fwd" + + def recompute_w_u_fwd( + self, + k, + v, + beta, + A, + g=None, + cu_seqlens=None, + chunk_indices=None, + ): + from fla.ops.gated_delta_rule.backends.triton_ascend.wy_fast import recompute_w_u_fwd_npu + return recompute_w_u_fwd_npu(k, v, beta, A, g, cu_seqlens, chunk_indices) + + def prepare_wy_repr_bwd_verifier(self, *args, **kwargs): + return True, None + + def prepare_wy_repr_bwd(self, *args, **kwargs): + from fla.ops.gated_delta_rule.backends.triton_ascend.wy_fast import prepare_wy_repr_bwd_npu + return prepare_wy_repr_bwd_npu(*args, **kwargs) + + def chunk_gated_delta_rule_fwd_intra_verifier(self, *args, **kwargs): + return True, None + + def chunk_gated_delta_rule_fwd_intra(self, *args, **kwargs): + from fla.ops.gated_delta_rule.backends.triton_ascend.chunk_fwd import chunk_gated_delta_rule_fwd_intra_npu + return chunk_gated_delta_rule_fwd_intra_npu(*args, **kwargs) diff --git a/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/chunk_fwd.py b/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/chunk_fwd.py new file mode 100644 index 00000000..040e5c97 --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/chunk_fwd.py @@ -0,0 +1,62 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""chunk_gated_delta_rule_fwd_intra adapted for triton-ascend on Ascend NPU.""" + +from __future__ import annotations + +import torch + +from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from fla.ops.gated_delta_rule.wy_fast import recompute_w_u_fwd +from fla.ops.utils import prepare_chunk_indices, solve_tril +from fla.utils import input_guard + + +@input_guard +def chunk_gated_delta_rule_fwd_intra_npu( + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if chunk_size not in (16, 32, 64): + raise ValueError(f'`chunk_size` must be 16, 32, or 64, got {chunk_size}.') + + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + + # Unfused kkt + solve_tril path to stay within UB budget. + A = chunk_scaled_dot_kkt_fwd( + k=k, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=BT, + 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=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return w, u, A diff --git a/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/gate.py b/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/gate.py new file mode 100644 index 00000000..dbd1ac73 --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/gate.py @@ -0,0 +1,375 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""GDN gate kernels adapted for triton-ascend on Ascend NPU.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.index import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.ops.utils.softplus import softplus +from fla.utils import input_guard +from fla.utils.ascend_ub_manager import ( + ASCEND_MAX_GRID_DIM, + compute_ub_block_size, + max_grid_axis_chunks, +) + +_NUM_WARPS = 4 +# Peak live fp32 vectors: input + output (+ bias path). +_GATE_FWD_MEM_MULT = 3.0 +_GATE_BWD_MEM_MULT = 5.0 +_SAFETY_MARGIN = 0.85 +_FALLBACK_BT = 32 +_FALLBACK_BT_FWD = 64 + + +def _get_gate_fwd_bt(T: int) -> int: + return compute_ub_block_size( + T, + _GATE_FWD_MEM_MULT, + safety_margin=_SAFETY_MARGIN, + dtype_size=4, + fallback=_FALLBACK_BT_FWD, + desired=min(triton.next_power_of_2(T), _FALLBACK_BT_FWD), + ) + + +def _get_gate_bwd_bt(T: int) -> int: + return compute_ub_block_size( + T, + _GATE_BWD_MEM_MULT, + safety_margin=_SAFETY_MARGIN, + dtype_size=4, + fallback=_FALLBACK_BT, + desired=min(triton.next_power_of_2(T), _FALLBACK_BT), + ) + + +@triton.heuristics({ + 'HAS_BIAS': lambda args: args['dt_bias'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def gdn_gate_fwd_kernel_npu( + g, + A_log, + dt_bias, + yg, + T, + H: tl.constexpr, + BT: tl.constexpr, + HAS_BIAS: tl.constexpr, + NT_OFFSET: tl.constexpr, + H_OFFSET: tl.constexpr, +): + i_t = tl.program_id(0) + NT_OFFSET + i_h = tl.program_id(1) + H_OFFSET + + b_A = tl.load(A_log + i_h).to(tl.float32) + + p_g = tl.make_block_ptr(g + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_yg = tl.make_block_ptr(yg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + if HAS_BIAS: + b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) + b_yg = -exp(b_A) * softplus(b_g) + tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), boundary_check=(0,)) + + +def _launch_gate_fwd( + *, + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + yg: torch.Tensor, + T: int, + H: int, + BT: int, +) -> None: + NT = triton.cdiv(T, BT) + kernel_kwargs = dict( + g=g, + A_log=A_log, + dt_bias=dt_bias, + yg=yg, + T=T, + H=H, + BT=BT, + num_warps=_NUM_WARPS, + ) + max_nt = max_grid_axis_chunks(NT, H, max_grid=ASCEND_MAX_GRID_DIM) + for nt_off in range(0, NT, max_nt): + nt_len = min(max_nt, NT - nt_off) + max_h = max_grid_axis_chunks(H, nt_len, max_grid=ASCEND_MAX_GRID_DIM) + for h_off in range(0, H, max_h): + h_len = min(max_h, H - h_off) + gdn_gate_fwd_kernel_npu[(nt_len, h_len)]( + **kernel_kwargs, + NT_OFFSET=nt_off, + H_OFFSET=h_off, + ) + + +@triton.heuristics({ + 'HAS_BIAS': lambda args: args['dt_bias'] is not None, + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def gdn_gate_chunk_cumsum_scalar_kernel_npu( + g, + A_log, + dt_bias, + o, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + REVERSE: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + NT_OFFSET: tl.constexpr, + BH_OFFSET: tl.constexpr, +): + i_t = tl.program_id(0) + NT_OFFSET + i_bh = tl.program_id(1) + BH_OFFSET + 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 + + p_g = tl.make_block_ptr(g + 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,)) + + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + if HAS_BIAS: + b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) + b_A = tl.load(A_log + i_h).to(tl.float32) + b_gate = -exp(b_A) * softplus(b_g) + + b_o = tl.cumsum(b_gate, axis=0) + if REVERSE: + b_z = tl.sum(b_gate, axis=0) + b_o = -b_o + b_z[None] + b_gate + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + + +def _launch_gate_chunk_cumsum( + *, + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + o: torch.Tensor, + scale: float | None, + cu_seqlens: torch.LongTensor | None, + chunk_indices: torch.LongTensor | None, + T: int, + B: int, + H: int, + BT: int, + NT: int, + reverse: bool, +) -> None: + bh_total = B * H + kernel_kwargs = dict( + g=g, + A_log=A_log, + dt_bias=dt_bias, + o=o, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + REVERSE=reverse, + num_warps=_NUM_WARPS, + ) + max_nt = max_grid_axis_chunks(NT, bh_total, max_grid=ASCEND_MAX_GRID_DIM) + for nt_off in range(0, NT, max_nt): + nt_len = min(max_nt, NT - nt_off) + max_bh = max_grid_axis_chunks(bh_total, nt_len, max_grid=ASCEND_MAX_GRID_DIM) + for bh_off in range(0, bh_total, max_bh): + bh_len = min(max_bh, bh_total - bh_off) + gdn_gate_chunk_cumsum_scalar_kernel_npu[(nt_len, bh_len)]( + **kernel_kwargs, + NT_OFFSET=nt_off, + BH_OFFSET=bh_off, + ) + + +@triton.heuristics({ + 'HAS_BIAS': lambda args: args['dt_bias'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def gdn_gate_bwd_kernel_npu( + g, + A_log, + dt_bias, + dyg, + dg, + dA, + T, + H: tl.constexpr, + BT: tl.constexpr, + HAS_BIAS: tl.constexpr, + NT_OFFSET: tl.constexpr, + H_OFFSET: tl.constexpr, +): + i_t = tl.program_id(0) + NT_OFFSET + i_h = tl.program_id(1) + H_OFFSET + + b_A = tl.load(A_log + i_h).to(tl.float32) + + p_g = tl.make_block_ptr(g + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dg = tl.make_block_ptr(dg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_dyg = tl.make_block_ptr(dyg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + b_dyg = tl.load(p_dyg, boundary_check=(0,)).to(tl.float32) + + if HAS_BIAS: + b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) + + b_neg_expA = -exp(b_A) + b_yg = b_neg_expA * softplus(b_g) + b_dg = b_neg_expA * (b_dyg * tl.sigmoid(b_g)) + b_dA = tl.sum(b_dyg * b_yg, 0) + + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + tl.store(dA + i_t * H + i_h, b_dA) + + +def _launch_gate_bwd( + *, + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + dyg: torch.Tensor, + dg: torch.Tensor, + dA: torch.Tensor, + T: int, + H: int, + BT: int, +) -> None: + NT = triton.cdiv(T, BT) + kernel_kwargs = dict( + g=g, + A_log=A_log, + dt_bias=dt_bias, + dyg=dyg, + dg=dg, + dA=dA, + T=T, + H=H, + BT=BT, + num_warps=_NUM_WARPS, + ) + max_nt = max_grid_axis_chunks(NT, H, max_grid=ASCEND_MAX_GRID_DIM) + for nt_off in range(0, NT, max_nt): + nt_len = min(max_nt, NT - nt_off) + max_h = max_grid_axis_chunks(H, nt_len, max_grid=ASCEND_MAX_GRID_DIM) + for h_off in range(0, H, max_h): + h_len = min(max_h, H - h_off) + gdn_gate_bwd_kernel_npu[(nt_len, h_len)]( + **kernel_kwargs, + NT_OFFSET=nt_off, + H_OFFSET=h_off, + ) + + +@input_guard +def gdn_gate_fwd_npu( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + H = g.shape[-1] + T = g.numel() // H + BT = _get_gate_fwd_bt(T) + yg = torch.empty_like(g, dtype=output_dtype) + _launch_gate_fwd(g=g, A_log=A_log, dt_bias=dt_bias, yg=yg, T=T, H=H, BT=BT) + return yg + + +@input_guard +def gdn_gate_chunk_cumsum_npu( + g: torch.Tensor, + A_log: torch.Tensor, + chunk_size: int, + scale: float = None, + dt_bias: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + output_dtype: torch.dtype | None = torch.float, +) -> torch.Tensor: + B, T, H = g.shape + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), "chunk_size must be a power of 2" + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + o = torch.empty_like(g, dtype=output_dtype or g.dtype) + _launch_gate_chunk_cumsum( + g=g, + A_log=A_log, + dt_bias=dt_bias, + o=o, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + B=B, + H=H, + BT=BT, + NT=NT, + reverse=False, + ) + return o + + +def gdn_gate_bwd_npu( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + dyg: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + H = g.shape[-1] + T = g.numel() // H + BT = _get_gate_bwd_bt(T) + dg = torch.empty_like(g, dtype=torch.float32) + NT = triton.cdiv(T, BT) + dA = A_log.new_empty(NT, H, dtype=torch.float32) + _launch_gate_bwd( + g=g, + A_log=A_log, + dt_bias=dt_bias, + dyg=dyg, + dg=dg, + dA=dA, + T=T, + H=H, + BT=BT, + ) + dg = dg.view_as(g).type_as(g) + dA = dA.sum(0).view_as(A_log).type_as(A_log) + dbias = dg.view(-1, H).sum(0).to(dt_bias) if dt_bias is not None else None + return dg, dA, dbias diff --git a/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py b/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py new file mode 100644 index 00000000..fe6625bf --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py @@ -0,0 +1,823 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""WY-representation kernels adapted for triton-ascend on Ascend NPU.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +import triton.runtime.driver as driver + +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.op import exp2 +from fla.utils import input_guard +from fla.utils.ascend_ub_manager import ( + ASCEND_MAX_GRID_DIM, + compute_row_tile_block_size, + max_grid_axis_chunks, +) + + +def get_npu_properties(): + device = torch.npu.current_device() + return driver.active.utils.get_device_properties(device) + + +# prepare_wy_repr_bwd stage-specific UB models +_PREPARE_BWD_K_MEM_MULT = 4.5 +_PREPARE_BWD_V_MEM_MULT = 8.0 +_SAFETY_MARGIN = 0.75 +_FALLBACK_TILE = 8 +_MAX_TILE_BWD = 128 + + +def _g_npu_arg(g: torch.Tensor | None, HV: int) -> tuple[torch.Tensor | None, bool]: + if g is None or HV == 1: + return g, False + return g.transpose(1, 2).contiguous(), True + + +def _beta_npu_arg(beta: torch.Tensor, HV: int) -> tuple[torch.Tensor, bool]: + if HV == 1: + return beta, False + return beta.transpose(1, 2).contiguous(), True + + +def _t_npu_buf( + B: int, T: int, HV: int, *, dtype: torch.dtype, device: torch.device, +) -> tuple[torch.Tensor, bool]: + if HV == 1: + return torch.empty(B, T, HV, dtype=dtype, device=device), False + return torch.empty(B, HV, T, dtype=dtype, device=device), True + + +def _get_bwd_k_tile(BT: int, K: int) -> int: + return compute_row_tile_block_size( + BT, K, _PREPARE_BWD_K_MEM_MULT, + tiling_row=False, + safety_margin=_SAFETY_MARGIN, + fallback=_FALLBACK_TILE, + min_block=8, + max_block=min(_MAX_TILE_BWD, triton.next_power_of_2(K)), + ) + + +def _get_bwd_v_tile(BT: int, V: int) -> int: + return compute_row_tile_block_size( + BT, V, _PREPARE_BWD_V_MEM_MULT, + tiling_row=False, + safety_margin=_SAFETY_MARGIN, + fallback=_FALLBACK_TILE, + min_block=8, + max_block=min(_MAX_TILE_BWD, triton.next_power_of_2(V)), + ) + + +def _get_bwd_tiles(BT: int, K: int, V: int) -> tuple[int, int]: + return _get_bwd_k_tile(BT, K), _get_bwd_v_tile(BT, V) + + +@triton.jit +def _g_contig_base(g, bos, i_b, i_h, T_seq, HV, IS_VARLEN: tl.constexpr): + if IS_VARLEN: + return g + bos + i_h * T_seq + return g + i_b * HV * T_seq + i_h * T_seq + + +@triton.jit +def _t_block_ptr(base, T, offset, BLK, CONTIG: tl.constexpr, HV: tl.constexpr): + if CONTIG: + return tl.make_block_ptr(base, (T,), (1,), (offset,), (BLK,), (0,)) + return tl.make_block_ptr(base, (T,), (HV,), (offset,), (BLK,), (0,)) + + +def _launch_wy_kernel(kernel, *, NT: int, bh_total: int, kernel_kwargs: dict) -> None: + max_nt = max_grid_axis_chunks(NT, bh_total, max_grid=ASCEND_MAX_GRID_DIM) + chunk_indices = kernel_kwargs.get('chunk_indices') + cu_seqlens = kernel_kwargs.get('cu_seqlens') + for nt_off in range(0, NT, max_nt): + nt_len = min(max_nt, NT - nt_off) + if cu_seqlens is not None and chunk_indices is not None: + kernel_kwargs['chunk_indices'] = chunk_indices[nt_off:nt_off + nt_len] + kernel_kwargs['NT_OFFSET'] = 0 + else: + kernel_kwargs['NT_OFFSET'] = nt_off + max_bh = max_grid_axis_chunks(bh_total, nt_len, max_grid=ASCEND_MAX_GRID_DIM) + for bh_off in range(0, bh_total, max_bh): + bh_len = min(max_bh, bh_total - bh_off) + kernel_kwargs['BH_OFFSET'] = bh_off + kernel[(nt_len, bh_len)](**kernel_kwargs) + + +def _launch_wy_core_grid(kernel, *, task_num: int, kernel_kwargs: dict) -> None: + num_core = get_npu_properties()["num_aicore"] + kernel[(num_core,)](task_num=task_num, num_core=num_core, **kernel_kwargs) + + +@triton.heuristics({ + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + "USE_G": lambda args: args["g"] is not None, +}) +@triton.jit(do_not_specialize=["T", "B", "task_num", "num_core"]) +def recompute_w_u_fwd_kernel_npu( + k, + v, + beta, + w, + u, + A, + g, + cu_seqlens, + chunk_indices, + T, + B, + task_num, + num_core, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + T_max = T + core_id = tl.program_id(0) + + for task_id in tl.range(core_id, task_num, num_core): + i_t_o = task_id // (B * HV) + i_bh = task_id % (B * HV) + i_b, i_h = i_bh // HV, i_bh % HV + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t_o * 2).to(tl.int32), tl.load( + chunk_indices + i_t_o * 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 + bos_bh = bos + else: + i_t = i_t_o + bos, eos = i_b * T, i_b * T + T + bos_bh = i_b * HV * T_max + + offs_t = tl.arange(0, BT) + global_offs_t = i_t * BT + offs_t + mask_t = global_offs_t < T + + offs_t_2d = global_offs_t[:, None] + offs_bt = tl.arange(0, BT)[None, :] + ptr_A = A + (bos * HV + i_h) * BT + offs_t_2d * (HV * BT) + offs_bt * 1 + mask_A = mask_t[:, None] + b_A = tl.load(ptr_A, mask=mask_A, other=0.0).to(tl.float32) + + ptr_beta = beta + bos_bh + i_h * T_max + global_offs_t + b_beta = tl.load(ptr_beta, mask=mask_t, other=0.0).to(tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + offs_v = i_v * BV + tl.arange(0, BV)[None, :] + mask_v = (mask_t[:, None]) & (offs_v < V) + + ptr_v = v + (bos * HV + i_h) * V + offs_t_2d * (HV * V) + offs_v * 1 + b_v = tl.load(ptr_v, mask=mask_v, other=0.0).to(tl.float32) + + b_vb = b_v * b_beta[:, None] + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + + ptr_u = u + (bos * HV + i_h) * V + offs_t_2d * (HV * V) + offs_v * 1 + tl.store(ptr_u, b_u.to(ptr_u.dtype.element_ty), mask=mask_v) + + if USE_G: + ptr_g = g + bos_bh + i_h * T_max + global_offs_t + b_g = exp2(tl.load(ptr_g, mask=mask_t, other=0.0)).to(tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + offs_k = i_k * BK + tl.arange(0, BK)[None, :] + mask_k = (mask_t[:, None]) & (offs_k < K) + ptr_k = ( + k + + (bos * H + i_h // (HV // H)) * K + + offs_t_2d * (H * K) + + offs_k * 1 + ) + b_k = tl.load(ptr_k, mask=mask_k, other=0.0).to(tl.float32) + + b_kb = b_k * b_beta[:, None] + if USE_G: + b_kb = b_kb * b_g[:, None] + b_w = tl.dot(b_A, b_kb, allow_tf32=False) + + ptr_w = w + (bos * HV + i_h) * K + offs_t_2d * (HV * K) + offs_k * 1 + tl.store(ptr_w, b_w.to(ptr_w.dtype.element_ty), mask=mask_k) + + +@triton.heuristics({ + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, +}) +@triton.jit(do_not_specialize=["T", "B", "task_num", "num_core"]) +def prepare_wy_repr_bwd_kv_npu( + k, v, beta, g, A, dw, du, dk, dv, dA_scr, db, dg, + cu_seqlens, chunk_indices, T, B, + task_num, num_core, + H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, + BT: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + USE_G: tl.constexpr, IS_VARLEN: tl.constexpr, + G_T_CONTIG: tl.constexpr, BETA_T_CONTIG: tl.constexpr, + DG_T_CONTIG: tl.constexpr, DB_T_CONTIG: tl.constexpr, + G_EXP_PRECOMP: tl.constexpr, +): + T_seq = T + core_id = tl.program_id(0) + for task_id in tl.range(core_id, task_num, num_core): + i_t_o = task_id // (B * HV) + i_bh = task_id % (B * HV) + i_b, i_h = i_bh // HV, i_bh % HV + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t_o * 2).to(tl.int32), tl.load( + chunk_indices + i_t_o * 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: + i_t = i_t_o + bos, eos = i_b * T_seq, i_b * T_seq + T_seq + + if BETA_T_CONTIG: + beta_ptr = _g_contig_base(beta, bos, i_b, i_h, T_seq, HV, IS_VARLEN) + p_b = _t_block_ptr(beta_ptr, T, i_t * BT, BT, True, HV) + else: + p_b = tl.make_block_ptr(beta + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) + if DB_T_CONTIG: + db_ptr = _g_contig_base(db, bos, i_b, i_h, T_seq, HV, IS_VARLEN) + p_db = _t_block_ptr(db_ptr, T, i_t * BT, BT, True, HV) + else: + p_db = tl.make_block_ptr(db + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr( + A + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1), + ) + p_dA = tl.make_block_ptr( + dA_scr + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1), + ) + + b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + b_db = tl.zeros([BT], dtype=tl.float32) + b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + + if USE_G: + if G_T_CONTIG: + g_ptr = _g_contig_base(g, bos, i_b, i_h, T_seq, HV, IS_VARLEN) + p_g = _t_block_ptr(g_ptr, T, i_t * BT, BT, True, HV) + else: + p_g = tl.make_block_ptr(g + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + b_g_exp = b_g if G_EXP_PRECOMP else exp2(b_g) + b_dg = tl.zeros([BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * H + i_h // (HV // H)) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0), + ) + p_dk = tl.make_block_ptr( + dk + (bos * HV + i_h) * K, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0), + ) + p_dw = tl.make_block_ptr( + dw + (bos * HV + i_h) * K, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) + if USE_G: + b_kbg = b_k * (b_b * b_g_exp)[:, None] + else: + b_kbg = b_k * b_b[:, None] + b_dw = tl.load(p_dw, boundary_check=(0, 1)).to(tl.float32) + b_dA += tl.dot(b_dw, tl.trans(b_kbg), allow_tf32=False) + b_dkbg = tl.dot(b_A, b_dw, allow_tf32=False) + if USE_G: + b_dk = b_dkbg * (b_g_exp * b_b)[:, None] + b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1) + b_dg += tl.sum(b_dkbg * b_kbg, 1) + else: + b_dk = b_dkbg * b_b[:, None] + b_db += tl.sum(b_dkbg * b_k, 1) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * HV + i_h) * V, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0), + ) + p_dv = tl.make_block_ptr( + dv + (bos * HV + i_h) * V, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0), + ) + p_du = tl.make_block_ptr( + du + (bos * HV + i_h) * V, (T, V), (HV * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) + b_du = tl.load(p_du, boundary_check=(0, 1)).to(tl.float32) + b_vb = b_v * b_b[:, None] + b_dA += tl.dot(b_du, tl.trans(b_vb), allow_tf32=False) + b_dvb = tl.dot(b_A, b_du, allow_tf32=False) + b_dv = b_dvb * b_b[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + if USE_G: + if DG_T_CONTIG: + dg_ptr = _g_contig_base(dg, bos, i_b, i_h, T_seq, HV, IS_VARLEN) + p_dg = _t_block_ptr(dg_ptr, T, i_t * BT, BT, True, HV) + else: + p_dg = tl.make_block_ptr(dg + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_da_mask_dot1_npu( + A, dA_scr, dA_mid, + cu_seqlens, chunk_indices, T, + HV: tl.constexpr, BT: tl.constexpr, + IS_VARLEN: tl.constexpr, + NT_OFFSET: tl.constexpr, BH_OFFSET: tl.constexpr, +): + i_t = tl.program_id(0) + NT_OFFSET + i_bh = tl.program_id(1) + BH_OFFSET + i_b, i_h = i_bh // HV, i_bh % HV + 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 + + p_A = tl.make_block_ptr( + A + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1), + ) + p_in = tl.make_block_ptr( + dA_scr + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1), + ) + p_out = tl.make_block_ptr( + dA_mid + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1), + ) + b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_dA = tl.load(p_in, boundary_check=(0, 1)).to(tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA, 0) + b_out = tl.dot(b_dA, b_A, allow_tf32=False) + tl.store(p_out, b_out.to(p_out.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_da_dot2_npu( + A, dA_mid, dA_out, + cu_seqlens, chunk_indices, T, + HV: tl.constexpr, BT: tl.constexpr, + IS_VARLEN: tl.constexpr, + NT_OFFSET: tl.constexpr, BH_OFFSET: tl.constexpr, +): + i_t = tl.program_id(0) + NT_OFFSET + i_bh = tl.program_id(1) + BH_OFFSET + i_b, i_h = i_bh // HV, i_bh % HV + 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 + + p_A = tl.make_block_ptr(A + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1)) + p_in = tl.make_block_ptr(dA_mid + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1)) + p_out = tl.make_block_ptr(dA_out + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1)) + b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_dA = tl.load(p_in, boundary_check=(0, 1)).to(tl.float32) + b_dA = tl.dot(b_A, b_dA, allow_tf32=False) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, -b_dA, 0) + tl.store(p_out, b_dA.to(p_out.dtype.element_ty), boundary_check=(0, 1)) + + +_DG_BLK = 16 + + +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_da_gate_npu( + g, dA_out, + cu_seqlens, chunk_indices, T, + HV: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, + IS_VARLEN: tl.constexpr, G_T_CONTIG: tl.constexpr, + NT_OFFSET: tl.constexpr, BH_OFFSET: tl.constexpr, +): + i_t = tl.program_id(0) + NT_OFFSET + i_bh = tl.program_id(1) + BH_OFFSET + i_b, i_h = i_bh // HV, i_bh % HV + T_seq = T + 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 + + n_sub = BT // BC + if G_T_CONTIG: + g_ptr = _g_contig_base(g, bos, i_b, i_h, T_seq, HV, IS_VARLEN) + else: + g_ptr = g + (bos * HV + i_h) + + for r in range(n_sub): + i_tr = i_t * BT + r * BC + p_gr = _t_block_ptr(g_ptr, T, i_tr, BC, G_T_CONTIG, HV) + b_gr = tl.load(p_gr, boundary_check=(0,)).to(tl.float32) + for c in range(n_sub): + i_tc = i_t * BT + c * BC + p_dA = tl.make_block_ptr( + dA_out + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), + (r * BC, i_t * BT + c * BC), (BC, BC), (0, 1), + ) + b_dA = tl.load(p_dA, boundary_check=(0, 1)).to(tl.float32) + p_gc = _t_block_ptr(g_ptr, T, i_tc, BC, G_T_CONTIG, HV) + b_gc = tl.load(p_gc, boundary_check=(0,)).to(tl.float32) + b_gate = exp2(b_gr[:, None] - b_gc[None, :]) + b_prod = b_dA * b_gate + b_dA = tl.where(b_prod == b_prod, b_prod, 0.0) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.heuristics({ + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, +}) +@triton.jit(do_not_specialize=["T", "B", "task_num", "num_core"]) +def prepare_wy_repr_bwd_finalize_k_npu( + k, beta, dA_out, dk, db, + cu_seqlens, chunk_indices, T, B, + task_num, num_core, + H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, + BT: tl.constexpr, BK: tl.constexpr, + IS_VARLEN: tl.constexpr, BETA_T_CONTIG: tl.constexpr, DB_T_CONTIG: tl.constexpr, +): + T_seq = T + core_id = tl.program_id(0) + for task_id in tl.range(core_id, task_num, num_core): + i_t_o = task_id // (B * HV) + i_bh = task_id % (B * HV) + i_b, i_h = i_bh // HV, i_bh % HV + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t_o * 2).to(tl.int32), tl.load( + chunk_indices + i_t_o * 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: + i_t = i_t_o + bos, eos = i_b * T_seq, i_b * T_seq + T_seq + + if BETA_T_CONTIG: + beta_ptr = _g_contig_base(beta, bos, i_b, i_h, T_seq, HV, IS_VARLEN) + p_b = _t_block_ptr(beta_ptr, T, i_t * BT, BT, True, HV) + else: + p_b = tl.make_block_ptr(beta + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) + if DB_T_CONTIG: + db_ptr = _g_contig_base(db, bos, i_b, i_h, T_seq, HV, IS_VARLEN) + p_db = _t_block_ptr(db_ptr, T, i_t * BT, BT, True, HV) + else: + p_db = tl.make_block_ptr(db + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) + p_dA = tl.make_block_ptr( + dA_out + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1), + ) + + b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + b_db = tl.load(p_db, boundary_check=(0,)).to(tl.float32) + b_dA = tl.load(p_dA, boundary_check=(0, 1)).to(tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * H + i_h // (HV // H)) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0), + ) + p_dk = tl.make_block_ptr( + dk + (bos * HV + i_h) * K, (T, K), (HV * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) + b_kb = b_k * b_b[:, None] + b_dkb = tl.dot(b_dA, b_k, allow_tf32=False) + b_db += tl.sum(b_dkb * b_k, 1) + b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(tl.trans(b_kb), b_dA, allow_tf32=False)) + b_dk += tl.load(p_dk, boundary_check=(0, 1)).to(tl.float32) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + + tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + + +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_finalize_a2_npu( + k, beta, a2_scr, + cu_seqlens, chunk_indices, T, + H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, + BT: tl.constexpr, BK: tl.constexpr, + IS_VARLEN: tl.constexpr, BETA_T_CONTIG: tl.constexpr, + NT_OFFSET: tl.constexpr, BH_OFFSET: tl.constexpr, +): + i_t = tl.program_id(0) + NT_OFFSET + i_bh = tl.program_id(1) + BH_OFFSET + i_b, i_h = i_bh // HV, i_bh % HV + T_seq = T + 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 BETA_T_CONTIG: + beta_ptr = _g_contig_base(beta, bos, i_b, i_h, T_seq, HV, IS_VARLEN) + p_b = _t_block_ptr(beta_ptr, T, i_t * BT, BT, True, HV) + else: + p_b = tl.make_block_ptr(beta + (bos * HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) + p_a2 = tl.make_block_ptr(a2_scr + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1)) + b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + b_A2 = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * H + i_h // (HV // H)) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) + b_A2 += tl.dot(b_k, tl.trans(b_k), allow_tf32=False) + b_A2 *= b_b[:, None] + tl.store(p_a2, b_A2.to(p_a2.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_finalize_dg_npu( + dA_out, a2_scr, dg, col_acc_scr, + cu_seqlens, chunk_indices, T, + HV: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, + IS_VARLEN: tl.constexpr, DG_T_CONTIG: tl.constexpr, + NT_OFFSET: tl.constexpr, BH_OFFSET: tl.constexpr, +): + i_t = tl.program_id(0) + NT_OFFSET + i_bh = tl.program_id(1) + BH_OFFSET + i_b, i_h = i_bh // HV, i_bh % HV + T_seq = T + if IS_VARLEN: + i_tg = i_t + 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: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + n_sub = BT // BC + col_off = (i_tg * HV + i_h) * BT + p_col0 = tl.make_block_ptr(col_acc_scr + col_off, (BT,), (1,), (0,), (BT,), (0,)) + tl.store(p_col0, tl.zeros([BT], dtype=tl.float32), boundary_check=(0,)) + + if DG_T_CONTIG: + dg_ptr = _g_contig_base(dg, bos, i_b, i_h, T_seq, HV, IS_VARLEN) + else: + dg_ptr = dg + (bos * HV + i_h) + + for r in range(n_sub): + i_tr = i_t * BT + r * BC + p_dg_r = _t_block_ptr(dg_ptr, T, i_tr, BC, DG_T_CONTIG, HV) + b_dg_r = tl.load(p_dg_r, boundary_check=(0,)).to(tl.float32) + for c in range(n_sub): + p_dA = tl.make_block_ptr( + dA_out + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), + (r * BC, i_t * BT + c * BC), (BC, BC), (0, 1), + ) + p_a2 = tl.make_block_ptr( + a2_scr + (bos * HV + i_h) * BT, (BT, T), (1, HV * BT), + (r * BC, i_t * BT + c * BC), (BC, BC), (0, 1), + ) + b_dA = tl.load(p_dA, boundary_check=(0, 1)).to(tl.float32) + b_a2 = tl.load(p_a2, boundary_check=(0, 1)).to(tl.float32) + prod = b_dA * b_a2 + b_dg_r += tl.sum(prod, axis=1) + p_col = tl.make_block_ptr( + col_acc_scr + col_off, (BT,), (1,), (c * BC,), (BC,), (0,), + ) + b_col = tl.load(p_col, boundary_check=(0,)).to(tl.float32) + b_col += tl.sum(prod, axis=0) + tl.store(p_col, b_col.to(p_col.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dg_r, b_dg_r.to(p_dg_r.dtype.element_ty), boundary_check=(0,)) + + p_dg = _t_block_ptr(dg_ptr, T, i_t * BT, BT, DG_T_CONTIG, HV) + p_col = tl.make_block_ptr(col_acc_scr + col_off, (BT,), (1,), (0,), (BT,), (0,)) + b_dg = tl.load(p_dg, boundary_check=(0,)).to(tl.float32) + b_col = tl.load(p_col, boundary_check=(0,)).to(tl.float32) + b_dg -= b_col + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + + +@input_guard +def recompute_w_u_fwd_npu( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + g: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + BT = A.shape[-1] + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + BK = 64 + BV = 64 + + u = torch.empty_like(v) + w = k.new_empty(B, T, HV, K) + beta = beta.transpose(1, 2).contiguous() + if g is not None: + g = g.transpose(1, 2).contiguous() + + num_core = get_npu_properties()["num_aicore"] + task_num = NT * B * HV + recompute_w_u_fwd_kernel_npu[(num_core,)]( + k=k, + v=v, + beta=beta, + w=w, + u=u, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + B=B, + task_num=task_num, + num_core=num_core, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u + + +def prepare_wy_repr_bwd_npu( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + g: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] + BT = A.shape[-1] + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BK, BV = _get_bwd_tiles(BT, K, V) + use_g = g is not None + is_varlen = cu_seqlens is not None + + dk = k.new_empty(B, T, HV, K) + dv = torch.empty_like(v) + dg, dg_t_contig = None, False + if use_g: + dg, dg_t_contig = _t_npu_buf(B, T, HV, dtype=g.dtype, device=k.device) + db, db_t_contig = _t_npu_buf(B, T, HV, dtype=beta.dtype, device=k.device) + beta_arg, beta_t_contig = _beta_npu_arg(beta, HV) + g_gate, g_t_contig = None, False + g_exp_precomp = False + if use_g: + g_gate, g_t_contig = _g_npu_arg(g, HV) + g_k_arg = g_gate + if not is_varlen: + g_k_arg = torch.exp2(g_gate.float()).to(g_gate.dtype) + g_exp_precomp = True + dg_arg = dg if use_g else beta + dA_scr = torch.zeros_like(A, dtype=torch.float32) + dA_mid = torch.zeros_like(A, dtype=torch.float32) + dA_out = torch.zeros_like(A, dtype=torch.float32) + a2_scr = torch.zeros_like(A, dtype=torch.float32) + col_acc_scr = torch.zeros(B, triton.cdiv(T, BT) if cu_seqlens is None else len( + chunk_indices), HV, BT, dtype=torch.float32, device=k.device) + + base = dict( + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + BT=BT, + IS_VARLEN=is_varlen, + ) + core_base = dict(B=B, **base) + task_num = NT * B * HV + _launch_wy_core_grid( + prepare_wy_repr_bwd_kv_npu, + task_num=task_num, + kernel_kwargs=dict( + k=k, v=v, beta=beta_arg, g=g_k_arg if use_g else k, A=A, dw=dw, du=du, + dk=dk, dv=dv, dA_scr=dA_scr, db=db, dg=dg_arg, + H=H, HV=HV, K=K, V=V, BK=BK, BV=BV, USE_G=use_g, + G_T_CONTIG=g_t_contig, BETA_T_CONTIG=beta_t_contig, + DG_T_CONTIG=dg_t_contig, DB_T_CONTIG=db_t_contig, + G_EXP_PRECOMP=g_exp_precomp, + **core_base, + ), + ) + _launch_wy_kernel( + prepare_wy_repr_bwd_da_mask_dot1_npu, + NT=NT, + bh_total=B * HV, + kernel_kwargs=dict( + A=A, dA_scr=dA_scr, dA_mid=dA_mid, + HV=HV, + **base, + ), + ) + _launch_wy_kernel( + prepare_wy_repr_bwd_da_dot2_npu, + NT=NT, + bh_total=B * HV, + kernel_kwargs=dict( + A=A, dA_mid=dA_mid, dA_out=dA_out, + HV=HV, + **base, + ), + ) + if use_g: + _launch_wy_kernel( + prepare_wy_repr_bwd_da_gate_npu, + NT=NT, + bh_total=B * HV, + kernel_kwargs=dict( + g=g_gate, dA_out=dA_out, + HV=HV, BC=_DG_BLK, G_T_CONTIG=g_t_contig, + **base, + ), + ) + _launch_wy_core_grid( + prepare_wy_repr_bwd_finalize_k_npu, + task_num=task_num, + kernel_kwargs=dict( + k=k, beta=beta_arg, dA_out=dA_out, dk=dk, db=db, + H=H, HV=HV, K=K, BK=BK, BETA_T_CONTIG=beta_t_contig, DB_T_CONTIG=db_t_contig, + **core_base, + ), + ) + if use_g: + _launch_wy_kernel( + prepare_wy_repr_bwd_finalize_a2_npu, + NT=NT, + bh_total=B * HV, + kernel_kwargs=dict( + k=k, beta=beta_arg, a2_scr=a2_scr, + H=H, HV=HV, K=K, BK=BK, BETA_T_CONTIG=beta_t_contig, + **base, + ), + ) + _launch_wy_kernel( + prepare_wy_repr_bwd_finalize_dg_npu, + NT=NT, + bh_total=B * HV, + kernel_kwargs=dict( + dA_out=dA_out, a2_scr=a2_scr, dg=dg_arg, col_acc_scr=col_acc_scr, + HV=HV, BC=_DG_BLK, DG_T_CONTIG=dg_t_contig, + **base, + ), + ) + if H != HV: + dk = dk.view(B, T, H, HV // H, K).sum(3) + if db_t_contig: + db = db.transpose(1, 2).contiguous() + if use_g and dg_t_contig: + dg = dg.transpose(1, 2).contiguous() + return dk, dv, db, dg diff --git a/upstream_ref/fla/ops/gated_delta_rule/chunk.py b/upstream_ref/fla/ops/gated_delta_rule/chunk.py new file mode 100644 index 00000000..576278e9 --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/chunk.py @@ -0,0 +1,591 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import warnings + +import torch + +from fla.modules.l2norm import l2norm_bwd, l2norm_fwd +from fla.ops.backends import dispatch +from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o +from fla.ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd +from fla.ops.cp import FLACPContext +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu_pre_process, + chunk_gated_delta_rule_fwd_h_pre_process, + compress_h0, + expand_h0, +) +from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra +from fla.ops.gated_delta_rule.gate import gdn_gate_bwd, gdn_gate_chunk_cumsum +from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd +from fla.ops.utils import chunk_local_cumsum +from fla.ops.utils.constant import RCP_LN2 +from fla.ops.utils.index import prepare_chunk_indices +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +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, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, +): + g_input = g if use_gate_in_kernel else None + if use_gate_in_kernel: + g = gdn_gate_chunk_cumsum( + g=g, + A_log=A_log, + chunk_size=chunk_size, + scale=RCP_LN2, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + else: + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + # obtain WY representation. u is actually the new v. + # fused kkt + solve_tril + recompute_w_u + w, u, A = chunk_gated_delta_rule_fwd_intra( + k=k, + v=v, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, + w=w, + u=u, + g=g, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + 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, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + if cp_context is not None: + initial_state = compress_h0(initial_state, context=cp_context) + + 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, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + return g, o, A, final_state, initial_state, g_input + + +def chunk_gated_delta_rule_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + use_gate_in_kernel: bool = False, + g_input: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + chunk_size: int = 64, +): + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + if cp_context is not None: + initial_state = expand_h0(initial_state, context=cp_context) + + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k, + g=g, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ) + + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=q, + k=k, + w=w, + do=do, + dv=dv, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + context=cp_context, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k, + w=w, + g=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dq, dk, dw, dg = chunk_bwd_dqkwg( + q=q, + k=k, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + chunk_size=chunk_size, + ) + dk2, dv, db, dg2 = prepare_wy_repr_bwd( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + dk.add_(dk2) + dg.add_(dg2) + dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + dA_log, ddt_bias = None, None + if use_gate_in_kernel: + dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg) + return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias + + +class ChunkGatedDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + 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, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + cp_context: FLACPContext | None = None, + chunk_size: int = 64, + ): + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + beta_raw = beta + if use_beta_sigmoid_in_kernel: + beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) + + chunk_indices = None + if cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) + g, o, A, final_state, initial_state, g_input = 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, + cp_context=cp_context, + chunk_indices=chunk_indices, + state_v_first=state_v_first, + use_gate_in_kernel=use_gate_in_kernel, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=chunk_size, + ) + ctx.save_for_backward( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + ) + ctx.scale = scale + ctx.chunk_size = chunk_size + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel + ctx.allow_neg_eigval = allow_neg_eigval + ctx.cp_context = cp_context + ctx.state_v_first = state_v_first + ctx.use_gate_in_kernel = use_gate_in_kernel + return o.to(q.dtype), final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dht: torch.Tensor, + ): + ( + q, + q_rstd, + k, + k_rstd, + v, + g, + beta_raw, + beta, + A, + initial_state, + cu_seqlens, + chunk_indices, + g_input, + A_log, + dt_bias, + ) = ctx.saved_tensors + dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + cp_context=ctx.cp_context, + chunk_indices=chunk_indices, + state_v_first=ctx.state_v_first, + use_gate_in_kernel=ctx.use_gate_in_kernel, + g_input=g_input, + A_log=A_log, + dt_bias=dt_bias, + chunk_size=ctx.chunk_size, + ) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + if ctx.use_beta_sigmoid_in_kernel: + db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) + return ( + dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta_raw), + None, dh0, None, None, None, None, None, None, dA_log, ddt_bias, + None, None, None, None, + ) + + +@dispatch('gated_delta_rule') +@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 = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, +): + 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 (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. + g (torch.Tensor): + (forget) gating tensor of shape `[B, T, HV]`. + When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay). + When `use_gate_in_kernel=True`, `g` is the raw input before gate activation; + the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + 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, K, V]` 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, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q/k tensor internally. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space GDN decay internally. + When `True`, the passed `g` is the raw input, and `A_log` must be provided. + The kernel fuses gate activation + chunk cumsum in a single pass. + Default: `False`. + A_log (Optional[torch.Tensor]): + Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + dt_bias (Optional[torch.Tensor]): + Bias added to `g` before activation, of shape `[HV]`. + Only used when `use_gate_in_kernel=True`. + use_beta_sigmoid_in_kernel (bool): + Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel. + - If `True`, the passed `beta` acts as the raw beta logits. + - If `False`, `beta` is expected to already be in post-sigmoid space. + Default: `False`. + allow_neg_eigval (bool): + Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. + Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case + the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. + state_v_first (Optional[bool]): + Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + cp_context (Optional[FLACPContext]): + Context parallel context for distributed training across multiple devices. + When provided, `initial_state` and `output_final_state` are not supported, + and `cu_seqlens` will be overridden by the context. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` 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, HV, K, V = 4, 2048, 4, 8, 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, HV, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, HV, K, V, 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.long) + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if 'transpose_state_layout' in kwargs: + if state_v_first: + raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") + warnings.warn( + "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", + DeprecationWarning, + stacklevel=2, + ) + state_v_first = kwargs.pop('transpose_state_layout') + + # Validate head dimensions + if q.shape[2] != k.shape[2]: + raise ValueError( + f"q and k must have the same number of heads, " + f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}" + ) + H, HV = q.shape[2], v.shape[2] + if HV % H != 0: + raise ValueError( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by " + f"num_heads (H={H}), but got HV % H = {HV % H}" + ) + + if 'head_first' in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + + chunk_size = kwargs.pop('chunk_size', 64) + if chunk_size not in (16, 32, 64): + raise ValueError(f"`chunk_size` must be 16, 32, or 64 for Gated Delta Rule, got {chunk_size}.") + + if cp_context is not None: + assert initial_state is None, "Initial state is not supported for CP" + assert output_final_state is False, "Output final state is not supported for CP" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + + 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]}.", + ) + use_gate_in_kernel = kwargs.get('use_gate_in_kernel', False) + A_log = kwargs.get('A_log') + dt_bias = kwargs.get('dt_bias') + if use_gate_in_kernel: + assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True." + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") + + 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, + state_v_first, + cu_seqlens, + cu_seqlens_cpu, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + A_log, + dt_bias, + use_beta_sigmoid_in_kernel, + allow_neg_eigval, + cp_context, + chunk_size, + ) + return o, final_state + + +chunk_gdn = chunk_gated_delta_rule diff --git a/upstream_ref/fla/ops/gated_delta_rule/chunk_fwd.py b/upstream_ref/fla/ops/gated_delta_rule/chunk_fwd.py new file mode 100644 index 00000000..875bdde9 --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/chunk_fwd.py @@ -0,0 +1,425 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import torch +import triton +import triton.language as tl + +from fla.ops.backends import dispatch +from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from fla.ops.gated_delta_rule.wy_fast import recompute_w_u_fwd +from fla.ops.utils import prepare_chunk_indices, solve_tril +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.op import exp2 +from fla.utils import IS_TF32_SUPPORTED, autotune_cache_kwargs + +if IS_TF32_SUPPORTED: + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') +else: + SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps) + for BK in [32, 64] + for num_warps in [1, 2, 4] + ], + key=['H', 'HV', 'K', 'BC'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gated_delta_rule_fwd_kkt_solve_kernel( + k, + g, + beta, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Fused kernel: compute beta * K @ K^T (lower triangular) + solve_tril (I+A)^{-1} in one pass. + + This kernel fuses chunk_scaled_dot_kkt_fwd and solve_tril into a single kernel, + avoiding the HBM round-trip for the intermediate A matrix. + + Steps: + 1. Compute all 10 lower-triangular [BC, BC] blocks of beta * K @ K^T in registers + 2. Apply gate and beta scaling + 3. Forward substitution on diagonal blocks + 4. Block merge to get full (I+A)^{-1} + 5. Write result to A (output) + """ + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV + + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + k += (bos * H + i_h // (HV // H)) * K + A += (bos * HV + i_h) * BT + + o_i = tl.arange(0, BC) + m_tc0 = (i_tc0 + o_i) < T + m_tc1 = (i_tc1 + o_i) < T + m_tc2 = (i_tc2 + o_i) < T + m_tc3 = (i_tc3 + o_i) < T + + # load beta for each sub-chunk + p_b0 = beta + bos * HV + i_h + (i_tc0 + o_i) * HV + p_b1 = beta + bos * HV + i_h + (i_tc1 + o_i) * HV + p_b2 = beta + bos * HV + i_h + (i_tc2 + o_i) * HV + p_b3 = beta + bos * HV + i_h + (i_tc3 + o_i) * HV + b_b0 = tl.load(p_b0, mask=m_tc0, other=0.0).to(tl.float32) + b_b1 = tl.load(p_b1, mask=m_tc1, other=0.0).to(tl.float32) + b_b2 = tl.load(p_b2, mask=m_tc2, other=0.0).to(tl.float32) + b_b3 = tl.load(p_b3, mask=m_tc3, other=0.0).to(tl.float32) + + # load gate if used + if USE_G: + p_g0 = g + bos * HV + i_h + (i_tc0 + o_i) * HV + p_g1 = g + bos * HV + i_h + (i_tc1 + o_i) * HV + p_g2 = g + bos * HV + i_h + (i_tc2 + o_i) * HV + p_g3 = g + bos * HV + i_h + (i_tc3 + o_i) * HV + + b_g0 = tl.load(p_g0, mask=m_tc0, other=0.0).to(tl.float32) + b_g1 = tl.load(p_g1, mask=m_tc1, other=0.0).to(tl.float32) + b_g2 = tl.load(p_g2, mask=m_tc2, other=0.0).to(tl.float32) + b_g3 = tl.load(p_g3, mask=m_tc3, other=0.0).to(tl.float32) + + ############################################################################ + # Step 1: compute all 10 lower-triangular [BC, BC] blocks of K @ K^T + ############################################################################ + + # 4 diagonal blocks + b_A00 = tl.zeros([BC, BC], dtype=tl.float32) + b_A11 = tl.zeros([BC, BC], dtype=tl.float32) + b_A22 = tl.zeros([BC, BC], dtype=tl.float32) + b_A33 = tl.zeros([BC, BC], dtype=tl.float32) + + # 6 off-diagonal blocks + b_A10 = tl.zeros([BC, BC], dtype=tl.float32) + b_A20 = tl.zeros([BC, BC], dtype=tl.float32) + b_A21 = tl.zeros([BC, BC], dtype=tl.float32) + b_A30 = tl.zeros([BC, BC], dtype=tl.float32) + b_A31 = tl.zeros([BC, BC], dtype=tl.float32) + b_A32 = tl.zeros([BC, BC], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + p_k0 = k + (i_tc0 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k0 = tl.load(p_k0, mask=m_tc0[:, None] & (o_k[None, :] < K), other=0.0) + # diagonal block 0 + b_A00 += tl.dot(b_k0, tl.trans(b_k0)) + + if i_tc1 < T: + p_k1 = k + (i_tc1 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k1 = tl.load(p_k1, mask=m_tc1[:, None] & (o_k[None, :] < K), other=0.0) + # diagonal block 1 + b_A11 += tl.dot(b_k1, tl.trans(b_k1)) + # off-diagonal (1,0) + b_A10 += tl.dot(b_k1, tl.trans(b_k0)) + + if i_tc2 < T: + p_k2 = k + (i_tc2 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k2 = tl.load(p_k2, mask=m_tc2[:, None] & (o_k[None, :] < K), other=0.0) + # diagonal block 2 + b_A22 += tl.dot(b_k2, tl.trans(b_k2)) + # off-diagonal (2,0), (2,1) + b_A20 += tl.dot(b_k2, tl.trans(b_k0)) + b_A21 += tl.dot(b_k2, tl.trans(b_k1)) + + if i_tc3 < T: + p_k3 = k + (i_tc3 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k3 = tl.load(p_k3, mask=m_tc3[:, None] & (o_k[None, :] < K), other=0.0) + # diagonal block 3 + b_A33 += tl.dot(b_k3, tl.trans(b_k3)) + # off-diagonal (3,0), (3,1), (3,2) + b_A30 += tl.dot(b_k3, tl.trans(b_k0)) + b_A31 += tl.dot(b_k3, tl.trans(b_k1)) + b_A32 += tl.dot(b_k3, tl.trans(b_k2)) + + ############################################################################ + # Step 2: apply gate and beta scaling + ############################################################################ + + # apply gate, beta scaling, and masking + # m_d: strictly lower triangular mask for diagonal blocks + # m_tc: boundary mask to prevent NaN from 0 * inf (IEEE 754) when + # out-of-bounds g loads as 0 via boundary_check and exp2(0 - g_inbounds) overflows + m_d = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + if USE_G: + b_A00 *= tl.where(m_d & m_tc0[:, None] & m_tc0[None, :], exp2(b_g0[:, None] - b_g0[None, :]), 0.) + b_A11 *= tl.where(m_d & m_tc1[:, None] & m_tc1[None, :], exp2(b_g1[:, None] - b_g1[None, :]), 0.) + b_A22 *= tl.where(m_d & m_tc2[:, None] & m_tc2[None, :], exp2(b_g2[:, None] - b_g2[None, :]), 0.) + b_A33 *= tl.where(m_d & m_tc3[:, None] & m_tc3[None, :], exp2(b_g3[:, None] - b_g3[None, :]), 0.) + + b_A10 *= tl.where(m_tc1[:, None] & m_tc0[None, :], exp2(b_g1[:, None] - b_g0[None, :]), 0.) + b_A20 *= tl.where(m_tc2[:, None] & m_tc0[None, :], exp2(b_g2[:, None] - b_g0[None, :]), 0.) + b_A21 *= tl.where(m_tc2[:, None] & m_tc1[None, :], exp2(b_g2[:, None] - b_g1[None, :]), 0.) + b_A30 *= tl.where(m_tc3[:, None] & m_tc0[None, :], exp2(b_g3[:, None] - b_g0[None, :]), 0.) + b_A31 *= tl.where(m_tc3[:, None] & m_tc1[None, :], exp2(b_g3[:, None] - b_g1[None, :]), 0.) + b_A32 *= tl.where(m_tc3[:, None] & m_tc2[None, :], exp2(b_g3[:, None] - b_g2[None, :]), 0.) + else: + b_A00 = tl.where(m_d, b_A00, 0.) + b_A11 = tl.where(m_d, b_A11, 0.) + b_A22 = tl.where(m_d, b_A22, 0.) + b_A33 = tl.where(m_d, b_A33, 0.) + + # diagonal blocks: scaled by beta + b_A00 = b_A00 * b_b0[:, None] + b_A11 = b_A11 * b_b1[:, None] + b_A22 = b_A22 * b_b2[:, None] + b_A33 = b_A33 * b_b3[:, None] + + # off-diagonal blocks: full block, scaled by beta + b_A10 = b_A10 * b_b1[:, None] + b_A20 = b_A20 * b_b2[:, None] + b_A21 = b_A21 * b_b2[:, None] + b_A30 = b_A30 * b_b3[:, None] + b_A31 = b_A31 * b_b3[:, None] + b_A32 = b_A32 * b_b3[:, None] + + ############################################################################ + # Step 3: forward substitution on diagonal blocks -> (I + A_diag)^{-1} + # + # Same algorithm as solve_tril, but rows are extracted from in-register + # [BC, BC] tensor via tl.sum(tl.where(mask, tensor, 0), 0) instead of + # tl.load from HBM. + ############################################################################ + + b_Ai00 = -b_A00 + b_Ai11 = -b_A11 + b_Ai22 = -b_A22 + b_Ai33 = -b_A33 + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = tl.sum(tl.where((o_i == i)[:, None], -b_A00, 0.), 0) + b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 = b_a00 + tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(2, min(BC, T - i_tc1)): + b_a11 = tl.sum(tl.where((o_i == i)[:, None], -b_A11, 0.), 0) + b_a11 = tl.where(o_i < i, b_a11, 0.) + b_a11 = b_a11 + tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i)[:, None], b_a11, b_Ai11) + for i in range(2, min(BC, T - i_tc2)): + b_a22 = tl.sum(tl.where((o_i == i)[:, None], -b_A22, 0.), 0) + b_a22 = tl.where(o_i < i, b_a22, 0.) + b_a22 = b_a22 + tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i)[:, None], b_a22, b_Ai22) + for i in range(2, min(BC, T - i_tc3)): + b_a33 = tl.sum(tl.where((o_i == i)[:, None], -b_A33, 0.), 0) + b_a33 = tl.where(o_i < i, b_a33, 0.) + b_a33 = b_a33 + tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ############################################################################ + # Step 4: block merge -> full (I + A)^{-1} + ############################################################################ + + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_A10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_A21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_A32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_A20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_A31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_A30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_A32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION, + ) + + ############################################################################ + # Step 5: store full (I + A)^{-1} to output A + ############################################################################ + + p_A00 = A + (i_tc0 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A10 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A11 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :] + p_A20 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A21 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :] + p_A22 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :] + p_A30 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A31 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :] + p_A32 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :] + p_A33 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (3*BC + o_i)[None, :] + + m_A0 = m_tc0[:, None] & (o_i[None, :] < BT) + m_A1 = m_tc1[:, None] & (o_i[None, :] < BT) + m_A2 = m_tc2[:, None] & (o_i[None, :] < BT) + m_A3 = m_tc3[:, None] & (o_i[None, :] < BT) + m_A11 = m_tc1[:, None] & ((BC + o_i)[None, :] < BT) + m_A21 = m_tc2[:, None] & ((BC + o_i)[None, :] < BT) + m_A22 = m_tc2[:, None] & ((2*BC + o_i)[None, :] < BT) + m_A31 = m_tc3[:, None] & ((BC + o_i)[None, :] < BT) + m_A32 = m_tc3[:, None] & ((2*BC + o_i)[None, :] < BT) + m_A33 = m_tc3[:, None] & ((3*BC + o_i)[None, :] < BT) + + tl.store(p_A00, b_Ai00.to(A.dtype.element_ty), mask=m_A0) + tl.store(p_A10, b_Ai10.to(A.dtype.element_ty), mask=m_A1) + tl.store(p_A11, b_Ai11.to(A.dtype.element_ty), mask=m_A11) + tl.store(p_A20, b_Ai20.to(A.dtype.element_ty), mask=m_A2) + tl.store(p_A21, b_Ai21.to(A.dtype.element_ty), mask=m_A21) + tl.store(p_A22, b_Ai22.to(A.dtype.element_ty), mask=m_A22) + tl.store(p_A30, b_Ai30.to(A.dtype.element_ty), mask=m_A3) + tl.store(p_A31, b_Ai31.to(A.dtype.element_ty), mask=m_A31) + tl.store(p_A32, b_Ai32.to(A.dtype.element_ty), mask=m_A32) + tl.store(p_A33, b_Ai33.to(A.dtype.element_ty), mask=m_A33) + + +@dispatch('gated_delta_rule') +def chunk_gated_delta_rule_fwd_intra( + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + r""" + GDN intra-chunk forward: fused or unfused kkt + solve_tril + recompute_w_u. + + For ``chunk_size == 64``, this uses the fused kkt + solve_tril path. For + other supported chunk sizes, it computes the mathematically equivalent + representation with ``chunk_scaled_dot_kkt_fwd`` followed by ``solve_tril``. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + v (torch.Tensor): + The value tensor of shape `[B, T, HV, V]`. + g (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, HV]`. Default: `None`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, HV]`. + cu_seqlens (torch.LongTensor): + The cumulative sequence lengths. Default: `None`. + chunk_size (int): + The chunk size. Default: 64. + chunk_indices (torch.LongTensor): + Precomputed chunk indices. Default: `None`. + + Returns: + w (torch.Tensor): shape `[B, T, HV, K]` + u (torch.Tensor): shape `[B, T, HV, V]` + A (torch.Tensor): shape `[B, T, HV, BT]`, the solved (I+A)^{-1} matrix + """ + if chunk_size not in (16, 32, 64): + raise ValueError(f"`chunk_size` must be 16, 32, or 64, got {chunk_size}.") + + B, T, H, K, HV = *k.shape, beta.shape[2] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + + if BT == 64: + # Step 1: fused kkt + solve_tril + BC = 16 + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + A = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype) + chunk_gated_delta_rule_fwd_kkt_solve_kernel[(NT, B * HV)]( + k=k, + g=g, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + BC=BC, + ) + else: + # Step 1: mathematically equivalent unfused kkt + solve_tril for non-64 chunks + A = chunk_scaled_dot_kkt_fwd( + k=k, + g=g, + beta=beta, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=BT, + output_dtype=torch.float32, + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=k.dtype, + ) + + # Step 2: recompute_w_u + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return w, u, A diff --git a/upstream_ref/fla/ops/gated_delta_rule/fused_recurrent.py b/upstream_ref/fla/ops/gated_delta_rule/fused_recurrent.py new file mode 100644 index 00000000..0207dc9e --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/fused_recurrent.py @@ -0,0 +1,478 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import warnings + +import torch +import triton +import triton.language as tl + +from fla.ops.utils.op import exp +from fla.ops.utils.softplus import softplus +from fla.utils import input_guard + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'USE_GK': lambda args: args['gk'] is not None, + 'USE_GV': lambda args: args['gv'] is not None, + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_GATE_IN_KERNEL': lambda args: args['A_log'] is not None, + 'HAS_DT_BIAS': lambda args: args['dt_bias'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def fused_recurrent_gated_delta_rule_fwd_kernel( + q, + k, + v, + g, + gk, + gv, + beta, + A_log, + dt_bias, + o, + h0, + ht, + cu_seqlens, + scale, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_GV: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_BETA_HEADWISE: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + STATE_V_FIRST: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATE_IN_KERNEL: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + APPLY_BETA_SIGMOID: tl.constexpr, + ALLOW_NEG_EIGVAL: tl.constexpr, +): + pid = tl.program_id(0) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) + 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) + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + o_k = 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 USE_G: + p_g = g + bos * HV + i_hv + if USE_GK: + p_gk = gk + (bos * HV + i_hv) * K + o_k + if USE_GV: + p_gv = gv + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + bos * HV + i_hv + else: + p_beta = beta + (bos * HV + i_hv) * V + o_v + + p_o = o + (bos * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + if STATE_V_FIRST: + mask_h = mask_v[:, None] & mask_k[None, :] + else: + mask_h = mask_k[:, None] & mask_v[None, :] + + if STATE_V_FIRST: + b_h = tl.zeros([BV, BK], dtype=tl.float32) + else: + b_h = tl.zeros([BK, BV], dtype=tl.float32) + if USE_INITIAL_STATE: + if STATE_V_FIRST: + p_h0 = h0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :] + else: + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for _ in tl.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 + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta).to(tl.float32) + else: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + if APPLY_BETA_SIGMOID: + b_beta = tl.sigmoid(b_beta) + if ALLOW_NEG_EIGVAL: + b_beta = b_beta * 2 + + if USE_G: + b_g = tl.load(p_g).to(tl.float32) + if USE_GATE_IN_KERNEL: + b_A = tl.load(A_log + i_hv).to(tl.float32) + if HAS_DT_BIAS: + b_g = b_g + tl.load(dt_bias + i_hv).to(tl.float32) + b_g = -exp(b_A) * softplus(b_g) + b_h *= exp(b_g) + + if USE_GK: + b_gk = tl.load(p_gk).to(tl.float32) + if STATE_V_FIRST: + b_h *= exp(b_gk[None, :]) + else: + b_h *= exp(b_gk[:, None]) + + if USE_GV: + b_gv = tl.load(p_gv).to(tl.float32) + if STATE_V_FIRST: + b_h *= exp(b_gv[:, None]) + else: + b_h *= exp(b_gv[None, :]) + + if STATE_V_FIRST: + b_v = b_beta * (b_v - tl.sum(b_h * b_k[None, :], 1)) + b_h += b_v[:, None] * b_k[None, :] + b_o = tl.sum(b_h * b_q[None, :], 1) + else: + b_v = b_beta * (b_v - tl.sum(b_h * b_k[:, None], 0)) + b_h += b_k[:, None] * b_v + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + p_q += H*K + p_k += H*K + p_v += HV*V + if USE_G: + p_g += HV + if USE_GK: + p_gk += HV*K + if USE_GV: + p_gv += HV*V + p_beta += HV * (1 if IS_BETA_HEADWISE else V) + p_o += HV*V + + if STORE_FINAL_STATE: + if STATE_V_FIRST: + p_ht = ht + i_nh * K*V + o_v[:, None] * K + o_k[None, :] + else: + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + +def fused_recurrent_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> 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 = triton.next_power_of_2(K) + BV = min(8, triton.next_power_of_2(V)) if gv is None else triton.next_power_of_2(V) + NV = triton.cdiv(V, BV) + + o = torch.empty_like(v) + if output_final_state: + if state_v_first: + final_state = q.new_empty(N, HV, V, K, dtype=torch.float32) + else: + final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) + else: + final_state = None + + grid = (NV * N * HV,) + fused_recurrent_gated_delta_rule_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + gk=gk, + gv=gv, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + o=o, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + IS_BETA_HEADWISE=beta.ndim != v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel, + ALLOW_NEG_EIGVAL=allow_neg_eigval, + STATE_V_FIRST=state_v_first, + num_warps=1, + num_stages=3, + ) + return o, final_state + + +class FusedRecurrentFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + ): + o, final_state = fused_recurrent_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + gk=gk, + gv=gv, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, + allow_neg_eigval=allow_neg_eigval, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + ) + + return o, final_state + + @staticmethod + @input_guard + def backward(ctx, do, dht): + raise NotImplementedError( + "Backward pass is not implemented yet and we do not have plans to implement it " + "because we haven't figured out how to compute dg without materializing the full " + "hidden states for all time steps.", + ) + + +def fused_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + gv: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + **kwargs, +) -> 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 (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. + g (torch.Tensor): + g (decays) of shape `[B, T, HV]`. Default: `None`. + When `use_gate_in_kernel=False` (default), `g` must be in log space (pre-computed decay). + When `use_gate_in_kernel=True`, `g` is the raw pre-activation input; the kernel fuses + `-exp(A_log) * softplus(g + dt_bias)` internally per step. + gk (torch.Tensor): + gk (decays) of shape `[B, T, HV, K]`. Default: `None`. + gv (torch.Tensor): + gv (decays) of shape `[B, T, HV, V]`. Default: `None`. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + 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, K, V]` 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, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + Whether to use L2 normalization in the kernel. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space GDN decay internally. + When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses + gate activation into the recurrence. Default: `False`. + A_log (Optional[torch.Tensor]): + Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + dt_bias (Optional[torch.Tensor]): + Bias added to `g` before activation, of shape `[HV]`. + Only used when `use_gate_in_kernel=True`. + use_beta_sigmoid_in_kernel (Optional[bool]): + Whether to apply `torch.sigmoid(beta)` inside the kernel. + - If `True`, the passed `beta` acts as the raw beta logits. + - If `False`, `beta` is expected to already be in post-sigmoid space. + Default: `False`. + allow_neg_eigval (Optional[bool]): + Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. + Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case + the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. + state_v_first (Optional[bool]): + Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + cu_seqlens (torch.LongTensor): + 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, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, K, V]` 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 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, K, V, device='cuda') + >>> o, ht = fused_gated_recurrent_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, 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.long) + >>> o, ht = fused_gated_recurrent_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + if 'transpose_state_layout' in kwargs: + if state_v_first: + raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") + warnings.warn( + "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", + DeprecationWarning, + stacklevel=2, + ) + state_v_first = kwargs.pop('transpose_state_layout') + + 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 + if beta is None: + beta = torch.ones_like(q[..., 0]) + if use_gate_in_kernel: + if A_log is None: + raise ValueError("`A_log` must be provided when `use_gate_in_kernel=True`.") + if g is None: + raise ValueError("`g` (raw pre-activation) must be provided when `use_gate_in_kernel=True`.") + else: + A_log = None + dt_bias = None + if allow_neg_eigval and not use_beta_sigmoid_in_kernel: + raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") + + o, final_state = FusedRecurrentFunction.apply( + q, + k, + v, + g, + gk, + gv, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + use_beta_sigmoid_in_kernel, + allow_neg_eigval, + state_v_first, + cu_seqlens, + ) + return o, final_state + + +fused_recurrent_gdn = fused_recurrent_gated_delta_rule diff --git a/upstream_ref/fla/ops/gated_delta_rule/gate.py b/upstream_ref/fla/ops/gated_delta_rule/gate.py new file mode 100644 index 00000000..564177e1 --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/gate.py @@ -0,0 +1,344 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from fla.ops.backends import dispatch +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.index import prepare_chunk_indices +from fla.ops.utils.op import exp +from fla.ops.utils.softplus import softplus +from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + + +def naive_gdn_gate( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Torch reference implementation for GDN gate computation. + + Computes: ``g = -A_log.exp() * softplus(g + dt_bias)`` + + Args: + g (torch.Tensor): + Input tensor of shape `[..., HV]`. + A_log (torch.Tensor): + Decay parameter tensor with `HV` elements. + dt_bias (torch.Tensor | None): + Optional bias tensor added to `g` before activation, shape `[HV]`. + + Returns: + Output tensor of shape `[..., HV]`. + """ + g = g.float() + if dt_bias is not None: + g = g + dt_bias.float() + return (-A_log.float().exp() * F.softplus(g)).to(output_dtype) + + +@triton.heuristics({ + 'HAS_BIAS': lambda args: args['dt_bias'] is not None, + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['H', 'BT', 'IS_VARLEN', 'REVERSE'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def gdn_gate_chunk_cumsum_scalar_kernel( + g, + A_log, + dt_bias, + o, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + REVERSE: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_g = g + bos * H + i_h + o_t * H + p_o = o + bos * H + i_h + o_t * H + + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + if HAS_BIAS: + b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) + b_A = tl.load(A_log + i_h).to(tl.float32) + b_gate = -exp(b_A) * softplus(b_g) + + b_o = tl.cumsum(b_gate, axis=0) + if REVERSE: + b_z = tl.sum(b_gate, axis=0) + b_o = -b_o + b_z[None] + b_gate + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t) + + +@triton.heuristics({ + 'HAS_BIAS': lambda args: args['dt_bias'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8] + ], + key=['H', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def gdn_gate_bwd_kernel( + g, + A_log, + dt_bias, + dyg, + dg, + dA, + T, + H: tl.constexpr, + BT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) + + b_A = tl.load(A_log + i_h).to(tl.float32) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_g = g + i_h + o_t * H + p_dg = dg + i_h + o_t * H + p_dyg = dyg + i_h + o_t * H + + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + b_dyg = tl.load(p_dyg, mask=m_t, other=0.0).to(tl.float32) + + if HAS_BIAS: + b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) + + # gate = -exp(A_log) * softplus(g + bias) + # d(gate)/d(g) = -exp(A_log) * sigmoid(g + bias) (softplus' = sigmoid) + # d(gate)/d(A_log) = -exp(A_log) * softplus(g + bias) = gate + b_neg_expA = -exp(b_A) + b_yg = b_neg_expA * softplus(b_g) + b_dg = b_neg_expA * (b_dyg * tl.sigmoid(b_g)) + b_dA = tl.sum(b_dyg * b_yg, 0) + + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) + tl.store(dA + i_t * H + i_h, b_dA) + + +@input_guard +@dispatch('gated_delta_rule') +def gdn_gate_chunk_cumsum( + g: torch.Tensor, + A_log: torch.Tensor, + chunk_size: int, + scale: float = None, + dt_bias: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + output_dtype: torch.dtype | None = torch.float, +) -> torch.Tensor: + B, T, H = g.shape + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + o = torch.empty_like(g, dtype=output_dtype or g.dtype) + gdn_gate_chunk_cumsum_scalar_kernel[(NT, B * H)]( + g=g, + A_log=A_log, + dt_bias=dt_bias, + o=o, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + REVERSE=False, + ) + return o + + +@dispatch('gated_delta_rule') +def gdn_gate_bwd( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + dyg: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + H = g.shape[-1] + T = g.numel() // H + BT = 32 + NT = triton.cdiv(T, BT) + + dg = torch.empty_like(g, dtype=torch.float32) + dA = A_log.new_empty(NT, H, dtype=torch.float32) + + gdn_gate_bwd_kernel[(NT, H)]( + g=g, + A_log=A_log, + dt_bias=dt_bias, + dyg=dyg, + dg=dg, + dA=dA, + T=T, + H=H, + BT=BT, + ) + + dg = dg.view_as(g).type_as(g) + dA = dA.sum(0).view_as(A_log).type_as(A_log) + dbias = dg.view(-1, H).sum(0).to(dt_bias) if dt_bias is not None else None + + return dg, dA, dbias + + +@triton.heuristics({ + 'HAS_BIAS': lambda args: args['dt_bias'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) + for BT in [32, 64, 128] + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3] + ], + key=['H'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def gdn_gate_fwd_kernel( + g, + A_log, + dt_bias, + yg, + T, + H: tl.constexpr, + BT: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1) + + b_A = tl.load(A_log + i_h).to(tl.float32) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_g = g + i_h + o_t * H + p_yg = yg + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + if HAS_BIAS: + b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) + b_yg = -exp(b_A) * softplus(b_g) + tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), mask=m_t) + + +@dispatch('gated_delta_rule') +def gdn_gate_fwd( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + H = g.shape[-1] + T = g.numel() // H + + yg = torch.empty_like(g, dtype=output_dtype) + + def grid(meta): + return (triton.cdiv(T, meta['BT']), H) + + gdn_gate_fwd_kernel[grid]( + g=g, + A_log=A_log, + dt_bias=dt_bias, + yg=yg, + T=T, + H=H, + ) + return yg + + +class GDNGateFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + yg = gdn_gate_fwd(g=g, A_log=A_log, dt_bias=dt_bias, output_dtype=output_dtype) + ctx.save_for_backward(g, A_log, dt_bias) + return yg + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, dyg: torch.Tensor): + g, A_log, dt_bias = ctx.saved_tensors + dg, dA, dbias = gdn_gate_bwd(g=g, A_log=A_log, dt_bias=dt_bias, dyg=dyg) + return dg, dA, dbias, None + + +@torch.compiler.disable +def fused_gdn_gate( + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + r""" + Fused GDN gate computation with autograd support. + + Computes: ``g = -A_log.exp() * softplus(g + dt_bias)`` + + Args: + g (torch.Tensor): + Input tensor of shape `[..., HV]`. + A_log (torch.Tensor): + Decay parameter tensor with `HV` elements. + dt_bias (torch.Tensor | None): + Optional bias tensor added to `g` before activation, shape `[HV]`. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32`. + + Returns: + Output tensor of shape `[..., HV]`. + """ + return GDNGateFunction.apply(g, A_log, dt_bias, output_dtype) diff --git a/upstream_ref/fla/ops/gated_delta_rule/naive.py b/upstream_ref/fla/ops/gated_delta_rule/naive.py new file mode 100644 index 00000000..cd0cf0d1 --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/naive.py @@ -0,0 +1,161 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import torch +import torch.nn.functional as F +from einops import rearrange + + +def naive_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + """ + Reference PyTorch implementation of recurrent gated delta rule. + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, H, V] + beta: [B, T, H] + g: [B, T, H] + scale: float, optional + initial_state: [B, H, K, V], optional + output_final_state: bool + + Returns: + o: [B, T, H, V] + final_state: [B, H, K, V] if output_final_state else None + """ + q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]) + B, H, T, K, V = *k.shape, v.shape[-1] + o = torch.zeros(B, H, T, V).to(v) + h = torch.zeros(B, H, K, V).to(v) + if initial_state is not None: + h = initial_state.to(torch.float32) + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + q = q * scale + + for i in range(T): + b_q = q[:, :, i] + b_k = k[:, :, i] + b_v = v[:, :, i].clone() + h = h.clone() * g[:, :, i].exp()[..., None, None] + b_beta = beta[:, :, i] + b_v = b_v - (h.clone() * b_k[..., None]).sum(-2) + b_v = b_v * b_beta[..., None] + h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2) + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h) + + if not output_final_state: + h = None + o = o.transpose(1, 2).contiguous() + return o, h + + +def naive_chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = 64, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, +): + """ + Reference PyTorch implementation of chunk gated delta rule. + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, H, V] + g: [B, T, H] + beta: [B, T, H] + chunk_size: int + scale: float, optional + initial_state: [B, H, K, V], optional + output_final_state: bool + + Returns: + o: [B, T, H, V] + final_state: [B, H, K, V] if output_final_state else None + """ + BT = chunk_size + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + + q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]) + + T = q.shape[-2] + pad_len = (BT - (T % BT)) % BT + if pad_len > 0: + q = F.pad(q, (0, 0, 0, pad_len)) + k = F.pad(k, (0, 0, 0, pad_len)) + v = F.pad(v, (0, 0, 0, pad_len)) + beta = F.pad(beta, (0, pad_len)) + g = F.pad(g, (0, pad_len)) + + q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g]) + decay = g + chunk_size = BT + b, h, l, d_k = q.shape + d_v = v.shape[-1] + q = q * scale + v = v * beta[..., None] + k_beta = k * beta[..., None] + assert l % chunk_size == 0 + + # note that diagonal is masked. + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) + q, k, v, k_beta, decay = map( + lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), + [q, k, v, k_beta, decay.unsqueeze(-1)], + ) + decay = decay.squeeze(-1).cumsum(-1) + decay_exp = decay.exp()[..., None] + L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() + attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0) + for i in range(1, chunk_size): + attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) + attn = attn + k_cumsum = attn @ v + k_cumdecay = attn @ (k_beta * decay_exp) + v = k_cumsum + + S = k.new_zeros(b, h, d_k, d_v) + if initial_state is not None: + S = initial_state.to(torch.float32) + + o = torch.zeros_like(v) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) + for i in range(0, l // chunk_size): + q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = (k_cumdecay[:, :, i]) @ S + v_new = v_i - v_prime + o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S + o[:, :, i] = o_inter + attn @ v_new + S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp() + [..., None]).transpose(-1, -2) @ v_new + if not output_final_state: + S = None + + # unpad + o = rearrange(o, 'b h n c d -> b h (n c) d') + o = o[:, :, :T] + o = o.transpose(1, 2) + return o, S diff --git a/upstream_ref/fla/ops/gated_delta_rule/wy_fast.py b/upstream_ref/fla/ops/gated_delta_rule/wy_fast.py new file mode 100644 index 00000000..63821cd7 --- /dev/null +++ b/upstream_ref/fla/ops/gated_delta_rule/wy_fast.py @@ -0,0 +1,347 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import torch +import triton +import triton.language as tl + +from fla.ops.backends import dispatch +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.cache import fla_cache_autotune +from fla.ops.utils.op import exp2 +from fla.utils import IS_NVIDIA_BLACKWELL, autotune_cache_kwargs, check_shared_mem + +# Blackwell can select unstable Triton configs for prepare_wy_repr_bwd_kernel +# during autotuning (see #913). Restrict it to the config that has been +# validated on B200 until the wider config space is re-validated. +PREPARE_WY_REPR_BWD_NUM_WARPS = [2] if IS_NVIDIA_BLACKWELL else [2, 4] +PREPARE_WY_REPR_BWD_NUM_STAGES = [4] if IS_NVIDIA_BLACKWELL else [2, 3, 4] + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kernel( + k, + v, + beta, + w, + u, + A, + g, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_b = beta + bos*HV + i_h + o_t * HV + b_b = tl.load(p_b, mask=m_t, other=0.0) + + p_A = A + (bos*HV + i_h) * BT + o_t[:, None] * (HV*BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) + + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_u = u + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) + + if USE_G: + p_g = g + (bos*HV + i_h) + o_t * HV + b_g = exp2(tl.load(p_g, mask=m_t, other=0.0)) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_w = w + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_kb = b_k * b_b[:, None] + if USE_G: + b_kb *= b_g[:, None] + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in PREPARE_WY_REPR_BWD_NUM_WARPS + for num_stages in PREPARE_WY_REPR_BWD_NUM_STAGES + ], + key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_bwd_kernel( + k, + v, + beta, + g, + A, + dw, + du, + dk, + dv, + db, + dg, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_AT = (o_A[:, None] < BT) & m_t[None, :] + p_b = beta + (bos*HV + i_h) + o_t * HV + p_db = db + (bos*HV + i_h) + o_t * HV + p_A = A + (bos*HV + i_h) * BT + o_A[:, None] + o_t[None, :] * (HV*BT) + + b_b = tl.load(p_b, mask=m_t, other=0.0) + b_db = tl.zeros([BT], dtype=tl.float32) + b_A = tl.load(p_A, mask=m_AT, other=0.0) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + + if USE_G: + p_g = g + (bos*HV + i_h) + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_g_exp = exp2(b_g) + b_dg = tl.zeros([BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + p_dw = dw + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + # [BT, BK] + b_k = tl.load(p_k, mask=m_k, other=0.0) + if USE_G: + b_kbg = b_k * (b_b * b_g_exp)[:, None] + else: + b_kbg = b_k * b_b[:, None] + b_dw = tl.load(p_dw, mask=m_k, other=0.0) + + b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) + b_dkbg = tl.dot(b_A, b_dw) + if USE_G: + b_dk = b_dkbg * (b_g_exp * b_b)[:, None] + b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1) + b_dg += tl.sum(b_dkbg * b_kbg, 1) + else: + b_dk = b_dkbg * b_b[:, None] + b_db += tl.sum(b_dkbg * b_k, 1) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv = dv + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_du = du + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_du = tl.load(p_du, mask=m_v, other=0.0) + b_dA += tl.dot(b_du, tl.trans(b_vb)) + b_dvb = tl.dot(b_A, b_du) + b_dv = b_dvb * b_b[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA, 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + + if USE_G: + b_dA *= exp2(b_g[:, None] - b_g[None, :]) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + b_dA = tl.where(m_A, -b_dA, 0).to(k.dtype.element_ty) + + tl.debug_barrier() + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_kt = tl.trans(b_k) + b_kb = b_k * b_b[:, None] + + b_A += tl.dot(b_k, b_kt) + b_dkb = tl.dot(b_dA, b_k) + b_db += tl.sum(b_dkb * b_k, 1) + b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(tl.trans(b_kb).to(b_dA.dtype), b_dA)) + b_dk += tl.load(p_dk, mask=m_k, other=0.0) + + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_t) + + b_A *= b_b[:, None] + if USE_G: + b_AdA = b_dA * b_A + p_dg = dg + (bos*HV + i_h) + o_t * HV + b_dg += tl.sum(b_AdA, axis=1) - tl.sum(b_AdA, axis=0) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) + + +@dispatch('gated_delta_rule') +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + g: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = k.new_empty(B, T, HV, K) + u = torch.empty_like(v) + recompute_w_u_fwd_kernel[(NT, B*HV)]( + k=k, + v=v, + beta=beta, + w=w, + u=u, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u + + +@dispatch('gated_delta_rule') +def prepare_wy_repr_bwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + g: torch.Tensor = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] + BT = A.shape[-1] + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + dk = k.new_empty(B, T, HV, K) + dv = torch.empty_like(v) + dg = torch.empty_like(g) if g is not None else None + db = torch.empty_like(beta) + prepare_wy_repr_bwd_kernel[(NT, B * HV)]( + k=k, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=du, + dk=dk, + dv=dv, + db=db, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + if H != HV: + dk = dk.view(B, T, H, HV // H, K).sum(3) + return dk, dv, db, dg + + +fwd_recompute_w_u = recompute_w_u_fwd +bwd_prepare_wy_repr = prepare_wy_repr_bwd diff --git a/upstream_ref/vllm_gdn/gdn/__init__.py b/upstream_ref/vllm_gdn/gdn/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/upstream_ref/vllm_gdn/gdn/base.py b/upstream_ref/vllm_gdn/gdn/base.py new file mode 100644 index 00000000..95ced2e0 --- /dev/null +++ b/upstream_ref/vllm_gdn/gdn/base.py @@ -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, + ) diff --git a/upstream_ref/vllm_gdn/gdn/kimi_gdn_linear_attn.py b/upstream_ref/vllm_gdn/gdn/kimi_gdn_linear_attn.py new file mode 100644 index 00000000..bc49226c --- /dev/null +++ b/upstream_ref/vllm_gdn/gdn/kimi_gdn_linear_attn.py @@ -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)) diff --git a/upstream_ref/vllm_gdn/gdn/olmo_gdn_linear_attn.py b/upstream_ref/vllm_gdn/gdn/olmo_gdn_linear_attn.py new file mode 100644 index 00000000..a72707c7 --- /dev/null +++ b/upstream_ref/vllm_gdn/gdn/olmo_gdn_linear_attn.py @@ -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 diff --git a/upstream_ref/vllm_gdn/gdn/qwen_gdn_linear_attn.py b/upstream_ref/vllm_gdn/gdn/qwen_gdn_linear_attn.py new file mode 100644 index 00000000..60879f8f --- /dev/null +++ b/upstream_ref/vllm_gdn/gdn/qwen_gdn_linear_attn.py @@ -0,0 +1,1751 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Qwen3-Next/Qwen3.5 model.""" + +from typing import Literal + +import torch +from einops import rearrange +from torch import nn + +from vllm import envs +from vllm._aiter_ops import rocm_aiter_ops +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.logger import init_logger +from vllm.model_executor.custom_op import CustomOp, 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_mixer2 import mamba_v2_sharded_weight_loader +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.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig +from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig +from vllm.model_executor.layers.quantization.inc import INCConfig +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 as fla_chunk_gated_delta_rule, +) +from vllm.third_party.flash_linear_attention.ops import ( + fused_post_conv_prep, + fused_recurrent_gated_delta_rule_packed_decode, + fused_sigmoid_gating_delta_rule_update, +) +from vllm.third_party.flash_linear_attention.ops.chunk import l2norm_fwd +from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE +from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata + +# Optional ROCm AITER Triton kernels for the GDN decode path. +# Availability is checked centrally via rocm_aiter_ops; the actual function +# references are imported here so that they can be called without per-call +# import overhead. +GDN_AITER_TRITON_AVAILABLE = ( + rocm_aiter_ops.are_gdn_triton_kernels_available() + or rocm_aiter_ops.is_rdna_gdn_triton_kernels_available() +) + +if GDN_AITER_TRITON_AVAILABLE: + from aiter.ops.triton.causal_conv1d_update_single_token import ( + fused_reshape_causal_conv1d_update_single_token as gdn_aiter_fused_reshape_causal_conv1d_update_single_token, # noqa: E501 + ) + from aiter.ops.triton.gated_delta_net.fused_rearrange_sigmoid_gdr import ( + fused_rearrange_sigmoid_gated_delta_rule as gdn_aiter_fused_rearrange_sigmoid_gated_delta_rule, # noqa: E501 + ) + +logger = init_logger(__name__) + + +def _resolve_gdn_prefill_backend( + vllm_config: VllmConfig, +) -> tuple[str, Literal["triton", "flashinfer", "cutedsl"]]: + """Resolve GDN prefill backend. + + FlashInfer's GDN prefill kernel is chosen when: + * ``requested in ["flashinfer", "auto"]``; + * ``platform == cuda``; + * one of the following: + - Hopper (SM90) — no further constraints; + - Blackwell (SM10.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``. + + In-tree CuteDSL GDN prefill kernel is chosen when: + * "cutedsl" is requested; (opt-in only) + * Blackwell (SM10.x) with ``head_k_dim == 128``; + """ + additional_config = vllm_config.additional_config + backend_cfg = ( + additional_config.get("gdn_prefill_backend", "auto") + if isinstance(additional_config, dict) + else "auto" + ) + backend = str(backend_cfg).strip().lower() + + if not current_platform.is_cuda(): + return backend, "triton" + + head_k_dim = getattr( + vllm_config.model_config.hf_text_config, "linear_key_head_dim", None + ) + + supports_flashinfer = False + supports_cutedsl = False + + if current_platform.is_device_capability(90): + supports_flashinfer = True + elif ( + current_platform.is_device_capability_family(100) + and head_k_dim == 128 + and current_platform.get_cuda_runtime_major() >= 13 + ): + supports_flashinfer = True + supports_cutedsl = True + + if backend in ["flashinfer", "auto"] and supports_flashinfer: + return backend, "flashinfer" + if backend == "cutedsl" and supports_cutedsl: + return backend, "cutedsl" + return backend, "triton" + + +def _log_gdn_backend_decision( + vllm_config: VllmConfig, + requested_backend: str, + active_backend: str, +) -> None: + """Log the GDN prefill backend choice in the attention-selector style.""" + head_k_dim = getattr( + vllm_config.model_config.hf_text_config, "linear_key_head_dim", None + ) + chosen = { + "flashinfer": "FlashInfer", + "cutedsl": "CuteDSL", + "triton": "Triton/FLA", + }[active_backend] + logger.info_once( + "Using %s GDN prefill kernel (requested=%s, head_k_dim=%s).", + chosen, + requested_backend, + head_k_dim, + ) + if active_backend == "flashinfer" and current_platform.is_device_capability(90): + logger.warning_once( + "FlashInfer GDN prefill is JIT-compiled; first run may take a " + "while. Set --gdn-prefill-backend triton to skip JIT.", + ) + + +def fi_chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = True, +): + from flashinfer.gdn_prefill import ( + chunk_gated_delta_rule as chunk_gated_delta_rule_fi, + ) + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + + # use flashinfer implementation + q = q.squeeze(0).contiguous() + k = k.squeeze(0).contiguous() + v = v.squeeze(0).contiguous() + + g = g.squeeze(0).contiguous() + beta = beta.squeeze(0).contiguous() + fi_state = initial_state.to(torch.float32) + fi_g = g.to(torch.float32) + fi_beta = beta.to(torch.float32) + result = chunk_gated_delta_rule_fi( + q=q, + k=k, + v=v, + g=torch.exp(fi_g), + beta=fi_beta, + initial_state=fi_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + # FlashInfer returns (output, state) when output_final_state=True, + # or just output when output_final_state=False. + # Unsqueeze back to 4D (1, L, H, D) to match fla output format + if output_final_state: + output, final_state = result + return output.unsqueeze(0), final_state + else: + return result.unsqueeze(0), None + + +@CustomOp.register("chunk_gated_delta_rule") +class ChunkGatedDeltaRule(CustomOp): + def __init__(self) -> None: + super().__init__() + vllm_config = get_current_vllm_config() + backend, active_backend = _resolve_gdn_prefill_backend(vllm_config) + self.gdn_prefill_backend = active_backend + + if backend in ("flashinfer", "cutedsl") and active_backend != backend: + logger.warning_once( + "GDN prefill backend '%s' is selected but cannot use this " + "kernel on the current platform. Falling back to Triton/FLA.", + backend, + ) + _log_gdn_backend_decision(vllm_config, backend, active_backend) + + if active_backend == "flashinfer": + self._forward_method = self.forward_cuda + elif active_backend == "cutedsl": + self._forward_method = self.forward_cutedsl + else: + self._forward_method = self.forward_native + + def forward_cuda( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + 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 = True, + core_attn_out: torch.Tensor | None = None, + ): + o, final_state = fi_chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + if core_attn_out is not None: + o_flat = o.squeeze(0).reshape(-1) + co_flat = core_attn_out.reshape(-1) + co_flat[: o_flat.numel()].copy_(o_flat) + return o, final_state + + def forward_native( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + 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 = True, + core_attn_out: torch.Tensor | None = None, + ): + return fla_chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + core_attn_out=core_attn_out, + ) + + def forward_cutedsl( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + 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 = True, + core_attn_out: torch.Tensor | None = None, + ): + from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( + chunk_gated_delta_rule_cutedsl, + ) + + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + + assert cu_seqlens is not None + assert chunk_indices is not None + assert chunk_offsets is not None + + o, final_state = chunk_gated_delta_rule_cutedsl( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + core_attn_out=core_attn_out, + ) + if not output_final_state: + final_state = None + return o, final_state + + +@PluggableLayer.register("qwen_gated_delta_net_attention") +class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): + 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: Qwen3NextConfig, + vllm_config: VllmConfig, + prefix: str = "", + gqa_interleaved_layout=False, + reduce_results: bool = True, + ) -> None: + super().__init__(config, vllm_config, prefix) + + 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.gqa_interleaved_layout = gqa_interleaved_layout + if current_platform.is_xpu(): + self._forward_method = self.forward_xpu + elif current_platform.is_cpu(): + from vllm.model_executor.layers.mamba.ops.cpu.gdn_attention import ( + register_cpu_gdn_attention_ops, + ) + + register_cpu_gdn_attention_ops() + self._forward_method = self.forward_cpu + elif current_platform.is_rocm(): + self._forward_method = self.forward_hip + else: + self._forward_method = self.forward_cuda + + # QKV + 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) + + # projection of the input hidden states + # Qwen3-Next and Qwen3.5 has a different qkv_proj layout, + # we need to create qkvz_proj adaptively here. + # When create_in_proj_qkvz is False (e.g. LoRA enabled in Qwen3.5), + # in_proj_qkv and in_proj_z are created separately instead. + self.in_proj_qkvz = self.create_qkvz_proj( + hidden_size=self.hidden_size, + key_dim=self.key_dim, + value_dim=self.value_dim, + quant_config=self.quant_config, + prefix=f"{prefix}.in_proj_qkvz", + ) + + # ba_proj doesn't support blockwise fp8 quantization. + # Qwen3-Next and Qwen3.5 have different in_proj_ba checkpoint + # layouts, so we use a factory method to create the projection. + self.in_proj_ba = self.create_ba_proj( + hidden_size=self.hidden_size, + num_v_heads=self.num_v_heads, + quant_config=self.quant_config, + prefix=f"{prefix}.in_proj_ba", + ) + self.disable_tp_for_ba_proj = self.maybe_disable_tp(self.quant_config) + + query_key_settings = (self.key_dim, 0, False) + value_settings = (self.value_dim, 0, False) + + self.conv1d.weight.weight_loader = mamba_v2_sharded_weight_loader( + [ + query_key_settings, + query_key_settings, + value_settings, + ], + self.tp_size, + self.tp_rank, + ) + + # selective projection used to make dt, B and C input dependent + + # time step projection (discretization) + # instantiate once and copy inv_dt in init_weights of PretrainedModel + 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), + dtype=torch.float32, + ) + ) + + set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)}) + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) + + output_gate_type = getattr(config, "output_gate_type", "silu") + if output_gate_type == "swish": + output_gate_type = "silu" + assert output_gate_type in ["silu", "swish", "sigmoid"], ( + f"unsupported {output_gate_type=}" + ) + + self.norm = RMSNormGated( + self.head_v_dim, + eps=self.layer_norm_epsilon, + group_size=None, + norm_before_gate=True, + activation=output_gate_type, + device=current_platform.current_device(), + ) + + self.out_proj = RowParallelLinear( + self.value_dim, + self.hidden_size, + bias=False, + input_is_parallel=True, + reduce_results=reduce_results, + quant_config=self.quant_config, + prefix=f"{prefix}.out_proj", + ) + + self.chunk_gated_delta_rule = ChunkGatedDeltaRule() + self.gdn_prefill_backend = self.chunk_gated_delta_rule.gdn_prefill_backend + self._prefill_kernels_warmed_up = False + self.enable_packed_recurrent_decode = ( + envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE + ) + + 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 create_qkvz_proj( + self, + hidden_size: int, + key_dim: int, + value_dim: int, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> MergedColumnParallelLinear: + # When gqa_interleaved_layout=True (Qwen3-Next), qkvz weights are + # stored as a single fused tensor with interleaved GQA layout, so we + # use one output shard to preserve the interleaving across TP ranks. + # When gqa_interleaved_layout=False (Qwen3.5), the checkpoint has + # separate q, k, v, z weights, so we use 4 independent output sizes. + output_sizes = ( + [sum((key_dim, key_dim, value_dim, value_dim))] + if self.gqa_interleaved_layout + else [key_dim, key_dim, value_dim, value_dim] + ) + return MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=output_sizes, + bias=False, + quant_config=quant_config, + prefix=prefix, + ) + + def create_ba_proj( + self, + hidden_size: int, + num_v_heads: int, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> MergedColumnParallelLinear: + # When gqa_interleaved_layout=True (Qwen3-Next), in_proj_ba is stored + # as a single fused weight [b_g0, a_g0, b_g1, a_g1, ...] interleaved + # by key-head group; a single output shard preserves this across TP. + # When gqa_interleaved_layout=False (Qwen3.5), in_proj_b and in_proj_a + # are separate checkpoint weights, so we use 2 independent output sizes. + output_sizes = ( + [num_v_heads * 2] if self.gqa_interleaved_layout else [num_v_heads] * 2 + ) + return MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=output_sizes, + bias=False, + quant_config=quant_config, + prefix=prefix, + disable_tp=self.maybe_disable_tp(quant_config), + ) + + def maybe_disable_tp(self, quant_config: QuantizationConfig | None) -> bool: + """Whether to replicate ba_proj instead of TP-sharding it. + + Marlin requires output_size_per_partition >= MIN_THREAD_N=64, which + the Qwen3.5 non-interleaved [num_v_heads]*2 layout violates at TP>=2 + (e.g. num_v_heads=64, TP=4 -> 16). Replicating the projection keeps + each rank above the Marlin threshold; forward() then slices b/a to + the local TP partition. Qwen3-Next's interleaved [num_v_heads*2] + layout is unaffected and stays TP-sharded. + + See https://github.com/vllm-project/vllm/issues/35924 + """ + return ( + current_platform.is_cuda() + and not self.gqa_interleaved_layout + and isinstance(quant_config, (AutoAWQConfig, AutoGPTQConfig, INCConfig)) + ) + + def split_ba(self, ba: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + b, a = ba.chunk(2, dim=-1) + if self.disable_tp_for_ba_proj and self.tp_size > 1: + # ba_proj is replicated for Marlin; slice b/a to local TP rank. + ba_chunk = self.num_v_heads // self.tp_size + ba_start = self.tp_rank * ba_chunk + b = b[:, ba_start : ba_start + ba_chunk] + a = a[:, ba_start : ba_start + ba_chunk] + return b, a + + def fix_query_key_value_ordering( + self, + mixed_qkvz: torch.Tensor, + mixed_ba: torch.Tensor, + ): + """ + Derives `query`, `key` and `value` tensors from `mixed_qkvzba`. + """ + new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + ( + self.num_k_heads // self.tp_size, + ( + self.head_k_dim + + self.head_k_dim + + (self.head_v_dim + self.head_v_dim) + * self.num_v_heads + // self.num_k_heads + ), + ) + new_tensor_shape_ba = mixed_ba.size()[:-1] + ( + self.num_k_heads // self.tp_size, + 2 * self.num_v_heads // self.num_k_heads, + ) + + mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz) + mixed_ba = mixed_ba.view(*new_tensor_shape_ba) + + split_arg_list_qkvz = [ + self.head_k_dim, + self.head_k_dim, + (self.num_v_heads // self.num_k_heads * self.head_v_dim), + (self.num_v_heads // self.num_k_heads * self.head_v_dim), + ] + split_arg_list_ba = [ + self.num_v_heads // self.num_k_heads, + self.num_v_heads // self.num_k_heads, + ] + + # [b, sq, ng, (hn + hn + np/ng * hn + np/ng + np/ng)] + # --> [b, sq, ng, hn], [b, sq, ng, hn], [b, sq, ng, np/ng * hn], + # [b, sq, ng, np/ng * hn], [b, sq, ng, np/ng], [b, sq, ng, np/ng] + (query, key, value, z) = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=2) + (b, a) = torch.split(mixed_ba, split_arg_list_ba, dim=2) + + # [b, sq, ng, np/ng * hn] -> [b, sq, np, hn] + value = value.reshape(value.size(0), -1, self.head_v_dim) + z = z.reshape(z.size(0), -1, self.head_v_dim) + b = b.reshape(b.size(0), self.num_v_heads // self.tp_size) + a = a.reshape(a.size(0), self.num_v_heads // self.tp_size) + + return query, key, value, z, b, a + + @torch.compile(fullgraph=True) + def prepare_gdn_attention_core_inputs( + self, + mixed_qkvz: torch.Tensor, + mixed_ba: torch.Tensor, + num_tokens: int, + ): + """ + Derives mixed_qkv, z, b, a from projected qkvz/ba for the GDN custom op. + + For gqa_interleaved_layout (Qwen3-Next): unpack the interleaved + [ng, (hk + hk + np/ng*hv + np/ng*hv)] layout into contiguous qkv. + For non-interleaved layout (Qwen3.5): simple split along last dim. + """ + if not self.gqa_interleaved_layout: + # Qwen3.5: weights are in [q, k, v, z] order + assert num_tokens == mixed_qkvz.shape[0] + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size + z_size = self.value_dim // self.tp_size + mixed_qkv, z_flat = mixed_qkvz.split([qkv_size, z_size], dim=-1) + n = mixed_qkvz.shape[0] + z_out = z_flat.reshape(n, -1, self.head_v_dim) + b, a = mixed_ba.chunk(2, dim=-1) + return mixed_qkv, z_out, b, a + + # Qwen3-Next: interleaved GQA layout + base_shape_qkvz = mixed_qkvz.size()[:-1] + base_shape_ba = mixed_ba.size()[:-1] + ng = self.num_k_heads // self.tp_size + + new_tensor_shape_qkvz = base_shape_qkvz + ( + ng, + ( + self.head_k_dim + + self.head_k_dim + + (self.head_v_dim + self.head_v_dim) + * self.num_v_heads + // self.num_k_heads + ), + ) + new_tensor_shape_ba = base_shape_ba + ( + ng, + 2 * self.num_v_heads // self.num_k_heads, + ) + + mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz) + mixed_ba = mixed_ba.view(*new_tensor_shape_ba) + + split_arg_list_qkvz = [ + self.head_k_dim, + self.head_k_dim, + (self.num_v_heads // self.num_k_heads * self.head_v_dim), + (self.num_v_heads // self.num_k_heads * self.head_v_dim), + ] + split_arg_list_ba = [ + self.num_v_heads // self.num_k_heads, + self.num_v_heads // self.num_k_heads, + ] + + (query, key, value, z) = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=-1) + (b, a) = torch.split(mixed_ba, split_arg_list_ba, dim=-1) + + mixed_qkv_logical = torch.cat( + [ + query.reshape(num_tokens, -1), + key.reshape(num_tokens, -1), + value.reshape(num_tokens, -1), + ], + dim=-1, + ) + + # The split above produces non-contiguous views into the interleaved + # buffer. Concatenating everything into a single flat tensor forces a + # contiguous copy, then slicing back out gives contiguous q/k/v/z/b/a + # tensors that downstream kernels require. Doing this in one cat+slice + # keeps torch.compile in a single Triton graph instead of emitting + # separate copy kernels per tensor. The original code used + # rearrange(...).contiguous() on each tensor individually. + fused = torch.cat( + [ + mixed_qkv_logical.reshape(-1), + z.reshape(-1), + b.reshape(-1), + a.reshape(-1), + ], + dim=0, + ) + + curr = 0 + qkv_numel = mixed_qkv_logical.numel() + z_numel = z.numel() + b_numel = b.numel() + a_numel = a.numel() + + mixed_qkv_out = fused[curr : curr + qkv_numel].view(num_tokens, -1) + curr += qkv_numel + + z_out = fused[curr : curr + z_numel].view( + num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim + ) + curr += z_numel + + b_out = fused[curr : curr + b_numel].view( + num_tokens, self.num_v_heads // self.tp_size + ) + curr += b_numel + + a_out = fused[curr : curr + a_numel].view( + num_tokens, self.num_v_heads // self.tp_size + ) + + return mixed_qkv_out, z_out, b_out, a_out + + def rearrange_mixed_qkv(self, mixed_qkv): + """Split packed qkv into contiguous (1, seq, heads, dim) tensors. + + The original code used ``rearrange(x, "l (h d) -> 1 l h d", d=...)`` + followed by ``.contiguous()`` on each tensor. This version flattens + all three splits into a single buffer via ``torch.cat`` so that + torch.compile emits one Triton copy kernel instead of three separate + contiguous() calls. + """ + if mixed_qkv is None: + return None, None, None + + seq_len = mixed_qkv.shape[0] + q_dim = self.key_dim // self.tp_size + k_dim = self.key_dim // self.tp_size + v_dim = self.value_dim // self.tp_size + + query, key, value = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) + + fused = torch.cat( + [query.reshape(-1), key.reshape(-1), value.reshape(-1)], dim=0 + ) + + q_size = seq_len * q_dim + k_size = seq_len * k_dim + + q_contig = fused[0:q_size] + k_contig = fused[q_size : q_size + k_size] + v_contig = fused[q_size + k_size :] + + query = q_contig.view(1, seq_len, -1, self.head_k_dim) + key = k_contig.view(1, seq_len, -1, self.head_k_dim) + value = v_contig.view(1, seq_len, -1, self.head_v_dim) + + return query, key, value + + def forward( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + return self._forward_method(hidden_states) + + def _output_projection( + self, + core_attn_out: torch.Tensor, + z: torch.Tensor, + ) -> torch.Tensor: + """Part 3: RMSNormGated + output linear projection. + + The RMSNormGated + quant sequence is eligible for fusion + by the compilation pass when fuse_norm_quant is enabled. + """ + z_shape_og = z.shape + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z = z.reshape(-1, z.shape[-1]) + core_attn_out = self.norm(core_attn_out, z) + core_attn_out = core_attn_out.reshape(z_shape_og) + core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) + output, _ = self.out_proj(core_attn_out) + return output + + def forward_hip( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """ROCm forward using AITER Triton fused projection+attention when + available, otherwise falling back to the generic CUDA path.""" + if GDN_AITER_TRITON_AVAILABLE: + num_tokens = hidden_states.size(0) + projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states) + projected_states_ba, _ = self.in_proj_ba(hidden_states) + projected_states_qkvz = projected_states_qkvz.view(num_tokens, -1) + projected_states_ba = projected_states_ba.view(num_tokens, -1) + core_attn_out = torch.empty( + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + z = torch.empty( + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim), + dtype=projected_states_qkvz.dtype, + device=projected_states_qkvz.device, + ) + + torch.ops.vllm.qwen_gdn_attention_core( + projected_states_qkvz, + projected_states_ba, + z, + core_attn_out, + layer_name=_encode_layer_name(self.prefix), + use_aiter=True, + ) + + return self._output_projection(core_attn_out, z) + else: + return self.forward_cuda(hidden_states) + + def forward_cuda( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """ + Forward pass with three parts: + 1. Input projection + 2. Core attention (custom op) + 3. Output projection + """ + num_tokens = hidden_states.size(0) + # ============================================================ + # Part 1: Input Projection + # ============================================================ + mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) + ba, _ = self.in_proj_ba(hidden_states) + + if self.gqa_interleaved_layout: + # Qwen3-Next: unpack the interleaved GQA layout + query, key, value, z, b, a = self.fix_query_key_value_ordering( + mixed_qkvz, ba + ) + query, key, value = map( + lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value) + ) + mixed_qkv = torch.cat((query, key, value), dim=-1) + else: + # Qwen3.5: weights are already in [q, k, v, z] and [b, a] order + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size + z_size = self.value_dim // self.tp_size + mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1) + z = z.reshape(z.size(0), -1, self.head_v_dim) + b, a = self.split_ba(ba) + b = b.contiguous() + a = a.contiguous() + + # ============================================================ + # Part 2: Core Attention (Custom Op) + # ============================================================ + # Note: we should not use torch.empty here like other attention backends, + # see discussions in https://github.com/vllm-project/vllm/pull/28182 + 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, + ) + + torch.ops.vllm.qwen_gdn_attention_core( + mixed_qkv, + b, + a, + core_attn_out, + layer_name=_encode_layer_name(self.prefix), + ) + + # ============================================================ + # Part 3: Output Projection + # ============================================================ + return self._output_projection(core_attn_out, z) + + def forward_xpu( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """ + Forward pass with three parts: + 1. Input projection + 2. Core attention (custom op) + 3. Output projection + """ + num_tokens = hidden_states.size(0) + + # ============================================================ + # Part 1: Input Projection + # ============================================================ + projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states) + projected_states_ba, _ = self.in_proj_ba(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, + ) + z = torch.empty_like(core_attn_out) + + torch.ops.vllm.gdn_attention_core_xpu( + core_attn_out, + z, + projected_states_qkvz, + projected_states_ba, + self.prefix, + ) + + # ============================================================ + # Part 3: Output Projection + # ============================================================ + z_shape_og = z.shape + # Reshape input data into 2D tensor + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z = z.reshape(-1, z.shape[-1]) + core_attn_out = self.norm(core_attn_out, z) + core_attn_out = core_attn_out.reshape(z_shape_og) + core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) + out, _ = self.out_proj(core_attn_out) + return out + + def forward_cpu( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on CPU." + + mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) + ba, _ = self.in_proj_ba(hidden_states) + + if self.gqa_interleaved_layout: + # Qwen3-Next: unpack the interleaved GQA layout + query, key, value, z, b, a = self.fix_query_key_value_ordering( + mixed_qkvz, ba + ) + query, key, value = map( + lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value) + ) + mixed_qkv = torch.cat((query, key, value), dim=-1) + else: + # Qwen3.5: weights are already in [q, k, v, z] and [b, a] order + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size + z_size = self.value_dim // self.tp_size + mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1) + z = z.reshape(z.size(0), -1, self.head_v_dim) + b, a = ba.chunk(2, dim=-1) + + num_tokens = hidden_states.size(0) + 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, + ) + + torch.ops.vllm.cpu_gdn_attention_core( + mixed_qkv, + b, + a, + core_attn_out, + _encode_layer_name(self.prefix), + ) + + z_shape_og = z.shape + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z = z.reshape(-1, z.shape[-1]) + core_attn_out = self.norm(core_attn_out, z) + core_attn_out = core_attn_out.reshape(z_shape_og) + core_attn_out = core_attn_out.flatten(-2) # ... h d -> ... (h d) + out, _ = self.out_proj(core_attn_out) + return out + + def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None: + """Warm up GDN prefill kernels during V1 profiling. + + During V1 profile runs, ``_forward_core`` returns early because + ``attn_metadata`` is ``None``, so the autotuned kernels used by + ``chunk_gated_delta_rule`` (e.g. ``solve_tril``, + ``chunk_scaled_dot_kkt``) are never invoked. After profiling, + vLLM allocates KV cache using most of the remaining GPU memory. + When the first real inference triggers the autotuner it OOMs + because there is not enough memory left for benchmarking. + + This method runs minimal forward passes through + ``chunk_gated_delta_rule`` with small dummy tensors to force + autotuning while GPU memory is still plentiful. The autotuner + results are cached globally, so only the first layer incurs + actual benchmarking cost. + + All kernels including ``chunk_fwd_kernel_o`` now use a fixed + ``BT = chunk_size`` (64). A single warmup pass with T = 64 + is sufficient to populate the autotuner cache. + + The decode path uses ``gdn_aiter_fused_rearrange_sigmoid_gated_delta_rule`` + which has fixed kernel parameters (no autotuning), so only the + prefill (chunked) path needs warming up. + """ + if self._prefill_kernels_warmed_up: + return + self._prefill_kernels_warmed_up = True + + device = qkv_or_qkvz.device + dtype = qkv_or_qkvz.dtype + num_k_heads = self.num_k_heads // self.tp_size + num_v_heads = self.num_v_heads // self.tp_size + _, state_dtype = self.get_state_dtype() + + # All kernels use BT = chunk_size, so a single pass with T = chunk_size + # is sufficient to populate every autotuner cache. Mirror the real + # prefill path here: build q/k/v/g/beta via fused_post_conv_prep and + # then run chunk_gated_delta_rule with in-kernel L2 norm disabled. + T = FLA_CHUNK_SIZE + dummy_mixed_qkv = torch.randn( + T, qkv_or_qkvz.shape[-1] - v_dim, device=device, dtype=dtype + ) + dummy_a = torch.randn(T, num_v_heads, device=device, dtype=dtype) + dummy_b = torch.randn(T, num_v_heads, device=device, dtype=dtype) + q, k, v, g, beta = fused_post_conv_prep( + conv_output=dummy_mixed_qkv, + a=dummy_a, + b=dummy_b, + A_log=self.A_log, + dt_bias=self.dt_bias, + num_k_heads=num_k_heads, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + apply_l2norm=True, + output_g_exp=False, + ) + q = q.unsqueeze(0) + k = k.unsqueeze(0) + v = v.unsqueeze(0) + g = g.unsqueeze(0) + beta = beta.unsqueeze(0) + state = torch.zeros( + 1, + num_v_heads, + self.head_v_dim, + self.head_k_dim, + device=device, + dtype=state_dtype, + ) + cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.int32) + + # CuteDSL kernels require metadata + chunk_indices = None + chunk_offsets = None + if self.gdn_prefill_backend == "cutedsl": + from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( + prepare_metadata_cutedsl, + ) + + chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, T) + + try: + self.chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=state, + output_final_state=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + use_qk_l2norm_in_kernel=False, + ) + except Exception: + logger.warning( + "GDN prefill kernel warmup (T=%d) failed for " + "layer %s. First inference may OOM due to " + "autotuner.", + T, + self.prefix, + exc_info=True, + ) + else: + logger.debug( + "GDN prefill kernel warmup (T=%d) completed for layer %s", + T, + self.prefix, + ) + finally: + del ( + dummy_mixed_qkv, + q, + k, + v, + dummy_a, + dummy_b, + g, + beta, + state, + cu_seqlens, + chunk_indices, + chunk_offsets, + ) + + torch.accelerator.empty_cache() + + def _forward_core_rocm( + self, + qkvz: torch.Tensor, + ba: torch.Tensor, + z_out: torch.Tensor, + core_attn_out: torch.Tensor, + ): + """ROCm AITER fast path: conv1d + recurrent attention from packed + qkvz/ba layout. + + For decode-only (no spec, no prefill) interleaved-GQA layouts, + dispatches directly to ``_forward_core_decode_aiter``. Otherwise unpacks + the packed layout and falls through to ``_forward_core``. + + Args: + qkvz: packed [q, k, v, z] projection (num_tokens, qkvz_dim) + ba: packed [b, a] gating vectors (num_tokens, 2*num_heads) + z_out: **output** buffer for z (num_tokens, num_heads, + head_dim); mutated in-place. + core_attn_out: Pre-allocated output buffer for attention results. + """ + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + + if attn_metadata_raw is None: + v_dim = core_attn_out.shape[-1] * core_attn_out.shape[-2] + self._warmup_prefill_kernels(qkvz, v_dim) + return + + assert isinstance(attn_metadata_raw, dict) + attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index] + assert isinstance(attn_metadata, GDNAttentionMetadata) + + # The AITER fused reshape/conv kernel expects Qwen3-Next's interleaved + # GQA layout. Qwen3.5 uses a non-interleaved q/k/v/z layout and must use + # the generic path below to split/rearrange inputs correctly. + if ( + self.gqa_interleaved_layout + and attn_metadata.spec_sequence_masks is None + and attn_metadata.num_prefills == 0 + and attn_metadata.num_decodes > 0 + ): + return self._forward_core_decode_aiter( + qkvz=qkvz, + ba=ba, + z_out=z_out, + core_attn_out=core_attn_out, + attn_metadata=attn_metadata, + ) + + core_attn_out.zero_() + num_tokens_all = qkvz.shape[0] + mixed_qkv, z, b, a = self.prepare_gdn_attention_core_inputs( + qkvz, ba, num_tokens_all + ) + z_out[:] = z + self._forward_core( + mixed_qkv=mixed_qkv, + b=b, + a=a, + core_attn_out=core_attn_out, + ) + + def _forward_core( + self, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + core_attn_out: torch.Tensor, + ): + """Core conv1d + recurrent attention (standard path). + + Args: + mixed_qkv: packed [q, k, v] projection (num_tokens, qkv_dim) + b: beta gating vector (num_tokens, num_heads) + a: alpha gating vector (num_tokens, num_heads) + core_attn_out: Pre-allocated output buffer for attention results. + """ + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + + if attn_metadata_raw is None: + self._warmup_prefill_kernels(mixed_qkv, 0) + return + + assert isinstance(attn_metadata_raw, dict) + attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index] + assert isinstance(attn_metadata, GDNAttentionMetadata) + + if ( + self.enable_packed_recurrent_decode + and attn_metadata.spec_sequence_masks is None + and attn_metadata.num_prefills == 0 + and attn_metadata.num_decodes > 0 + ): + return self._forward_core_decode_non_spec( + mixed_qkv=mixed_qkv, + b=b, + a=a, + core_attn_out=core_attn_out, + attn_metadata=attn_metadata, + ) + + 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 # noqa: E501 + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501 + 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] + + # 1. Convolution sequence transformation + 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 + + # 1.1: Process the multi-query part + if spec_sequence_masks is not None: + # spec_state_indices_tensor is always set when spec_sequence_masks is set + assert spec_state_indices_tensor is not None + mixed_qkv_spec = causal_conv1d_update( + mixed_qkv_spec, + conv_state, + conv_weights, + self.conv1d.bias, + self.activation, + conv_state_indices=spec_state_indices_tensor[:, 0][ # type: ignore[index] + : attn_metadata.num_spec_decodes # type: ignore[attr-defined] + ], + 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, + ) + + # 1.2: Process the remaining part + 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) + # - "cache_indices" updates the conv_state cache in positions + # pointed to by "state_indices_tensor" + mixed_qkv_non_spec = causal_conv1d_fn( + mixed_qkv_non_spec_T, + conv_weights, + self.conv1d.bias, + 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 mixed_qkv_non_spec is not None + mixed_qkv_non_spec = causal_conv1d_update( + mixed_qkv_non_spec, + conv_state, + conv_weights, + self.conv1d.bias, + self.activation, + conv_state_indices=non_spec_state_indices_tensor[ # type: ignore[index] + : attn_metadata.num_actual_tokens # type: ignore[attr-defined] + ], + validate_data=True, + ) + else: + mixed_qkv_non_spec = None + + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec) + + # Split mixed non-spec-decode+prefill to process independently + split_non_spec = ( + spec_sequence_masks is None + and attn_metadata.num_prefills > 0 + and attn_metadata.num_decodes > 0 + ) + num_decode_tokens = attn_metadata.num_decode_tokens + + if attn_metadata.num_prefills > 0: + assert mixed_qkv_non_spec is not None, ( + "mixed_qkv_non_spec must be provided for prefill path" + ) + if spec_sequence_masks is not None: + a_non_spec = a.index_select(0, non_spec_token_indx) + b_non_spec = b.index_select(0, non_spec_token_indx) + else: + a_non_spec = a + b_non_spec = b + + if split_non_spec: + conv_output_prefill = mixed_qkv_non_spec[num_decode_tokens:] + a_prefill = a_non_spec[num_decode_tokens:] + b_prefill = b_non_spec[num_decode_tokens:] + else: + conv_output_prefill = mixed_qkv_non_spec + a_prefill = a_non_spec + b_prefill = b_non_spec + + ( + query_non_spec, + key_non_spec, + value_non_spec, + g_non_spec, + beta_non_spec, + ) = fused_post_conv_prep( + conv_output=conv_output_prefill, + a=a_prefill, + b=b_prefill, + A_log=self.A_log, + dt_bias=self.dt_bias, + num_k_heads=self.num_k_heads // self.tp_size, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + apply_l2norm=True, + output_g_exp=False, + ) + query_non_spec = query_non_spec.unsqueeze(0) + key_non_spec = key_non_spec.unsqueeze(0) + value_non_spec = value_non_spec.unsqueeze(0) + g_non_spec = g_non_spec.unsqueeze(0) + beta_non_spec = beta_non_spec.unsqueeze(0) + else: + query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv( + mixed_qkv_non_spec + ) + g_non_spec = None + beta_non_spec = None + + # 2. Recurrent attention + + # 2.1: Process the multi-query part + if spec_sequence_masks is not None: + core_attn_out_spec, last_recurrent_state = ( + fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, + a=a, + b=b, + dt_bias=self.dt_bias, + q=query_spec, + k=key_spec, + v=value_spec, + initial_state=ssm_state, + inplace_final_state=True, + cu_seqlens=spec_query_start_loc[ # type: ignore[index] + : attn_metadata.num_spec_decodes + + 1 # type: ignore[attr-defined] + ], + 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 + + # 2.2: Process non-spec-decode part + if split_non_spec: + query_decode, key_decode, value_decode = self.rearrange_mixed_qkv( + mixed_qkv_non_spec[:num_decode_tokens] # type: ignore[index] + ) + core_attn_out_decode, _ = fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, + a=a[:num_decode_tokens], + b=b[:num_decode_tokens], + dt_bias=self.dt_bias, + q=query_decode, + k=key_decode, + v=value_decode, + initial_state=ssm_state, + inplace_final_state=True, + cu_seqlens=non_spec_query_start_loc[ # type: ignore[index] + : attn_metadata.num_decodes + 1 + ], + ssm_state_indices=non_spec_state_indices_tensor, + use_qk_l2norm_in_kernel=True, + ) + else: + core_attn_out_decode = None + + # 2.3: Process the remaining part (prefill chunk, or non-spec decode-only) + if attn_metadata.num_prefills > 0: + # State indices, initial-state mask and cu_seqlens for the chunk + # kernel are precomputed by the metadata builder (the prefill tail + # when decodes are peeled off, else the full non-spec batch), so they + # don't need to be re-derived per layer. + prefill_state_indices = attn_metadata.prefill_state_indices + prefill_has_initial_state = attn_metadata.prefill_has_initial_state + assert prefill_state_indices is not None + assert prefill_has_initial_state is not None + initial_state = ssm_state[prefill_state_indices] + initial_state[~prefill_has_initial_state, ...] = 0 + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = self.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=attn_metadata.prefill_query_start_loc, + chunk_indices=attn_metadata.chunk_indices, + chunk_offsets=attn_metadata.chunk_offsets, + use_qk_l2norm_in_kernel=False, + ) + # Init cache + ssm_state[prefill_state_indices] = last_recurrent_state.to(ssm_state.dtype) + + if split_non_spec: + # Stitch the peeled decode outputs in front of the prefill + # outputs (decode-first order). + core_attn_out_non_spec = torch.cat( + [core_attn_out_decode, core_attn_out_non_spec], dim=1 + ) + elif attn_metadata.num_decodes > 0: + core_attn_out_non_spec, last_recurrent_state = ( + fused_sigmoid_gating_delta_rule_update( + A_log=self.A_log, + a=a, + b=b, + dt_bias=self.dt_bias, + q=query_non_spec, + k=key_non_spec, + v=value_non_spec, + initial_state=ssm_state, + inplace_final_state=True, + cu_seqlens=non_spec_query_start_loc[ # type: ignore[index] + : attn_metadata.num_decodes + + 1 # type: ignore[attr-defined] + ], + 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 + + # 3. Merge core attention output + 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 _forward_core_decode_aiter( + self, + qkvz: torch.Tensor, + ba: torch.Tensor, + z_out: torch.Tensor, + core_attn_out: torch.Tensor, + attn_metadata: GDNAttentionMetadata, + ): + non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501 + 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] + + # 1. Convolution sequence transformation + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + + mixed_qkv_non_spec, b, a = ( + gdn_aiter_fused_reshape_causal_conv1d_update_single_token( + qkvz, + attn_metadata.num_actual_tokens, + self.num_k_heads // self.tp_size, + self.num_v_heads // self.tp_size, + self.head_k_dim, + self.head_v_dim, + ba, + z_out, + core_attn_out, + conv_state, + conv_weights, + self.conv1d.bias, + self.activation, + conv_state_indices=non_spec_state_indices_tensor[ # type: ignore[index] + : attn_metadata.num_actual_tokens + ], + validate_data=True, + ) + ) + + # 2. Recurrent attention + gdn_aiter_fused_rearrange_sigmoid_gated_delta_rule( + A_log=self.A_log, + a=a, + b=b, + dt_bias=self.dt_bias, + qkv=mixed_qkv_non_spec, + key_dim=self.key_dim // self.tp_size, + value_dim=self.value_dim // self.tp_size, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + initial_state=ssm_state, + inplace_final_state=True, + cu_seqlens=non_spec_query_start_loc[: attn_metadata.num_decodes + 1], # type: ignore[index] + ssm_state_indices=non_spec_state_indices_tensor, + use_qk_l2norm_in_kernel=True, + core_attn_out=core_attn_out.reshape(-1), + ) + + def _forward_core_decode_non_spec( + self, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + core_attn_out: torch.Tensor, + attn_metadata: GDNAttentionMetadata, + ): + """ + Core attention computation with a packed non-spec decode fast path. + """ + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501 + 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 + + 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) + ) + mixed_qkv_non_spec = causal_conv1d_update( + mixed_qkv, + conv_state, + conv_weights, + self.conv1d.bias, + self.activation, + conv_state_indices=non_spec_state_indices_tensor[:num_actual_tokens], # type: ignore[index] + validate_data=False, + ) + out_buf = core_attn_out[:num_actual_tokens].unsqueeze(1) + fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=mixed_qkv_non_spec, + a=a, + b=b, + A_log=self.A_log, + dt_bias=self.dt_bias, + scale=self.head_k_dim**-0.5, + initial_state=ssm_state, + out=out_buf, + ssm_state_indices=non_spec_state_indices_tensor[:num_actual_tokens], # type: ignore[index] + use_qk_l2norm_in_kernel=True, + ) + return + + +def qwen_gdn_attention_core( + qkv_or_qkvz: torch.Tensor, + b_or_ba: torch.Tensor, + a_or_z_out: torch.Tensor, + core_attn_out: torch.Tensor, + layer_name: LayerNameType, + use_aiter: bool = False, +) -> None: + """Custom op dispatching to _forward_core or _forward_core_rocm. + + Handles conv1d + recurrent attention only; input/output projections + are performed by the caller. + + When ``use_aiter=False`` (standard path): + qkv_or_qkvz is [q, k, v], b_or_ba is b, a_or_z_out is a (read-only). + When ``use_aiter=True`` (AITER Triton path, ROCm only): + qkv_or_qkvz is [q, k, v, z], b_or_ba is [b, a], a_or_z_out is the + z output buffer (mutated in-place). + + ``core_attn_out`` is always mutated in-place. + """ + layer_name = _resolve_layer_name(layer_name) + forward_context: ForwardContext = get_forward_context() + self = forward_context.no_compile_layers[layer_name] + if use_aiter: + self._forward_core_rocm( + qkvz=qkv_or_qkvz, + ba=b_or_ba, + z_out=a_or_z_out, + core_attn_out=core_attn_out, + ) + else: + self._forward_core( + mixed_qkv=qkv_or_qkvz, + b=b_or_ba, + a=a_or_z_out, + core_attn_out=core_attn_out, + ) + + +def gdn_attention_core_fake( + qkv_or_qkvz: torch.Tensor, + b_or_ba: torch.Tensor, + a_or_z_out: torch.Tensor, + core_attn_out: torch.Tensor, + layer_name: LayerNameType, + use_aiter: bool = False, +) -> None: + """Fake implementation for torch.compile.""" + return + + +direct_register_custom_op( + op_name="qwen_gdn_attention_core", + op_func=qwen_gdn_attention_core, + mutates_args=["a_or_z_out", "core_attn_out"], + fake_impl=gdn_attention_core_fake, +) + + +@triton.jit +def fused_gdn_gating_kernel( + g, + beta_output, + A_log, + a, + b, + dt_bias, + seq_len, + 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) + # If the model is loaded in fp16, without the .float() here, A might be -inf + 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) + # compute beta_output = sigmoid(b) + blk_beta_output = tl.sigmoid(blk_b.to(tl.float32)) + tl.store( + beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask + ) + + +def fused_gdn_gating( + A_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, + beta: float = 1.0, + threshold: float = 20.0, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Fused computation of g and beta for Gated Delta Net. + g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) + beta_output = b.sigmoid() + TODO maybe use torch.compile to replace this triton kernel + """ + 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=b.dtype, device=b.device) + fused_gdn_gating_kernel[grid]( + g, + beta_output, + A_log, + a, + b, + dt_bias, + seq_len, + num_heads, + beta, + threshold, + 8, + num_warps=1, + ) + return g, beta_output diff --git a/upstream_ref/vllm_gdn/models/qwen3_5.py b/upstream_ref/vllm_gdn/models/qwen3_5.py new file mode 100644 index 00000000..6ebf26c8 --- /dev/null +++ b/upstream_ref/vllm_gdn/models/qwen3_5.py @@ -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() diff --git a/upstream_ref/vllm_gdn/models/qwen3_next.py b/upstream_ref/vllm_gdn/models/qwen3_next.py new file mode 100644 index 00000000..b43aea87 --- /dev/null +++ b/upstream_ref/vllm_gdn/models/qwen3_next.py @@ -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) diff --git a/upstream_ref/vllm_gdn/ops/causal_conv1d.py b/upstream_ref/vllm_gdn/ops/causal_conv1d.py new file mode 100644 index 00000000..8335c849 --- /dev/null +++ b/upstream_ref/vllm_gdn/ops/causal_conv1d.py @@ -0,0 +1,1289 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2024, Tri Dao. +# Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py + + +import numpy as np +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID + + +@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"]) +def _causal_conv1d_fwd_kernel( # continuous batching + # Pointers to matrices + x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences + w_ptr, # (dim, width) + bias_ptr, + initial_states_ptr, # conv_states_ptr + cache_indices_ptr, # (batch, n_blocks + padding) The second dimension contains + # the block indices relevant for each sequence + # plus potential 0-padding at the beginning and at the end + has_initial_states_ptr, + query_start_loc_ptr, + batch_ptr, + token_chunk_offset_ptr, + block_idx_first_scheduled_token, # (batch,) + block_idx_last_scheduled_token, # (batch,) + initial_state_idx, # (batch,) + num_computed_tokens, # (batch,) + o_ptr, # (dim, seqlen) - actually pointing to x_ptr + # Matrix dimensions + dim: tl.constexpr, + num_cache_lines, # added to support vLLM larger cache lines + # Strides + stride_x_dim: tl.constexpr, # stride to get to next feature-value, + stride_x_token: tl.int64, # stride to get to next token (same feature-index, same sequence-index) + stride_w_dim: tl.constexpr, # stride to get to next dim-axis value + stride_w_width: tl.constexpr, # stride to get to next width-axis value + stride_istate_seq: tl.constexpr, + stride_istate_dim: tl.constexpr, + stride_istate_token: tl.constexpr, + stride_cache_indices: tl.constexpr, + stride_o_dim: tl.constexpr, + stride_o_token: tl.int64, + stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M + # others + pad_slot_id: tl.constexpr, + null_block_id: tl.constexpr, + # Meta-parameters + HAS_BIAS: tl.constexpr, + KERNEL_WIDTH: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + IS_APC_ENABLED: tl.constexpr, + HAS_NULL_BLOCK: tl.constexpr, + NP2_STATELEN: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + launch_pdl: tl.constexpr, +): + conv_states_ptr = initial_states_ptr + conv_state_indices_ptr = cache_indices_ptr + stride_conv_state_seq = stride_istate_seq + stride_conv_state_dim = stride_istate_dim + stride_conv_state_tok = stride_istate_token + state_len = ( + KERNEL_WIDTH - 1 + ) # can be passed via argument if it's not the same as this value + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + # one program handles one chunk in a single sequence + # rather than mixing sequences - to make updating initial_states across sequences efficiently + + # single-sequence id + idx_seq = tl.load(batch_ptr + tl.program_id(0)).to(tl.int64) + chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(0)) + + # BLOCK_N elements along the feature-dimension (channel) + idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + + if idx_seq == pad_slot_id: + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + return + + sequence_start_index = tl.load(query_start_loc_ptr + idx_seq) + sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1) + # find the actual sequence length + seqlen = sequence_end_index - sequence_start_index + + B_size: tl.constexpr = stride_block_m * BLOCK_M + + if IS_APC_ENABLED: + # Handle the case if prefix caching is enabled. + # In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr" + + # Get the length of the completed sequence so far and compute the offset. + current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq) + current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) + sequence_completed_index = tl.load(num_computed_tokens + idx_seq) + + # Compute the offset where the first stride_block_m-aligned first full block is + # Value in "token-space" + sequence_completed_offset_token = sequence_completed_index % B_size + seq_completed_offset = B_size - sequence_completed_offset_token + seq_end_offset = (seqlen - seq_completed_offset) % B_size + last_full_block_token_index = sequence_end_index - seq_end_offset + # If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one + if seq_end_offset == 0: + last_full_block_token_index = last_full_block_token_index - B_size + + # Get the number of blocks to be filled for the current sequence + # If n_block_to_fill = 0, then only the state at the sequence end is stored + n_block_to_fill = current_last_index - current_first_index + + # Get the index of the init block + conv_state_init_index = tl.load(initial_state_idx + idx_seq) + else: + n_block_to_fill = 0 + current_last_index = 0 + conv_state_init_index = 0 + current_first_index = 0 + last_full_block_token_index = 0 + + token_offset = BLOCK_M * chunk_offset + segment_len = min(BLOCK_M, seqlen - token_offset) + + # base of the sequence + x_base = ( + x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim + ) # [BLOCK_N,] + + # cache_idx + conv_states_input_coord = tl.load( + conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index + ).to(tl.int64) + + if HAS_NULL_BLOCK: # noqa + if conv_states_input_coord == null_block_id: + # not processing as this is a null block (padding) + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + return + conv_states_base = ( + conv_states_ptr + + (conv_states_input_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim) + ) # [BLOCK_N,] + + w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] + + # Does 2 things: + # 1. READ prior-block init-state data - [done by every Triton programs] + # 2. update conv_state with new data [only by the Triton program handles chunk_offset=0] + if chunk_offset == 0: + # read from conv_states + load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1) + if load_init_state: + # load from conv_states + prior_tokens = conv_states_base + (state_len - 1) * stride_conv_state_tok + mask_w = idx_feats < dim + if KERNEL_WIDTH == 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 3: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 4: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH == 5: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + else: + # prior-tokens are zeros + if KERNEL_WIDTH >= 2: # STRATEGY1 + # first chunk and does not have prior-token, so just set to 0 + col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 3: # STRATEGY1 + col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 4: # STRATEGY1 + col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + if KERNEL_WIDTH >= 5: # STRATEGY1 + col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) + + # STEP 2: + # here prepare data for updating conv_state + if ( + state_len <= seqlen + ): # SMALL_CACHE=True (only move part of 'x' into conv_state cache) + # just read from 'x' + # copy 'x' data to conv_state + # load only 'x' data (and set 0 before 'x' if seqlen < state_len) + idx_tokens_last = (seqlen - state_len) + tl.arange( + 0, NP2_STATELEN + ) # [BLOCK_M] + x_ptrs = ( + x_ptr + + ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None] + + (idx_feats * stride_x_dim)[None, :] + ) # [BLOCK_M,BLOCK_N,] + mask_x = ( + (idx_tokens_last >= 0)[:, None] + & (idx_tokens_last < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + # Compute the offset where the last block should be written in the conv_states + conv_states_output_coord = tl.load( + conv_state_indices_ptr + + idx_seq * stride_cache_indices + + current_last_index + ).to(tl.int64) + + conv_states_ptrs_target = ( + conv_states_ptr + + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq + + (idx_feats * stride_conv_state_dim) + )[None, :] + ( # [BLOCK_N,] + idx_tokens_conv * stride_conv_state_tok + )[:, None] + + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.debug_barrier() # NOTE: use this due to bug in Triton compiler + tl.store(conv_states_ptrs_target, loaded_x, mask) + + else: + if load_init_state: + # update conv_state by shifting left, i.e. take last few cols from conv_state + cols from 'x' + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + conv_states_ptrs_source = ( + conv_states_ptr + + (conv_states_input_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim)[None, :] + + ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = ( + (conv_states_input_coord < num_cache_lines) + & ((idx_tokens_conv + seqlen) < state_len)[:, None] + & (idx_feats < dim)[None, :] + ) + conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0) + + VAL = state_len - seqlen + + x_ptrs = ( + x_base[None, :] + + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens_conv - VAL >= 0)[:, None] + & (idx_tokens_conv - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + + tl.debug_barrier() # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load + new_conv_state = tl.where( + mask, conv_state, loaded_x + ) # BUG in 'tl.where' which requires a barrier before this + conv_states_ptrs_target = ( + conv_states_base + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[ + None, : + ] + tl.store(conv_states_ptrs_target, new_conv_state, mask) + else: # load_init_state == False + # update conv_state by shifting left, BUT + # set cols prior to 'x' as zeros + cols from 'x' + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + VAL = state_len - seqlen + + x_ptrs = ( + x_base[None, :] + + ((idx_tokens_conv - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens_conv - VAL >= 0)[:, None] + & (idx_tokens_conv - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + new_conv_state = tl.load(x_ptrs, mask_x, 0.0) + + conv_states_ptrs_target = ( + conv_states_base + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[ + None, : + ] + tl.store(conv_states_ptrs_target, new_conv_state, mask) + + else: # chunk_offset > 0 + # read prior-token data from `x` + load_init_state = True + prior_tokens = x_base + (token_offset - 1) * stride_x_token + mask_w = idx_feats < dim + if KERNEL_WIDTH == 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 3: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 4: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + if KERNEL_WIDTH == 5: + # ruff: noqa: F841 + conv_states_ptrs = prior_tokens # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") + + # Store intermediate states aligned with stride_block_m + # The additional states are cached starting from the last stride_block_m. + # For example: + # If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved. + # If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last + # stride_block_m are cached. + # For example chunk_offset = n_block_to_fill stores the state at last_full_block + if (chunk_offset - 1) < n_block_to_fill: + # Store the states at the chunk boundaries from the start of the sequence + idx_tokens_last = ( + last_full_block_token_index + - (n_block_to_fill - chunk_offset) * B_size + - state_len + ) + tl.arange(0, NP2_STATELEN) # [BLOCK_M] + x_ptrs = ( + x_ptr + + (idx_tokens_last * stride_x_token)[:, None] + + (idx_feats * stride_x_dim)[None, :] + ) # [BLOCK_M,BLOCK_N,] + + mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[ + None, : + ] # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + # cache_idx + conv_states_output_coord = tl.load( + conv_state_indices_ptr + + idx_seq * stride_cache_indices + + current_first_index + + (chunk_offset - 1) + ).to(tl.int64) + + conv_states_ptrs_target = ( + conv_states_ptr + + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq + + (idx_feats * stride_conv_state_dim) + )[None, :] + ( # [BLOCK_N,] + idx_tokens_conv * stride_conv_state_tok + )[:, None] + + mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.debug_barrier() # NOTE: use this due to bug in Triton compiler + tl.store(conv_states_ptrs_target, loaded_x, mask) + + if HAS_BIAS: + bias = bias_ptr + idx_feats + mask_bias = idx_feats < dim + acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to( + tl.float32 + ) # [BLOCK_N] + else: + acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) + + x_base_1d = x_base + token_offset * stride_x_token # starting of chunk + + # PRE-LOAD WEIGHTS + mask_w = idx_feats < dim + if KERNEL_WIDTH >= 2: + w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor + w_col0 = tl.load(w_ptrs, mask_w, other=0.0) + w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor + w_col1 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 3: + w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor + w_col2 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 4: + w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor + w_col3 = tl.load(w_ptrs, mask_w, other=0.0) + mask_x_1d = idx_feats < dim + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + + for idx_token in range(segment_len): + acc = acc_preload + + matrix_w = w_col0 + matrix_x = col0 + for j in tl.static_range(KERNEL_WIDTH): + if KERNEL_WIDTH == 2: + if j == 1: # KERNEL_WIDTH-1: + matrix_w = w_col1 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 3: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 4: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + + acc += matrix_x * matrix_w # [BLOCK_N] + + if KERNEL_WIDTH == 2: + col0 = matrix_x + elif KERNEL_WIDTH == 3: + col0 = col1 + col1 = matrix_x + elif KERNEL_WIDTH == 4: + col0 = col1 + col1 = col2 + col2 = matrix_x + + if SILU_ACTIVATION: + acc = acc / (1 + tl.exp(-acc)) + mask_1d = (idx_token < segment_len) & ( + idx_feats < dim + ) # token-index # feature-index + o_ptrs = ( + o_ptr + + (sequence_start_index + token_offset + idx_token) * stride_o_token + + (idx_feats * stride_o_dim) + ) + + tl.store(o_ptrs, acc, mask=mask_1d) + + +def causal_conv1d_fn( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor | None = None, + has_initial_state: torch.Tensor | None = None, + activation: str | None = "silu", + pad_slot_id: int = PAD_SLOT_ID, + null_block_id: int = NULL_BLOCK_ID, + block_idx_first_scheduled_token: torch.Tensor | None = None, + block_idx_last_scheduled_token: torch.Tensor | None = None, + initial_state_idx: torch.Tensor | None = None, + num_computed_tokens: torch.Tensor | None = None, + block_size_to_align=0, + metadata=None, + validate_data=False, +): + """support varlen + continuous batching when x is 2D tensor + + x: (dim,cu_seq_len) + cu_seq_len = total tokens of all seqs in that batch + sequences are concatenated from left to right for varlen + weight: (dim, width) + conv_states: (...,dim,width - 1) itype + updated inplace if cache_indices are not provided + [it use `cache_indices` to get the index to the cache of conv_state for that sequence + + conv_state[cache_indices[i]] for seq-i - to be used as initial_state when has_initial_state[i] = True + and after that conv_state[cache_indices[i]] need to be shift-left and updated with values from 'x' + ] + query_start_loc: (batch + 1) int32 + The cumulative sequence lengths of the sequences in + the batch, used to index into sequence. prepended by 0. + if + x = [5, 1, 1, 1] <- continuous batching (batch=4) + then + query_start_loc = [0, 5, 6, 7, 8] <- the starting index of the next sequence; while the last value is + the ending index of the last sequence + [length(query_start_loc)-1 == batch] + for example: query_start_loc = torch.Tensor([0,10,16,17]), + x.shape=(dim,17) + cache_indices: (batch) int32 + indicates the corresponding state index, + like so: conv_state = conv_states[cache_indices[batch_id]] + has_initial_state: (batch) bool + indicates whether should the kernel take the current state as initial + state for the calculations + [single boolean for each sequence in the batch: True or False] + bias: (dim,) + activation: either None or "silu" or "swish" or True + pad_slot_id: int + if cache_indices is passed, lets the kernel identify padded + entries that will not be processed, + for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] + in this case, the kernel will not process entries at + indices 0 and 3 + block_idx_first_scheduled_token: (batch,), dtype int32 + The pointer into cache_indices, where the first cache block to be filled is located. + block_idx_last_scheduled_token: (batch,), dtype int32 + The pointer into cache_indices, where the last cache block to be filled is located. + initial_state_idx: (batch,), dtype int32 + The pointer into cache_indices, where the cache block containing the initial state is located. + num_computed_tokens: (batch,), dtype int32 + The number of tokens already completed for each sequence + block_size_to_align: int + The block size to align the cached states to + out: same shape as `x` + """ + if isinstance(activation, bool) and activation: + activation = "silu" + + args = None + # Store original dtype to cast back at the end + original_x_dtype = x.dtype + x = x.to(conv_states.dtype) + out = torch.empty_like(x) + if metadata is not None: + nums_dict = metadata.nums_dict + args = nums_dict + batch_ptr = metadata.batch_ptr + token_chunk_offset_ptr = metadata.token_chunk_offset_ptr + else: + seqlens = query_start_loc.diff().to("cpu") + args = seqlens + MAX_NUM_PROGRAMS = 1024 + + batch_ptr = torch.full( + (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device + ) # tracking which seq-idx the Triton program is handling + token_chunk_offset_ptr = torch.full( + (MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device + ) # tracking BLOCK_M-based index in the sequence the Triton program is handling + + is_channel_last = (x.stride(0) == 1) & (x.stride(1) > 1) + dim, cu_seqlen = x.shape + _, width = weight.shape + state_len = width - 1 + np2_statelen = triton.next_power_of_2(state_len) + + padded_batch = query_start_loc.size(0) - 1 + stride_x_dim = x.stride(0) + stride_x_token = x.stride(1) + stride_w_dim = weight.stride(0) + stride_w_width = weight.stride(1) + stride_istate_seq = 0 + stride_istate_dim = 0 + stride_istate_token = 0 + num_cache_lines = 0 + BLOCK_M = 8 + if conv_states is not None: + # extensions to support vLLM: + # 1. conv_states is used to replaced initial_states + # 2. conv_states serve as a cache with num cache lines can be larger than batch size + # 3. mapping from sequence x[idx] to a cache line at index as specified via cache_indices[idx] + # 4. computation can be skipped if cache_indices[idx] == pad_slot_id + num_cache_lines = conv_states.size(0) + assert ( + num_cache_lines == conv_states.shape[0] + and dim == conv_states.shape[1] + and width - 1 <= conv_states.shape[2] + ) + stride_istate_seq = conv_states.stride(0) + stride_istate_dim = conv_states.stride(1) + stride_istate_token = conv_states.stride(2) + if out.dim() == 2: + stride_o_dim = out.stride(0) + stride_o_token = out.stride(1) + else: + stride_o_dim = out.stride(1) + stride_o_token = out.stride(2) + stride_cache_indices = cache_indices.stride(0) if cache_indices is not None else 0 + + if validate_data: + assert x.dim() == 2 + assert query_start_loc is not None + assert query_start_loc.dim() == 1 + assert x.stride(0) == 1 or x.stride(1) == 1 + if bias is not None: + assert bias.dim() == 1 + assert dim == bias.size(0) + if cache_indices is not None: + assert cache_indices.dim() == 1 + assert padded_batch == cache_indices.size(0) + if has_initial_state is not None: + assert has_initial_state.size() == (padded_batch,) + assert conv_states is not None, ( + "ERROR: `has_initial_state` is used, which needs also `conv_states`" + ) + assert weight.stride(1) == 1 + assert (dim, width) == weight.shape + assert is_channel_last, "Need to run in channel-last layout" + if block_size_to_align is not None and block_size_to_align > 0: + assert (block_size_to_align % BLOCK_M) == 0, ( + "The mamba block size needs to be divisible by the BLOCK_M" + ) + else: + block_size_to_align = BLOCK_M + + if metadata is None: + + def num_program(META, seqlens): + tot = 0 + + mlist = [] + offsetlist = [] # type: ignore + + nums = -(-seqlens // META["BLOCK_M"]) + + tot = nums.sum().item() + mlist = np.repeat(np.arange(len(nums)), nums) + for idx, num in enumerate(nums): + offsetlist.extend( + range(num) + ) # chunk-idx if a sequence is split into multiple chunks + + if META["batch_ptr"].nelement() < len(mlist): + newlen = len(mlist) + 1 + META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + + if META["batch_ptr"].nelement() >= len(mlist): + META["batch_ptr"][0 : len(mlist)].copy_( + torch.from_numpy(np.array(mlist)) + ) + META["token_chunk_offset_ptr"][0 : len(mlist)].copy_( + torch.from_numpy(np.array(offsetlist)) + ) + + META["batch_ptr"] = META["batch_ptr"].to(META["x_ptr"].device) + META["token_chunk_offset_ptr"] = META["token_chunk_offset_ptr"].to( + META["x_ptr"].device + ) + return tot + else: + + def num_program(META, nums_dict): + tot = nums_dict[META["BLOCK_M"]]["tot"] + + mlist = nums_dict[META["BLOCK_M"]]["mlist"] + mlist_len = nums_dict[META["BLOCK_M"]]["mlist_len"] + + offsetlist = nums_dict[META["BLOCK_M"]]["offsetlist"] + + if nums_dict[META["BLOCK_M"]]["batch_ptr"] is not None: + META["batch_ptr"] = nums_dict[META["BLOCK_M"]]["batch_ptr"] + META["token_chunk_offset_ptr"] = nums_dict[META["BLOCK_M"]][ + "token_chunk_offset_ptr" + ] + else: + if META["batch_ptr"].nelement() < mlist_len: + newlen = mlist_len + 1 + META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID) + + if META["batch_ptr"].nelement() >= mlist_len: + META["batch_ptr"][0:mlist_len].copy_(mlist) + META["token_chunk_offset_ptr"][0:mlist_len].copy_(offsetlist) + return tot + + def grid(META): + return ( + num_program(META, args), + triton.cdiv(dim, META["BLOCK_N"]), + ) + + if batch_ptr.device != x.device: + batch_ptr = batch_ptr.to(x.device) + token_chunk_offset_ptr = token_chunk_offset_ptr.to(x.device) + + _causal_conv1d_fwd_kernel[grid]( + # Pointers to matrices + x, + weight, + bias, + conv_states, + cache_indices, + has_initial_state, + query_start_loc, + batch_ptr, + token_chunk_offset_ptr, + block_idx_first_scheduled_token, + block_idx_last_scheduled_token, + initial_state_idx, + num_computed_tokens, + out, + # Matrix dimensions + dim, + num_cache_lines, + # stride + stride_x_dim, + stride_x_token, + stride_w_dim, + stride_w_width, + stride_istate_seq, + stride_istate_dim, + stride_istate_token, + stride_cache_indices, + stride_o_dim, + stride_o_token, + block_size_to_align // BLOCK_M, + # others + pad_slot_id, + null_block_id, + # META + HAS_BIAS=bias is not None, + KERNEL_WIDTH=width, + SILU_ACTIVATION=activation in ["silu", "swish"], + IS_APC_ENABLED=block_idx_last_scheduled_token is not None, + HAS_NULL_BLOCK=null_block_id is not None, + NP2_STATELEN=np2_statelen, + # launch_cooperative_grid=True + BLOCK_M=BLOCK_M, + BLOCK_N=256, + num_stages=2, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out.to(original_x_dtype) + + +@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"]) +def _causal_conv1d_update_kernel( + # Pointers to matrices + x_ptr, # (batch, dim, seqlen) + w_ptr, # (dim, width) + bias_ptr, + conv_state_ptr, + conv_state_indices_ptr, + num_accepted_tokens_ptr, + query_start_loc_ptr, # (batch + 1) + block_idx_last_scheduled_token, # (batch,) + initial_state_idx, # (batch,) + o_ptr, # (batch, dim, seqlen) + # Matrix dimensions + batch: int, + dim: tl.constexpr, + seqlen: tl.constexpr, + state_len: tl.constexpr, + num_cache_lines, # added to support vLLM larger cache lines + # Strides + stride_x_seq: tl.constexpr, + stride_x_dim: tl.constexpr, + stride_x_token: tl.int64, + stride_w_dim: tl.constexpr, + stride_w_width: tl.constexpr, + stride_conv_state_seq: tl.constexpr, + stride_conv_state_dim: tl.constexpr, + stride_conv_state_tok: tl.constexpr, + stride_state_indices: tl.constexpr, + stride_o_seq: tl.constexpr, + stride_o_dim: tl.constexpr, + stride_o_token: tl.int64, + # others + null_block_id: tl.constexpr, + # Meta-parameters + HAS_BIAS: tl.constexpr, + KERNEL_WIDTH: tl.constexpr, + SILU_ACTIVATION: tl.constexpr, + IS_VARLEN: tl.constexpr, + IS_APC_ENABLED: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + NP2_STATELEN: tl.constexpr, + HAS_NULL_BLOCK: tl.constexpr, + BLOCK_N: tl.constexpr, + launch_pdl: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + + # ruff: noqa: E501 + idx_seq = tl.program_id(0) + if idx_seq >= batch: + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + return + + # [BLOCK_N,] elements along the feature-dimension (channel) + idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + + if IS_APC_ENABLED: + # Get the state from the initial_state_idx + conv_state_init = tl.load(initial_state_idx + idx_seq) + current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) + else: + conv_state_init = 0 + current_last_index = 0 + + # cache_idx + conv_states_input_coord = tl.load( + conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init + ).to(tl.int64) + + if HAS_NULL_BLOCK: # noqa + if conv_states_input_coord == null_block_id: + # not processing as this is not the actual sequence + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + return + + if IS_VARLEN: + query_start_index = tl.load(query_start_loc_ptr + idx_seq).to(tl.int64) + query_end_index = tl.load(query_start_loc_ptr + (idx_seq + 1)).to(tl.int64) + # revise state_len and seqlen + state_len = state_len - (seqlen - (query_end_index - query_start_index)) + seqlen = query_end_index - query_start_index + x_offset = query_start_index * stride_x_token + o_offset = query_start_index * stride_o_token + else: + query_start_index = idx_seq * seqlen + query_end_index = query_start_index + seqlen + x_offset = idx_seq * stride_x_seq + o_offset = idx_seq * stride_o_seq + + if query_start_index == query_end_index: + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + return + + if IS_SPEC_DECODING: + # The rolling of conv state: + # + # Before forward, the conv_state is: + # [history1, history2, ..., historyM]. + # + # After forward, the conv_state becomes: + # [history2, ..., historyM, draft1, draft2, ..., draftN]. + # + # After acceptance, it becomes: + # + # - accept 1 tokens: [history2, ..., historyM, draft1] + # - accept 2 tokens: [history3, ..., historyM, draft1, draft2] + # - and so on. + conv_state_token_offset = ( + tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1 + ) + else: + conv_state_token_offset = 0 + + # STEP 1: READ init_state data + conv_states_base = ( + conv_state_ptr + + (conv_states_input_coord * stride_conv_state_seq) + + (idx_feats * stride_conv_state_dim) + ) + mask_w = idx_feats < dim + + prior_tokens = conv_states_base + conv_state_token_offset * stride_conv_state_tok + if KERNEL_WIDTH >= 2: + conv_states_ptrs = prior_tokens # [BLOCK_N] + col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 3: + conv_states_ptrs = prior_tokens + 1 * stride_conv_state_tok # [BLOCK_N] + col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 4: + conv_states_ptrs = prior_tokens + 2 * stride_conv_state_tok # [BLOCK_N] + col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 5: + conv_states_ptrs = prior_tokens + 3 * stride_conv_state_tok # [BLOCK_N] + col3 = tl.load(conv_states_ptrs, mask_w, 0.0) + if KERNEL_WIDTH >= 6: + conv_states_ptrs = prior_tokens + 4 * stride_conv_state_tok # [BLOCK_N] + col4 = tl.load(conv_states_ptrs, mask_w, 0.0) + + # STEP 2: assume state_len > seqlen + idx_tokens = tl.arange(0, NP2_STATELEN) # [BLOCK_M] + + # With speculative decoding, the conv_state updates works in a sliding + # window manner, at each forward pass, the tokens are shift by 1, so we + # load since idx_tokens + 1. + conv_state_ptrs_source = ( + conv_state_ptr + + (conv_states_input_coord * stride_conv_state_seq) + + conv_state_token_offset * stride_conv_state_tok + + (idx_feats * stride_conv_state_dim)[None, :] + + ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[ + :, None + ] + ) # [BLOCK_M, BLOCK_N] + mask = ( + (conv_states_input_coord < num_cache_lines) + & ((idx_tokens + seqlen) < state_len)[:, None] + & (idx_feats < dim)[None, :] + ) + conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0) + + VAL = state_len - seqlen + x_base = x_ptr + x_offset + (idx_feats * stride_x_dim) # [BLOCK_N] + + x_ptrs = ( + x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None] + ) # [BLOCK_M, BLOCK_N] + + mask_x = ( + (idx_tokens - VAL >= 0)[:, None] + & (idx_tokens - VAL < seqlen)[:, None] + & (idx_feats < dim)[None, :] + ) # token-index # token-index # feature-index + loaded_x = tl.load(x_ptrs, mask_x, 0.0) + tl.debug_barrier() + + new_conv_state = tl.where(mask, conv_state, loaded_x) + + # Get the state from the initial_state_idx + # cache_idx + conv_states_offset = tl.load( + conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index + ).to(tl.int64) + conv_state_ptrs_target = ( + conv_state_ptr + + (conv_states_offset * stride_conv_state_seq) # Offset from seq + + (idx_feats * stride_conv_state_dim) + )[None, :] + ( # [BLOCK_N,] + idx_tokens * stride_conv_state_tok + )[:, None] + mask = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :] + tl.store(conv_state_ptrs_target, new_conv_state, mask) + + # STEP 3: init accumulator + if HAS_BIAS: + bias = bias_ptr + idx_feats + mask_bias = idx_feats < dim + acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to( + tl.float32 + ) # [BLOCK_N] + else: + acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32) + + # STEP 4: + # PRE-LOAD WEIGHTS + # first kernel column, configured for weights to handle BLOCK_N features in range + w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,] + mask_w = idx_feats < dim + if KERNEL_WIDTH >= 2: + w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor + w_col0 = tl.load(w_ptrs, mask_w, other=0.0) + w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor + w_col1 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 3: + w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor + w_col2 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 4: + w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor + w_col3 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 5: + w_ptrs = w_base + (4 * stride_w_width) # [BLOCK_N] tensor + w_col4 = tl.load(w_ptrs, mask_w, other=0.0) + if KERNEL_WIDTH >= 6: + w_ptrs = w_base + (5 * stride_w_width) # [BLOCK_N] tensor + w_col5 = tl.load(w_ptrs, mask_w, other=0.0) + + x_base_1d = x_base # starting of chunk [BLOCK_N] + mask_x_1d = idx_feats < dim + + # STEP 5: compute each token + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + + for idx_token in tl.range(seqlen): + acc = acc_preload + + matrix_w = w_col0 + matrix_x = col0 + for j in tl.static_range(KERNEL_WIDTH): + if KERNEL_WIDTH == 2: + if j == 1: # KERNEL_WIDTH-1: + matrix_w = w_col1 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 3: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 4: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 5: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + matrix_x = col3 + elif j == 4: + matrix_w = w_col4 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + elif KERNEL_WIDTH == 6: + if j == 1: + matrix_w = w_col1 + matrix_x = col1 + elif j == 2: + matrix_w = w_col2 + matrix_x = col2 + elif j == 3: + matrix_w = w_col3 + matrix_x = col3 + elif j == 4: + matrix_w = w_col4 + matrix_x = col4 + elif j == 5: + matrix_w = w_col5 + x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N] + matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d) + + acc += matrix_x * matrix_w # [BLOCK_N] + + if KERNEL_WIDTH == 2: + col0 = matrix_x + elif KERNEL_WIDTH == 3: + col0 = col1 + col1 = matrix_x + elif KERNEL_WIDTH == 4: + col0 = col1 + col1 = col2 + col2 = matrix_x + elif KERNEL_WIDTH == 5: + col0 = col1 + col1 = col2 + col2 = col3 + col3 = matrix_x + elif KERNEL_WIDTH == 6: + col0 = col1 + col1 = col2 + col2 = col3 + col3 = col4 + col4 = matrix_x + + if SILU_ACTIVATION: + acc = acc / (1 + tl.exp(-acc)) + mask_1d = (idx_token < seqlen) & ( + idx_feats < dim + ) # token-index # feature-index + o_ptrs = ( + o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim) + ) + + tl.store(o_ptrs, acc, mask=mask_1d) + + +def causal_conv1d_update( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: bool | str | None = None, + conv_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + max_query_len: int = -1, + null_block_id: int = NULL_BLOCK_ID, + block_idx_last_scheduled_token: torch.Tensor | None = None, + initial_state_idx: torch.Tensor | None = None, + validate_data=False, + out: torch.Tensor | None = None, +): + """ + x: Input tensor which can take the following shapes: + + - `[batch, dim]` - single token prediction + - `[batch, dim, seqlen]` - single or multiple tokens prediction + - `[num_tokens, dim]` - continuous batching, where num_tokens is + the total tokens of all sequences in that batch + + conv_state: (..., dim, state_len), where state_len >= width - 1 + weight: (dim, width) + bias: (dim,) + conv_state_indices: (batch,), dtype int32 + If not None, the conv_state is a larger tensor along the batch dim, + and we are selecting the batch coords specified by conv_state_indices. + Useful for a continuous batching scenario. + block_idx_last_scheduled_token: (batch,), dtype int32 + The pointer into conv_state_indices, where the last cache block to be filled is located. + initial_state_idx: (batch,), dtype int32 + The pointer into conv_state_indices, where the cache block containing the initial state is located. + num_accepted_tokens: (batch,), dtype int32 + If not None, it indicates the number of accepted tokens for each + sequence in the batch. + This is used in speculative decoding, where the conv_state is updated + in a sliding window manner. + query_start_loc: (batch + 1,) int32 + If not None, the inputs is given in a varlen fashion and this indicates + the starting index of each sequence in the batch. + max_query_len: int + If query_start_loc is not None, this indicates the maximum query + length in the batch. + null_block_id: int + Block ID used to identify padded entries in + conv_state_indices. Block 0 is the null block. + for example: conv_state_indices = [null_block_id, 1, 20, null_block_id] + in this case, the kernel will not process entries at + indices 0 and 3 + out: optional output tensor with the same shape as `x`. When omitted, + the input is overwritten. + """ + if validate_data: + assert null_block_id is not None + assert x.stride(1) == 1 + if isinstance(activation, bool): + activation = "silu" if activation is True else None + elif activation is not None: + assert activation in ["silu", "swish"] + + original_x_dtype = x.dtype + x = x.to(conv_state.dtype) + if out is None: + out = x + else: + if out.shape != x.shape: + raise ValueError( + f"`out` shape {tuple(out.shape)} must match `x` shape {tuple(x.shape)}." + ) + if out.dtype != original_x_dtype or out.device != x.device: + raise ValueError( + "`out` must have the same dtype and device as the input `x`." + ) + unsqueeze = query_start_loc is None and x.dim() == 2 + if unsqueeze: + # make it (batch, dim, seqlen) with seqlen == 1 + x = x.unsqueeze(-1) + out = out.unsqueeze(-1) + if query_start_loc is None: + batch, dim, seqlen = x.shape + else: + assert conv_state_indices is not None + batch = conv_state_indices.size(0) + dim = x.size(1) + seqlen = max_query_len + _, width = weight.shape + # conv_state: (..., dim, state_len), where state_len >= width - 1 + num_cache_lines, _, state_len = conv_state.size() + + if validate_data: + assert dim == weight.size(0) + assert state_len >= width - 1 + # when above happens, we don't shift-left to keep any records in conv_state + assert dim == conv_state.size(1) + if conv_state_indices is None: + assert conv_state.size(0) >= batch + else: + assert batch == conv_state_indices.shape[0], ( + f"ERROR: conv_state_indices should have shape ({batch},*) but got {conv_state_indices.shape}" + ) + + assert num_cache_lines >= batch + assert weight.stride(1) == 1 # Need this + + stride_w_dim, stride_w_width = weight.stride() + + if query_start_loc is None: + # X (batch, dim, seqlen) + stride_x_seq, stride_x_dim, stride_x_token = x.stride() + stride_o_seq, stride_o_dim, stride_o_token = out.stride() + else: + # X (dim, cu_seqlen) + stride_x_token, stride_x_dim = x.stride() + stride_x_seq = 0 + stride_o_token, stride_o_dim = out.stride() + stride_o_seq = 0 + + stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride() + stride_state_indices = ( + conv_state_indices.stride(0) if conv_state_indices is not None else 0 + ) + if num_accepted_tokens is not None: + state_len = width - 1 + (seqlen - 1) # effective state_len needed + else: + state_len = width - 1 + np2_statelen = triton.next_power_of_2(state_len) + + def grid(META): + return ( + batch, + triton.cdiv(dim, META["BLOCK_N"]), + ) + + _causal_conv1d_update_kernel[grid]( + # Pointers to matrices + x, + weight, + bias, + conv_state, + conv_state_indices, + num_accepted_tokens, + query_start_loc, + block_idx_last_scheduled_token, + initial_state_idx, + out, + # Matrix dimensions + batch, + dim, + seqlen, + state_len, + num_cache_lines, + # stride + stride_x_seq, + stride_x_dim, + stride_x_token, + stride_w_dim, + stride_w_width, + stride_istate_seq, + stride_istate_dim, + stride_istate_token, + stride_state_indices, + stride_o_seq, + stride_o_dim, + stride_o_token, + # others + null_block_id, + # META + HAS_BIAS=bias is not None, + KERNEL_WIDTH=width, + SILU_ACTIVATION=activation in ["silu", "swish"], + IS_VARLEN=query_start_loc is not None, + IS_APC_ENABLED=block_idx_last_scheduled_token is not None, + IS_SPEC_DECODING=num_accepted_tokens is not None, + NP2_STATELEN=np2_statelen, + HAS_NULL_BLOCK=null_block_id is not None, + BLOCK_N=256, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + if unsqueeze: + out = out.squeeze(-1) + return out.to(original_x_dtype) + + +if current_platform.is_cpu(): + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_fn_cpu, + causal_conv1d_update_cpu, + ) + + causal_conv1d_fn = causal_conv1d_fn_cpu # type: ignore + causal_conv1d_update = causal_conv1d_update_cpu # type: ignore diff --git a/upstream_ref/vllm_gdn/third_party/__init__.py b/upstream_ref/vllm_gdn/third_party/__init__.py new file mode 100644 index 00000000..0e89cf9f --- /dev/null +++ b/upstream_ref/vllm_gdn/third_party/__init__.py @@ -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 diff --git a/upstream_ref/vllm_gdn/third_party/ops/__init__.py b/upstream_ref/vllm_gdn/third_party/ops/__init__.py new file mode 100644 index 00000000..1942d898 --- /dev/null +++ b/upstream_ref/vllm_gdn/third_party/ops/__init__.py @@ -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", +] diff --git a/upstream_ref/vllm_gdn/third_party/ops/chunk.py b/upstream_ref/vllm_gdn/third_party/ops/chunk.py new file mode 100644 index 00000000..caf8b0c9 --- /dev/null +++ b/upstream_ref/vllm_gdn/third_party/ops/chunk.py @@ -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 diff --git a/upstream_ref/vllm_gdn/third_party/ops/cumsum.py b/upstream_ref/vllm_gdn/third_party/ops/cumsum.py new file mode 100644 index 00000000..b0820104 --- /dev/null +++ b/upstream_ref/vllm_gdn/third_party/ops/cumsum.py @@ -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" + ) diff --git a/upstream_ref/vllm_gdn/third_party/ops/fused_gdn_prefill_post_conv.py b/upstream_ref/vllm_gdn/third_party/ops/fused_gdn_prefill_post_conv.py new file mode 100644 index 00000000..4807c78e --- /dev/null +++ b/upstream_ref/vllm_gdn/third_party/ops/fused_gdn_prefill_post_conv.py @@ -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 diff --git a/upstream_ref/vllm_gdn/third_party/ops/fused_recurrent.py b/upstream_ref/vllm_gdn/third_party/ops/fused_recurrent.py new file mode 100644 index 00000000..c8004cb5 --- /dev/null +++ b/upstream_ref/vllm_gdn/third_party/ops/fused_recurrent.py @@ -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 diff --git a/upstream_ref/vllm_gdn/third_party/ops/l2norm.py b/upstream_ref/vllm_gdn/third_party/ops/l2norm.py new file mode 100644 index 00000000..2eb137a2 --- /dev/null +++ b/upstream_ref/vllm_gdn/third_party/ops/l2norm.py @@ -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) diff --git a/upstream_ref/vllm_gdn/third_party/ops/utils.py b/upstream_ref/vllm_gdn/third_party/ops/utils.py new file mode 100644 index 00000000..d8e7db93 --- /dev/null +++ b/upstream_ref/vllm_gdn/third_party/ops/utils.py @@ -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 diff --git a/upstream_ref/vllm_gdn/v1/attention/backends/gdn_attn.py b/upstream_ref/vllm_gdn/v1/attention/backends/gdn_attn.py new file mode 100644 index 00000000..0b32dfcf --- /dev/null +++ b/upstream_ref/vllm_gdn/v1/attention/backends/gdn_attn.py @@ -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) diff --git a/upstream_ref/xllm_latest/core/kernels/npu/npu_causal_conv1d.cpp b/upstream_ref/xllm_latest/core/kernels/npu/npu_causal_conv1d.cpp new file mode 100644 index 00000000..dc8274be --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/npu/npu_causal_conv1d.cpp @@ -0,0 +1,59 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp" +#include "core/kernels/npu/utils.h" +#include "core/kernels/npu/xllm_ops/xllm_ops_api.h" + +namespace xllm::kernel::npu { + +torch::Tensor causal_conv1d(const torch::Tensor& x, + const torch::Tensor& weight, + const torch::Tensor& conv_state, + const std::optional& bias_opt, + const torch::IntArrayRef query_start_loc_opt, + const torch::IntArrayRef cache_indices_opt, + const torch::IntArrayRef initial_state_mode_opt, + const torch::IntArrayRef num_accepted_tokens_opt, + int64_t activation_mode, + int64_t pad_slot_id, + int64_t run_mode) { + check_tensor(x, "x", "causal_conv1d"); + check_tensor(weight, "weight", "causal_conv1d"); + check_tensor(conv_state, "conv_state", "causal_conv1d"); + + c10::optional bias_tensor = c10::nullopt; + if (bias_opt.has_value() && bias_opt.value().defined()) { + bias_tensor = bias_opt.value(); + } + + torch::Tensor output = torch::empty(x.sizes(), x.options()); + EXEC_NPU_CMD(aclnnCausalConv1d, + x, + weight, + bias_tensor, + conv_state, + query_start_loc_opt, + cache_indices_opt, + initial_state_mode_opt, + num_accepted_tokens_opt, + activation_mode, + pad_slot_id, + run_mode, + output); + return output; +} + +} // namespace xllm::kernel::npu diff --git a/upstream_ref/xllm_latest/core/kernels/npu/npu_recurrent_gated_delta_rule.cpp b/upstream_ref/xllm_latest/core/kernels/npu/npu_recurrent_gated_delta_rule.cpp new file mode 100644 index 00000000..d75c4c04 --- /dev/null +++ b/upstream_ref/xllm_latest/core/kernels/npu/npu_recurrent_gated_delta_rule.cpp @@ -0,0 +1,83 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include + +#include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp" +#include "core/kernels/npu/npu_ops_api.h" +#include "core/kernels/npu/utils.h" + +namespace { + +c10::optional to_c10_optional_tensor( + const std::optional& tensor_opt) { + if (tensor_opt.has_value() && tensor_opt.value().defined()) { + return tensor_opt.value(); + } + return c10::nullopt; +} + +} // namespace + +namespace xllm::kernel::npu { + +torch::Tensor npu_recurrent_gated_delta_rule( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + torch::Tensor& state, + const std::optional& beta, + const std::optional scale, + const std::optional& actual_seq_lengths, + const std::optional& ssm_state_indices, + const std::optional& num_accepted_tokens, + const std::optional& g, + const std::optional& gk) { + check_tensor(query, "query", "recurrent_gated_delta_rule"); + check_tensor(key, "key", "recurrent_gated_delta_rule"); + check_tensor(value, "value", "recurrent_gated_delta_rule"); + check_tensor(state, "state", "recurrent_gated_delta_rule"); + CHECK(scale.has_value()) + << "recurrent_gated_delta_rule requires a valid scale value"; + + c10::optional beta_tensor = to_c10_optional_tensor(beta); + c10::optional actual_seq_lengths_tensor = + to_c10_optional_tensor(actual_seq_lengths); + c10::optional ssm_state_indices_tensor = + to_c10_optional_tensor(ssm_state_indices); + c10::optional num_accepted_tokens_tensor = + to_c10_optional_tensor(num_accepted_tokens); + c10::optional g_tensor = to_c10_optional_tensor(g); + c10::optional gk_tensor = to_c10_optional_tensor(gk); + float scale_value = static_cast(scale.value()); + torch::Tensor output = torch::empty_like(value); + + EXEC_NPU_CMD(aclnnRecurrentGatedDeltaRule, + query, + key, + value, + beta_tensor, + state, + actual_seq_lengths_tensor, + ssm_state_indices_tensor, + g_tensor, + gk_tensor, + num_accepted_tokens_tensor, + scale_value, + output); + return output; +} + +} // namespace xllm::kernel::npu diff --git a/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.cpp b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.cpp new file mode 100644 index 00000000..da8f3abf --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.cpp @@ -0,0 +1,236 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "qwen3_5_attention.h" + +#include + +#include + +#include "kernels/ops_api.h" +namespace xllm { +namespace layer { + +Qwen3_5AttentionImpl::Qwen3_5AttentionImpl(const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options, + int32_t layer_id) { + const int64_t tp_size = parallel_args.tp_group_->world_size(); + const int64_t total_num_heads = args.n_heads(); + const int64_t total_num_kv_heads = args.n_kv_heads().value_or(args.n_heads()); + layer_id_ = layer_id; + rank_ = parallel_args.tp_group_->rank(); + CHECK(total_num_heads % tp_size == 0); + num_heads_ = total_num_heads / tp_size; + + if (total_num_kv_heads >= tp_size) { + CHECK(total_num_kv_heads % tp_size == 0); + num_kv_heads_ = total_num_kv_heads / tp_size; + num_kv_head_replicas_ = 1; + } else { + CHECK(tp_size % total_num_kv_heads == 0); + num_kv_heads_ = 1; + num_kv_head_replicas_ = tp_size / total_num_kv_heads; + } + + head_dim_ = args.head_dim(); + q_size_ = num_heads_ * head_dim_; + kv_size_ = num_kv_heads_ * head_dim_; + scaling_ = 1.0f / std::sqrt(static_cast(head_dim_)); + attn_output_gate_ = args.attn_output_gate(); + mrope_cu_seq_lens_ = torch::zeros(2, torch::kInt32).to(options.device()); + // 1. QKV linear + qkv_proj_ = register_module( + "qkv_proj", + QKVParallelLinear(args.hidden_size(), + attn_output_gate_ ? num_heads_ * 2 : num_heads_, + num_kv_heads_, + args.head_dim(), + num_kv_head_replicas_, + /*bias=*/args.attention_bias(), + /*gather_output=*/false, + parallel_args, + options)); + + // 2. O proj + o_proj_ = register_module("o_proj", + RowParallelLinear(total_num_heads * head_dim_, + args.hidden_size(), + /*bias=*/false, + /*input_is_parallelized=*/true, + /*if_reduce_results=*/true, + quant_args, + parallel_args.tp_group_, + options)); + + // 3. Q norm + q_norm_ = register_module( + "q_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options)); + + // 4. K norm + k_norm_ = register_module( + "k_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options)); + + // 5. Attention + attn_ = register_module("attn", + Attention(num_heads_, + head_dim_, + scaling_, + num_kv_heads_, + args.sliding_window())); + + // 6. Rotary embedding + const int32_t rotary_dim = + static_cast(head_dim_ * args.partial_rotary_factor()); + rotary_emb_ = + register_module("rope", + MRotaryEmbedding(rotary_dim, + args.max_position_embeddings(), + args.rope_theta(), + /*interleaved=*/false, + args.rope_scaling_mrope_section(), + options)); +} + +void Qwen3_5AttentionImpl::rotary_emb_forward( + torch::Tensor& q, + torch::Tensor& k, + const torch::Tensor& positions, + const AttentionMetadata& attn_metadata) { + auto q_shape = q.sizes(); + auto k_shape = k.sizes(); + auto num_tokens = positions.size(-1); + mrope_cu_seq_lens_[1] = num_tokens; + + xllm::kernel::RotaryParams rotary_params; + bool only_prefill = + (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill); + if (only_prefill) { + rotary_params.sin = attn_metadata.mrope_sin; + rotary_params.cos = attn_metadata.mrope_cos; + rotary_params.position_ids = std::nullopt; + rotary_params.cu_query_lens = mrope_cu_seq_lens_; + rotary_params.interleaved = false; + rotary_params.discrete = false; + rotary_params.max_query_len = num_tokens; + + rotary_params.q = q.view({num_tokens, -1, head_dim_}); + xllm::kernel::apply_rotary(rotary_params); + q = rotary_params.q.reshape(q_shape); + + rotary_params.q = k.view({num_tokens, -1, head_dim_}); + xllm::kernel::apply_rotary(rotary_params); + k = rotary_params.q.reshape(k_shape); + } else { + if (positions.dim() == 2) { + rotary_params.position_ids = positions[0]; + } else { + rotary_params.position_ids = positions; + } + rotary_params.sin = rotary_emb_->get_sin_cache(); + rotary_params.cos = rotary_emb_->get_cos_cache(); + + rotary_params.interleaved = false; + rotary_params.discrete = true; + rotary_params.max_query_len = num_tokens; + rotary_params.q = q.view({1, num_tokens, -1, head_dim_}); + xllm::kernel::apply_rotary(rotary_params); + q = rotary_params.q.reshape(q_shape); + + rotary_params.q = k.view({1, num_tokens, -1, head_dim_}); + xllm::kernel::apply_rotary(rotary_params); + k = rotary_params.q.reshape(k_shape); + } +} + +torch::Tensor Qwen3_5AttentionImpl::forward( + const torch::Tensor& positions, + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache) { + // 1. qkv projection + auto qkv = qkv_proj_->forward(hidden_states); + torch::Tensor q, k, v; + torch::Tensor gate; + + if (attn_output_gate_) { + // Split qkv for attn_output_gate case: [q_size*2, kv_size, kv_size] + auto q_gate = qkv.slice(/*dim=*/-1, 0, q_size_ * 2); + k = qkv.slice(/*dim=*/-1, q_size_ * 2, q_size_ * 2 + kv_size_); + v = qkv.slice( + /*dim=*/-1, q_size_ * 2 + kv_size_, q_size_ * 2 + kv_size_ * 2); + v = v.contiguous(); + + std::vector orig_shape; + for (int64_t i = 0; i < q_gate.dim() - 1; i++) { + orig_shape.push_back(q_gate.size(i)); + } + std::vector new_shape = orig_shape; + new_shape.push_back(num_heads_); + new_shape.push_back(-1); + torch::Tensor q_gate_reshaped = q_gate.reshape(new_shape); + auto chunks = torch::chunk(q_gate_reshaped, 2, /*dim=*/-1); + q = chunks[0]; + gate = chunks[1]; + + std::vector q_new_shape = orig_shape; + q_new_shape.push_back(-1); + q = q.reshape(q_new_shape); + + std::vector gate_new_shape = orig_shape; + gate_new_shape.push_back(-1); + gate = gate.reshape(gate_new_shape); + } else { + // Normal case: [q_size, kv_size, kv_size] + q = qkv.slice(/*dim=*/-1, 0, q_size_); + k = qkv.slice(/*dim=*/-1, q_size_, q_size_ + kv_size_); + v = qkv.slice(/*dim=*/-1, q_size_ + kv_size_, q_size_ + 2 * kv_size_); + } + + const int64_t T = q.size(0); + + auto q_reshaped = q.reshape({T, num_heads_, head_dim_}); + auto q_normed = std::get<0>(q_norm_->forward(q_reshaped)); + auto k_reshaped = k.reshape({T, num_kv_heads_, head_dim_}); + auto k_normed = std::get<0>(k_norm_->forward(k_reshaped)); + + q = q_normed.view({T, q_size_}); + k = k_normed.view({T, kv_size_}); + rotary_emb_forward(q, k, positions, attn_metadata); + auto out = std::get<0>(attn_->forward(attn_metadata, q, k, v, kv_cache)); + + if (attn_output_gate_) { + gate = torch::sigmoid(gate); + out = out * gate; + } + + out = o_proj_->forward(out); + return out; +} + +void Qwen3_5AttentionImpl::load_state_dict(const StateDict& state_dict) { + qkv_proj_->load_state_dict(state_dict, {"q_proj.", "k_proj.", "v_proj."}); + o_proj_->load_state_dict(state_dict.get_dict_with_prefix("o_proj.")); + if (auto w = state_dict.get_tensor("q_norm.weight"); w.defined()) { + q_norm_->load_state_dict(StateDict({{"weight", w}})); + } + if (auto w = state_dict.get_tensor("k_norm.weight"); w.defined()) { + k_norm_->load_state_dict(StateDict({{"weight", w}})); + } +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.h b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.h new file mode 100644 index 00000000..72fd2334 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_attention.h @@ -0,0 +1,79 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include "attention.h" +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_args.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "layers/common/linear.h" +#include "layers/common/partial_rotary_embedding.h" +#include "layers/common/qwen3_next_rms_norm.h" +#include "layers/common/rotary_embedding.h" + +namespace xllm { +namespace layer { + +class Qwen3_5AttentionImpl : public torch::nn::Module { + public: + Qwen3_5AttentionImpl() = default; + Qwen3_5AttentionImpl(const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options, + int32_t layer_id); + + torch::Tensor forward(const torch::Tensor& positions, + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache); + + void load_state_dict(const StateDict& state_dict); + void rotary_emb_forward(torch::Tensor& q, + torch::Tensor& k, + const torch::Tensor& positions, + const AttentionMetadata& attn_metadata); + + private: + int64_t num_heads_; + int64_t num_kv_heads_; + int64_t num_kv_head_replicas_; + int64_t head_dim_; + int64_t q_size_; + int64_t kv_size_; + float scaling_; + bool attn_output_gate_; + int32_t layer_id_; + int32_t rank_; + + QKVParallelLinear qkv_proj_{nullptr}; + RowParallelLinear o_proj_{nullptr}; + + Qwen3NextRMSNorm q_norm_{nullptr}; + Qwen3NextRMSNorm k_norm_{nullptr}; + + Attention attn_{nullptr}; + MRotaryEmbedding rotary_emb_{nullptr}; + torch::Tensor mrope_cu_seq_lens_; +}; +TORCH_MODULE(Qwen3_5Attention); + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_decoder_layer.cpp b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_decoder_layer.cpp new file mode 100644 index 00000000..1a6021e0 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_decoder_layer.cpp @@ -0,0 +1,193 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "qwen3_5_decoder_layer.h" + +#include + +#include "common/global_flags.h" +#include "layers/common/dp_utils.h" + +namespace xllm { +namespace layer { +namespace { +bool use_moe_all2all(bool enable_deep_ep, + const ModelInputParams& input_params) { + return enable_deep_ep && all_dp_ranks_are_decode(input_params); +} + +bool is_moe_layer(const ModelArgs& model_args, int32_t layer_id) { + const auto& mlp_only_layers = model_args.mlp_only_layers(); + return std::count(mlp_only_layers.begin(), mlp_only_layers.end(), layer_id) == + 0 && + model_args.n_routed_experts() > 0 && + (layer_id + 1) % model_args.decoder_sparse_step() == 0; +} +} // namespace + +Qwen3_5DecoderLayerImpl::Qwen3_5DecoderLayerImpl(const ModelContext& context, + int32_t layer_id) + : parallel_args_(context.get_parallel_args()) { + const auto& model_args = context.get_model_args(); + const auto& quant_args = context.get_quant_args(); + const auto& options = context.get_tensor_options(); + + const bool use_moe = is_moe_layer(model_args, layer_id); + + enable_deep_ep_ = use_moe && FLAGS_expert_parallel_degree == 2; + if (enable_deep_ep_) { + CHECK_EQ(parallel_args_.dp_size(), parallel_args_.world_size()) + << "Qwen3.5 MoE only support deep ep all2all when dp_size == " + "world_size"; + CHECK_EQ(parallel_args_.dp_size(), parallel_args_.ep_size()) + << "Qwen3.5 MoE only support deep ep all2all when dp_size == ep_size"; + } + + auto layer_types = model_args.layer_types(); + if (layer_types.empty()) { + int32_t interval = model_args.full_attention_interval(); + for (int32_t i = 0; i < model_args.n_layers(); i++) { + layer_types.push_back((i + 1) % interval == 0 ? "full_attention" + : "linear_attention"); + } + } + + if (layer_id >= 0 && layer_id < static_cast(layer_types.size())) { + layer_type_ = layer_types[layer_id]; + } else { + layer_type_ = "full_attention"; + } + + if (layer_type_ == "linear_attention") { + // TODO: support linear attention + } else { + full_attention_ = register_module( + "self_attn", + Qwen3_5Attention( + model_args, quant_args, parallel_args_, options, layer_id)); + } + + input_norm_ = register_module( + "input_layernorm", + Qwen3NextRMSNorm( + model_args.hidden_size(), model_args.rms_norm_eps(), options)); + + post_norm_ = register_module( + "post_attention_layernorm", + Qwen3NextRMSNorm( + model_args.hidden_size(), model_args.rms_norm_eps(), options)); + + if (use_moe) { + moe_mlp_ = register_module("mlp", + Qwen3_5FusedMoE(model_args, + FusedMoEArgs{.is_gated = true}, + quant_args, + parallel_args_, + options)); + } else { + mlp_ = register_module("mlp", + DenseMLP(model_args.hidden_size(), + model_args.intermediate_size(), + true, + false, + model_args.hidden_act(), + /*enable_result_reduction=*/true, + quant_args, + parallel_args_.tp_group_, + options)); + } +} + +void Qwen3_5DecoderLayerImpl::load_state_dict(const StateDict& state_dict) { + if (layer_type_ == "linear_attention") { + // TODO: support linear attention + } else { + full_attention_->load_state_dict( + state_dict.get_dict_with_prefix("self_attn.")); + } + input_norm_->load_state_dict( + state_dict.get_dict_with_prefix("input_layernorm.")); + post_norm_->load_state_dict( + state_dict.get_dict_with_prefix("post_attention_layernorm.")); + if (moe_mlp_) { + moe_mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp.")); + } else { + mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp.")); + } +} + +torch::Tensor Qwen3_5DecoderLayerImpl::run_moe( + torch::Tensor x, + const ModelInputParams& input_params) { + const bool enable_moe_all2all = + use_moe_all2all(enable_deep_ep_, input_params); + if (need_dp_moe_gather(parallel_args_, enable_moe_all2all)) { + x = gather_dp_tokens(x, input_params, parallel_args_); + x = moe_mlp_->forward_experts(x, enable_moe_all2all); + return get_dp_local_slice(x, input_params, parallel_args_); + } + return moe_mlp_->forward_experts(x, enable_moe_all2all); +} + +std::tuple> +Qwen3_5DecoderLayerImpl::apply_norm(Qwen3NextRMSNorm& norm, + torch::Tensor& input, + std::optional& residual) { + if (!residual.has_value()) { + auto new_residual = input; + auto output = std::get<0>(norm->forward(input)); + return {output, new_residual}; + } + auto orig_dtype = input.dtype(); + input = input + residual.value(); + auto new_residual = input; + input = input.to(orig_dtype); + auto output = std::get<0>(norm->forward(input)); + return {output, new_residual}; +} + +torch::Tensor Qwen3_5DecoderLayerImpl::forward( + torch::Tensor& x, + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params) { + // Pre-attention norm + std::tie(x, residual) = apply_norm(input_norm_, x, residual); + + // Attention + if (full_attention_) { + x = full_attention_->forward(positions, x, attn_metadata, kv_cache); + } else { + // TODO: support linear attention + } + + auto orig_dtype = x.dtype(); + // Post-attention norm + std::tie(x, residual) = apply_norm(post_norm_, x, residual); + + // MLP/MoE + if (moe_mlp_) { + x = run_moe(x, input_params); + } else { + x = mlp_->forward(x); + } + x = x.to(orig_dtype); + return x; +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_decoder_layer.h b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_decoder_layer.h new file mode 100644 index 00000000..efe56d1b --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_decoder_layer.h @@ -0,0 +1,73 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include + +#include +#include + +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_args.h" +#include "framework/model/model_input_params.h" +#include "framework/model_context.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/state_dict/state_dict.h" +#include "layers/common/dense_mlp.h" +#include "layers/common/qwen3_next_rms_norm.h" +#include "layers/mlu/qwen3_5_attention.h" +#include "layers/mlu/qwen3_5_fused_moe.h" + +namespace xllm { +namespace layer { + +class Qwen3_5DecoderLayerImpl final : public torch::nn::Module { + public: + Qwen3_5DecoderLayerImpl(const ModelContext& context, int32_t layer_id); + + void load_state_dict(const StateDict& state_dict); + + torch::Tensor forward(torch::Tensor& x, + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params); + + private: + std::tuple> apply_norm( + Qwen3NextRMSNorm& norm, + torch::Tensor& input, + std::optional& residual); + + torch::Tensor run_moe(torch::Tensor x, const ModelInputParams& input_params); + + std::string layer_type_; + Qwen3_5Attention full_attention_{nullptr}; + // TODO: support linear attention + // Qwen3_5GatedDeltaNet linear_attention_{nullptr}; + DenseMLP mlp_{nullptr}; + Qwen3_5FusedMoE moe_mlp_{nullptr}; + Qwen3NextRMSNorm input_norm_{nullptr}; + Qwen3NextRMSNorm post_norm_{nullptr}; + ParallelArgs parallel_args_; + bool enable_deep_ep_ = false; +}; + +TORCH_MODULE(Qwen3_5DecoderLayer); + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.cpp b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.cpp new file mode 100644 index 00000000..a32ba794 --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.cpp @@ -0,0 +1,209 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "qwen3_5_fused_moe.h" + +#include + +#include "framework/parallel_state/parallel_state.h" +#include "framework/state_dict/utils.h" + +namespace xllm { +namespace layer { +namespace { +torch::Tensor get_tensor_with_weight_suffix(const StateDict& state_dict, + const std::string& tensor_name) { + auto tensor = state_dict.get_tensor(tensor_name); + if (!tensor.defined()) { + tensor = state_dict.get_tensor(tensor_name + ".weight"); + } + return tensor; +} + +torch::Tensor slice_expert_weights(const torch::Tensor& weight, + int64_t start_expert_id, + int64_t num_experts_per_rank) { + return weight + .slice(0, start_expert_id, start_expert_id + num_experts_per_rank) + .contiguous(); +} + +bool load_fused_gate_up_fallback(const StateDict& state_dict, + int64_t rank, + int64_t world_size, + int64_t start_expert_id, + int64_t num_experts_per_rank, + torch::Tensor& w13) { + auto fused_gate_up = + get_tensor_with_weight_suffix(state_dict, "gate_up_proj"); + if (!fused_gate_up.defined()) { + return false; + } + + if (world_size > 1) { + CHECK_EQ(fused_gate_up.size(1) % 2, 0) + << "gate_up_proj dim1 must be even, got " << fused_gate_up.size(1); + const int64_t full_intermediate = fused_gate_up.size(1) / 2; + CHECK_EQ(full_intermediate % world_size, 0) + << "gate_up_proj intermediate dim is not divisible by world_size"; + const int64_t inter_shard = full_intermediate / world_size; + + auto gate_full = fused_gate_up.slice(1, 0, full_intermediate); + auto up_full = + fused_gate_up.slice(1, full_intermediate, full_intermediate * 2); + auto gate_shard = + gate_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard); + auto up_shard = + up_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard); + fused_gate_up = torch::cat({gate_shard, up_shard}, 1); + } + + auto gate_up_slice = slice_expert_weights( + fused_gate_up, start_expert_id, num_experts_per_rank); + CHECK_EQ(w13.sizes(), gate_up_slice.sizes()) + << "weight size mismatch for " << state_dict.prefix() + << "experts.gate_up_proj"; + w13.copy_(gate_up_slice); + return true; +} + +bool load_fused_down_fallback(const StateDict& state_dict, + int64_t rank, + int64_t world_size, + int64_t start_expert_id, + int64_t num_experts_per_rank, + torch::Tensor& w2) { + auto fused_down = get_tensor_with_weight_suffix(state_dict, "down_proj"); + if (!fused_down.defined()) { + return false; + } + + if (world_size > 1) { + CHECK_EQ(fused_down.size(2) % world_size, 0) + << "down_proj dim2 is not divisible by world_size"; + const int64_t down_shard = fused_down.size(2) / world_size; + fused_down = + fused_down.slice(2, rank * down_shard, (rank + 1) * down_shard); + } + + auto down_slice = + slice_expert_weights(fused_down, start_expert_id, num_experts_per_rank); + CHECK_EQ(w2.sizes(), down_slice.sizes()) + << "weight size mismatch for " << state_dict.prefix() + << "experts.down_proj"; + w2.copy_(down_slice); + return true; +} +} // namespace + +Qwen3_5FusedMoEImpl::Qwen3_5FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options) + : FusedMoEImpl(model_args, moe_args, quant_args, parallel_args, options) { + if (n_shared_experts_ > 0) { + shared_expert_gate_ = register_module( + "shared_expert_gate", + torch::nn::Linear( + torch::nn::LinearOptions(hidden_size_, 1).bias(false))); + shared_expert_gate_->weight.set_data( + shared_expert_gate_->weight.to(options)); + } +} + +void Qwen3_5FusedMoEImpl::load_experts(const StateDict& state_dict) { + FusedMoEImpl::load_experts(state_dict); + + if (!is_smoothquant_) { + if (!w13_is_loaded_) { + w13_is_loaded_ = load_fused_gate_up_fallback(state_dict, + tp_pg_->rank(), + tp_pg_->world_size(), + start_expert_id_, + num_experts_per_rank_, + w13_); + } + + if (!w2_is_loaded_) { + w2_is_loaded_ = load_fused_down_fallback(state_dict, + tp_pg_->rank(), + tp_pg_->world_size(), + start_expert_id_, + num_experts_per_rank_, + w2_); + } + } +} + +void Qwen3_5FusedMoEImpl::load_state_dict(const StateDict& state_dict) { + if (state_dict.size() == 0) { + return; + } + + if (n_shared_experts_ > 0) { + shared_experts_->load_state_dict( + state_dict.get_dict_with_prefix("shared_expert.")); + auto weight = state_dict.get_tensor("shared_expert_gate.weight"); + if (weight.defined()) { + weight = weight.reshape({weight.size(0), -1}); + DCHECK_EQ(shared_expert_gate_->weight.sizes(), weight.sizes()) + << "proj weight size mismatch for " << name(); + shared_expert_gate_->weight.data().copy_(weight); + } + } + gate_->load_state_dict(state_dict.get_dict_with_prefix("gate.")); + load_experts(state_dict.get_dict_with_prefix("experts.")); +} + +void Qwen3_5FusedMoEImpl::final_comm_allreduce( + torch::Tensor& final_hidden_states, + const torch::Tensor& hidden_states, + torch::Tensor& shared_expert_output) { + auto current_stream = device_.current_stream(); + routed_stream_->wait_stream(*current_stream); + { + torch::StreamGuard stream_guard = routed_stream_->set_stream_guard(); + if (tp_pg_->world_size() > 1) { + final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_); + } + if (parallel_args_.ep_size() > 1) { + final_hidden_states = parallel_state::reduce( + final_hidden_states, parallel_args_.moe_ep_group_); + } + } + + if (n_shared_experts_ > 0) { + shared_stream_->wait_stream(*current_stream); + torch::StreamGuard stream_guard = shared_stream_->set_stream_guard(); + shared_expert_output = shared_experts_(hidden_states); + if (shared_expert_gate_) { + auto gate = torch::sigmoid(shared_expert_gate_->forward(hidden_states)); + shared_expert_output = gate * shared_expert_output; + } + shared_expert_output = + shared_expert_output.reshape({-1, shared_expert_output.size(-1)}); + } + + // join for parallelization + current_stream->wait_stream(*routed_stream_); + if (n_shared_experts_ > 0) { + current_stream->wait_stream(*shared_stream_); + final_hidden_states += shared_expert_output; + } +} + +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.h b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.h new file mode 100644 index 00000000..150ae93c --- /dev/null +++ b/upstream_ref/xllm_latest/core/layers/mlu/qwen3_5_fused_moe.h @@ -0,0 +1,47 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#pragma once + +#include "layers/mlu/fused_moe.h" + +namespace xllm { +namespace layer { + +class Qwen3_5FusedMoEImpl final : public FusedMoEImpl { + public: + Qwen3_5FusedMoEImpl() = default; + + Qwen3_5FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + void load_state_dict(const StateDict& state_dict) override; + + protected: + void final_comm_allreduce(torch::Tensor& final_hidden_states, + const torch::Tensor& hidden_states, + torch::Tensor& shared_expert_output) override; + + private: + void load_experts(const StateDict& state_dict); + torch::nn::Linear shared_expert_gate_{nullptr}; +}; + +TORCH_MODULE(Qwen3_5FusedMoE); +} // namespace layer +} // namespace xllm diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp index 0b97be6d..7d572476 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.cpp @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2026 The xLLM Authors. All Rights Reserved. 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 @@ -123,53 +123,19 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations( } std::pair -Qwen3_5GatedDeltaNetImpl::project_decode_inputs( - const torch::Tensor& hidden_states) { - const auto reshape_projection = [](const torch::Tensor& projection) { - return projection.view({projection.size(0), -1, projection.size(-1)}); - }; - auto qkv = reshape_projection(in_proj_qkv_->forward(hidden_states)); - auto z_proj = reshape_projection(in_proj_z_->forward(hidden_states)); - auto b_proj = reshape_projection(in_proj_b_->forward(hidden_states)); - auto a_proj = reshape_projection(in_proj_a_->forward(hidden_states)); - return {merge_qkvz_from_split_activations(qkv, z_proj), - merge_ba_from_split_activations(b_proj, a_proj)}; -} - -std::pair -Qwen3_5GatedDeltaNetImpl::project_flat_inputs( - const torch::Tensor& hidden_states) { - auto qkv = in_proj_qkv_->forward(hidden_states).unsqueeze(0); - auto z_proj = in_proj_z_->forward(hidden_states).unsqueeze(0); - auto b_proj = in_proj_b_->forward(hidden_states).unsqueeze(0); - auto a_proj = in_proj_a_->forward(hidden_states).unsqueeze(0); - auto qkvz = merge_qkvz_from_split_activations(qkv, z_proj); - auto ba = merge_ba_from_split_activations(b_proj, a_proj); - return {qkvz.view({hidden_states.size(0), qkvz.size(-1)}).contiguous(), - ba.view({hidden_states.size(0), ba.size(-1)}).contiguous()}; -} - -std::optional< - std::tuple> -Qwen3_5GatedDeltaNetImpl::project_split_inputs( +Qwen3_5GatedDeltaNetImpl::project_padded_inputs( const torch::Tensor& hidden_states, const AttentionMetadata& attn_metadata) { - auto qkv = reshape_projected_tokens_with_pad( - attn_metadata, in_proj_qkv_->forward(hidden_states)); - auto z_proj = reshape_projected_tokens_with_pad( - attn_metadata, in_proj_z_->forward(hidden_states)); - auto b_proj = reshape_projected_tokens_with_pad( - attn_metadata, in_proj_b_->forward(hidden_states)); - auto a_proj = reshape_projected_tokens_with_pad( - attn_metadata, in_proj_a_->forward(hidden_states)); - - const int64_t batch_size = qkv.size(0); - const int64_t seq_len = qkv.size(1); - auto z = - z_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); - auto b = b_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_}); - auto a = a_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_}); - return std::make_tuple(qkv, z, b, a); + auto qkv = reshape_qkvz_with_pad(attn_metadata, + in_proj_qkv_->forward(hidden_states)); + auto z_proj = + reshape_qkvz_with_pad(attn_metadata, in_proj_z_->forward(hidden_states)); + auto b_proj = + reshape_qkvz_with_pad(attn_metadata, in_proj_b_->forward(hidden_states)); + auto a_proj = + reshape_qkvz_with_pad(attn_metadata, in_proj_a_->forward(hidden_states)); + return {merge_qkvz_from_split_activations(qkv, z_proj), + merge_ba_from_split_activations(b_proj, a_proj)}; } void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict( diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h index 7c782e3f..bec6c1c6 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_gated_delta_net.h @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2026 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -17,9 +17,7 @@ limitations under the License. #include -#include #include -#include #include #include "qwen3_next_gated_delta_net.h" @@ -36,15 +34,9 @@ class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl { const torch::TensorOptions& options); protected: - std::pair project_decode_inputs( - const torch::Tensor& hidden_states) override; - std::pair project_flat_inputs( - const torch::Tensor& hidden_states) override; - std::optional< - std::tuple> - project_split_inputs(const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata) override; - bool use_fla_ssm_state_layout() const override { return true; } + std::pair project_padded_inputs( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata) override; void load_projection_state_dict(const StateDict& state_dict) override; void verify_projection_weights(const std::string& prefix) const override; diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp index 7f8b4b5c..cec9a95e 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2026 The xLLM Authors. All Rights Reserved. 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 @@ -15,12 +15,9 @@ limitations under the License. #include #include -#include #include -#include "xllm/core/kernels/npu/npu_ops_api.h" #include "xllm/core/kernels/ops_api.h" -#include "xllm/core/platform/npu/acl_graph_task_update_context.h" namespace xllm { namespace layer { @@ -31,31 +28,6 @@ torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) { return x / norm; } -torch::Tensor repeat_tensor_heads(const torch::Tensor& tensor, - int64_t target_heads, - int64_t head_dim) { - const int64_t current_heads = tensor.size(head_dim); - if (current_heads == target_heads) { - return tensor; - } - CHECK_GT(current_heads, 0) << "current heads must be positive"; - CHECK_EQ(target_heads % current_heads, 0) - << "target heads must be divisible by current heads, target_heads=" - << target_heads << ", current_heads=" << current_heads; - - const int64_t repeats = target_heads / current_heads; - std::vector view_shape = tensor.sizes().vec(); - view_shape.insert(view_shape.begin() + head_dim + 1, 1); - std::vector expand_shape = view_shape; - expand_shape[head_dim + 1] = repeats; - std::vector output_shape = tensor.sizes().vec(); - output_shape[head_dim] = target_heads; - return tensor.unsqueeze(head_dim + 1) - .expand(expand_shape) - .reshape(output_shape) - .contiguous(); -} - std::tuple torch_recurrent_gated_delta_rule( torch::Tensor query, torch::Tensor key, @@ -80,9 +52,6 @@ std::tuple torch_recurrent_gated_delta_rule( value = to_float32_and_transpose(value); beta = to_float32_and_transpose(beta); g = to_float32_and_transpose(g); - const int64_t value_num_heads = value.size(1); - query = repeat_tensor_heads(query, value_num_heads, 1); - key = repeat_tensor_heads(key, value_num_heads, 1); int64_t batch_size = key.size(0); int64_t num_heads = key.size(1); @@ -150,15 +119,12 @@ std::tuple torch_chunk_gated_delta_rule( value = to_float32(value); beta = to_float32(beta); g = to_float32(g); - const int64_t value_num_heads = value.size(1); - query = repeat_tensor_heads(query, value_num_heads, 1); - key = repeat_tensor_heads(key, value_num_heads, 1); - int64_t batch_size = query.size(0); - int64_t num_heads = query.size(1); - int64_t sequence_length = query.size(2); - int64_t k_head_dim = key.size(-1); - int64_t v_head_dim = value.size(-1); + auto batch_size = query.size(0); + auto num_heads = query.size(1); + auto sequence_length = query.size(2); + auto k_head_dim = key.size(-1); + auto v_head_dim = value.size(-1); int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size; query = torch::nn::functional::pad( @@ -276,164 +242,6 @@ std::tuple torch_chunk_gated_delta_rule( core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype); return std::make_tuple(core_attn_out, last_recurrent_state); } - -int64_t get_checkpoint_stride(const torch::Tensor& conv_cache, - const torch::Tensor& ssm_cache) { - if (!conv_cache.defined() || !ssm_cache.defined() || - conv_cache.numel() == 0 || ssm_cache.numel() == 0) { - return 1; - } - CHECK_GT(conv_cache.size(0), 0) << "conv cache must have positive batch dim"; - CHECK_EQ(ssm_cache.size(0) % conv_cache.size(0), 0) - << "ssm cache checkpoint layout mismatch, ssm_rows=" << ssm_cache.size(0) - << ", conv_rows=" << conv_cache.size(0); - return ssm_cache.size(0) / conv_cache.size(0); -} - -torch::Tensor build_linear_state_base_indices( - const torch::Tensor& logical_state_indices, - int64_t checkpoint_stride) { - if (checkpoint_stride == 1) { - return logical_state_indices; - } - return logical_state_indices * checkpoint_stride; -} - -torch::Tensor expand_sequence_tensor_to_batch(const torch::Tensor& tensor, - int64_t target_batch, - const char* tensor_name) { - CHECK(tensor.defined()) << tensor_name << " must be defined"; - CHECK_EQ(tensor.dim(), 1) << tensor_name << " must be a 1D tensor."; - const int64_t source_batch = tensor.size(0); - if (source_batch == target_batch) { - return tensor.contiguous(); - } - CHECK_GT(source_batch, 0) << tensor_name << " must not be empty."; - CHECK_EQ(target_batch % source_batch, 0) - << tensor_name << " cannot be expanded from " << source_batch << " to " - << target_batch; - const int64_t repeat_count = target_batch / source_batch; - return tensor.unsqueeze(1) - .expand({source_batch, repeat_count}) - .reshape({target_batch}) - .contiguous(); -} - -torch::Tensor run_causal_conv1d_graph_update( - const std::shared_ptr& graph_context, - const torch::Tensor& x, - const torch::Tensor& weight, - const torch::Tensor& conv_state, - const std::optional& bias, - const std::vector& query_start_loc, - const std::vector& cache_indices, - const std::vector& num_accepted_tokens, - xllm::npu::CausalConv1dGraphBranch branch) { - CHECK(graph_context != nullptr && graph_context->capturing) - << "causal_conv1d graph update can only be registered during capture"; - - c10_npu::NPUStream stream = c10_npu::getCurrentNPUStream(); - auto event = std::make_shared(ACL_EVENT_EXTERNAL); - event->block(stream); - event->reset(stream); - - torch::Tensor output; - c10_npu::graph_task_group_begin(stream); - const std::vector empty_host_args; - CHECK(!query_start_loc.empty()) - << "query_start_loc must be populated for causal_conv1d graph update"; - CHECK_EQ(query_start_loc.back(), x.size(0)) - << "query_start_loc must be padded to x.shape[0] during graph capture"; - CHECK_EQ(cache_indices.size() + 1, query_start_loc.size()) - << "cache_indices must be sequence-scoped"; - if (branch == xllm::npu::CausalConv1dGraphBranch::kSpecVerify) { - CHECK_EQ(num_accepted_tokens.size(), cache_indices.size()) - << "num_accepted_tokens must be sequence-scoped for spec verify"; - } - - output = torch::empty_like(x); - xllm::kernel::causal_conv1d_out(output, - x, - weight, - conv_state, - bias, - torch::IntArrayRef(query_start_loc), - torch::IntArrayRef(cache_indices), - torch::IntArrayRef(empty_host_args), - torch::IntArrayRef(num_accepted_tokens), - xllm::npu::kCausalConv1dActivationSilu, - xllm::npu::kCausalConv1dGraphPadSlotId, - xllm::npu::kCausalConv1dRunModeUpdate); - c10_npu::NPUTaskGroupHandle handle = c10_npu::graph_task_group_end(stream); - - xllm::npu::CausalConv1dGraphTask task; - task.output = output; - task.x = x; - task.weight = weight; - task.conv_state = conv_state; - task.bias = bias; - task.activation_mode = xllm::npu::kCausalConv1dActivationSilu; - task.pad_slot_id = xllm::npu::kCausalConv1dGraphPadSlotId; - task.run_mode = xllm::npu::kCausalConv1dRunModeUpdate; - task.branch = branch; - task.handle = handle; - task.event = std::move(event); - graph_context->causal_conv1d_tasks.emplace_back(std::move(task)); - return output; -} - -torch::Tensor run_spec_verify_gated_delta_rule( - torch::Tensor query, - torch::Tensor key, - torch::Tensor value, - torch::Tensor g, - torch::Tensor beta, - torch::Tensor& ssm_cache, - const torch::Tensor& checkpoint_indices, - const torch::Tensor& num_accepted_tokens, - const torch::Tensor& cu_seq_lens, - const std::vector& q_seq_lens_vec, - double scale) { - const auto device = value.device(); - const int64_t batch_size = value.size(0); - const int64_t seq_len = value.size(1); - const int64_t total_seq_len = batch_size * seq_len; - CHECK_EQ(cu_seq_lens.numel(), batch_size + 1) - << "GDN spec verify cu_seq_lens must be cumulative."; - CHECK_EQ(q_seq_lens_vec.size(), static_cast(batch_size)) - << "GDN spec verify q_seq_lens_vec must be per sequence."; - for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - CHECK_EQ(q_seq_lens_vec[batch_idx], seq_len) - << "Qwen3.5 spec verify fused recurrent path expects dense " - "same-length validate tokens."; - } - - xllm::kernel::FusedRecurrentGatedDeltaRuleParams params; - params.q = query.reshape({1, total_seq_len, query.size(-2), query.size(-1)}) - .contiguous(); - params.k = - key.reshape({1, total_seq_len, key.size(-2), key.size(-1)}).contiguous(); - params.v = value.reshape({1, total_seq_len, value.size(-2), value.size(-1)}) - .contiguous(); - params.g = g.to(torch::kFloat32) - .reshape({1, total_seq_len, g.size(-1)}) - .contiguous(); - params.beta = beta.reshape({1, total_seq_len, beta.size(-1)}).contiguous(); - params.scale = static_cast(scale); - params.initial_state = ssm_cache; - params.inplace_final_state = true; - params.cu_seqlens = cu_seq_lens.to(torch::kLong).contiguous(); - params.ssm_state_indices = checkpoint_indices.contiguous(); - params.num_accepted_tokens = - num_accepted_tokens.to(device, torch::kInt32).contiguous(); - params.use_qk_l2norm_in_kernel = true; - - auto output_and_state = - xllm::kernel::fused_recurrent_gated_delta_rule(params); - return output_and_state.first.view( - {batch_size, seq_len, value.size(-2), value.size(-1)}); -} - } // namespace Qwen3GatedDeltaNetBaseImpl::Qwen3GatedDeltaNetBaseImpl( @@ -495,11 +303,7 @@ void Qwen3GatedDeltaNetBaseImpl::load_common_state_dict( if (auto w = state_dict.get_tensor("conv1d.weight"); w.defined()) { conv1d_->load_state_dict( - StateDict({{"weight", w.squeeze(1)}}, - static_cast(state_dict.prefix()) + "conv1d."), - shard_tensor_count, - shard_sizes); - conv1d_->weight().set_(conv1d_->weight().transpose(0, 1).contiguous()); + StateDict({{"weight", w.squeeze(1)}}), shard_tensor_count, shard_sizes); } o_proj_->load_state_dict(state_dict.get_dict_with_prefix("out_proj.")); if (auto w = state_dict.get_tensor("norm.weight"); w.defined()) { @@ -518,279 +322,87 @@ void Qwen3GatedDeltaNetBaseImpl::verify_common_loaded_weights( << prefix << "A_log"; } -std::pair -Qwen3GatedDeltaNetBaseImpl::project_padded_inputs( - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata) { - if (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill) { - auto [qkvz_flat, ba_flat] = project_flat_inputs(hidden_states); - return {reshape_projected_tokens_with_pad(attn_metadata, qkvz_flat), - reshape_projected_tokens_with_pad(attn_metadata, ba_flat)}; - } - return project_decode_inputs(hidden_states); -} - torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( const torch::Tensor& hidden_states, const AttentionMetadata& attn_metadata, KVCache& kv_cache, const ModelInputParams& input_params) { - // Early-return on dummy shards. Under dp>1, an empty shard is padded with a - // fake token by worker_impl but its GDN state tensors (kv_cache_tokens_nums, - // linear_state_ids etc.) are left undefined. This mirrors the is_dummy - // early-return in Attention::forward (npu_torch/attention.cpp). Uses - // zeros_like rather than empty_like so downstream post-norm / mlp do not - // read uninitialized data. Placed before FlashComm1 sequence gather so - // dummy shards do not enter the collective and waste bandwidth. - if (attn_metadata.is_dummy) { - return torch::zeros_like(hidden_states); - } - const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); - torch::Tensor h = hidden_states; - if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { - h = gather_sequence(hidden_states, *fc1_ctx); - } + auto [qkvz_padded, ba_padded] = + project_padded_inputs(hidden_states, attn_metadata); + int64_t batch_size = qkvz_padded.size(0); + int64_t seq_len = qkvz_padded.size(1); + + torch::Tensor qkvz_flat = + qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)}); + torch::Tensor ba_flat = + ba_padded.view({batch_size * seq_len, ba_padded.size(-1)}); + xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params; + fused_params.mixed_qkvz = qkvz_flat; + fused_params.mixed_ba = ba_flat; + fused_params.num_heads_qk = static_cast(num_k_heads_ / tp_size_); + fused_params.num_heads_v = static_cast(num_v_heads_ / tp_size_); + fused_params.head_qk = static_cast(head_k_dim_); + fused_params.head_v = static_cast(head_v_dim_); - // Save the gathered hidden-state size for potential padding later. - const int64_t original_num_tokens = h.size(0); - const bool use_spec_verify = input_params.is_spec_verify; - const bool is_any_prefill = - attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; torch::Tensor mixed_qkv, z, b, a; - torch::Tensor processed_q, processed_k, processed_v; - int64_t batch_size = 0; - int64_t seq_len = 0; + std::tie(mixed_qkv, z, b, a) = + xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params); - // Qwen3.5 stores qkv, z, b, and a as separate projection weights, so it can - // use their outputs directly in every forward mode. Qwen3Next stores qkvz - // and ba as packed weights and uses the fused-split fallback below. - auto split_inputs = project_split_inputs(h, attn_metadata); - if (split_inputs.has_value()) { - std::tie(mixed_qkv, z, b, a) = split_inputs.value(); - batch_size = mixed_qkv.size(0); - seq_len = mixed_qkv.size(1); - } else { - auto [qkvz_padded, ba_padded] = project_padded_inputs(h, attn_metadata); - batch_size = qkvz_padded.size(0); - seq_len = qkvz_padded.size(1); - - torch::Tensor qkvz_flat = - qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)}); - torch::Tensor ba_flat = - ba_padded.view({batch_size * seq_len, ba_padded.size(-1)}); - xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params; - fused_params.mixed_qkvz = qkvz_flat; - fused_params.mixed_ba = ba_flat; - fused_params.num_heads_qk = static_cast(num_k_heads_ / tp_size_); - fused_params.num_heads_v = static_cast(num_v_heads_ / tp_size_); - fused_params.head_qk = static_cast(head_k_dim_); - fused_params.head_v = static_cast(head_v_dim_); - - std::tie(mixed_qkv, z, b, a) = - xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params); - - mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)}); - z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); - b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_}); - a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_}); - } - - const bool fla_ssm_state_layout = use_fla_ssm_state_layout(); - const int64_t local_q_heads = num_k_heads_ / tp_size_; - const int64_t local_v_heads = num_v_heads_ / tp_size_; - const int64_t local_conv_dim = - 2 * local_q_heads * head_k_dim_ + local_v_heads * head_v_dim_; - bool used_direct_prefill_qkv = false; + mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)}); + z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); + b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_}); + a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_}); torch::Tensor conv_cache = kv_cache.get_conv_cache(); torch::Tensor ssm_cache = kv_cache.get_ssm_cache(); - torch::Device device = mixed_qkv.device(); - torch::Tensor conv_weight = conv1d_->weight(); - torch::Tensor logical_state_indices = - get_linear_state_indices(input_params, device); - const int64_t checkpoint_stride = - get_checkpoint_stride(conv_cache, ssm_cache); - torch::Tensor linear_state_base_indices = - build_linear_state_base_indices(logical_state_indices, checkpoint_stride); - auto graph_context = input_params.graph.acl_graph_task_update_context; - const bool register_conv1d_graph_update = - graph_context != nullptr && graph_context->capturing; + torch::Tensor g, beta, core_attn_out, last_recurrent_state; + auto device = mixed_qkv.device(); + auto conv_weight = conv1d_->weight(); + auto linear_state_indices = get_linear_state_indices(input_params, device); - if (!use_spec_verify && is_any_prefill) { - torch::IntArrayRef num_accepted_tokens_opt; - std::vector linear_state_indices_vec( - input_params.embedding.linear_state_ids.begin(), - input_params.embedding.linear_state_ids.end()); - torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv); + if (attn_metadata.is_prefill) { + mixed_qkv = mixed_qkv.transpose(1, 2); + torch::Tensor conv_state = + (seq_len < conv_kernel_size_ - 1) + ? torch::pad(mixed_qkv, {0, conv_kernel_size_ - 1 - seq_len}) + : (seq_len > conv_kernel_size_ - 1) + ? mixed_qkv.narrow( + -1, seq_len - conv_kernel_size_ + 1, conv_kernel_size_ - 1) + : mixed_qkv; + conv_state = conv_state.transpose(1, 2).contiguous(); + conv_cache.index_put_({linear_state_indices}, + conv_state.to(conv_cache.dtype())); + torch::Tensor bias; + auto conv_output = + torch::conv1d(mixed_qkv, + conv_weight.unsqueeze(1).to(device), + bias, + /*stride=*/std::vector{1}, + /*padding=*/std::vector{3}, + /*dilation=*/std::vector{1}, + /*groups=*/static_cast(mixed_qkv.size(1))); + mixed_qkv = torch::silu(conv_output.slice(2, 0, seq_len)); - const bool direct_qkv_model_supported = - fla_ssm_state_layout && num_k_heads_ % tp_size_ == 0 && - num_v_heads_ % tp_size_ == 0 && local_q_heads > 0 && - local_v_heads > 0 && head_k_dim_ == 128 && head_v_dim_ == 128; - const bool direct_qkv_metadata_available = - attn_metadata.q_seq_lens_vec.size() == - static_cast(batch_size) && - input_params.parallel.query_start_loc.size() == - static_cast(batch_size + 1) && - input_params.embedding.linear_state_ids.size() == - static_cast(batch_size) && - input_params.linear_state_validity_mask.size() == - static_cast(batch_size); - int64_t total_valid_tokens = 0; - bool direct_qkv_lengths_valid = direct_qkv_metadata_available; - if (direct_qkv_metadata_available) { - for (const int32_t valid_len : attn_metadata.q_seq_lens_vec) { - direct_qkv_lengths_valid = - direct_qkv_lengths_valid && valid_len >= 0 && valid_len <= seq_len; - total_valid_tokens += valid_len; - } - } - const bool direct_qkv_sequence_supported = - direct_qkv_model_supported && direct_qkv_lengths_valid && - conv_input.dim() == 2 && total_valid_tokens == conv_input.size(0); - const bool direct_qkv_shape_supported = - direct_qkv_sequence_supported && conv_input.size(1) == local_conv_dim && - conv_weight.dim() == 2 && conv_weight.size(0) == 4 && - conv_weight.size(1) == local_conv_dim && conv_cache.dim() == 3 && - conv_cache.size(1) >= 3 && conv_cache.size(2) == local_conv_dim; - const bool direct_qkv_dtype_supported = - direct_qkv_shape_supported && - conv_input.scalar_type() == torch::kBFloat16 && - conv_weight.scalar_type() == torch::kBFloat16 && - conv_cache.scalar_type() == torch::kBFloat16; - const bool use_direct_prefill_qkv = - direct_qkv_dtype_supported && conv_input.is_contiguous() && - conv_weight.is_contiguous() && conv_cache.is_contiguous(); - if (use_direct_prefill_qkv) { - std::tie(processed_q, processed_k, processed_v) = - xllm::kernel::npu::causal_conv1d_qkv( - conv_input, - conv_weight, - conv_cache, - torch::IntArrayRef(input_params.parallel.query_start_loc), - torch::IntArrayRef(linear_state_indices_vec), - torch::IntArrayRef(input_params.linear_state_validity_mask), - local_q_heads, - local_v_heads, - head_k_dim_, - head_v_dim_); - used_direct_prefill_qkv = true; - } else { - mixed_qkv = xllm::kernel::causal_conv1d( - conv_input, - conv_weight, - conv_cache, - std::optional(), // bias (no bias for qwen3) - torch::IntArrayRef(input_params.parallel.query_start_loc), - torch::IntArrayRef(linear_state_indices_vec), - torch::IntArrayRef(input_params.linear_state_validity_mask), - num_accepted_tokens_opt, - xllm::npu::kCausalConv1dActivationSilu, - xllm::npu::kCausalConv1dGraphPadSlotId, - xllm::npu::kCausalConv1dRunModeForward); - - mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv); - mixed_qkv = mixed_qkv.transpose(1, 2); - } } else { - if (use_spec_verify) { - CHECK(input_params.num_accepted_tokens.defined()) - << "num_accepted_tokens must be populated for Qwen3.5 spec verify"; - } - torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv); - const auto& num_accepted = use_spec_verify - ? input_params.num_accepted_tokens_host - : std::vector(); - const std::vector linear_state_indices_host( - input_params.embedding.linear_state_ids.begin(), - input_params.embedding.linear_state_ids.end()); - if (register_conv1d_graph_update) { - if (use_spec_verify) { - const auto conv1d_branch = - xllm::npu::CausalConv1dGraphBranch::kSpecVerify; - mixed_qkv = run_causal_conv1d_graph_update( - graph_context, - conv_input, - conv_weight, - conv_cache, - std::optional(), - input_params.parallel.query_start_loc, - linear_state_indices_host, - num_accepted, - conv1d_branch); - } else { - auto conv_input_2d = conv_input.dim() == 3 - ? conv_input.reshape({-1, conv_input.size(-1)}) - : conv_input; - xllm::kernel::CausalConv1dUpdateParams conv1d_params; - conv1d_params.x = conv_input_2d; - conv1d_params.conv_state = conv_cache; - conv1d_params.weight = conv_weight; - conv1d_params.conv_state_indices = logical_state_indices; - conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens; - conv1d_params.max_query_len = attn_metadata.max_query_len; - mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params); - if (conv_input.dim() == 3) { - mixed_qkv = - mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)}); - } - } - } else { - if (use_spec_verify) { - torch::Tensor output = torch::empty_like(conv_input); - xllm::kernel::causal_conv1d_out( - output, - conv_input, - conv_weight, - conv_cache, - std::optional(), - torch::IntArrayRef(input_params.parallel.query_start_loc), - torch::IntArrayRef(linear_state_indices_host), - torch::IntArrayRef(std::vector()), - torch::IntArrayRef(num_accepted), - xllm::npu::kCausalConv1dActivationSilu, - xllm::npu::kCausalConv1dGraphPadSlotId, - xllm::npu::kCausalConv1dRunModeUpdate); - mixed_qkv = output; - } else { - auto conv_input_2d = conv_input.dim() == 3 - ? conv_input.reshape({-1, conv_input.size(-1)}) - : conv_input; - xllm::kernel::CausalConv1dUpdateParams conv1d_params; - conv1d_params.x = conv_input_2d; - conv1d_params.conv_state = conv_cache; - conv1d_params.weight = conv_weight; - conv1d_params.conv_state_indices = logical_state_indices; - conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens; - conv1d_params.max_query_len = attn_metadata.max_query_len; - mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params); - if (conv_input.dim() == 3) { - mixed_qkv = - mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)}); - } - } - } - mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv); + xllm::kernel::CausalConv1dUpdateParams conv1d_params; + conv1d_params.x = mixed_qkv.reshape({-1, mixed_qkv.size(-1)}); + conv1d_params.conv_state = conv_cache; + conv1d_params.weight = conv_weight; + conv1d_params.conv_state_indices = linear_state_indices; + conv1d_params.block_idx_last_scheduled_token = + std::optional(); + conv1d_params.initial_state_idx = std::optional(); + conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens; + conv1d_params.max_query_len = attn_metadata.max_query_len; + mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params); + // Reshape back to 3D [batch_size, dim, seq_len] + mixed_qkv = + mixed_qkv.view({batch_size, -1, mixed_qkv.size(-1)}).contiguous(); mixed_qkv = mixed_qkv.transpose(1, 2); } - const bool use_fused_sigmoid_gdn_decode = - fla_ssm_state_layout && !use_spec_verify && !is_any_prefill && - checkpoint_stride == 1; - torch::Tensor g; - torch::Tensor beta; + // Compute gated delta net decay and beta terms. - if (use_spec_verify || attn_metadata.is_chunked_prefill || - checkpoint_stride > 1) { - beta = torch::sigmoid(b); - torch::Tensor A_log_exp = A_log_.exp(); - torch::Tensor a_float = a.to(torch::kFloat32); - torch::Tensor a_plus_dt = a_float + dt_bias_; - torch::Tensor softplus_out = torch::nn::functional::softplus( - a_plus_dt, - torch::nn::functional::SoftplusFuncOptions().beta(1.0).threshold(20.0)); - g = -A_log_exp * softplus_out; - g = g.to(a.dtype()).contiguous(); - } else if (attn_metadata.is_prefill) { + if (attn_metadata.is_prefill) { xllm::kernel::FusedGdnGatingParams gdn_params; gdn_params.A_log = A_log_; gdn_params.a = a.contiguous().view({-1, a.size(-1)}); @@ -801,7 +413,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params); g = g.squeeze(0).contiguous().view({batch_size, seq_len, a.size(-1)}); beta = beta.squeeze(0).contiguous().view({batch_size, seq_len, b.size(-1)}); - } else if (!use_fused_sigmoid_gdn_decode) { + } else { xllm::kernel::FusedGdnGatingParams gdn_params; gdn_params.A_log = A_log_; gdn_params.a = a.view({-1, a.size(-1)}); @@ -811,216 +423,57 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( gdn_params.threshold = 20.0f; std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params); } - if (!used_direct_prefill_qkv) { - std::tie(processed_q, processed_k, processed_v) = - process_mixed_qkv(mixed_qkv); - } - torch::Tensor core_attn_out; - torch::Tensor last_recurrent_state; + auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv); // Apply chunked or recurrent gated-delta attention and update caches. - if (use_spec_verify) { - torch::Tensor spec_num_accepted_tokens = expand_sequence_tensor_to_batch( - input_params.num_accepted_tokens.to(device, torch::kInt32), - batch_size, - "num_accepted_tokens"); - torch::Tensor spec_linear_state_base_indices = - expand_sequence_tensor_to_batch( - linear_state_base_indices, batch_size, "linear_state_base_indices"); - torch::Tensor step_offsets = - torch::arange(seq_len, - torch::TensorOptions() - .dtype(spec_linear_state_base_indices.dtype()) - .device(device)); - torch::Tensor checkpoint_indices = - spec_linear_state_base_indices.unsqueeze(1) + step_offsets; - double scale = 1.0 / std::sqrt(static_cast(processed_q.size(-1))); - core_attn_out = - run_spec_verify_gated_delta_rule(processed_q, - processed_k, - processed_v, - g, - beta, - ssm_cache, - checkpoint_indices, - spec_num_accepted_tokens, - attn_metadata.q_cu_seq_lens, - attn_metadata.q_seq_lens_vec, - scale); - } else if (is_any_prefill) { - CHECK_GE(attn_metadata.q_seq_lens_vec.size(), - static_cast(batch_size)) - << "q_seq_lens_vec must be populated for Qwen3.5 prefill."; - const bool use_single_prefill_pack = - batch_size == 1 && attn_metadata.q_seq_lens_vec.size() == 1 && - attn_metadata.q_seq_lens_vec[0] == seq_len; - torch::Tensor packed_processed_q; - torch::Tensor packed_processed_k; - torch::Tensor packed_processed_v; - torch::Tensor packed_g_tensor; - torch::Tensor packed_beta_tensor; - if (use_single_prefill_pack) { - packed_processed_q = processed_q; - packed_processed_k = processed_k; - packed_processed_v = processed_v; - packed_g_tensor = g; - packed_beta_tensor = beta; - } else { - std::vector packed_q; - std::vector packed_k; - std::vector packed_v; - std::vector packed_g; - std::vector packed_beta; - packed_q.reserve(batch_size); - packed_k.reserve(batch_size); - packed_v.reserve(batch_size); - packed_g.reserve(batch_size); - packed_beta.reserve(batch_size); - for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx]; - if (!used_direct_prefill_qkv) { - packed_q.emplace_back(processed_q[batch_idx].narrow( - /*dim=*/0, /*start=*/0, valid_len)); - packed_k.emplace_back(processed_k[batch_idx].narrow( - /*dim=*/0, /*start=*/0, valid_len)); - packed_v.emplace_back(processed_v[batch_idx].narrow( - /*dim=*/0, /*start=*/0, valid_len)); - } - packed_g.emplace_back( - g[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len)); - packed_beta.emplace_back( - beta[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len)); - } - if (used_direct_prefill_qkv) { - packed_processed_q = processed_q; - packed_processed_k = processed_k; - packed_processed_v = processed_v; - } else { - packed_processed_q = torch::cat(packed_q, 0).unsqueeze(0); - packed_processed_k = torch::cat(packed_k, 0).unsqueeze(0); - packed_processed_v = torch::cat(packed_v, 0).unsqueeze(0); - } - packed_g_tensor = torch::cat(packed_g, 0).unsqueeze(0); - packed_beta_tensor = torch::cat(packed_beta, 0).unsqueeze(0); - } - - xllm::kernel::MegaChunkGdnParams mega_chunk_gdn_params; - mega_chunk_gdn_params.q = packed_processed_q; - mega_chunk_gdn_params.k = packed_processed_k; - mega_chunk_gdn_params.v = packed_processed_v; - mega_chunk_gdn_params.g = packed_g_tensor; - mega_chunk_gdn_params.beta = packed_beta_tensor; + if (attn_metadata.is_prefill) { + xllm::kernel::ChunkGatedDeltaRuleParams chunk_gated_delta_params; + chunk_gated_delta_params.q = processed_q; + chunk_gated_delta_params.k = processed_k; + chunk_gated_delta_params.v = processed_v; + chunk_gated_delta_params.g = g; + chunk_gated_delta_params.beta = beta; // Get initial state from ssm_cache for sequences with previous state // Shape: [batch_size, num_heads, head_k_dim, head_v_dim] torch::Tensor initial_state_tensor = - torch::index_select(ssm_cache, 0, linear_state_base_indices); - CHECK_EQ(input_params.linear_state_validity_mask.size(), - input_params.embedding.linear_state_ids.size()) - << "linear state validity mask must be sequence-scoped."; - for (size_t i = 0; i < input_params.linear_state_validity_mask.size(); - ++i) { - if (input_params.linear_state_validity_mask[i] == 0) { - initial_state_tensor.select(0, static_cast(i)).fill_(0.0); - } - } - if (!fla_ssm_state_layout && attn_metadata.is_chunked_prefill) { - initial_state_tensor = - initial_state_tensor.transpose(-1, -2).contiguous(); - } - mega_chunk_gdn_params.initial_state = initial_state_tensor; - mega_chunk_gdn_params.output_final_state = true; - mega_chunk_gdn_params.cu_seqlens = attn_metadata.q_cu_seq_lens; - mega_chunk_gdn_params.q_seq_lens = c10::ArrayRef( - attn_metadata.q_seq_lens_vec.data(), static_cast(batch_size)); - mega_chunk_gdn_params.use_qk_l2norm_in_kernel = !used_direct_prefill_qkv; - torch::Tensor packed_core_attn_out; - std::tie(packed_core_attn_out, last_recurrent_state) = - xllm::kernel::mega_chunk_gdn(mega_chunk_gdn_params); - if (use_single_prefill_pack) { - core_attn_out = packed_core_attn_out; - if (core_attn_out.scalar_type() != processed_v.scalar_type()) { - core_attn_out = core_attn_out.to(processed_v.scalar_type()); - } - } else { - core_attn_out = - used_direct_prefill_qkv - ? torch::zeros({batch_size, seq_len, local_v_heads, head_v_dim_}, - z.options()) - : torch::zeros_like(processed_v); - int64_t packed_offset = 0; - for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx]; - core_attn_out[batch_idx] - .narrow(/*dim=*/0, /*start=*/0, valid_len) - .copy_(packed_core_attn_out[0].narrow( - /*dim=*/0, packed_offset, valid_len)); - packed_offset += valid_len; - } - } - torch::Tensor state_to_store = fla_ssm_state_layout - ? last_recurrent_state - : last_recurrent_state.transpose(-1, -2); - ssm_cache.index_put_({linear_state_base_indices}, - state_to_store.to(ssm_cache.dtype())); - } else if (checkpoint_stride > 1) { - auto ssm_state = - torch::index_select(ssm_cache, 0, linear_state_base_indices); - if (!fla_ssm_state_layout) { - ssm_state = ssm_state.transpose(-1, -2); - } - ssm_state = ssm_state.contiguous(); + torch::index_select(ssm_cache, 0, linear_state_indices); + // Todo: chunked-prefill/prefix-cache use initial_state + initial_state_tensor.fill_(0.0); + chunk_gated_delta_params.initial_state = initial_state_tensor; + chunk_gated_delta_params.output_final_state = true; + chunk_gated_delta_params.cu_seqlens = attn_metadata.q_cu_seq_lens; + chunk_gated_delta_params.head_first = false; + chunk_gated_delta_params.use_qk_l2norm_in_kernel = true; std::tie(core_attn_out, last_recurrent_state) = - torch_recurrent_gated_delta_rule( - processed_q, processed_k, processed_v, g, beta, ssm_state); - torch::Tensor state_to_store = fla_ssm_state_layout - ? last_recurrent_state - : last_recurrent_state.transpose(-1, -2); - ssm_cache.index_put_({linear_state_base_indices}, - state_to_store.to(ssm_cache.dtype())); + xllm::kernel::chunk_gated_delta_rule(chunk_gated_delta_params); + ssm_cache.index_put_( + {linear_state_indices}, + last_recurrent_state.transpose(-1, -2).to(ssm_cache.dtype())); } else { + processed_q = xllm::kernel::l2_norm(processed_q, 1e-6); + processed_k = xllm::kernel::l2_norm(processed_k, 1e-6); + auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options()); + torch::Tensor actual_seq_lengths = + torch::cat({zero, attn_metadata.q_seq_lens}, 0); double scale = 1.0 / std::sqrt(static_cast(processed_q.size(-1))); - if (fla_ssm_state_layout) { - xllm::kernel::FusedSigmoidGatingDeltaRuleUpdateParams params; - params.A_log = A_log_.contiguous(); - params.a = a.contiguous(); - params.dt_bias = dt_bias_.contiguous(); - params.q = processed_q.contiguous(); - params.k = processed_k.contiguous(); - params.v = processed_v.contiguous(); - params.b = b.contiguous(); - params.initial_state_source = ssm_cache; - params.initial_state_indices = linear_state_base_indices.contiguous(); - params.cu_seqlens = attn_metadata.q_cu_seq_lens.contiguous(); - params.scale = static_cast(scale); - params.use_qk_l2norm_in_kernel = true; - params.softplus_beta = 1.0f; - params.softplus_threshold = 20.0f; - core_attn_out = - xllm::kernel::fused_sigmoid_gating_delta_rule_update(params); - } else { - processed_q = xllm::kernel::l2_norm(processed_q, /*eps=*/1e-6); - processed_k = xllm::kernel::l2_norm(processed_k, /*eps=*/1e-6); - auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options()); - torch::Tensor actual_seq_lengths = - torch::cat({zero, attn_metadata.q_seq_lens}, 0); - core_attn_out = xllm::kernel::recurrent_gated_delta_rule( - processed_q.reshape( - {-1, processed_q.size(-2), processed_q.size(-1)}), - processed_k.reshape( - {-1, processed_k.size(-2), processed_k.size(-1)}), - processed_v.reshape( - {-1, processed_v.size(-2), processed_v.size(-1)}), - ssm_cache, - beta.squeeze(0).contiguous(), - scale, - actual_seq_lengths, - logical_state_indices, - c10::nullopt, - g.squeeze(0).contiguous(), - c10::nullopt) - .unsqueeze(0) - .contiguous(); - } + core_attn_out = xllm::kernel::recurrent_gated_delta_rule( + processed_q.reshape( + {-1, processed_q.size(-2), processed_q.size(-1)}), + processed_k.reshape( + {-1, processed_k.size(-2), processed_k.size(-1)}), + processed_v.reshape( + {-1, processed_v.size(-2), processed_v.size(-1)}), + ssm_cache, + beta.squeeze(0).contiguous(), + scale, + actual_seq_lengths, + linear_state_indices, + c10::nullopt, + g.squeeze(0).contiguous(), + c10::nullopt) + .unsqueeze(0) + .contiguous(); } + auto z_reshaped = z.view({-1, z.size(-1)}); auto core_attn_out_reshaped = core_attn_out.view({-1, core_attn_out.size(-1)}); @@ -1033,47 +486,25 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( auto rearranged_norm = norm_out.reshape({norm_out.size(0), norm_out.size(1) * norm_out.size(2)}); rearranged_norm = reshape_qkvz_unpad(attn_metadata, rearranged_norm); - // For chunked prefill or spec verify, reshape_projected_tokens_with_pad may - // pad each batch to max_len, causing output tokens > original_num_tokens. We - // need to slice back to original_num_tokens to match the residual shape. - if (rearranged_norm.size(0) > original_num_tokens) { - // Slice excess padding tokens - rearranged_norm = - rearranged_norm.slice(0, 0, original_num_tokens).contiguous(); - } - if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { - return o_proj_->forward(rearranged_norm, - row_parallel_reduce_mode_for_fc1(*fc1_ctx)); - } - return o_proj_->forward(rearranged_norm); + auto attn_output = o_proj_->forward(rearranged_norm); + return attn_output; } torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad( const AttentionMetadata& attn_metadata, const torch::Tensor& padded_qkvz) const { - const bool has_padded_queries = - attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; - if (!has_padded_queries) { + if (!attn_metadata.is_prefill) { return padded_qkvz; } std::vector valid_batches; - const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty(); - int64_t bs = has_host_lens - ? static_cast(attn_metadata.q_seq_lens_vec.size()) - : attn_metadata.q_seq_lens.size(0); - valid_batches.reserve(bs); + int64_t bs = attn_metadata.q_seq_lens.size(0); int64_t max_len = attn_metadata.max_query_len; const auto& ori_seq_lens = attn_metadata.q_seq_lens; auto reshaped_qkvz = padded_qkvz.view({bs, max_len, -1}); for (int64_t b = 0; b < bs; ++b) { - int64_t ori_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b] - : ori_seq_lens[b].template item(); - torch::Tensor valid_batch = - reshaped_qkvz[b].slice(/*dim=*/0, /*start=*/0, ori_len); - valid_batches.emplace_back(valid_batch); - } - if (valid_batches.size() == 1) { - return valid_batches[0].contiguous(); + int64_t ori_len = ori_seq_lens[b].template item(); + torch::Tensor valid_batch = reshaped_qkvz[b].slice(0, 0, ori_len); + valid_batches.push_back(valid_batch); } return torch::cat(valid_batches, 0).contiguous(); } @@ -1081,60 +512,41 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad( torch::Tensor Qwen3GatedDeltaNetBaseImpl::get_linear_state_indices( const ModelInputParams& input_params, const torch::Device& device) const { - CHECK(!input_params.embedding.linear_state_ids.empty()) + CHECK(!input_params.linear_state_ids.empty()) << "linear_state_ids must be populated for gated delta net"; - if (input_params.embedding.linear_state_indices.defined()) { - auto indices = input_params.embedding.linear_state_indices; - if (indices.device() != device || indices.scalar_type() != torch::kInt) { - indices = - indices.to(torch::TensorOptions().dtype(torch::kInt).device(device), - /*non_blocking=*/true, - /*copy=*/true); - } - return indices.contiguous(); + if (input_params.linear_state_indices.defined()) { + return input_params.linear_state_indices; } return torch::tensor( - input_params.embedding.linear_state_ids, + input_params.linear_state_ids, torch::TensorOptions().dtype(torch::kInt).device(device)); } -torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_projected_tokens_with_pad( +torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad( const AttentionMetadata& attn_metadata, - const torch::Tensor& projected_tokens) const { - const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty(); - int64_t bs = has_host_lens - ? static_cast(attn_metadata.q_seq_lens_vec.size()) - : attn_metadata.q_seq_lens.size(0); + const torch::Tensor& qkvz) const { + int64_t bs = attn_metadata.q_seq_lens.size(0); int64_t max_len = attn_metadata.max_query_len; const auto& start_loc = attn_metadata.q_seq_lens; - const bool need_padding = - attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; - if (!need_padding) { - return projected_tokens.view({bs, -1, projected_tokens.size(-1)}); - } - if (has_host_lens && bs == 1 && attn_metadata.q_seq_lens_vec[0] == max_len && - projected_tokens.dim() == 2 && projected_tokens.size(0) == max_len) { - return projected_tokens.view({1, max_len, projected_tokens.size(-1)}); + if (!attn_metadata.is_prefill) { + return qkvz.view({qkvz.size(0), -1, qkvz.size(-1)}); } std::vector batches; - batches.reserve(bs); int64_t idx = 0; for (int64_t b = 0; b < bs; ++b) { - int64_t cur_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b] - : start_loc[b].template item(); - torch::Tensor batch = - projected_tokens.slice(/*dim=*/0, idx, idx + cur_len).contiguous(); + int64_t cur_len = start_loc[b].template item(); + torch::Tensor batch = qkvz.slice(0, idx, idx + cur_len).contiguous(); idx = idx + cur_len; if (batch.size(0) != max_len) { batch = batch.size(0) > max_len - ? batch.slice(/*dim=*/0, /*start=*/0, max_len).contiguous() + ? batch.slice(0, 0, max_len).contiguous() : torch::nn::functional::pad( batch, torch::nn::functional::PadFuncOptions( {0, 0, 0, max_len - batch.size(0)})) .contiguous(); } - batches.emplace_back(batch); + batches.push_back(batch); } auto ret = torch::stack(batches, 0).contiguous(); return ret; diff --git a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.h b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.h index fdc82b4d..2994f329 100644 --- a/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.h +++ b/upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.h @@ -1,4 +1,4 @@ -/* Copyright 2025-2026 The xLLM Authors. +/* Copyright 2026 The xLLM Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -17,7 +17,6 @@ limitations under the License. #include -#include #include #include #include @@ -52,40 +51,19 @@ class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module { const ModelInputParams& input_params); protected: - virtual std::pair project_decode_inputs( - const torch::Tensor& hidden_states) = 0; - virtual std::pair project_flat_inputs( - const torch::Tensor& hidden_states) = 0; - // Qwen3.5 overrides this to project and reshape its separate qkv/z/b/a - // weights in every forward mode. Qwen3Next keeps qkvz/ba packed and returns - // nullopt to select the fused-split fallback. - virtual std::optional< - std::tuple> - project_split_inputs(const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata) { - return std::nullopt; - } - virtual bool use_fla_ssm_state_layout() const { return false; } + virtual std::pair project_padded_inputs( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata) = 0; void load_common_state_dict(const StateDict& state_dict); void verify_common_loaded_weights(const std::string& prefix) const; - torch::Tensor get_linear_state_indices(const ModelInputParams& input_params, - const torch::Device& device) const; - - std::pair project_padded_inputs( - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata); - + torch::Tensor reshape_qkvz_with_pad(const AttentionMetadata& attn_metadata, + const torch::Tensor& qkvz) const; torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata, const torch::Tensor& padded_qkvz) const; - - // Projection outputs are packed as [total_tokens, dim], while GDN kernels - // consume dense [batch, max_query_len, dim] tensors. Split the packed tokens - // by query length and pad each sequence before entering the kernels. - torch::Tensor reshape_projected_tokens_with_pad( - const AttentionMetadata& attn_metadata, - const torch::Tensor& projected_tokens) const; + torch::Tensor get_linear_state_indices(const ModelInputParams& input_params, + const torch::Device& device) const; std::tuple process_mixed_qkv( torch::Tensor& mixed_qkv) const;