minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CPU] Add Qwen3.5 model optimization for CPU (#19484)

Browse Source 
Co-authored-by: Zheng, Beilei <beilei.zheng@intel.com>
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
This commit is contained in:
jianan-gu
2026-04-26 10:12:36 -07:00
committed by GitHub
co-authored by Zheng, Beilei Ma Mingfei Xinyuan Tong
parent 7d49564431
commit 10fd0faccd
20 changed files with 768 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/update_config.py
+179 -76
View File
@@ -1,7 +1,13 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from sglang.srt.utils import (
log_debug_on_rank0,
)
logger = logging.getLogger(__name__)
DEFAULT_MOE_PADDING_SIZE = 32
@@ -40,7 +46,14 @@ def get_moe_padding_size(weight_block_size):
return DEFAULT_MOE_PADDING_SIZE
def get_num_heads_padding_size(tp_size, weight_block_size, head_dim):
def get_num_heads_padding_size(tp_size, weight_block_size, head_dim=None):
if head_dim is None:
pad_size = (
tp_size * 2
if tp_size % 2 == 1 and weight_block_size is not None
else tp_size
)
return pad_size
pad_size = tp_size
if weight_block_size is not None and head_dim % weight_block_size[0] != 0:
@@ -53,6 +66,25 @@ def get_num_heads_padding_size(tp_size, weight_block_size, head_dim):
return pad_size
def resolve_head_dim(cfg, num_heads, is_text_config):
# default getting head_dim by hidden_size and num_heads
hidden_size = getattr(cfg, "hidden_size", getattr(cfg, "d_model", None))
head_dim = hidden_size // num_heads if hidden_size else None
# update head_dim if specified in model config
if is_text_config:
if hasattr(cfg.hf_config, "qk_head_dim"):
head_dim = cfg.hf_config.qk_head_dim
elif hasattr(cfg.hf_text_config, "head_dim"):
head_dim = cfg.hf_text_config.head_dim
elif hasattr(cfg.hf_config, "head_dim"):
head_dim = cfg.hf_config.head_dim
else:
if hasattr(cfg, "head_dim"):
head_dim = cfg.head_dim
return head_dim
def adjust_tp_num_heads_if_necessary(model_config, tp_size, is_post_update):
# is_post_update: whether to update an existing config
from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size
@@ -75,25 +107,45 @@ def adjust_tp_num_heads_if_necessary(model_config, tp_size, is_post_update):
// model_config.linear_num_key_heads
)
if is_post_update:
model_config.linear_num_key_heads_cpu = linear_num_key_heads_cpu
model_config.linear_num_value_heads_cpu = linear_num_value_heads_cpu
update_config(
model_config, "linear_num_key_heads_cpu", linear_num_key_heads_cpu
)
update_config(
model_config,
"linear_num_value_heads_cpu",
linear_num_value_heads_cpu,
)
else:
model_config.linear_num_key_heads = linear_num_key_heads_cpu
model_config.linear_num_value_heads = linear_num_value_heads_cpu
update_config(
model_config, "linear_num_key_heads", linear_num_key_heads_cpu
)
update_config(
model_config, "linear_num_value_heads", linear_num_value_heads_cpu
)
else:
if is_post_update:
model_config.linear_num_key_heads_cpu = (
model_config.linear_num_key_heads
update_config(
model_config,
"linear_num_key_heads_cpu",
model_config.linear_num_key_heads,
)
model_config.linear_num_value_heads_cpu = (
model_config.linear_num_value_heads
update_config(
model_config,
"linear_num_value_heads_cpu",
model_config.linear_num_value_heads,
)
def update_intermediate_size(model_config, attr_name, intermediate_padding_size):
attr_value = intermediate_padding_size
if hasattr(model_config, "hf_config") and hasattr(
if (
hasattr(model_config, "hf_config")
and hasattr(model_config.hf_config, "text_config")
and hasattr(model_config.hf_config.text_config, attr_name)
):
attr_value = getattr(model_config.hf_config.text_config, attr_name)
elif hasattr(model_config, "hf_config") and hasattr(
model_config.hf_config, attr_name
):
attr_value = getattr(model_config.hf_config, attr_name)
@@ -105,50 +157,62 @@ def update_intermediate_size(model_config, attr_name, intermediate_padding_size)
attr_value = pad_vocab_size(attr_value, intermediate_padding_size)
if hasattr(model_config, "hf_config"):
setattr(model_config.hf_config, attr_name, attr_value)
update_config(model_config.hf_config, attr_name, attr_value)
if hasattr(model_config, "hf_text_config"):
setattr(model_config.hf_text_config, attr_name, attr_value)
update_config(model_config.hf_text_config, attr_name, attr_value)
if hasattr(model_config.hf_config, "text_config"):
update_config(model_config.hf_config.text_config, attr_name, attr_value)
else:
setattr(model_config, attr_name, attr_value)
update_config(model_config, attr_name, attr_value)
return model_config
def update_config(model_config, attr_name, new_value):
config_name = model_config.__class__.__name__
if hasattr(model_config, attr_name):
old_value = getattr(model_config, attr_name)
if old_value != new_value:
log_debug_on_rank0(
logger,
f"Updating {config_name}.{attr_name} from {old_value} to {new_value}",
)
else:
log_debug_on_rank0(logger, f"Setting {config_name}.{attr_name} to {new_value}")
setattr(model_config, attr_name, new_value)
def adjust_config_with_unaligned_cpu_tp(
model_config: ModelConfig, load_config: LoadConfig, tp_size: int
) -> ModelConfig:
# Support the case where the num_attention_heads is not divisible by the TP size.
weight_block_size = may_get_weight_block_size(model_config, load_config)
model_config.hf_config.original_num_attention_heads = (
model_config.num_attention_heads
)
model_config.hf_text_config.original_num_attention_heads = (
model_config.num_attention_heads
)
model_config.hf_config.original_total_num_kv_heads = (
model_config.get_total_num_kv_heads()
)
model_config.hf_text_config.original_total_num_kv_heads = (
model_config.get_total_num_kv_heads()
)
for config in [model_config.hf_config, model_config.hf_text_config]:
update_config(
config,
"original_num_attention_heads",
model_config.num_attention_heads,
)
update_config(
config,
"original_total_num_kv_heads",
model_config.get_total_num_kv_heads(),
)
if (
model_config.num_attention_heads % tp_size != 0
or model_config.get_total_num_kv_heads() % tp_size != 0
):
# Compute the head_dim using the model_config.num_attention_heads before padding
if not hasattr(model_config.hf_config, "head_dim"):
model_config.hf_config.head_dim = (
model_config.hidden_size // model_config.num_attention_heads
)
if hasattr(model_config.hf_config, "qk_nope_head_dim") and hasattr(
model_config.hf_config, "qk_rope_head_dim"
):
model_config.hf_config.qk_head_dim = (
update_config(
model_config.hf_config,
"qk_head_dim",
model_config.hf_config.qk_nope_head_dim
+ model_config.hf_config.qk_rope_head_dim
+ model_config.hf_config.qk_rope_head_dim,
)
query_heads_per_kv = (
@@ -157,60 +221,99 @@ def adjust_config_with_unaligned_cpu_tp(
total_kv_heads = model_config.get_total_num_kv_heads()
from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size
head_dim = (
model_config.hf_config.qk_head_dim
if hasattr(model_config.hf_config, "qk_head_dim")
else model_config.hf_config.head_dim
head_dim = resolve_head_dim(
model_config, model_config.num_attention_heads, True
)
pad_size = get_num_heads_padding_size(tp_size, weight_block_size, head_dim)
num_key_value_heads = pad_vocab_size(total_kv_heads, pad_size)
model_config.num_key_value_heads = num_key_value_heads
model_config.hf_config.num_key_value_heads = num_key_value_heads
model_config.hf_text_config.num_key_value_heads = num_key_value_heads
num_attention_heads = num_key_value_heads * query_heads_per_kv
model_config.num_attention_heads = num_attention_heads
model_config.hf_config.num_attention_heads = num_attention_heads
model_config.hf_text_config.num_attention_heads = num_attention_heads
for config in [
model_config,
model_config.hf_config,
model_config.hf_text_config,
]:
update_config(config, "num_key_value_heads", num_key_value_heads)
update_config(config, "num_attention_heads", num_attention_heads)
adjust_tp_num_heads_if_necessary(model_config.hf_config, tp_size, True)
if hasattr(model_config.hf_config, "text_config"):
adjust_tp_num_heads_if_necessary(
model_config.hf_config.text_config, tp_size, True
)
intermediate_padding_size = tp_size * get_moe_padding_size(weight_block_size)
model_config = update_intermediate_size(
model_config, "moe_intermediate_size", intermediate_padding_size
)
model_config = update_intermediate_size(
model_config, "intermediate_size", intermediate_padding_size
)
model_config = update_intermediate_size(
model_config, "intermediate_size_mlp", intermediate_padding_size
)
model_config = update_intermediate_size(
model_config, "shared_expert_intermediate_size", intermediate_padding_size
)
if (
hasattr(model_config.hf_config, "vision_config")
and model_config.hf_config.vision_config.model_type == "siglip_vision_model"
):
model_config.hf_config.vision_config.original_num_attention_heads = (
model_config.num_attention_heads
for moe_intermediate_attr in [
"moe_intermediate_size",
"intermediate_size",
"intermediate_size_mlp",
"shared_expert_intermediate_size",
]:
model_config = update_intermediate_size(
model_config, moe_intermediate_attr, intermediate_padding_size
)
if model_config.hf_config.vision_config.num_attention_heads % tp_size != 0:
model_config.hf_config.vision_config.head_dim = (
model_config.hf_config.vision_config.hidden_size
// model_config.hf_config.vision_config.num_attention_heads
)
from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size
pad_size = get_num_heads_padding_size(tp_size, weight_block_size)
model_config.hf_config.vision_config.num_attention_heads = pad_vocab_size(
model_config.hf_config.vision_config.num_attention_heads, pad_size
)
model_config.hf_config.vision_config = update_intermediate_size(
model_config.hf_config.vision_config,
"intermediate_size",
intermediate_padding_size,
multimodal_config = [
[
model_config.hf_config,
"vision_config",
"siglip_vision_model",
"num_attention_heads",
],
[model_config.hf_config, "vision_config", "qwen3_vl_moe", "num_heads"],
[model_config.hf_config, "vision_config", "qwen3_vl", "num_heads"],
[model_config.hf_config, "vision_config", "qwen3_5_moe", "num_heads"],
[model_config.hf_config, "vision_config", "qwen3_5", "num_heads"],
]
if hasattr(model_config.hf_config, "thinker_config"):
multimodal_config.append(
[
model_config.hf_config.thinker_config,
"vision_config",
"qwen3_omni_moe_vision_encoder",
"num_heads",
]
)
multimodal_config.append(
[
model_config.hf_config.thinker_config,
"audio_config",
"qwen3_omni_moe_audio_encoder",
"encoder_attention_heads",
]
)
for m_config, config_name, model_type, num_head_str in multimodal_config:
if (
hasattr(m_config, config_name)
and getattr(m_config, config_name).model_type == model_type
):
num_heads = getattr(getattr(m_config, config_name), num_head_str)
update_config(
getattr(m_config, config_name), "original_" + num_head_str, num_heads
)
if num_heads % tp_size != 0:
from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size
multimodal_head_dim = resolve_head_dim(
getattr(m_config, config_name), num_heads, False
)
pad_size = get_num_heads_padding_size(
tp_size, weight_block_size, multimodal_head_dim
)
new_num_heads = pad_vocab_size(num_heads, pad_size)
update_config(
getattr(m_config, config_name), num_head_str, new_num_heads
)
setattr(
m_config,
config_name,
update_intermediate_size(
getattr(m_config, config_name),
"intermediate_size",
intermediate_padding_size,
),
)
return model_config
python/sglang/srt/layers/attention/fla/fused_norm_gate.py
+19 -11
View File
@@ -375,14 +375,22 @@ class FusedRMSNormGated(nn.Module):
prenorm: bool = False,
residual_in_fp32: bool = False,
) -> torch.Tensor:
return rms_norm_gated(
x,
g,
self.weight,
self.bias,
self.activation,
residual=residual,
eps=self.eps,
prenorm=prenorm,
residual_in_fp32=residual_in_fp32,
)
if _use_cpu:
assert (
self.activation == "silu"
), "CPU rmsnorm_gated currently only supports activation silu"
return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(
x, self.weight, g, self.eps
)
else:
return rms_norm_gated(
x,
g,
self.weight,
self.bias,
self.activation,
residual=residual,
eps=self.eps,
prenorm=prenorm,
residual_in_fp32=residual_in_fp32,
)
python/sglang/srt/layers/attention/mamba/mamba.py
+18 -2
View File
@@ -1,3 +1,4 @@
import logging
from typing import Callable, List, Optional, Tuple
import torch
@@ -29,7 +30,12 @@ from sglang.srt.model_loader.weight_utils import (
composed_weight_loader,
sharded_weight_loader,
)
from sglang.srt.utils import is_cpu, is_cuda, is_npu, set_weight_attrs
from sglang.srt.utils import (
is_cpu,
is_cuda,
is_npu,
set_weight_attrs,
)
if is_cuda():
from sglang.srt.layers.attention.mamba.causal_conv1d import (
@@ -52,6 +58,8 @@ elif is_npu():
LoaderFunction = Callable[[torch.Tensor, torch.Tensor], None]
logger = logging.getLogger(__name__)
def mamba_v2_sharded_weight_loader(
shard_spec: List[Tuple[int, int, float]],
@@ -81,6 +89,14 @@ def mamba_v2_sharded_weight_loader(
weight_full_dim_list.append(
int(full_dim / full_dim_sum * loaded_weight.size(0))
)
assert sum(weight_full_dim_list) == loaded_weight.size(
0
), f"Padding the loaded weight failed due to sizes are not divisible cleanly from {weight_full_dim_list} to {loaded_weight.size(0)}"
if loaded_weight.size(0) < full_dim_sum and tp_rank == 0:
logger.warning(
f"[ZERO-PADDING] Loaded_weight.dim(0) size:{loaded_weight.size(0)} is padding to {full_dim_sum}"
f", where original sizes of {weight_full_dim_list} will be updated to {full_dim_list}",
)
# - iterate over the shard specs
for full_dim, extra, duplicate_groups in shard_spec:
@@ -110,7 +126,7 @@ def mamba_v2_sharded_weight_loader(
# CPU logic of padding size for qwen3-next
# TODO : make this common for all mamba.
if is_cpu() and loaded_weight.size(0) % tp_size != 0:
if is_cpu() and (loaded_weight.size(0) < full_dim_sum):
import copy
loaded_weight_ = copy.deepcopy(loaded_weight)
python/sglang/srt/layers/attention/vision.py
+2 -1
View File
@@ -749,6 +749,7 @@ class VisionAttention(nn.Module):
num_heads: int,
projection_size: int,
use_qkv_parallel: bool,
head_size: Optional[int] = None,
qkv_backend: Optional[str] = None,
quant_config: Optional[QuantizationConfig] = None,
dropout: float = 0.0,
@@ -775,7 +776,7 @@ class VisionAttention(nn.Module):
self.tp_size = 1 if use_data_parallel else get_attention_tp_size()
self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank()
self.dropout = dropout
self.head_size = embed_dim // num_heads
self.head_size = head_size if head_size is not None else embed_dim // num_heads
self.hidden_size_per_attention_head = dist_utils.divide(
projection_size, num_heads
)
python/sglang/srt/layers/linear.py
+17 -1
View File
@@ -742,6 +742,14 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
for i, output_size in enumerate(output_sizes):
shard_offsets.append((i, current_shard_offset, output_size))
current_shard_offset += output_size
if _is_cpu:
from sglang.srt.model_loader.weight_utils import (
pad_loaded_weight,
)
loaded_weight = pad_loaded_weight(
loaded_weight, param.output_dim, output_sizes
)
for shard_id, shard_offset, shard_size in shard_offsets:
# Special case for Quantization.
@@ -754,7 +762,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
shard_size, shard_offset = param.adjust_shard_indexes_for_packing(
shard_size=shard_size, shard_offset=shard_offset
)
loaded_weight_shard = loaded_weight.narrow(
param.output_dim, shard_offset, shard_size
)
@@ -781,6 +788,15 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
shard_block_offsets.append(current_block_offset)
current_block_offset += shard_block_size
if _is_cpu:
from sglang.srt.model_loader.weight_utils import (
pad_loaded_weight,
)
loaded_weight = pad_loaded_weight(
loaded_weight, param.output_dim, shard_block_sizes
)
# Load each shard
for shard_id, (shard_block_offset, shard_block_size) in enumerate(
zip(shard_block_offsets, shard_block_sizes)
python/sglang/srt/model_executor/cpu_graph_runner.py
+12
View File
@@ -408,6 +408,18 @@ def register_fake_ops():
a = mixed_ba.new_empty(batch, num_heads_v)
return mixed_qkv, z, b, a
@torch.library.register_fake(
"sgl_kernel::fused_qkvzba_split_reshape_cat_contiguous_cpu"
)
def _(mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_qk, head_v):
batch = mixed_qkvz.shape[0]
qkv_dim = num_heads_qk * head_qk * 2 + num_heads_v * head_v
mixed_qkv = mixed_qkvz.new_empty(batch, qkv_dim)
z = mixed_qkvz.new_empty(batch, num_heads_v, head_v)
b = mixed_ba.new_empty(batch, num_heads_v)
a = mixed_ba.new_empty(batch, num_heads_v)
return mixed_qkv, z, b, a
@torch.library.register_fake(
"sgl_kernel::fused_sigmoid_gating_delta_rule_update_cpu"
)
python/sglang/srt/model_loader/weight_utils.py
+35 -1
View File
@@ -1248,7 +1248,11 @@ def sharded_weight_loader(shard_axis: int) -> LoaderFunction:
if (
is_cpu()
and loaded_weight.size(0) % get_tensor_model_parallel_world_size() != 0
and (
loaded_weight.size(0) % get_tensor_model_parallel_world_size() != 0
or loaded_weight.size(0)
< get_tensor_model_parallel_world_size() * shard_size
)
and loaded_weight.dim() == 1
):
param_data = param.data # view copy on param for uneven padding
@@ -1623,3 +1627,33 @@ def narrow_padded_param_and_loaded_weight(
param_data = param_data.narrow(dim, param_data_start, actual_shard_size)
return param_data, loaded_weight
def pad_loaded_weight(loaded_weight, output_dim, output_sizes):
# This function is for padding zeros when loaded_weight is less than output_sizes.
# Most cases, sum(output_sizes) = loaded_weight.size(output_dim),
# while in some TP cases like TP6, output_sizes will be padded, thus loaded_weight needs padding.
total_output_size = sum(output_sizes)
raw_output_size = loaded_weight.size(output_dim)
if total_output_size > raw_output_size:
loaded_weight_pad = []
weight_split_size = [
int(output_size / total_output_size * raw_output_size)
for output_size in output_sizes
]
assert (
sum(weight_split_size) == raw_output_size
), f"Padding the loaded weight failed due to sizes are not divisible cleanly from {output_sizes} to {raw_output_size}"
split_weight = loaded_weight.split_with_sizes(weight_split_size, dim=output_dim)
for i, output_size in enumerate(output_sizes):
pad_size = output_size - weight_split_size[i]
target_pad_shape = list(loaded_weight.size())
target_pad_shape[output_dim] = pad_size
pad_tensor = torch.zeros(target_pad_shape).to(loaded_weight.dtype)
loaded_weight_pad.append(
torch.cat([split_weight[i], pad_tensor], dim=output_dim)
)
return torch.cat(loaded_weight_pad, dim=output_dim)
else:
return loaded_weight
python/sglang/srt/models/qwen3_5.py
+37 -4
View File
@@ -124,8 +124,16 @@ class Qwen3_5GatedDeltaNet(nn.Module):
self.attn_tp_rank = get_attention_tp_rank()
self.attn_tp_size = get_attention_tp_size()
self.hidden_size = config.hidden_size
self.num_v_heads = config.linear_num_value_heads
self.num_k_heads = config.linear_num_key_heads
self.num_v_heads = (
config.linear_num_value_heads
if not _is_cpu
else config.linear_num_value_heads_cpu
)
self.num_k_heads = (
config.linear_num_key_heads
if not _is_cpu
else config.linear_num_key_heads_cpu
)
self.head_k_dim = config.linear_key_head_dim
self.head_v_dim = config.linear_value_head_dim
self.key_dim = self.head_k_dim * self.num_k_heads
@@ -321,7 +329,20 @@ class Qwen3_5GatedDeltaNet(nn.Module):
chunks = [loaded_weight.reshape(1)]
else:
split_dim = getattr(param, "output_dim", 0)
chunks = loaded_weight.split(split_sizes, dim=split_dim)
if _is_cpu:
cpu_split_sizes = []
split_size_sum = sum(split_sizes)
target_size_sim = loaded_weight.size(split_dim)
for i in range(len(split_sizes)):
cpu_split_sizes.append(
int(target_size_sim * split_sizes[i] / split_size_sum)
)
assert (
sum(cpu_split_sizes) == target_size_sim
), f"Padding the loaded weight failed due to sizes are not divisible cleanly from {cpu_split_sizes} to {target_size_sim}"
chunks = loaded_weight.split(cpu_split_sizes, dim=split_dim)
else:
chunks = loaded_weight.split(split_sizes, dim=split_dim)
assert len(chunks) == len(loaded_shard_id), (
f"Chunk/shard mismatch: {len(chunks)=}, "
@@ -454,7 +475,7 @@ class Qwen3_5GatedDeltaNet(nn.Module):
)
elif _is_cpu and _is_amx_available:
mixed_qkv, z, b, a = (
torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_cpu(
torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_contiguous_cpu(
projected_states_qkvz,
projected_states_ba,
self.num_k_heads // self.attn_tp_size,
@@ -467,10 +488,12 @@ class Qwen3_5GatedDeltaNet(nn.Module):
query, key, value, z, b, a = self.fix_query_key_value_ordering(
projected_states_qkvz, projected_states_ba
)
query, key, value = map(
lambda x: x.reshape(x.shape[0], -1), (query, key, value)
)
mixed_qkv = torch.cat((query, key, value), dim=-1)
core_attn_out = self.attn(
forward_batch,
mixed_qkv=mixed_qkv,
@@ -1484,6 +1507,16 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration):
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
if (
self.config.tie_word_embeddings
and name == "model.embed_tokens.weight"
and (_is_cpu and _is_amx_available)
):
param_lm_head = params_dict["lm_head.weight"]
weight_loader = getattr(
param_lm_head, "weight_loader", default_weight_loader
)
weight_loader(param_lm_head, loaded_weight)
loaded_params.add(name)
return loaded_params
python/sglang/srt/models/qwen3_next.py
+14 -5
View File
@@ -245,14 +245,23 @@ class Qwen3GatedDeltaNet(nn.Module):
if output_dim is not None and module.tp_size > 1:
shard_size = param.data.shape[output_dim]
start_idx = module.tp_rank * shard_size
if (
_is_cpu and _is_amx_available
) and start_idx + shard_size > loaded_weight.shape[output_dim]:
shard_size = loaded_weight.shape[output_dim] - start_idx
loaded_weight = loaded_weight.narrow(
output_dim, start_idx, shard_size
)
assert param.data.shape == loaded_weight.shape, (
f"Shape mismatch: param {param.data.shape} vs "
f"loaded {loaded_weight.shape}"
)
param.data.copy_(loaded_weight)
if _is_cpu and _is_amx_available:
slices = tuple(slice(0, s) for s in loaded_weight.shape)
param.data.zero_()
param.data[slices].copy_(loaded_weight)
else:
assert param.data.shape == loaded_weight.shape, (
f"Shape mismatch: param {param.data.shape} vs "
f"loaded {loaded_weight.shape}"
)
param.data.copy_(loaded_weight)
else:
# Split checkpoint (int or tuple shard_id) → standard path
original_loader(param, loaded_weight, loaded_shard_id)
python/sglang/srt/models/qwen3_vl.py
+29 -6
View File
@@ -72,7 +72,13 @@ from sglang.srt.models.utils import (
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix, is_npu, round_up
from sglang.srt.utils import (
add_prefix,
cpu_has_amx_support,
is_cpu,
is_npu,
round_up,
)
from sglang.srt.utils.hf_transformers_utils import get_processor
_is_npu = is_npu()
@@ -87,6 +93,9 @@ if _is_npu:
logger = logging.getLogger(__name__)
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu()
class Qwen3_VisionMLP(nn.Module):
@@ -169,6 +178,7 @@ class Qwen3_VisionBlock(nn.Module):
dim: int,
num_heads: int,
intermediate_dim: int,
head_size: Optional[int] = None,
hidden_act="silu",
norm_layer: Optional[Callable[[int], nn.Module]] = None,
quant_config: Optional[QuantizationConfig] = None,
@@ -185,7 +195,8 @@ class Qwen3_VisionBlock(nn.Module):
self.attn = VisionAttention(
embed_dim=dim,
num_heads=num_heads,
projection_size=dim,
head_size=head_size,
projection_size=num_heads * head_size,
use_qkv_parallel=True,
proj_bias=True,
flatten_batch=True,
@@ -240,6 +251,7 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module):
self,
dim: int,
context_dim: int,
padded_context_dim: int,
norm_layer: Optional[Callable[[int], nn.Module]] = None,
spatial_merge_size: int = 2,
use_postshuffle_norm: bool = False,
@@ -249,6 +261,7 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module):
) -> None:
super().__init__()
self.hidden_size = context_dim * (spatial_merge_size**2)
self.padded_context_dim = padded_context_dim * (spatial_merge_size**2)
self.use_postshuffle_norm = use_postshuffle_norm
@@ -261,7 +274,7 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module):
self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank()
self.linear_fc1 = ColumnParallelLinear(
self.hidden_size,
self.hidden_size,
self.padded_context_dim,
bias=True,
quant_config=quant_config,
prefix=add_prefix("linear_fc1", prefix),
@@ -270,7 +283,7 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module):
)
self.act_fn = nn.GELU()
self.linear_fc2 = RowParallelLinear(
self.hidden_size,
self.padded_context_dim,
dim,
bias=True,
quant_config=quant_config,
@@ -336,7 +349,10 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
self.pos_embed = PPMissingLayer()
norm_layer = partial(nn.LayerNorm, eps=norm_eps)
head_dim = self.hidden_size // self.num_heads
if is_cpu() and hasattr(vision_config, "original_num_heads"):
head_dim = self.hidden_size // vision_config.original_num_heads
else:
head_dim = self.hidden_size // self.num_heads
self.rotary_pos_emb = get_rope(
head_size=head_dim,
rotary_dim=head_dim // 2,
@@ -363,6 +379,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
dim=self.hidden_size,
num_heads=self.num_heads,
intermediate_dim=vision_config.intermediate_size,
head_size=head_dim,
hidden_act=vision_config.hidden_act,
norm_layer=norm_layer,
quant_config=quant_config,
@@ -376,6 +393,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
self.merger = Qwen3VLMoeVisionPatchMerger(
dim=vision_config.out_hidden_size,
context_dim=self.hidden_size,
padded_context_dim=self.num_heads * head_dim,
norm_layer=norm_layer,
spatial_merge_size=self.spatial_merge_size,
quant_config=quant_config,
@@ -388,6 +406,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
Qwen3VLMoeVisionPatchMerger(
dim=vision_config.out_hidden_size,
context_dim=self.hidden_size,
padded_context_dim=self.num_heads * head_dim,
spatial_merge_size=self.spatial_merge_size,
use_postshuffle_norm=True,
norm_layer=norm_layer,
@@ -1108,7 +1127,11 @@ class Qwen3VLForConditionalGeneration(nn.Module):
prefix=add_prefix("model.language_model", prefix),
)
if self.pp_group.is_last_rank:
if self.pp_group.world_size == 1 and self.config.tie_word_embeddings:
if (
self.pp_group.world_size == 1
and self.config.tie_word_embeddings
and not (_is_cpu and _is_cpu_amx_available)
):
self.lm_head = self.model.embed_tokens
else:
self.lm_head = ParallelLMHead(
python/sglang/srt/utils/common.py
+24 -2
View File
@@ -2862,8 +2862,30 @@ def log_info_on_rank0(logger, msg):
try:
if torch.distributed.is_initialized() and get_tensor_model_parallel_rank() == 0:
logger.info(msg)
except:
logger.info(msg)
except Exception as e:
if torch.distributed.is_initialized():
if torch.distributed.get_rank() == 0:
logger.info(f"{msg} (rank-check failed: {e})")
else:
logger.info(f"{msg} (rank-check failed: {e})")
def log_debug_on_rank0(logger, msg):
"""
Log a debug message only on tensor model parallel rank 0.
Falls back to logging if distributed is not initialized or error occurs.
"""
from sglang.srt.distributed import get_tensor_model_parallel_rank
try:
if torch.distributed.is_initialized() and get_tensor_model_parallel_rank() == 0:
logger.debug(msg)
except Exception as e:
if torch.distributed.is_initialized():
if torch.distributed.get_rank() == 0:
logger.debug(f"{msg} (rank-check failed: {e})")
else:
logger.debug(f"{msg} (rank-check failed: {e})")
def load_json_config(data: str):
sgl-kernel/csrc/cpu/common.h
+50
View File
@@ -97,6 +97,43 @@ namespace {
TORCH_CHECK(false, "Unsupported floating data type."); \
}
// Helper MICRO for CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT:
// TYPE1: the primary dtype (input, output, weight);
// TYPE2: defined as PARAM_T input
#define CPU_DISPATCH_TYPE1_WITH_PARAM_REDUCED(TYPE1, PARAM_T, ...) \
switch (TYPE1) { \
case at::ScalarType::BFloat16: { \
using scalar_t = at::BFloat16; \
using param_t = PARAM_T; \
return __VA_ARGS__(); \
} \
case at::ScalarType::Half: { \
using scalar_t = at::Half; \
using param_t = PARAM_T; \
return __VA_ARGS__(); \
} \
default: \
TORCH_CHECK(false, "Unsupported floating data type."); \
}
// Helper MICRO for CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT:
// TYPE1: the dtype both for scalar_t and param_t
#define CPU_DISPATCH_TYPE1_WITH_SAME_PARAM_REDUCED(TYPE1, ...) \
switch (TYPE1) { \
case at::ScalarType::BFloat16: { \
using scalar_t = at::BFloat16; \
using param_t = at::BFloat16; \
return __VA_ARGS__(); \
} \
case at::ScalarType::Half: { \
using scalar_t = at::Half; \
using param_t = at::Half; \
return __VA_ARGS__(); \
} \
default: \
TORCH_CHECK(false, "Unsupported reduced floating data type."); \
}
// dispatch with mixed dtypes (TYPE1, TYPE2):
// TYPE1: the primary dtype (input, output, weight);
// TYPE2: the secondary dtype (bias, etc.).
@@ -113,6 +150,19 @@ namespace {
} \
}()
// dispatch with mixed dtypes (reduced one, no float for TYPE1) (TYPE1, TYPE2):
// TYPE1: the primary dtype (input, output, weight);
// TYPE2: the secondary dtype (bias, etc.).
#define CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(TYPE1, TYPE2, ...) \
[&] { \
if (TYPE2 == at::kFloat) { \
CPU_DISPATCH_TYPE1_WITH_PARAM_REDUCED(TYPE1, float, __VA_ARGS__) \
} else { \
TORCH_CHECK(TYPE1 == TYPE2); \
CPU_DISPATCH_TYPE1_WITH_SAME_PARAM_REDUCED(TYPE1, __VA_ARGS__) \
} \
}()
#define UNUSED(x) (void)(x)
#define CHECK_CPU(x) TORCH_CHECK(x.device().type() == at::kCPU, #x " must be a CPU tensor")
sgl-kernel/csrc/cpu/mamba/fla.cpp
+93 -37
View File
@@ -814,12 +814,12 @@ inline at::vec::Vectorized<float> softplus(const at::vec::Vectorized<float>& x,
return Vec::blendv(Vec::blendv(log1pex, expx, mask_lo), x, mask_hi);
}
template <typename scalar_t>
template <typename scalar_t, typename param_t>
void fused_sigmoid_gating_delta_rule_update_kernel_impl(
const scalar_t* __restrict__ q_ptr,
const scalar_t* __restrict__ k_ptr,
const scalar_t* __restrict__ v_ptr,
const float* __restrict__ A_log_ptr,
const param_t* __restrict__ A_log_ptr,
const scalar_t* __restrict__ a_ptr,
const scalar_t* __restrict__ dt_bias_ptr,
const scalar_t* __restrict__ b_ptr,
@@ -903,7 +903,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
for (int64_t i = begin; i < end; ++i) {
int64_t cache_index = indices_ptr[bi];
int64_t state_offset = (cache_index * v_num_heads + ni) * head_dim * v_head_dim;
float g_val = -std::exp(A_log_ptr[ni]) *
float g_val = -std::exp(float(A_log_ptr[ni])) *
softplus(float(a_ptr[bi * v_num_heads + ni]) + float(dt_bias_ptr[ni]), softplus_threshold);
float g_val_exp = std::exp(g_val);
fVec g_val_exp_vec = fVec(g_val_exp);
@@ -1021,6 +1021,55 @@ void fused_gdn_gating_kernel_impl(
});
}
template <typename scalar_t>
void fused_gdn_gating_kernel_impl(
scalar_t* __restrict__ A_log,
const scalar_t* __restrict__ a,
const scalar_t* __restrict__ b,
const scalar_t* __restrict__ dt_bias,
float* __restrict__ out,
scalar_t* __restrict__ beta,
int64_t batch,
int64_t num_heads) {
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int vec_size = bVec::size();
constexpr int fvec_size = fVec::size();
const fVec neg_one(-1.0f);
const fVec one(1.0f);
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
for (int64_t i = begin; i < end; ++i) {
int64_t j = 0;
for (; j < num_heads - (num_heads % vec_size); j += vec_size) {
bVec A_log_bvec = bVec::loadu(A_log + j);
fVec A_log_vec0, A_log_vec1;
std::tie(A_log_vec0, A_log_vec1) = at::vec::convert_to_float(A_log_bvec);
bVec dt_bias_vec = bVec::loadu(dt_bias + j);
bVec a_bvec = bVec::loadu(a + i * num_heads + j);
bVec b_bvec = bVec::loadu(b + i * num_heads + j);
fVec a0, a1, dt_bias_vec0, dt_bias_vec1, b0, b1;
std::tie(a0, a1) = at::vec::convert_to_float(a_bvec);
std::tie(b0, b1) = at::vec::convert_to_float(b_bvec);
std::tie(dt_bias_vec0, dt_bias_vec1) = at::vec::convert_to_float(dt_bias_vec);
fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0);
fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1);
fVec beta0 = one / (one + (neg_one * b0).exp_u20());
fVec beta1 = one / (one + (neg_one * b1).exp_u20());
g0.store(out + i * num_heads + j);
g1.store(out + i * num_heads + j + fvec_size);
bVec beta_vec = at::vec::convert_from_float<scalar_t>(beta0, beta1);
beta_vec.store(beta + i * num_heads + j);
}
for (; j < num_heads; ++j) {
out[i * num_heads + j] = -std::exp(float(A_log[j])) * softplus(float(a[i * num_heads + j]) + float(dt_bias[j]));
beta[i * num_heads + j] = 1 / (1 + std::exp(-b[i * num_heads + j]));
}
}
});
}
} // anonymous namespace
template <bool is_last_dim_contiguous>
@@ -1242,7 +1291,6 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
int64_t v_head_dim = v.size(3);
CHECK_INPUT_SHAPE_DTYPE<true>(k, {seq_len, batch_size, num_heads, head_dim}, q.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<true>(v, {seq_len, batch_size, v_num_heads, v_head_dim}, q.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<true>(A_log, {v_num_heads}, at::kFloat);
CHECK_INPUT_SHAPE_DTYPE<true>(a, {batch_size, v_num_heads}, q.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<true>(dt_bias, {v_num_heads}, q.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<true>(b, {batch_size, v_num_heads}, q.scalar_type());
@@ -1252,6 +1300,12 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
initial_state_source, {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}, at::kFloat);
CHECK(initial_state_source.size(0) >= batch_size);
CHECK_EQ(v_num_heads % num_heads, 0);
TORCH_CHECK(
A_log.sizes() == at::IntArrayRef({v_num_heads}),
"Input tensor shape mismatch: expected ",
at::IntArrayRef({v_num_heads}),
", got ",
A_log.sizes());
int64_t q_strideB = q.stride(1);
int64_t q_strideS = q.stride(0);
@@ -1264,37 +1318,39 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
int64_t v_strideH = v.stride(2);
at::Tensor core_attn_out = at::empty({batch_size, seq_len, v_num_heads, v_head_dim}, q.options());
at::Tensor qk_scale_buf = at::empty({2 * batch_size, seq_len, num_heads}, at::kFloat);
AT_DISPATCH_REDUCED_FLOATING_TYPES(q.scalar_type(), "fused_sigmoid_gating_delta_rule_update_kernel_impl", [&] {
fused_sigmoid_gating_delta_rule_update_kernel_impl<scalar_t>(
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
A_log.data_ptr<float>(),
a.data_ptr<scalar_t>(),
dt_bias.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
initial_state_indices.data_ptr<int32_t>(),
initial_state_source.data_ptr<float>(),
core_attn_out.data_ptr<scalar_t>(),
qk_scale_buf.data_ptr<float>(),
seq_len,
batch_size,
num_heads,
head_dim,
v_num_heads,
v_head_dim,
q_strideB,
q_strideS,
q_strideH,
k_strideB,
k_strideS,
k_strideH,
v_strideB,
v_strideS,
v_strideH,
use_qk_l2norm_in_kernel,
softplus_threshold);
});
CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
q.scalar_type(), A_log.scalar_type(), "fused_sigmoid_gating_delta_rule_update_kernel_impl", [&] {
fused_sigmoid_gating_delta_rule_update_kernel_impl<scalar_t, param_t>(
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
A_log.data_ptr<param_t>(),
a.data_ptr<scalar_t>(),
dt_bias.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
initial_state_indices.data_ptr<int32_t>(),
initial_state_source.data_ptr<float>(),
core_attn_out.data_ptr<scalar_t>(),
qk_scale_buf.data_ptr<float>(),
seq_len,
batch_size,
num_heads,
head_dim,
v_num_heads,
v_head_dim,
q_strideB,
q_strideS,
q_strideH,
k_strideB,
k_strideS,
k_strideH,
v_strideB,
v_strideS,
v_strideH,
use_qk_l2norm_in_kernel,
softplus_threshold);
});
return core_attn_out;
}
@@ -1318,9 +1374,9 @@ fused_gdn_gating_cpu(const at::Tensor& A_log, const at::Tensor& a, const at::Ten
CHECK_EQ(b.size(1), num_heads);
at::Tensor out = at::empty({1, batch, num_heads}, a.options().dtype(at::kFloat));
at::Tensor beta = at::empty({1, batch, num_heads}, b.options());
AT_DISPATCH_REDUCED_FLOATING_TYPES(a.scalar_type(), "fused_gdn_gating_kernel", [&] {
CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(a.scalar_type(), A_log.scalar_type(), "fused_gdn_gating_kernel", [&] {
fused_gdn_gating_kernel_impl<scalar_t>(
A_log.data_ptr<float>(),
A_log.data_ptr<param_t>(),
a.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
dt_bias.data_ptr<scalar_t>(),
sgl-kernel/csrc/cpu/model/qwen3.cpp
+86
View File
@@ -61,6 +61,41 @@ void fused_qkvzba_split_reshape_cat_impl(
}
});
}
template <typename scalar_t>
void fused_qkvzba_split_reshape_cat_contiguous_impl(
const scalar_t* __restrict__ mixed_qkvz,
const scalar_t* __restrict__ mixed_ba,
scalar_t* __restrict__ mixed_qkv,
scalar_t* __restrict__ z,
scalar_t* __restrict__ b,
scalar_t* __restrict__ a,
int64_t batch,
int64_t k_tp,
int64_t v_tp,
int64_t num_heads_v,
int64_t qkv_dim,
int64_t qkv_strideB,
int64_t qkvz_strideB,
int64_t ba_strideB) {
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
for (int64_t bi = begin; bi < end; ++bi) {
scalar_t* __restrict__ qkv_out_ptr = mixed_qkv + bi * qkv_strideB;
const scalar_t* __restrict__ qkv_in_ptr = mixed_qkvz + bi * qkvz_strideB;
scalar_t* __restrict__ z_out_ptr = z + bi * v_tp;
const scalar_t* __restrict__ z_in_ptr = qkv_in_ptr + qkv_dim;
copy_stub(qkv_out_ptr, qkv_in_ptr, qkv_dim);
copy_stub(z_out_ptr, z_in_ptr, v_tp);
scalar_t* __restrict__ b_out_ptr = b + bi * num_heads_v;
const scalar_t* __restrict__ b_in_ptr = mixed_ba + bi * ba_strideB;
scalar_t* __restrict__ a_out_ptr = a + bi * num_heads_v;
const scalar_t* __restrict__ a_in_ptr = b_in_ptr + num_heads_v;
copy_stub(b_out_ptr, b_in_ptr, num_heads_v);
copy_stub(a_out_ptr, a_in_ptr, num_heads_v);
}
});
}
} // anonymous namespace
// mixed_qkvz: [batch, num_heads_qk * head_qk * 2 + num_heads_v * head_v * 2]
@@ -83,6 +118,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_re
CHECK_EQ(mixed_qkvz.size(1), expected_dim);
CHECK_EQ(mixed_ba.size(0), batch);
CHECK_EQ(mixed_ba.size(1), ba_dim);
TORCH_CHECK(mixed_ba.scalar_type() == mixed_qkvz.scalar_type(), "mixed_ba and mixed_qkvz must share same dtype");
CHECK_EQ(num_heads_v % num_heads_qk, 0);
at::Tensor mixed_qkv = at::empty({batch, qkv_dim}, mixed_qkvz.options());
at::Tensor z = at::empty({batch, num_heads_v, head_v}, mixed_qkvz.options());
@@ -112,3 +148,53 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_re
});
return std::make_tuple(mixed_qkv, z, b, a);
}
// mixed_qkvz: [batch, num_heads_qk * head_qk * 2 + num_heads_v * head_v * 2]
// mixed_ba: [batch, num_heads_v * 2]
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_reshape_cat_contiguous_cpu(
const at::Tensor& mixed_qkvz,
const at::Tensor& mixed_ba,
int64_t num_heads_qk,
int64_t num_heads_v,
int64_t head_qk,
int64_t head_v) {
CHECK_DIM(2, mixed_qkvz);
CHECK_DIM(2, mixed_ba);
CHECK_INPUT(mixed_qkvz);
CHECK_INPUT(mixed_ba);
int64_t batch = mixed_qkvz.size(0);
int64_t k_tp = num_heads_qk * head_qk;
int64_t v_tp = num_heads_v * head_v;
int64_t qkv_dim = k_tp * 2 + v_tp;
int64_t ba_dim = num_heads_v * 2;
int64_t expected_dim = qkv_dim + v_tp;
CHECK_EQ(mixed_qkvz.size(1), expected_dim);
CHECK_EQ(mixed_ba.size(0), batch);
CHECK_EQ(mixed_ba.size(1), ba_dim);
TORCH_CHECK(mixed_ba.scalar_type() == mixed_qkvz.scalar_type(), "mixed_ba and mixed_qkvz must share same dtype");
at::Tensor mixed_qkv = at::empty({batch, qkv_dim}, mixed_qkvz.options());
at::Tensor z = at::empty({batch, num_heads_v, head_v}, mixed_qkvz.options());
at::Tensor b = at::empty({batch, num_heads_v}, mixed_ba.options());
at::Tensor a = at::empty({batch, num_heads_v}, mixed_ba.options());
int64_t qkvz_strideB = mixed_qkvz.size(1);
int64_t qkv_strideB = mixed_qkv.size(1);
int64_t ba_strideB = mixed_ba.size(1);
AT_DISPATCH_REDUCED_FLOATING_TYPES(mixed_qkvz.scalar_type(), "fused_qkvzba_split_reshape_cat_contiguous_impl", [&] {
fused_qkvzba_split_reshape_cat_contiguous_impl<scalar_t>(
mixed_qkvz.data_ptr<scalar_t>(),
mixed_ba.data_ptr<scalar_t>(),
mixed_qkv.data_ptr<scalar_t>(),
z.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
a.data_ptr<scalar_t>(),
batch,
k_tp,
v_tp,
num_heads_v,
qkv_dim,
qkv_strideB,
qkvz_strideB,
ba_strideB);
});
return std::make_tuple(mixed_qkv, z, b, a);
}
sgl-kernel/csrc/cpu/moe.cpp
+7 -5
View File
@@ -495,13 +495,15 @@ void fused_experts_kernel_impl(
const scalar_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb_upper * BLOCK_N * stride_n;
const scalar_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb_lower * BLOCK_N * stride_n;
// 1.a load A
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
int64_t m_size = offsets[mb + 1] - offsets[mb];
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m] / topk;
copy_stub(A + m * K, input + index * K, K);
if (nb_offset == 0) {
// 1.a load A
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m] / topk;
copy_stub(A + m * K, input + index * K, K);
}
}
if (use_brgemm) {
sgl-kernel/csrc/cpu/moe_fp8.cpp
+7 -5
View File
@@ -65,13 +65,15 @@ void fused_experts_fp8_kernel_impl(
int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
// 1.a load A
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
int64_t m_size = offsets[mb + 1] - offsets[mb];
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m] / topk;
copy_stub(A + m * K, input + index * K, K);
if (nb_offset == 0) {
// 1.a load A
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m] / topk;
copy_stub(A + m * K, input + index * K, K);
}
}
const int64_t offset = offsets[mb];
sgl-kernel/csrc/cpu/moe_int8.cpp
+8 -6
View File
@@ -550,14 +550,16 @@ void fused_experts_int8_kernel_impl(
const float* __restrict__ Bs0 = w1s + expert_id * 2 * N + nb_upper * BLOCK_N;
const float* __restrict__ Bs1 = w1s + expert_id * 2 * N + nb_lower * BLOCK_N;
// 1.a load A
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
int64_t m_size = offsets[mb + 1] - offsets[mb];
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m] / topk;
copy_stub(A + m * K, Aq_tmp + index * K, K);
As[m] = As_tmp[index];
if (nb_offset == 0) {
// 1.a load A
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
for (int64_t m = 0; m < m_size; ++m) {
int32_t index = A_ids[m] / topk;
copy_stub(A + m * K, Aq_tmp + index * K, K);
As[m] = As_tmp[index];
}
}
if (use_brgemm) {
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+15
View File
@@ -374,6 +374,15 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_re
int64_t head_qk,
int64_t head_v);
// fused_qkvzba_split_reshape_cat_cpu_contiguous
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_reshape_cat_contiguous_cpu(
const at::Tensor& mixed_qkvz,
const at::Tensor& mixed_ba,
int64_t num_heads_qk,
int64_t num_heads_v,
int64_t head_qk,
int64_t head_v);
// image preprocessor
std::tuple<at::Tensor, at::Tensor> image_preprocess_cpu(
at::TensorList images,
@@ -621,6 +630,12 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"fused_qkvzba_split_reshape_cat_cpu(Tensor mixed_qkvz, Tensor mixed_ba, int num_heads_qk, int num_heads_v, int "
"head_qk, int head_v) -> (Tensor, Tensor, Tensor, Tensor)");
m.impl("fused_qkvzba_split_reshape_cat_cpu", torch::kCPU, &fused_qkvzba_split_reshape_cat_cpu);
// fused_qkvzba_split_reshape_cat_contiguous_cpu
m.def(
"fused_qkvzba_split_reshape_cat_contiguous_cpu(Tensor mixed_qkvz, Tensor mixed_ba, int num_heads_qk, int "
"num_heads_v, int "
"head_qk, int head_v) -> (Tensor, Tensor, Tensor, Tensor)");
m.impl("fused_qkvzba_split_reshape_cat_contiguous_cpu", torch::kCPU, &fused_qkvzba_split_reshape_cat_contiguous_cpu);
// image preprocessor
m.def(
test/srt/cpu/test_mamba.py
+54 -47
View File
@@ -291,19 +291,20 @@ class TestMambaAttention(CustomTestCase):
def test_fused_gdn_gating(self):
dims = [6, 32]
for dim in dims:
A_log = torch.rand(dim)
a = torch.rand(1024, dim, dtype=torch.bfloat16)
b = torch.rand(1024, dim, dtype=torch.bfloat16)
dt_bias = torch.rand(dim, dtype=torch.bfloat16)
for A_log_dtype in [torch.float32, torch.bfloat16]:
A_log = torch.rand(dim, dtype=A_log_dtype)
a = torch.rand(1024, dim, dtype=torch.bfloat16)
b = torch.rand(1024, dim, dtype=torch.bfloat16)
dt_bias = torch.rand(dim, dtype=torch.bfloat16)
g, beta = torch_gdn_gating(A_log, a, b, dt_bias)
g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(
A_log, a, b, dt_bias
)
atol = rtol = precision[g.dtype]
atol2 = rtol2 = precision[beta.dtype]
torch.testing.assert_close(g, g_sgl, atol=atol, rtol=rtol)
torch.testing.assert_close(beta, beta_sgl, atol=atol2, rtol=rtol2)
g, beta = torch_gdn_gating(A_log, a, b, dt_bias)
g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(
A_log, a, b, dt_bias
)
atol = rtol = precision[g.dtype]
atol2 = rtol2 = precision[beta.dtype]
torch.testing.assert_close(g, g_sgl, atol=atol, rtol=rtol)
torch.testing.assert_close(beta, beta_sgl, atol=atol2, rtol=rtol2)
def test_fused_sigmoid_gating_delta_rule_update(self):
batch_size = 1
@@ -346,41 +347,47 @@ class TestMambaAttention(CustomTestCase):
if num_value_heads // num_heads > 1:
query_ref = query_ref.repeat_interleave(num_value_heads // num_heads, dim=2)
key_ref = key_ref.repeat_interleave(num_value_heads // num_heads, dim=2)
core_attn_out_ref, last_recurrent_state_ref = sigmoid_gating_delta_rule_update(
query_ref.transpose(0, 1),
key_ref.transpose(0, 1),
value.transpose(0, 1),
A_log,
a,
dt_bias,
b,
initial_state=ssm_states[cache_indices],
output_final_state=True,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
)
core_attn_out = torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
A_log=A_log,
dt_bias=dt_bias,
q=query,
k=key,
v=value,
a=a,
b=b,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
softplus_beta=1.0,
softplus_threshold=20.0,
)
last_recurrent_state = ssm_states[cache_indices]
atol = rtol = precision[core_attn_out.dtype]
torch.testing.assert_close(
core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol
)
torch.testing.assert_close(
last_recurrent_state, last_recurrent_state_ref, atol=atol, rtol=rtol
)
for A_log_dtype in [torch.float32, torch.bfloat16]:
A_log = A_log.to(A_log_dtype)
core_attn_out_ref, last_recurrent_state_ref = (
sigmoid_gating_delta_rule_update(
query_ref.transpose(0, 1),
key_ref.transpose(0, 1),
value.transpose(0, 1),
A_log,
a,
dt_bias,
b,
initial_state=ssm_states[cache_indices],
output_final_state=True,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
)
)
core_attn_out = (
torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
A_log=A_log,
dt_bias=dt_bias,
q=query,
k=key,
v=value,
a=a,
b=b,
initial_state_source=ssm_states,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
softplus_beta=1.0,
softplus_threshold=20.0,
)
)
last_recurrent_state = ssm_states[cache_indices]
atol = rtol = precision[core_attn_out.dtype]
torch.testing.assert_close(
core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol
)
torch.testing.assert_close(
last_recurrent_state, last_recurrent_state_ref, atol=atol, rtol=rtol
)
if __name__ == "__main__":
test/srt/cpu/test_qwen3.py
+62
View File
@@ -53,6 +53,34 @@ def fix_query_key_value_ordering_reshape_cat(
return mixed_qkv, z, b, a
def fix_query_key_value_ordering_reshape_cat_contiguous(
mixed_qkvz: torch.Tensor,
mixed_ba: torch.Tensor,
key_dim: int,
value_dim: int,
num_v_heads: int,
head_v_dim: int,
attn_tp_size: int,
):
"""
Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.
"""
k_tp = key_dim // attn_tp_size
v_tp = value_dim // attn_tp_size
nv_tp = num_v_heads // attn_tp_size
# Directly split, no head group reshape
query, key, value, z = mixed_qkvz.split([k_tp, k_tp, v_tp, v_tp], dim=-1)
b, a = mixed_ba.split([nv_tp, nv_tp], dim=-1)
# value / z reshape to (seq, num_v_heads/tp, head_v_dim)
value = value.reshape(value.size(0), -1, head_v_dim)
z = z.reshape(z.size(0), -1, head_v_dim)
query, key, value = map(lambda x: x.reshape(x.shape[0], -1), (query, key, value))
mixed_qkv = torch.cat((query, key, value), dim=-1)
return mixed_qkv, z, b, a
class TestQwen3(CustomTestCase):
def test_fused_qkvzba_split_reshape_cat(self):
mixed_qkvz = torch.rand(1024, 12288, dtype=torch.bfloat16)
@@ -82,6 +110,40 @@ class TestQwen3(CustomTestCase):
torch.testing.assert_close(b, b_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(a, a_ref, atol=atol, rtol=rtol)
def test_fused_qkvzba_split_reshape_cat_contiguous(self):
mixed_qkvz = torch.rand(1, 12288, dtype=torch.bfloat16)
mixed_ba = torch.rand(1, 64, dtype=torch.bfloat16)
head_k_dim = 128
head_v_dim = 128
num_v_heads = 32
num_k_heads = 16
attn_tp_size = 1
key_dim = head_k_dim * num_k_heads
value_dim = head_v_dim * num_v_heads
mixed_qkv_ref, z_ref, b_ref, a_ref = (
fix_query_key_value_ordering_reshape_cat_contiguous(
mixed_qkvz,
mixed_ba,
key_dim,
value_dim,
num_v_heads,
head_v_dim,
attn_tp_size,
)
)
num_heads_qk = num_k_heads // attn_tp_size
num_heads_v = num_v_heads // attn_tp_size
mixed_qkv, z, b, a = (
torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_contiguous_cpu(
mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_k_dim, head_v_dim
)
)
atol = rtol = precision[mixed_qkv.dtype]
torch.testing.assert_close(mixed_qkv, mixed_qkv_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(z, z_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(b, b_ref, atol=atol, rtol=rtol)
torch.testing.assert_close(a, a_ref, atol=atol, rtol=rtol)
if __name__ == "__main__":
unittest.main()