ref(upstream): 搬运 3 大 GDN 上游仓库 — FLA naive ops + vllm GDN 子树 + xllm C++ 参考

来源:
  1. fla-org/flash-linear-attention (5538 stars)
     → upstream_ref/fla/ops/gated_delta_rule/naive.py (正确的纯 PyTorch GDN)
     → upstream_ref/fla/ops/gated_delta_rule/chunk.py (Triton chunk kernel)
     → upstream_ref/fla/layers/gated_deltanet.py (层集成)

  2. vllm-project/vllm main (88717 stars)
     → upstream_ref/vllm_gdn/gdn/qwen_gdn_linear_attn.py (1751行, Qwen3.5 原生 GDN)
     → upstream_ref/vllm_gdn/ops/causal_conv1d.py (1289行, 正确的 Conv1d)
     → upstream_ref/vllm_gdn/third_party/ops/ (FLA Triton ops vendored)
     → upstream_ref/vllm_gdn/models/qwen3_5.py (vllm 最新 Qwen3.5 模型)

  3. Deep-Spark/xllm (BI-V100 硬件厂商)
     → upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp (576行)
     → upstream_ref/xllm_latest/core/kernels/npu/npu_causal_conv1d.cpp
     → upstream_ref/xllm_latest/core/kernels/npu/npu_recurrent_gated_delta_rule.cpp

目的: 修复 corex_gdn.py Conv1d groups 接口不匹配问题
  错误: conv1d_weight shape (2560,1,4) 被当成 (num_k_heads,1,4) 索引
  conv_dim = key_dim*2 + value_dim = 10240, TP=4 后 2560
  FLA naive.py 和 vllm qwen_gdn_linear_attn.py 有正确的实现可直接对接

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 <https://arxiv.org/abs/2412.06464>`_
+
+    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 <https://arxiv.org/abs/2411.12537>`_
+        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

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',
+]

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',
+]

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',
+]

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",
+]

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']

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,
+        )

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)

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

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

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

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

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

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

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)

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

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

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,
+        )

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))

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

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()

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)

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

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",
+]

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

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"
+        )

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

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

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)

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

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)

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<torch::Tensor>& 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<torch::Tensor> 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

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 <glog/logging.h>
+
+#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<torch::Tensor> to_c10_optional_tensor(
+    const std::optional<torch::Tensor>& 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<torch::Tensor>& beta,
+    const std::optional<double> scale,
+    const std::optional<torch::Tensor>& actual_seq_lengths,
+    const std::optional<torch::Tensor>& ssm_state_indices,
+    const std::optional<torch::Tensor>& num_accepted_tokens,
+    const std::optional<torch::Tensor>& g,
+    const std::optional<torch::Tensor>& 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<torch::Tensor> beta_tensor = to_c10_optional_tensor(beta);
+  c10::optional<torch::Tensor> actual_seq_lengths_tensor =
+      to_c10_optional_tensor(actual_seq_lengths);
+  c10::optional<torch::Tensor> ssm_state_indices_tensor =
+      to_c10_optional_tensor(ssm_state_indices);
+  c10::optional<torch::Tensor> num_accepted_tokens_tensor =
+      to_c10_optional_tensor(num_accepted_tokens);
+  c10::optional<torch::Tensor> g_tensor = to_c10_optional_tensor(g);
+  c10::optional<torch::Tensor> gk_tensor = to_c10_optional_tensor(gk);
+  float scale_value = static_cast<float>(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

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 <glog/logging.h>
+
+#include <tuple>
+
+#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<float>(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<int32_t>(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<int64_t> orig_shape;
+    for (int64_t i = 0; i < q_gate.dim() - 1; i++) {
+      orig_shape.push_back(q_gate.size(i));
+    }
+    std::vector<int64_t> 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<int64_t> q_new_shape = orig_shape;
+    q_new_shape.push_back(-1);
+    q = q.reshape(q_new_shape);
+
+    std::vector<int64_t> 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

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 <torch/torch.h>
+
+#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

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 <glog/logging.h>
+
+#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<int32_t>(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<torch::Tensor, std::optional<torch::Tensor>>
+Qwen3_5DecoderLayerImpl::apply_norm(Qwen3NextRMSNorm& norm,
+                                    torch::Tensor& input,
+                                    std::optional<torch::Tensor>& 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<torch::Tensor>& 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

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 <torch/torch.h>
+
+#include <optional>
+#include <string>
+
+#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<torch::Tensor>& residual,
+                        torch::Tensor& positions,
+                        const AttentionMetadata& attn_metadata,
+                        KVCache& kv_cache,
+                        const ModelInputParams& input_params);
+
+ private:
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> apply_norm(
+      Qwen3NextRMSNorm& norm,
+      torch::Tensor& input,
+      std::optional<torch::Tensor>& 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

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 <glog/logging.h>
+
+#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

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

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<torch::Tensor, torch::Tensor>
-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<torch::Tensor, torch::Tensor>
-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<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
-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(

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 <torch/torch.h>
 
-#include <optional>
 #include <string>
-#include <tuple>
 #include <utility>
 
 #include "qwen3_next_gated_delta_net.h"
@@ -36,15 +34,9 @@ class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl {
                            const torch::TensorOptions& options);
 
  protected:
-  std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
-      const torch::Tensor& hidden_states) override;
-  std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
-      const torch::Tensor& hidden_states) override;
-  std::optional<
-      std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
-  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<torch::Tensor, torch::Tensor> 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;

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 <glog/logging.h>
 #include <torch/torch.h>
 
-#include <optional>
 #include <tuple>
 
-#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<int64_t> view_shape = tensor.sizes().vec();
-  view_shape.insert(view_shape.begin() + head_dim + 1, 1);
-  std::vector<int64_t> expand_shape = view_shape;
-  expand_shape[head_dim + 1] = repeats;
-  std::vector<int64_t> 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::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
     torch::Tensor query,
     torch::Tensor key,
@@ -80,9 +52,6 @@ std::tuple<torch::Tensor, torch::Tensor> 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::Tensor, torch::Tensor> 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::Tensor, torch::Tensor> 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<xllm::npu::AclGraphTaskUpdateContext>& graph_context,
-    const torch::Tensor& x,
-    const torch::Tensor& weight,
-    const torch::Tensor& conv_state,
-    const std::optional<torch::Tensor>& bias,
-    const std::vector<int64_t>& query_start_loc,
-    const std::vector<int64_t>& cache_indices,
-    const std::vector<int64_t>& 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<c10_npu::NPUEvent>(ACL_EVENT_EXTERNAL);
-  event->block(stream);
-  event->reset(stream);
-
-  torch::Tensor output;
-  c10_npu::graph_task_group_begin(stream);
-  const std::vector<int64_t> 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<int32_t>& 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<size_t>(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<float>(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<std::string>(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<torch::Tensor, torch::Tensor>
-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<int32_t>(num_k_heads_ / tp_size_);
+  fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
+  fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
+  fused_params.head_v = static_cast<int32_t>(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<int32_t>(num_k_heads_ / tp_size_);
-    fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
-    fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
-    fused_params.head_v = static_cast<int32_t>(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<int64_t> 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<int64_t>{1},
+                      /*padding=*/std::vector<int64_t>{3},
+                      /*dilation=*/std::vector<int64_t>{1},
+                      /*groups=*/static_cast<int64_t>(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<size_t>(batch_size) &&
-        input_params.parallel.query_start_loc.size() ==
-            static_cast<size_t>(batch_size + 1) &&
-        input_params.embedding.linear_state_ids.size() ==
-            static_cast<size_t>(batch_size) &&
-        input_params.linear_state_validity_mask.size() ==
-            static_cast<size_t>(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<torch::Tensor>(),  // 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<int64_t>();
-    const std::vector<int64_t> 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<torch::Tensor>(),
-            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::Tensor>(),
-            torch::IntArrayRef(input_params.parallel.query_start_loc),
-            torch::IntArrayRef(linear_state_indices_host),
-            torch::IntArrayRef(std::vector<int64_t>()),
-            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<torch::Tensor>();
+    conv1d_params.initial_state_idx = std::optional<torch::Tensor>();
+    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<float>(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<size_t>(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<torch::Tensor> packed_q;
-      std::vector<torch::Tensor> packed_k;
-      std::vector<torch::Tensor> packed_v;
-      std::vector<torch::Tensor> packed_g;
-      std::vector<torch::Tensor> 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<int64_t>(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<int32_t>(
-        attn_metadata.q_seq_lens_vec.data(), static_cast<size_t>(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<float>(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<float>(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<torch::Tensor> valid_batches;
-  const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty();
-  int64_t bs = has_host_lens
-                   ? static_cast<int64_t>(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<int64_t>();
-    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<int64_t>();
+    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<int64_t>(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<torch::Tensor> 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<int64_t>();
-    torch::Tensor batch =
-        projected_tokens.slice(/*dim=*/0, idx, idx + cur_len).contiguous();
+    int64_t cur_len = start_loc[b].template item<int64_t>();
+    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;

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 <torch/torch.h>
 
-#include <optional>
 #include <string>
 #include <tuple>
 #include <utility>
@@ -52,40 +51,19 @@ class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module {
                         const ModelInputParams& input_params);
 
  protected:
-  virtual std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
-      const torch::Tensor& hidden_states) = 0;
-  virtual std::pair<torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
-  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<torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor, torch::Tensor> process_mixed_qkv(
       torch::Tensor& mixed_qkv) const;