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

[hybrid-model] clean up and consolidate redundant fields in RadixLinearAttention (#17660)

Browse Source
This commit is contained in:
Minglei Zhu
2026-01-27 10:37:58 -08:00
committed by GitHub
parent 8acd4d7d7e
commit d90c0837e5
4 changed files with 54 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -635,15 +635,7 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
(q_conv_weights, k_conv_weights, v_conv_weights) = layer.conv_weights
(q_conv_bias, k_conv_bias, v_conv_bias) = layer.bias
head_dim = layer.head_qk_dim
layer_id = layer.layer_id
beta = b
g = a
A_log = layer.A_log
dt_bias = layer.dt_bias
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)
q_conv_state, k_conv_state, v_conv_state = layer_cache.conv
ssm_states = layer_cache.temporal
query_start_loc = self.forward_metadata.query_start_loc
@@ -678,18 +670,18 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
conv_state_indices=cache_indices,
)
q, k, v = map(
lambda x: rearrange(x, "n (h d) -> 1 n h d", d=head_dim), (q, k, v)
)
q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
core_attn_out = fused_sigmoid_gating_delta_rule_update(
A_log=A_log,
dt_bias=dt_bias,
return fused_sigmoid_gating_delta_rule_update(
A_log=layer.A_log,
dt_bias=layer.dt_bias,
q=q,
k=k,
v=v,
a=g,
b=beta,
a=a,
b=b,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
@@ -699,8 +691,6 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
is_kda=True,
)
return core_attn_out
def forward_extend(
self,
layer: RadixLinearAttention,
@@ -719,18 +709,10 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
(q_conv_weights, k_conv_weights, v_conv_weights) = layer.conv_weights
(q_conv_bias, k_conv_bias, v_conv_bias) = layer.bias
head_dim = layer.head_qk_dim
layer_id = layer.layer_id
beta = b
g = a
A_log = layer.A_log
dt_bias = layer.dt_bias
query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices
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_state_q, conv_state_k, conv_state_v = mamba_cache_params.conv
# deal with strides
conv_state_q = conv_state_q.transpose(-1, -2)
@@ -781,16 +763,16 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)
q, k, v = map(
lambda x: rearrange(x, "n (h d) -> 1 n h d", d=head_dim), (q, k, v)
)
q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
core_attn_out = chunk_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
g=a,
beta=b,
initial_state=ssm_states,
initial_state_indices=cache_indices,
use_qk_l2norm_in_kernel=True,
@@ -829,15 +811,6 @@ class GDNAttnBackend(MambaAttnBackendBase):
b: torch.Tensor,
**kwargs, # Unused, for compatibility with HybridLinearAttnBackend
):
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.layer_id)
conv_states = layer_cache.conv[0]
ssm_states = layer_cache.temporal
@@ -848,27 +821,22 @@ class GDNAttnBackend(MambaAttnBackendBase):
mixed_qkv = causal_conv1d_update(
mixed_qkv,
conv_states,
conv_weights,
bias,
activation,
layer.conv_weights,
layer.bias,
layer.activation,
conv_state_indices=cache_indices,
)
query, key, value = torch.split(
mixed_qkv,
[
key_dim // attn_tp_size,
key_dim // attn_tp_size,
value_dim // attn_tp_size,
],
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
# Reshape from [l, h*d] to [1, l, h, d]
seq_len = query.shape[0]
num_heads = query.shape[1] // head_k_dim
query = query.view(1, seq_len, num_heads, head_k_dim)
key = key.view(1, seq_len, num_heads, head_k_dim)
value = value.view(1, seq_len, value.shape[1] // head_v_dim, head_v_dim)
# Reshape from [bs, h*d] to [1, bs, h, d]
bs = forward_batch.batch_size
query = query.view(1, bs, layer.num_q_heads, layer.head_q_dim)
key = key.view(1, bs, layer.num_k_heads, layer.head_k_dim)
value = value.view(1, bs, layer.num_v_heads, layer.head_v_dim)
core_attn_out = self._kernel_func(
A_log=layer.A_log,
@@ -904,15 +872,6 @@ class GDNAttnBackend(MambaAttnBackendBase):
assert isinstance(mixed_qkv, torch.Tensor)
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
@@ -951,9 +910,9 @@ class GDNAttnBackend(MambaAttnBackendBase):
mixed_qkv_processed = causal_conv1d_update(
mixed_qkv_reshaped,
conv_states,
conv_weights,
bias,
activation,
layer.conv_weights,
layer.bias,
layer.activation,
conv_state_indices=cache_indices[:batch_size],
intermediate_conv_window=intermediate_conv_window_cache,
intermediate_state_indices=intermediate_state_indices[:batch_size],
@@ -980,9 +939,9 @@ class GDNAttnBackend(MambaAttnBackendBase):
mixed_qkv = causal_conv1d_fn(
mixed_qkv,
conv_weights,
bias,
activation=activation,
layer.conv_weights,
layer.bias,
activation=layer.activation,
conv_states=conv_states,
has_initial_state=has_initial_states,
cache_indices=cache_indices,
@@ -990,22 +949,16 @@ class GDNAttnBackend(MambaAttnBackendBase):
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)[:seq_len]
key_split_dim = key_dim // attn_tp_size
value_split_dim = value_dim // attn_tp_size
query, key, value = torch.split(
mixed_qkv,
[key_split_dim, key_split_dim, value_split_dim],
[layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1,
)
actual_seq_len = query.shape[0]
num_heads = query.shape[1] // head_k_dim
num_value_heads = value.shape[1] // head_v_dim
query = query.view(1, actual_seq_len, num_heads, head_k_dim)
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)
query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim)
key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim)
value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim)
g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)

30
python/sglang/srt/layers/radix_linear_attention.py
View File

@@ -31,36 +31,30 @@ class RadixLinearAttention(nn.Module):
def __init__(
self,
layer_id: int,
num_qk_heads: int,
num_q_heads: int,
num_k_heads: int,
num_v_heads: int,
head_qk_dim: int,
head_q_dim: int,
head_k_dim: int,
head_v_dim: int,
attention_tp_size: int = 1,
# GDN KDA Shared Weights
conv_weights: Optional[Union[torch.Tensor, Tuple[torch.Tensor, ...]]] = None,
bias: Optional[torch.Tensor] = None,
bias: Optional[Union[torch.Tensor, Tuple[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_q_heads = num_q_heads
self.num_k_heads = num_k_heads
self.num_v_heads = num_v_heads
self.head_qk_dim = head_qk_dim
self.head_q_dim = head_q_dim
self.head_k_dim = head_k_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.q_dim = num_q_heads * head_q_dim
self.k_dim = num_k_heads * head_k_dim
self.v_dim = num_v_heads * head_v_dim
self.conv_weights = conv_weights
self.bias = bias

7
python/sglang/srt/models/kimi_linear.py
View File

@@ -315,11 +315,12 @@ class KimiDeltaAttention(nn.Module):
self.linear_attn = RadixLinearAttention(
layer_id=self.layer_idx,
num_qk_heads=self.num_k_heads // self.attn_tp_size,
num_q_heads=self.num_k_heads // self.attn_tp_size,
num_k_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_q_dim=self.head_k_dim,
head_k_dim=self.head_k_dim,
head_v_dim=self.head_v_dim,
attention_tp_size=self.attn_tp_size,
conv_weights=conv_weights,
bias=bias,
A_log=self.A_log,

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

@@ -306,11 +306,12 @@ class Qwen3GatedDeltaNet(nn.Module):
self.linear_attn = RadixLinearAttention(
layer_id=layer_id,
num_qk_heads=self.num_k_heads // self.attn_tp_size,
num_q_heads=self.num_k_heads // self.attn_tp_size,
num_k_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_q_dim=self.head_k_dim,
head_k_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,