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

fuse qkvbfg linear into one gemm and f_b g_b into batched gemm. (#17801)

Browse Source
This commit is contained in:
strgrb
2026-02-04 11:41:26 +08:00
committed by GitHub
parent c1d529c196
commit 37c33cc0aa
2 changed files with 244 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

106
python/sglang/srt/layers/linear.py
View File

@@ -7,6 +7,7 @@ import logging
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
import torch
from torch import nn
from torch.nn.parameter import Parameter, UninitializedParameter
from sglang.srt.distributed import (
@@ -1446,3 +1447,108 @@ class RowParallelLinear(LinearBase):
s += f", tp_size={self.tp_size}"
s += f", reduce_results={self.reduce_results}"
return s
class MergedColumnParallelRepeatedLinear(LinearBase):
"""Merged column parallel linear and repeated linear layer.
TODO: quantization is not supported yet.
Args:
input_size: input dimension of the linear layer.
column_output_sizes: output dimension of the column linear layers.
repeated_output_sizes: output dimension of the repeated linear layers.
skip_bias_add: If true, skip adding bias but instead return it.
params_dtype: Data type for the parameters.
quant_config: Quantization configure.
"""
def __init__(
self,
input_size: int,
column_output_sizes: List[int],
repeated_output_sizes: List[int],
skip_bias_add: bool = False,
params_dtype: Optional[torch.dtype] = None,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
output_size = sum(column_output_sizes) + sum(repeated_output_sizes)
super().__init__(
input_size=input_size,
output_size=output_size,
skip_bias_add=skip_bias_add,
params_dtype=params_dtype,
quant_config=quant_config,
prefix=prefix,
)
self.num_column_parallel = len(column_output_sizes)
self.tp_rank = get_tensor_model_parallel_rank()
self.tp_size = get_tensor_model_parallel_world_size()
self.output_partition_sizes = [
divide(x, self.tp_size) for x in column_output_sizes
] + repeated_output_sizes
self.quant_method.create_weights(
layer=self,
input_size_per_partition=self.input_size,
output_partition_sizes=self.output_partition_sizes,
input_size=self.input_size,
output_size=self.output_size,
params_dtype=self.params_dtype,
skip_block_quant_check=True,
weight_loader=self.weight_loader,
)
self.prefix = prefix
def forward(self, input_: torch.Tensor) -> torch.Tensor:
return self.quant_method.apply(self, input_)
def weight_loader(
self, param: Parameter, loaded_weight: torch.Tensor, loaded_shard_id: int
) -> torch.Tensor:
output_dim = param.output_dim
shard_offset = sum(self.output_partition_sizes[:loaded_shard_id])
shard_size = self.output_partition_sizes[loaded_shard_id]
param_data = param.data.narrow(output_dim, shard_offset, shard_size)
if loaded_shard_id < self.num_column_parallel:
start_idx = self.tp_rank * shard_size
loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
param_data.copy_(loaded_weight)
class ColumnParallelBatchedLinear(nn.Module):
"""Column parallel batched linear layer.
TODO: quantization is not supported yet.
Args:
batch: batch dimension of the linear layer.
input_size: input dimension of the linear layer.
output_size: output dimension of the linear layer.
dtype: Data type for the parameters.
"""
def __init__(
self, batch: int, input_size: int, output_size: int, dtype: torch.dtype
):
super().__init__()
self.tp_rank = get_tensor_model_parallel_rank()
self.tp_size = get_tensor_model_parallel_world_size()
self.weight = nn.Parameter(
torch.empty(batch, output_size // self.tp_size, input_size, dtype=dtype),
requires_grad=False,
)
setattr(self.weight, "weight_loader", self.weight_loader)
def forward(self, input: torch.Tensor) -> torch.Tensor:
return torch.bmm(input, self.weight.transpose(-1, -2))
def weight_loader(
self, param: Parameter, loaded_weight: torch.Tensor, loaded_shard_id: int
) -> torch.Tensor:
shard_size = self.weight.shape[-2]
start_idx = self.tp_rank * shard_size
loaded_weight = loaded_weight.narrow(0, start_idx, shard_size)
param.data[loaded_shard_id].copy_(loaded_weight)

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

@@ -19,7 +19,9 @@ from sglang.srt.layers.attention.fla.kda import FusedRMSNormGated, fused_kda_gat
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import (
ColumnParallelBatchedLinear,
ColumnParallelLinear,
MergedColumnParallelRepeatedLinear,
ReplicatedLinear,
RowParallelLinear,
)
@@ -190,57 +192,97 @@ class KimiDeltaAttention(nn.Module):
projection_size = self.head_dim * self.num_heads
self.conv_size = config.linear_attn_config["short_conv_kernel_size"]
self.q_proj = ColumnParallelLinear(
self.hidden_size,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.q_proj",
)
self.k_proj = ColumnParallelLinear(
self.hidden_size,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.k_proj",
)
self.v_proj = ColumnParallelLinear(
self.hidden_size,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.v_proj",
)
# TODO: support fusion with quant
self.do_fuse_qkvbfg = quant_config is None
if self.do_fuse_qkvbfg:
self.qkvb_sizes = [
projection_size,
projection_size,
projection_size,
self.num_heads,
]
self.fg_sizes = [self.head_dim, self.head_dim]
self.fused_qkvbfg_proj = MergedColumnParallelRepeatedLinear(
self.hidden_size,
self.qkvb_sizes,
self.fg_sizes,
quant_config=quant_config,
prefix=f"{prefix}.fused_qkvbfg_proj",
)
self.split_sizes = [x // self.tp_size for x in self.qkvb_sizes] + [
2 * self.head_dim
]
self.fused_fg_b_proj = ColumnParallelBatchedLinear(
2, self.head_dim, projection_size, dtype=config.dtype
)
else:
self.q_proj = ColumnParallelLinear(
self.hidden_size,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.q_proj",
)
self.k_proj = ColumnParallelLinear(
self.hidden_size,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.k_proj",
)
self.v_proj = ColumnParallelLinear(
self.hidden_size,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.v_proj",
)
self.f_a_proj = ReplicatedLinear(
self.hidden_size,
self.head_dim,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.f_a_proj",
)
self.f_a_proj = ReplicatedLinear(
self.hidden_size,
self.head_dim,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.f_a_proj",
)
self.f_b_proj = ColumnParallelLinear(
self.head_dim,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.f_b_proj",
)
self.b_proj = ColumnParallelLinear(
self.hidden_size,
self.num_heads,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.b_proj",
)
self.g_a_proj = ReplicatedLinear(
self.hidden_size,
self.head_dim,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.g_a_proj",
)
self.g_b_proj = ColumnParallelLinear(
self.head_dim,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.g_b_proj",
)
self.f_b_proj = ColumnParallelLinear(
self.head_dim,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.f_b_proj",
)
self.dt_bias = nn.Parameter(
torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32)
)
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
self.b_proj = ColumnParallelLinear(
self.hidden_size,
self.num_heads,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.b_proj",
)
self.q_conv1d = ColumnParallelLinear(
input_size=self.conv_size,
output_size=projection_size,
@@ -275,20 +317,6 @@ class KimiDeltaAttention(nn.Module):
)
set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)})
self.g_a_proj = ReplicatedLinear(
self.hidden_size,
self.head_dim,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.g_a_proj",
)
self.g_b_proj = ColumnParallelLinear(
self.head_dim,
projection_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.g_b_proj",
)
self.o_norm = FusedRMSNormGated(
self.head_dim, eps=rms_norm_eps, activation="sigmoid"
)
@@ -327,6 +355,36 @@ class KimiDeltaAttention(nn.Module):
dt_bias=self.dt_bias,
)
def forward_qkvbfg(self, hidden_states: torch.Tensor):
q_proj_states = self.q_proj(hidden_states)[0]
k_proj_states = self.k_proj(hidden_states)[0]
v_proj_states = self.v_proj(hidden_states)[0]
beta = self.b_proj(hidden_states)[0]
forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
return (
(q_proj_states, k_proj_states, v_proj_states),
beta,
forget_gate,
g_proj_states,
)
def forward_qkvbfg_fused(self, hidden_states: torch.Tensor):
fused_states = self.fused_qkvbfg_proj(hidden_states)
q_proj_states, k_proj_states, v_proj_states, beta, fg_a_states = torch.split(
fused_states, self.split_sizes, dim=-1
)
# use batch matmul to calculate forget_gate and g_proj_states
forget_gate, g_proj_states = self.fused_fg_b_proj(
fg_a_states.view(-1, 2, self.head_dim).transpose(0, 1)
)
return (
(q_proj_states, k_proj_states, v_proj_states),
beta,
forget_gate,
g_proj_states,
)
def forward(
self,
hidden_states: torch.Tensor,
@@ -334,15 +392,17 @@ class KimiDeltaAttention(nn.Module):
forward_batch: ForwardBatch,
zero_allocator: BumpAllocator,
) -> None:
q_proj_states = self.q_proj(hidden_states)[0]
k_proj_states = self.k_proj(hidden_states)[0]
v_proj_states = self.v_proj(hidden_states)[0]
mixed_qkv = (q_proj_states, k_proj_states, v_proj_states)
forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
if self.do_fuse_qkvbfg:
mixed_qkv, beta, forget_gate, g_proj_states = self.forward_qkvbfg_fused(
hidden_states
)
else:
mixed_qkv, beta, forget_gate, g_proj_states = self.forward_qkvbfg(
hidden_states
)
# fused_kda_gate is fused to KimiLinearAttentionBackend with decode
beta = self.b_proj(hidden_states)[0].float()
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
@@ -358,7 +418,6 @@ class KimiDeltaAttention(nn.Module):
b=beta,
)
g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
norm_gate = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
core_attn_out = self.o_norm(core_attn_out, norm_gate)
core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
@@ -623,7 +682,16 @@ class KimiLinearForCausalLM(nn.Module):
# (param_name, shard_name, shard_id)
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
(".fused_qkvbfg_proj", ".q_proj", 0),
(".fused_qkvbfg_proj", ".k_proj", 1),
(".fused_qkvbfg_proj", ".v_proj", 2),
(".fused_qkvbfg_proj", ".b_proj", 3),
(".fused_qkvbfg_proj", ".f_a_proj", 4),
(".fused_qkvbfg_proj", ".g_a_proj", 5),
(".fused_fg_b_proj", ".f_b_proj", 0),
(".fused_fg_b_proj", ".g_b_proj", 1),
]
fuse_qkvbfg_keys = {x[1] for x in stacked_params_mapping[2:]}
if self.config.is_moe:
# Params for weights, fp8 weight scales, fp8 activation scales
# (param_name, weight_name, expert_id, shard_id)
@@ -658,6 +726,11 @@ class KimiLinearForCausalLM(nn.Module):
# for mlp.experts[0].gate_gate_up_proj, which breaks load.
if ("mlp.experts." in name) and name not in params_dict:
continue
if weight_name in fuse_qkvbfg_keys:
layer_id = int(name.split(".")[2])
layer = self.model.layers[layer_id].self_attn
if not getattr(layer, "do_fuse_qkvbfg", False):
continue
name = name.replace(weight_name, param_name)
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict: