diff --git a/python/sglang/srt/configs/mamba_utils.py b/python/sglang/srt/configs/mamba_utils.py index c5cc000f8..96b2bca68 100644 --- a/python/sglang/srt/configs/mamba_utils.py +++ b/python/sglang/srt/configs/mamba_utils.py @@ -217,11 +217,13 @@ class KimiLinearStateShape: conv_state_k_shape = (divide(proj_k_size, tp_world_size), conv_kernel_size - 1) temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim) - conv_state_shape = conv_state_shape[1], conv_state_shape[0] - conv_state_k_shape = conv_state_k_shape[1], conv_state_k_shape[0] + conv_state_shape = ( + conv_state_shape[1], + conv_state_shape[0] + conv_state_k_shape[0] * 2, + ) return KimiLinearStateShape( - conv=[conv_state_shape, conv_state_k_shape, conv_state_k_shape], + conv=[conv_state_shape], temporal=temporal_state_shape, num_heads=num_heads, head_dim=head_dim, diff --git a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py index a38f9ce50..486537b0e 100644 --- a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py +++ b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py @@ -4,8 +4,6 @@ import torch import triton import triton.language as tl -from sglang.srt.layers.attention.fla.utils import input_guard - @triton.jit(do_not_specialize=["T"]) def fused_sigmoid_gating_delta_rule_update_kernel( @@ -24,6 +22,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel( cu_seqlens, scale, T, + stride_q, + stride_k, + stride_v, + stride_b, B: tl.constexpr, H: tl.constexpr, HV: tl.constexpr, @@ -57,10 +59,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel( o_k = i_k * BK + tl.arange(0, BK) o_v = i_v * BV + tl.arange(0, BV) - p_q = q + (bos * H + i_h) * K + o_k - p_k = k + (bos * H + i_h) * K + o_k - p_v = v + (bos * HV + i_hv) * V + o_v - p_b = b + bos * HV + i_hv + p_q = q + bos * stride_q + i_h * K + o_k + p_k = k + bos * stride_k + i_h * K + o_k + p_v = v + bos * stride_v + i_hv * V + o_v + p_b = b + bos * stride_b + i_hv p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v # Gating computation pointers @@ -164,7 +166,6 @@ def fused_sigmoid_gating_delta_rule_update_kernel( tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) -@input_guard def fused_sigmoid_gating_delta_rule_update( A_log: torch.Tensor, a: torch.Tensor, @@ -188,6 +189,10 @@ def fused_sigmoid_gating_delta_rule_update( and the recurrent delta rule update for better performance. """ B, T, H, K, V = *k.shape, v.shape[-1] + stride_q = q.stride()[1] + stride_k = k.stride()[1] + stride_v = v.stride()[1] + stride_b = b.stride()[-2] HV = v.shape[2] N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32) @@ -220,6 +225,10 @@ def fused_sigmoid_gating_delta_rule_update( cu_seqlens=cu_seqlens, scale=scale, T=T, + stride_q=stride_q, + stride_k=stride_k, + stride_v=stride_v, + stride_b=stride_b, B=B, H=H, HV=HV, diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index d005262b3..6e6b16056 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -136,49 +136,21 @@ class KDAAttnBackend(MambaAttnBackendBase): b: torch.Tensor, **kwargs, ): - q_proj_states, k_proj_states, v_proj_states = torch.split( - mixed_qkv, - [layer.q_dim, layer.k_dim, layer.v_dim], - dim=-1, - ) - q_conv_weights, k_conv_weights, v_conv_weights = layer.conv_weights - q_conv_bias, k_conv_bias, v_conv_bias = layer.bias - 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 + conv_states = layer_cache.conv[0] ssm_states = layer_cache.temporal query_start_loc = self.forward_metadata.query_start_loc cache_indices = self.forward_metadata.mamba_cache_indices - q_conv_state = q_conv_state.transpose(-1, -2) - k_conv_state = k_conv_state.transpose(-1, -2) - v_conv_state = v_conv_state.transpose(-1, -2) - - q = causal_conv1d_update( - q_proj_states, - q_conv_state, - q_conv_weights, - q_conv_bias, + qkv = causal_conv1d_update( + mixed_qkv, + conv_states.transpose(-1, -2), + layer.conv_weights, + layer.bias, activation="silu", conv_state_indices=cache_indices, ) - k = causal_conv1d_update( - k_proj_states, - k_conv_state, - k_conv_weights, - k_conv_bias, - activation="silu", - conv_state_indices=cache_indices, - ) - v = causal_conv1d_update( - v_proj_states, - v_conv_state, - v_conv_weights, - v_conv_bias, - activation="silu", - conv_state_indices=cache_indices, - ) - + q, k, v = qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1) 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) @@ -205,62 +177,55 @@ class KDAAttnBackend(MambaAttnBackendBase): b: torch.Tensor, **kwargs, ): - q_proj_states, k_proj_states, v_proj_states = torch.split( - mixed_qkv, - [layer.q_dim, layer.k_dim, layer.v_dim], - dim=-1, - ) - q_conv_weights, k_conv_weights, v_conv_weights = layer.conv_weights - q_conv_bias, k_conv_bias, v_conv_bias = layer.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.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) - conv_state_k = conv_state_k.transpose(-1, -2) - conv_state_v = conv_state_v.transpose(-1, -2) + conv_states = mamba_cache_params.conv[0].transpose(-1, -2) ssm_states = mamba_cache_params.temporal has_initial_state = forward_batch.extend_prefix_lens > 0 - q_proj_states = q_proj_states.transpose(0, 1) - k_proj_states = k_proj_states.transpose(0, 1) - v_proj_states = v_proj_states.transpose(0, 1) + splits = [layer.q_dim, layer.k_dim, layer.v_dim] + q, k, v = mixed_qkv.transpose(0, 1).split(splits, dim=0) + q_conv_weight, k_conv_weight, v_conv_weight = layer.conv_weights.split( + splits, dim=0 + ) + q_conv_state, k_conv_state, v_conv_state = conv_states.split(splits, dim=-2) + if layer.bias is not None: + q_bias, k_bias, v_bias = layer.bias.split(splits, dim=0) + else: + q_bias, k_bias, v_bias = None, None, None q = causal_conv1d_fn( - q_proj_states, - q_conv_weights, - q_conv_bias, + q, + q_conv_weight, + q_bias, activation="silu", - conv_states=conv_state_q, + conv_states=q_conv_state, has_initial_state=has_initial_state, cache_indices=cache_indices, query_start_loc=query_start_loc, seq_lens_cpu=forward_batch.extend_seq_lens_cpu, ).transpose(0, 1) - k = causal_conv1d_fn( - k_proj_states, - k_conv_weights, - k_conv_bias, + k, + k_conv_weight, + k_bias, activation="silu", - conv_states=conv_state_k, + conv_states=k_conv_state, has_initial_state=has_initial_state, cache_indices=cache_indices, query_start_loc=query_start_loc, seq_lens_cpu=forward_batch.extend_seq_lens_cpu, ).transpose(0, 1) - v = causal_conv1d_fn( - v_proj_states, - v_conv_weights, - v_conv_bias, + v, + v_conv_weight, + v_bias, activation="silu", - conv_states=conv_state_v, + conv_states=v_conv_state, has_initial_state=has_initial_state, cache_indices=cache_indices, query_start_loc=query_start_loc, diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 12544f028..0ba637c11 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -22,6 +22,7 @@ from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelBatchedLinear, ColumnParallelLinear, + MergedColumnParallelLinear, MergedColumnParallelRepeatedLinear, QKVParallelLinear, ReplicatedLinear, @@ -283,34 +284,18 @@ class KimiDeltaAttention(nn.Module): set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) - self.q_conv1d = ColumnParallelLinear( + self.qkv_conv1d = MergedColumnParallelLinear( input_size=self.conv_size, - output_size=projection_size, + output_sizes=[projection_size, projection_size, projection_size], bias=False, params_dtype=torch.float32, - prefix=f"{prefix}.q_conv1d", - ) - self.k_conv1d = ColumnParallelLinear( - input_size=self.conv_size, - output_size=projection_size, - bias=False, - params_dtype=torch.float32, - prefix=f"{prefix}.k_conv1d", - ) - self.v_conv1d = ColumnParallelLinear( - input_size=self.conv_size, - output_size=projection_size, - bias=False, - params_dtype=torch.float32, - prefix=f"{prefix}.v_conv1d", + prefix=f"{prefix}.qkv_conv1d", ) # unsqueeze to fit conv1d weights shape into the linear weights shape. # Can't do this in `weight_loader` since it already exists in # `ColumnParallelLinear` and `set_weight_attrs` # doesn't allow to override it - self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1) - self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1) - self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1) + self.qkv_conv1d.weight.data = self.qkv_conv1d.weight.data.unsqueeze(1) self.A_log = nn.Parameter( torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32) @@ -328,18 +313,8 @@ class KimiDeltaAttention(nn.Module): prefix=f"{prefix}.o_proj", ) - self.q_conv_weights = self.q_conv1d.weight.view( - self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2) - ) - self.k_conv_weights = self.k_conv1d.weight.view( - self.k_conv1d.weight.size(0), self.k_conv1d.weight.size(2) - ) - self.v_conv_weights = self.v_conv1d.weight.view( - self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2) - ) - - conv_weights = (self.q_conv_weights, self.k_conv_weights, self.v_conv_weights) - bias = (self.q_conv1d.bias, self.k_conv1d.bias, self.v_conv1d.bias) + conv_weights = self.qkv_conv1d.weight.squeeze(1) + bias = self.qkv_conv1d.bias self.attn = RadixLinearAttention( layer_id=self.layer_idx, @@ -409,12 +384,11 @@ class KimiDeltaAttention(nn.Module): ) # fused_kda_gate is fused to KimiLinearAttentionBackend with decode - beta = beta.float() if not forward_batch.forward_mode.is_decode(): forget_gate = fused_kda_gate( forget_gate, self.A_log, self.head_dim, g_bias=self.dt_bias ) - beta = beta.sigmoid() + beta = beta.float().sigmoid() forget_gate = forget_gate.unsqueeze(0) beta = beta.unsqueeze(0) @@ -703,6 +677,10 @@ class KimiLinearForCausalLM(nn.Module): (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), + # qkv conv fuse + (".qkv_conv1d", ".q_conv1d", 0), + (".qkv_conv1d", ".k_conv1d", 1), + (".qkv_conv1d", ".v_conv1d", 2), ] if self.config.is_moe: # Params for weights, fp8 weight scales, fp8 activation scales