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

python/sglang/srt/layers/linear.py

@@ -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)