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

support Kimi-K2.5-w4a8 on ascend

Browse Source
This commit is contained in:
khalilzhk
2026-03-10 14:43:27 +08:00
committed by GitHub
parent dadd4dde83
commit 5f717913a0
3 changed files with 91 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/layers/quantization/modelslim/modelslim.py
View File

@@ -152,6 +152,9 @@ class ModelSlimConfig(QuantizationConfig):
key = "vision_model"
elif "visual" in prefix:
key = "visual"
if "vision_tower" in prefix or "mm_projector" in prefix:
prefix = prefix.replace(r"attn.qkv_proj", r"wqkv")
prefix = prefix.replace(r"attn.proj", r"wo")
packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {})
prefix_in_quant_config = prefix
proj_name = prefix.split(".")[-1]

47
python/sglang/srt/models/kimi_k25.py
View File

@@ -26,6 +26,7 @@ except ImportError:
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
from sglang.srt.managers.schedule_batch import (
Modality,
MultimodalDataItem,
@@ -120,7 +121,13 @@ class MoonViTEncoderLayer(nn.Module):
self.norm0 = nn.LayerNorm(hidden_dim)
self.norm1 = nn.LayerNorm(hidden_dim)
self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation)
self.mlp = MLP2(
[hidden_dim, mlp_dim, hidden_dim],
activation,
quant_config=quant_config,
prefix=add_prefix("mlp", prefix),
)
self.attn = VisionAttention(
embed_dim=hidden_dim,
@@ -430,6 +437,8 @@ class MoonViT3dEncoder(nn.Module):
num_layers: int,
block_cfg: dict,
video_attn_type: str = "spatial_temporal",
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
@@ -441,7 +450,14 @@ class MoonViT3dEncoder(nn.Module):
block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512
)
self.blocks = nn.ModuleList(
[MoonViTEncoderLayer(**block_cfg) for _ in range(num_layers)]
[
MoonViTEncoderLayer(
**block_cfg,
quant_config=quant_config,
prefix=add_prefix(f"blocks.{layer_idx}", prefix),
)
for layer_idx in range(num_layers)
]
)
self.final_layernorm = nn.LayerNorm(hidden_dim)
@@ -480,7 +496,15 @@ class MoonViT3dPretrainedModel(nn.Module):
_supports_flash_attn_2 = True
_supports_sdpa = True
def __init__(self, config, *inputs, use_data_parallel: bool = False, **kwargs):
def __init__(
self,
config,
*inputs,
use_data_parallel: bool = False,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
**kwargs,
):
super().__init__()
config = deepcopy(config)
self.config = config
@@ -509,6 +533,8 @@ class MoonViT3dPretrainedModel(nn.Module):
"use_data_parallel": use_data_parallel,
},
video_attn_type=config.video_attn_type,
quant_config=quant_config,
prefix=add_prefix("encoder", prefix),
)
@property
@@ -671,12 +697,23 @@ class KimiK25ForConditionalGeneration(nn.Module):
self.use_data_parallel = get_global_server_args().mm_enable_dp_encoder
# Create vision tower
self.vision_tower = MoonViT3dPretrainedModel(
config.vision_config, use_data_parallel=self.use_data_parallel
config.vision_config,
use_data_parallel=self.use_data_parallel,
quant_config=(
quant_config if isinstance(quant_config, ModelSlimConfig) else None
),
prefix="vision_tower",
)
# Create mm projector
self.mm_projector = K2VLMultiModalProjector(config.vision_config)
self.language_model = DeepseekV3ForCausalLM(config.text_config, quant_config)
self.language_model = DeepseekV3ForCausalLM(
config.text_config,
quant_config,
prefix=(
"language_model" if isinstance(quant_config, ModelSlimConfig) else ""
),
)
# Ensure that the dtype of the vision_tower and mm_projector matches that of the language_model.
# This solves the dtype mismatch issue when using device_map="auto" and torch_dtype.

56
python/sglang/srt/models/kimi_vl_moonvit.py
View File

@@ -58,6 +58,10 @@ except ImportError:
flash_attn_varlen_func = None
from sglang.srt.configs import MoonViTConfig
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.quantization import QuantizationConfig
from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
from sglang.srt.utils import add_prefix
def multihead_attention(
@@ -393,21 +397,53 @@ class MLP2(nn.Module):
bias: whether to use bias in linear layer.
"""
def __init__(self, dims: list[int], activation, bias=True):
def __init__(
self,
dims: list[int],
activation,
bias: bool = True,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
):
super().__init__()
assert len(dims) == 3
self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)
self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)
self.quant_config = quant_config
if isinstance(self.quant_config, ModelSlimConfig):
self.fc0 = ReplicatedLinear(
dims[0],
dims[1],
bias=bias,
quant_config=quant_config,
prefix=add_prefix("fc0", prefix),
)
self.fc1 = ReplicatedLinear(
dims[1],
dims[2],
bias=bias,
quant_config=quant_config,
prefix=add_prefix("fc1", prefix),
)
else:
self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)
self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)
for m in [self.fc0, self.fc1]:
nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features))
if m.bias is not None:
nn.init.zeros_(m.bias)
self.activation = activation
for m in [self.fc0, self.fc1]:
nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features))
if m.bias is not None:
nn.init.zeros_(m.bias)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.fc0(x)
x = self.activation(x)
return self.fc1(x)
if isinstance(self.quant_config, ModelSlimConfig):
x = x.flatten(0, 1)
x, _ = self.fc0(x)
x = self.activation(x)
x, _ = self.fc1(x)
else:
x = self.fc0(x)
x = self.activation(x)
x = self.fc1(x)
return x
class MoonVitEncoderLayer(nn.Module):