From d90c0837e5b812864a7570234b1a498740a01e85 Mon Sep 17 00:00:00 2001 From: Minglei Zhu Date: Tue, 27 Jan 2026 10:37:58 -0800 Subject: [PATCH] [hybrid-model] clean up and consolidate redundant fields in RadixLinearAttention (#17660) --- .../attention/hybrid_linear_attn_backend.py | 115 ++++++------------ .../srt/layers/radix_linear_attention.py | 30 ++--- python/sglang/srt/models/kimi_linear.py | 7 +- python/sglang/srt/models/qwen3_next.py | 7 +- 4 files changed, 54 insertions(+), 105 deletions(-) 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 d8fc0103f..397427826 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -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) diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py index 43d8b582c..4cc44ff4f 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/python/sglang/srt/layers/radix_linear_attention.py @@ -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 diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index d9db0b120..a1447c5c3 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -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, diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 017c98a2c..dc85e9c52 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -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,