From 419bbcee10a5a2f76b8b2612c7c2fbde3d0be645 Mon Sep 17 00:00:00 2001 From: Minglei Zhu Date: Thu, 22 Jan 2026 01:42:06 -0800 Subject: [PATCH] refactor Qwen3-Next with a new RadixLinearAttention (#17373) --- .../attention/hybrid_linear_attn_backend.py | 163 ++++++++++-------- .../srt/layers/radix_linear_attention.py | 83 +++++++++ python/sglang/srt/models/qwen3_next.py | 58 +++---- 3 files changed, 199 insertions(+), 105 deletions(-) create mode 100644 python/sglang/srt/layers/radix_linear_attention.py diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index a0fb9d845..c148bc67e 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/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, - k: torch.Tensor, - v: torch.Tensor, - layer: RadixAttention, + layer: RadixLinearAttention, forward_batch: ForwardBatch, - save_kv_cache: bool = True, - **kwargs, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + **kwargs, # Unused, for compatibility with HybridLinearAttnBackend ): - mixed_qkv = kwargs["mixed_qkv"] - conv_weights = kwargs["conv_weights"] - bias = kwargs["bias"] - activation = kwargs["activation"] - key_dim = kwargs["key_dim"] - value_dim = kwargs["value_dim"] - attn_tp_size = kwargs["attention_tp_size"] - head_k_dim = kwargs["head_k_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"] + conv_weights = layer.conv_weights + bias = layer.bias + activation = layer.activation + key_dim = layer.key_dim + value_dim = layer.value_dim + attn_tp_size = layer.attention_tp_size + head_k_dim = layer.head_k_dim + head_v_dim = layer.head_v_dim - 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, - dt_bias=dt_bias, + A_log=layer.A_log, + 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, - k: torch.Tensor, - v: torch.Tensor, - layer: RadixAttention, + layer: RadixLinearAttention, forward_batch: ForwardBatch, - save_kv_cache: bool = True, - **kwargs, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + **kwargs, # Unused, for compatibility with HybridLinearAttnBackend ): - mixed_qkv = kwargs["mixed_qkv"] - conv_weights = kwargs["conv_weights"] - bias = kwargs["bias"] - activation = kwargs["activation"] - key_dim = kwargs["key_dim"] - value_dim = kwargs["value_dim"] - attn_tp_size = kwargs["attention_tp_size"] - head_k_dim = kwargs["head_k_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"] - seq_len = kwargs["seq_len"] + seq_len = mixed_qkv.shape[0] + + conv_weights = layer.conv_weights + bias = layer.bias + activation = layer.activation + key_dim = layer.key_dim + value_dim = layer.value_dim + attn_tp_size = layer.attention_tp_size + head_k_dim = layer.head_k_dim + head_v_dim = layer.head_v_dim 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, - k: torch.Tensor, - v: torch.Tensor, - layer: RadixAttention, - forward_batch: ForwardBatch, + q: Optional[torch.Tensor] = None, # For full attention + k: Optional[torch.Tensor] = None, # For full attention + v: Optional[torch.Tensor] = None, # For full attention + layer: RadixAttention = None, + 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: - return torch.empty_like(kwargs["z"]) + if is_linear_attn: + 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, - forward_batch, - save_kv_cache=save_kv_cache, + mixed_qkv, + a, + b, **kwargs, ) else: return self.forward_extend( + layer, + forward_batch, + save_kv_cache, q, k, v, - layer, - forward_batch, - save_kv_cache=save_kv_cache, + mixed_qkv, + a, + b, **kwargs, ) diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py new file mode 100644 index 000000000..2fe1dc749 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 6f11fdaed..6512afa08 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ 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 - conv_weights = self.conv1d.weight.view( - self.conv1d.weight.size(0), self.conv1d.weight.size(2) - ) - - 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, + core_attn_out = self.linear_attn( + forward_batch, + mixed_qkv=mixed_qkv, + a=a, + b=b, ) z_shape_og = z.shape