refactor Qwen3-Next with a new RadixLinearAttention (#17373)

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -29,6 +29,7 @@ from sglang.srt.layers.attention.mamba.mamba2_metadata import (
     Mamba2Metadata,
 )
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, MambaPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
@@ -833,30 +834,23 @@ class GDNAttnBackend(MambaAttnBackendBase):
 
     def forward_decode(
         self,
-        q: torch.Tensor,
+        layer: RadixLinearAttention,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        layer: RadixAttention,
         forward_batch: ForwardBatch,
-        save_kv_cache: bool = True,
+        mixed_qkv: torch.Tensor,
-        **kwargs,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        **kwargs,  # Unused, for compatibility with HybridLinearAttnBackend
     ):
-        mixed_qkv = kwargs["mixed_qkv"]
+        conv_weights = layer.conv_weights
-        conv_weights = kwargs["conv_weights"]
+        bias = layer.bias
-        bias = kwargs["bias"]
+        activation = layer.activation
-        activation = kwargs["activation"]
+        key_dim = layer.key_dim
-        key_dim = kwargs["key_dim"]
+        value_dim = layer.value_dim
-        value_dim = kwargs["value_dim"]
+        attn_tp_size = layer.attention_tp_size
-        attn_tp_size = kwargs["attention_tp_size"]
+        head_k_dim = layer.head_k_dim
-        head_k_dim = kwargs["head_k_dim"]
+        head_v_dim = layer.head_v_dim
-        head_v_dim = kwargs["head_v_dim"]
-        a = kwargs["a"]
-        b = kwargs["b"]
-        A_log = kwargs["A_log"]
-        dt_bias = kwargs["dt_bias"]
-        layer_id = kwargs["layer_id"]
 
-        layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer_id)
+        layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
         conv_states = layer_cache.conv[0]
         ssm_states = layer_cache.temporal
         query_start_loc = self.forward_metadata.query_start_loc
@@ -888,8 +882,8 @@ class GDNAttnBackend(MambaAttnBackendBase):
         value = value.view(1, seq_len, value.shape[1] // head_v_dim, head_v_dim)
 
         core_attn_out = self._kernel_func(
-            A_log=A_log,
+            A_log=layer.A_log,
-            dt_bias=dt_bias,
+            dt_bias=layer.dt_bias,
             q=query,
             k=key,
             v=value,
@@ -911,29 +905,23 @@ class GDNAttnBackend(MambaAttnBackendBase):
 
     def forward_extend(
         self,
-        q: torch.Tensor,
+        layer: RadixLinearAttention,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        layer: RadixAttention,
         forward_batch: ForwardBatch,
-        save_kv_cache: bool = True,
+        mixed_qkv: torch.Tensor,
-        **kwargs,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        **kwargs,  # Unused, for compatibility with HybridLinearAttnBackend
     ):
-        mixed_qkv = kwargs["mixed_qkv"]
+        seq_len = mixed_qkv.shape[0]
-        conv_weights = kwargs["conv_weights"]
+
-        bias = kwargs["bias"]
+        conv_weights = layer.conv_weights
-        activation = kwargs["activation"]
+        bias = layer.bias
-        key_dim = kwargs["key_dim"]
+        activation = layer.activation
-        value_dim = kwargs["value_dim"]
+        key_dim = layer.key_dim
-        attn_tp_size = kwargs["attention_tp_size"]
+        value_dim = layer.value_dim
-        head_k_dim = kwargs["head_k_dim"]
+        attn_tp_size = layer.attention_tp_size
-        head_v_dim = kwargs["head_v_dim"]
+        head_k_dim = layer.head_k_dim
-        a = kwargs["a"]
+        head_v_dim = layer.head_v_dim
-        b = kwargs["b"]
-        A_log = kwargs["A_log"]
-        dt_bias = kwargs["dt_bias"]
-        layer_id = kwargs["layer_id"]
-        seq_len = kwargs["seq_len"]
 
         is_target_verify = forward_batch.forward_mode.is_target_verify()
         forward_metadata = self.forward_metadata
@@ -944,7 +932,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
         retrieve_next_sibling = forward_metadata.retrieve_next_sibling
         retrieve_parent_token = forward_metadata.retrieve_parent_token
 
-        mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer_id)
+        mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
         conv_states = mamba_cache_params.conv[0]
         ssm_states = mamba_cache_params.temporal
         if is_target_verify:
@@ -1029,7 +1017,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
         key = key.view(1, actual_seq_len, num_heads, head_k_dim)
         value = value.view(1, actual_seq_len, num_value_heads, head_v_dim)
 
-        g, beta = fused_gdn_gating(A_log, a, b, dt_bias)
+        g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)
 
         if is_target_verify:
             core_attn_out = fused_recurrent_gated_delta_rule_update(
@@ -1240,12 +1228,15 @@ class HybridLinearAttnBackend(AttentionBackend):
 
     def forward_decode(
         self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
         layer: RadixAttention,
         forward_batch: ForwardBatch,
         save_kv_cache: bool = True,
+        q: Optional[torch.Tensor] = None,  # For full attention
+        k: Optional[torch.Tensor] = None,  # For full attention
+        v: Optional[torch.Tensor] = None,  # For full attention
+        mixed_qkv: Optional[torch.Tensor] = None,  # For GDN linear attention
+        a: Optional[torch.Tensor] = None,  # For GDN linear attention
+        b: Optional[torch.Tensor] = None,  # For GDN linear attention
         **kwargs,
     ):
         layer_id = layer.layer_id if layer else kwargs["layer_id"]
@@ -1253,18 +1244,31 @@ class HybridLinearAttnBackend(AttentionBackend):
             return self.full_attn_backend.forward_decode(
                 q, k, v, layer, forward_batch, save_kv_cache, **kwargs
             )
+        # Linear attention backend
         return self.linear_attn_backend.forward_decode(
-            q, k, v, layer, forward_batch, save_kv_cache, **kwargs
+            q=q,
+            k=k,
+            v=v,
+            layer=layer,
+            forward_batch=forward_batch,
+            save_kv_cache=save_kv_cache,
+            mixed_qkv=mixed_qkv,
+            a=a,
+            b=b,
+            **kwargs,
         )
 
     def forward_extend(
         self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
         layer: RadixAttention,
         forward_batch: ForwardBatch,
         save_kv_cache: bool = True,
+        q: Optional[torch.Tensor] = None,  # For full attention
+        k: Optional[torch.Tensor] = None,  # For full attention
+        v: Optional[torch.Tensor] = None,  # For full attention
+        mixed_qkv: Optional[torch.Tensor] = None,  # For GDN linear attention
+        a: Optional[torch.Tensor] = None,  # For GDN linear attention
+        b: Optional[torch.Tensor] = None,  # For GDN linear attention
         **kwargs,
     ):
         layer_id = layer.layer_id if layer else kwargs["layer_id"]
@@ -1272,43 +1276,66 @@ class HybridLinearAttnBackend(AttentionBackend):
             return self.full_attn_backend.forward_extend(
                 q, k, v, layer, forward_batch, save_kv_cache, **kwargs
             )
+        # Linear attention backend
         return self.linear_attn_backend.forward_extend(
-            q, k, v, layer, forward_batch, save_kv_cache, **kwargs
+            q=q,
+            k=k,
+            v=v,
+            layer=layer,
+            forward_batch=forward_batch,
+            save_kv_cache=save_kv_cache,
+            mixed_qkv=mixed_qkv,
+            a=a,
+            b=b,
+            **kwargs,
         )
 
     def forward(
         self,
-        q: torch.Tensor,
+        q: Optional[torch.Tensor] = None,  # For full attention
-        k: torch.Tensor,
+        k: Optional[torch.Tensor] = None,  # For full attention
-        v: torch.Tensor,
+        v: Optional[torch.Tensor] = None,  # For full attention
-        layer: RadixAttention,
+        layer: RadixAttention = None,
-        forward_batch: ForwardBatch,
+        forward_batch: ForwardBatch = None,
         save_kv_cache: bool = True,
+        mixed_qkv: Optional[torch.Tensor] = None,  # For GDN linear attention
+        a: Optional[torch.Tensor] = None,  # For GDN linear attention
+        b: Optional[torch.Tensor] = None,  # For GDN linear attention
         **kwargs,
     ):
-        """Run forward on an attention layer."""
+        layer_id = layer.layer_id if layer else kwargs["layer_id"]
+        is_linear_attn = layer_id not in self.full_attn_layers
+
         if forward_batch.forward_mode.is_idle():
-            if layer is None:
+            if is_linear_attn:
-                return torch.empty_like(kwargs["z"])
+                return mixed_qkv.new_empty(
+                    mixed_qkv.shape[0], layer.num_v_heads, layer.head_v_dim
+                )
             return q.new_empty(q.shape[0], layer.tp_q_head_num * layer.v_head_dim)
         elif forward_batch.forward_mode.is_decode():
             return self.forward_decode(
+                layer,
+                forward_batch,
+                save_kv_cache,
                 q,
                 k,
                 v,
-                layer,
+                mixed_qkv,
-                forward_batch,
+                a,
-                save_kv_cache=save_kv_cache,
+                b,
                 **kwargs,
             )
         else:
             return self.forward_extend(
+                layer,
+                forward_batch,
+                save_kv_cache,
                 q,
                 k,
                 v,
-                layer,
+                mixed_qkv,
-                forward_batch,
+                a,
-                save_kv_cache=save_kv_cache,
+                b,
                 **kwargs,
             )
 

python/sglang/srt/layers/radix_linear_attention.py

@@ -0,0 +1,83 @@
+# Copyright 2025-2026 SGLang Team
+# 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.
+# ==============================================================================
+"""Radix linear attention."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class RadixLinearAttention(nn.Module):
+    """
+    The Linear Attention Layer Implementation.
+    """
+
+    def __init__(
+        self,
+        layer_id: int,
+        num_qk_heads: int,
+        num_v_heads: int,
+        head_qk_dim: int,
+        head_v_dim: int,
+        attention_tp_size: int = 1,
+        conv_weights: Optional[torch.Tensor] = None,
+        bias: Optional[torch.Tensor] = None,
+        activation: str = "silu",
+        A_log: Optional[torch.Tensor] = None,
+        dt_bias: Optional[torch.Tensor] = None,
+    ):
+        super().__init__()
+        self.layer_id = layer_id
+        # Q and K share the same head count and dimension (per-TP values)
+        self.num_qk_heads = num_qk_heads
+        self.num_v_heads = num_v_heads
+        self.head_qk_dim = head_qk_dim
+        self.head_v_dim = head_v_dim
+        self.attention_tp_size = attention_tp_size
+
+        self.qk_dim_per_tp = num_qk_heads * head_qk_dim
+        self.value_dim_per_tp = num_v_heads * head_v_dim
+
+        self.key_dim = self.qk_dim_per_tp * attention_tp_size
+        self.value_dim = self.value_dim_per_tp * attention_tp_size
+
+        self.num_k_heads = num_qk_heads
+        self.num_q_heads = num_qk_heads
+        self.head_k_dim = head_qk_dim
+
+        self.conv_weights = conv_weights
+        self.bias = bias
+        self.activation = activation
+        self.A_log = A_log
+        self.dt_bias = dt_bias
+
+    def forward(
+        self,
+        forward_batch: ForwardBatch,
+        mixed_qkv: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+    ) -> torch.Tensor:
+        return forward_batch.attn_backend.forward(
+            layer=self,
+            forward_batch=forward_batch,
+            mixed_qkv=mixed_qkv,
+            a=a,
+            b=b,
+        )

python/sglang/srt/models/qwen3_next.py

@@ -29,6 +29,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.layers.rotary_embedding import get_rope
 from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
@@ -60,8 +61,6 @@ _is_npu = is_npu()
 import triton
 import triton.language as tl
 
-from sglang.srt.compilation.piecewise_context_manager import get_forward_context
-
 
 @triton.jit
 def fused_qkvzba_split_reshape_cat_kernel(
@@ -305,6 +304,20 @@ class Qwen3GatedDeltaNet(nn.Module):
             prefix=add_prefix("out_proj", prefix),
         )
 
+        self.linear_attn = RadixLinearAttention(
+            layer_id=layer_id,
+            num_qk_heads=self.num_k_heads // self.attn_tp_size,
+            num_v_heads=self.num_v_heads // self.attn_tp_size,
+            head_qk_dim=self.head_k_dim,
+            head_v_dim=self.head_v_dim,
+            attention_tp_size=self.attn_tp_size,
+            conv_weights=self.conv1d.weight.squeeze(1),
+            bias=self.conv1d.bias,
+            activation=self.activation,
+            A_log=self.A_log,
+            dt_bias=self.dt_bias,
+        )
+
     def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba):
         """
         Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.
@@ -379,8 +392,8 @@ class Qwen3GatedDeltaNet(nn.Module):
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
     ):
-        output = torch.empty_like(hidden_states)
         if forward_batch.forward_mode.is_extend() and get_forward_context() is not None:
+            output = torch.empty_like(hidden_states)
             gdn_with_output(
                 hidden_states,
                 output,
@@ -419,41 +432,12 @@ class Qwen3GatedDeltaNet(nn.Module):
                 lambda x: x.reshape(x.shape[0], -1), (query, key, value)
             )
             mixed_qkv = torch.cat((query, key, value), dim=-1)
-        # mixed_qkv = rearrange(mixed_qkv, "b l d -> b d l")
 
-        # 2. Convolution sequence transformation
+        core_attn_out = self.linear_attn(
-        conv_weights = self.conv1d.weight.view(
+            forward_batch,
-            self.conv1d.weight.size(0), self.conv1d.weight.size(2)
+            mixed_qkv=mixed_qkv,
-        )
+            a=a,
-
+            b=b,
-        kwargs = {
-            "mixed_qkv": mixed_qkv,
-            "conv_weights": conv_weights,
-            "bias": self.conv1d.bias,
-            "activation": self.activation,
-            "key_dim": self.key_dim,
-            "value_dim": self.value_dim,
-            "attention_tp_size": self.attn_tp_size,
-            "head_k_dim": self.head_k_dim,
-            "head_v_dim": self.head_v_dim,
-            "a": a,
-            "b": b,
-            "A_log": self.A_log,
-            "dt_bias": self.dt_bias,
-            "layer_id": self.layer_id,
-            "seq_len": seq_len,
-            "num_k_heads": self.num_k_heads,
-            "num_v_heads": self.num_v_heads,
-            "z": z,
-        }
-
-        core_attn_out = forward_batch.attn_backend.forward(
-            q=None,
-            k=None,
-            v=None,
-            layer=None,
-            forward_batch=forward_batch,
-            **kwargs,
         )
 
         z_shape_og = z.shape