minke.yu/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
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+34 -81
View File
@@ -635,15 +635,7 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
(q_conv_weights, k_conv_weights, v_conv_weights) = layer.conv_weights (q_conv_weights, k_conv_weights, v_conv_weights) = layer.conv_weights
(q_conv_bias, k_conv_bias, v_conv_bias) = layer.bias (q_conv_bias, k_conv_bias, v_conv_bias) = layer.bias
head_dim = layer.head_qk_dim layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
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)
q_conv_state, k_conv_state, v_conv_state = layer_cache.conv q_conv_state, k_conv_state, v_conv_state = layer_cache.conv
ssm_states = layer_cache.temporal ssm_states = layer_cache.temporal
query_start_loc = self.forward_metadata.query_start_loc query_start_loc = self.forward_metadata.query_start_loc
@@ -678,18 +670,18 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
conv_state_indices=cache_indices, conv_state_indices=cache_indices,
) )
q, k, v = map( q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
lambda x: rearrange(x, "n (h d) -> 1 n h d", d=head_dim), (q, k, v) 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( return fused_sigmoid_gating_delta_rule_update(
A_log=A_log, A_log=layer.A_log,
dt_bias=dt_bias, dt_bias=layer.dt_bias,
q=q, q=q,
k=k, k=k,
v=v, v=v,
a=g, a=a,
b=beta, b=b,
initial_state_source=ssm_states, initial_state_source=ssm_states,
initial_state_indices=cache_indices, initial_state_indices=cache_indices,
cu_seqlens=query_start_loc, cu_seqlens=query_start_loc,
@@ -699,8 +691,6 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
is_kda=True, is_kda=True,
) )
return core_attn_out
def forward_extend( def forward_extend(
self, self,
layer: RadixLinearAttention, layer: RadixLinearAttention,
@@ -719,18 +709,10 @@ class KimiLinearAttnBackend(MambaAttnBackendBase):
(q_conv_weights, k_conv_weights, v_conv_weights) = layer.conv_weights (q_conv_weights, k_conv_weights, v_conv_weights) = layer.conv_weights
(q_conv_bias, k_conv_bias, v_conv_bias) = layer.bias (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 query_start_loc = self.forward_metadata.query_start_loc
cache_indices = self.forward_metadata.mamba_cache_indices 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 conv_state_q, conv_state_k, conv_state_v = mamba_cache_params.conv
# deal with strides # deal with strides
conv_state_q = conv_state_q.transpose(-1, -2) 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, seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1) ).transpose(0, 1)
q, k, v = map( q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
lambda x: rearrange(x, "n (h d) -> 1 n h d", d=head_dim), (q, k, v) 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( core_attn_out = chunk_kda(
q=q, q=q,
k=k, k=k,
v=v, v=v,
g=g, g=a,
beta=beta, beta=b,
initial_state=ssm_states, initial_state=ssm_states,
initial_state_indices=cache_indices, initial_state_indices=cache_indices,
use_qk_l2norm_in_kernel=True, use_qk_l2norm_in_kernel=True,
@@ -829,15 +811,6 @@ class GDNAttnBackend(MambaAttnBackendBase):
b: torch.Tensor, b: torch.Tensor,
**kwargs, # Unused, for compatibility with HybridLinearAttnBackend **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) layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
conv_states = layer_cache.conv[0] conv_states = layer_cache.conv[0]
ssm_states = layer_cache.temporal ssm_states = layer_cache.temporal
@@ -848,27 +821,22 @@ class GDNAttnBackend(MambaAttnBackendBase):
mixed_qkv = causal_conv1d_update( mixed_qkv = causal_conv1d_update(
mixed_qkv, mixed_qkv,
conv_states, conv_states,
conv_weights, layer.conv_weights,
bias, layer.bias,
activation, layer.activation,
conv_state_indices=cache_indices, conv_state_indices=cache_indices,
) )
query, key, value = torch.split( query, key, value = torch.split(
mixed_qkv, mixed_qkv,
[ [layer.q_dim, layer.k_dim, layer.v_dim],
key_dim // attn_tp_size,
key_dim // attn_tp_size,
value_dim // attn_tp_size,
],
dim=-1, dim=-1,
) )
# Reshape from [l, h*d] to [1, l, h, d] # Reshape from [bs, h*d] to [1, bs, h, d]
seq_len = query.shape[0] bs = forward_batch.batch_size
num_heads = query.shape[1] // head_k_dim query = query.view(1, bs, layer.num_q_heads, layer.head_q_dim)
query = query.view(1, seq_len, num_heads, head_k_dim) key = key.view(1, bs, layer.num_k_heads, layer.head_k_dim)
key = key.view(1, seq_len, num_heads, head_k_dim) value = value.view(1, bs, layer.num_v_heads, layer.head_v_dim)
value = value.view(1, seq_len, value.shape[1] // head_v_dim, head_v_dim)
core_attn_out = self._kernel_func( core_attn_out = self._kernel_func(
A_log=layer.A_log, A_log=layer.A_log,
@@ -904,15 +872,6 @@ class GDNAttnBackend(MambaAttnBackendBase):
assert isinstance(mixed_qkv, torch.Tensor) assert isinstance(mixed_qkv, torch.Tensor)
seq_len = mixed_qkv.shape[0] 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() is_target_verify = forward_batch.forward_mode.is_target_verify()
forward_metadata = self.forward_metadata forward_metadata = self.forward_metadata
@@ -951,9 +910,9 @@ class GDNAttnBackend(MambaAttnBackendBase):
mixed_qkv_processed = causal_conv1d_update( mixed_qkv_processed = causal_conv1d_update(
mixed_qkv_reshaped, mixed_qkv_reshaped,
conv_states, conv_states,
conv_weights, layer.conv_weights,
bias, layer.bias,
activation, layer.activation,
conv_state_indices=cache_indices[:batch_size], conv_state_indices=cache_indices[:batch_size],
intermediate_conv_window=intermediate_conv_window_cache, intermediate_conv_window=intermediate_conv_window_cache,
intermediate_state_indices=intermediate_state_indices[:batch_size], intermediate_state_indices=intermediate_state_indices[:batch_size],
@@ -980,9 +939,9 @@ class GDNAttnBackend(MambaAttnBackendBase):
mixed_qkv = causal_conv1d_fn( mixed_qkv = causal_conv1d_fn(
mixed_qkv, mixed_qkv,
conv_weights, layer.conv_weights,
bias, layer.bias,
activation=activation, activation=layer.activation,
conv_states=conv_states, conv_states=conv_states,
has_initial_state=has_initial_states, has_initial_state=has_initial_states,
cache_indices=cache_indices, cache_indices=cache_indices,
@@ -990,22 +949,16 @@ class GDNAttnBackend(MambaAttnBackendBase):
seq_lens_cpu=forward_batch.extend_seq_lens_cpu, seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)[:seq_len] ).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( query, key, value = torch.split(
mixed_qkv, mixed_qkv,
[key_split_dim, key_split_dim, value_split_dim], [layer.q_dim, layer.k_dim, layer.v_dim],
dim=-1, dim=-1,
) )
actual_seq_len = query.shape[0] actual_seq_len = query.shape[0]
num_heads = query.shape[1] // head_k_dim query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim)
num_value_heads = value.shape[1] // head_v_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)
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)
g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)
python/sglang/srt/layers/radix_linear_attention.py
+12 -18
View File
@@ -31,36 +31,30 @@ class RadixLinearAttention(nn.Module):
def __init__( def __init__(
self, self,
layer_id: int, layer_id: int,
num_qk_heads: int, num_q_heads: int,
num_k_heads: int,
num_v_heads: int, num_v_heads: int,
head_qk_dim: int, head_q_dim: int,
head_k_dim: int,
head_v_dim: int, head_v_dim: int,
attention_tp_size: int = 1,
# GDN KDA Shared Weights # GDN KDA Shared Weights
conv_weights: Optional[Union[torch.Tensor, Tuple[torch.Tensor, ...]]] = None, 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", activation: str = "silu",
A_log: Optional[torch.Tensor] = None, A_log: Optional[torch.Tensor] = None,
dt_bias: Optional[torch.Tensor] = None, dt_bias: Optional[torch.Tensor] = None,
): ):
super().__init__() super().__init__()
self.layer_id = layer_id self.layer_id = layer_id
# Q and K share the same head count and dimension (per-TP values) self.num_q_heads = num_q_heads
self.num_qk_heads = num_qk_heads self.num_k_heads = num_k_heads
self.num_v_heads = num_v_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.head_v_dim = head_v_dim
self.attention_tp_size = attention_tp_size self.q_dim = num_q_heads * head_q_dim
self.k_dim = num_k_heads * head_k_dim
self.qk_dim_per_tp = num_qk_heads * head_qk_dim self.v_dim = num_v_heads * head_v_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.conv_weights = conv_weights
self.bias = bias self.bias = bias
python/sglang/srt/models/kimi_linear.py
+4 -3
View File
@@ -315,11 +315,12 @@ class KimiDeltaAttention(nn.Module):
self.linear_attn = RadixLinearAttention( self.linear_attn = RadixLinearAttention(
layer_id=self.layer_idx, 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, 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, head_v_dim=self.head_v_dim,
attention_tp_size=self.attn_tp_size,
conv_weights=conv_weights, conv_weights=conv_weights,
bias=bias, bias=bias,
A_log=self.A_log, A_log=self.A_log,
python/sglang/srt/models/qwen3_next.py
+4 -3
View File
@@ -306,11 +306,12 @@ class Qwen3GatedDeltaNet(nn.Module):
self.linear_attn = RadixLinearAttention( self.linear_attn = RadixLinearAttention(
layer_id=layer_id, 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, 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, head_v_dim=self.head_v_dim,
attention_tp_size=self.attn_tp_size,
conv_weights=self.conv1d.weight.squeeze(1), conv_weights=self.conv1d.weight.squeeze(1),
bias=self.conv1d.bias, bias=self.conv1d.bias,
activation=self.activation, activation=self.activation,