chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Minglei Zhu
2026-01-22 01:42:06 -08:00
committed by GitHub
parent f33022d039
commit 419bbcee10
3 changed files with 199 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

163
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
View File

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

83
python/sglang/srt/layers/radix_linear_attention.py Normal file
View File

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

58
python/sglang/srt/models/qwen3_next.py
View File

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