diff --git a/python/sglang/srt/configs/update_config.py b/python/sglang/srt/configs/update_config.py index b7b3a3c3b..7ee352f7d 100644 --- a/python/sglang/srt/configs/update_config.py +++ b/python/sglang/srt/configs/update_config.py @@ -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 diff --git a/python/sglang/srt/layers/attention/fla/fused_norm_gate.py b/python/sglang/srt/layers/attention/fla/fused_norm_gate.py index 68fcbf91d..ddaded752 100644 --- a/python/sglang/srt/layers/attention/fla/fused_norm_gate.py +++ b/python/sglang/srt/layers/attention/fla/fused_norm_gate.py @@ -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, + ) diff --git a/python/sglang/srt/layers/attention/mamba/mamba.py b/python/sglang/srt/layers/attention/mamba/mamba.py index d19be6f5b..1d48809ca 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba.py +++ b/python/sglang/srt/layers/attention/mamba/mamba.py @@ -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) diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index a624ad06e..0eb6de6f8 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -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 ) diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 26980e7de..a8e7bafda 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -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) diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 6628a516b..3dc6c11dc 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -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" ) diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index dc343bfdf..28d655ec3 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -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 diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 48b7f7e1d..516a3b89c 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -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 diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index ac63cae10..d550be33e 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -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) diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index f8abf9b96..1b6c185bc 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -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( diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 0094a5266..f68721f3d 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -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): diff --git a/sgl-kernel/csrc/cpu/common.h b/sgl-kernel/csrc/cpu/common.h index 6d2fb5424..139121859 100644 --- a/sgl-kernel/csrc/cpu/common.h +++ b/sgl-kernel/csrc/cpu/common.h @@ -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") diff --git a/sgl-kernel/csrc/cpu/mamba/fla.cpp b/sgl-kernel/csrc/cpu/mamba/fla.cpp index 38f38a091..1c551519d 100644 --- a/sgl-kernel/csrc/cpu/mamba/fla.cpp +++ b/sgl-kernel/csrc/cpu/mamba/fla.cpp @@ -814,12 +814,12 @@ inline at::vec::Vectorized softplus(const at::vec::Vectorized& x, return Vec::blendv(Vec::blendv(log1pex, expx, mask_lo), x, mask_hi); } -template +template 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 +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; + using fVec = at::vec::Vectorized; + 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(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 @@ -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(k, {seq_len, batch_size, num_heads, head_dim}, q.scalar_type()); CHECK_INPUT_SHAPE_DTYPE(v, {seq_len, batch_size, v_num_heads, v_head_dim}, q.scalar_type()); - CHECK_INPUT_SHAPE_DTYPE(A_log, {v_num_heads}, at::kFloat); CHECK_INPUT_SHAPE_DTYPE(a, {batch_size, v_num_heads}, q.scalar_type()); CHECK_INPUT_SHAPE_DTYPE(dt_bias, {v_num_heads}, q.scalar_type()); CHECK_INPUT_SHAPE_DTYPE(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( - q.data_ptr(), - k.data_ptr(), - v.data_ptr(), - A_log.data_ptr(), - a.data_ptr(), - dt_bias.data_ptr(), - b.data_ptr(), - initial_state_indices.data_ptr(), - initial_state_source.data_ptr(), - core_attn_out.data_ptr(), - qk_scale_buf.data_ptr(), - 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( + q.data_ptr(), + k.data_ptr(), + v.data_ptr(), + A_log.data_ptr(), + a.data_ptr(), + dt_bias.data_ptr(), + b.data_ptr(), + initial_state_indices.data_ptr(), + initial_state_source.data_ptr(), + core_attn_out.data_ptr(), + qk_scale_buf.data_ptr(), + 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( - A_log.data_ptr(), + A_log.data_ptr(), a.data_ptr(), b.data_ptr(), dt_bias.data_ptr(), diff --git a/sgl-kernel/csrc/cpu/model/qwen3.cpp b/sgl-kernel/csrc/cpu/model/qwen3.cpp index 05ec00c43..1095d5da2 100644 --- a/sgl-kernel/csrc/cpu/model/qwen3.cpp +++ b/sgl-kernel/csrc/cpu/model/qwen3.cpp @@ -61,6 +61,41 @@ void fused_qkvzba_split_reshape_cat_impl( } }); } + +template +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 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 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 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( + mixed_qkvz.data_ptr(), + mixed_ba.data_ptr(), + mixed_qkv.data_ptr(), + z.data_ptr(), + b.data_ptr(), + a.data_ptr(), + 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); +} diff --git a/sgl-kernel/csrc/cpu/moe.cpp b/sgl-kernel/csrc/cpu/moe.cpp index 940253e7f..8202a08e0 100644 --- a/sgl-kernel/csrc/cpu/moe.cpp +++ b/sgl-kernel/csrc/cpu/moe.cpp @@ -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) { diff --git a/sgl-kernel/csrc/cpu/moe_fp8.cpp b/sgl-kernel/csrc/cpu/moe_fp8.cpp index fb476b78f..ecd4e2adb 100644 --- a/sgl-kernel/csrc/cpu/moe_fp8.cpp +++ b/sgl-kernel/csrc/cpu/moe_fp8.cpp @@ -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]; diff --git a/sgl-kernel/csrc/cpu/moe_int8.cpp b/sgl-kernel/csrc/cpu/moe_int8.cpp index 95edd4562..b147b65f6 100644 --- a/sgl-kernel/csrc/cpu/moe_int8.cpp +++ b/sgl-kernel/csrc/cpu/moe_int8.cpp @@ -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) { diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 921e5b7d1..51700a281 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -374,6 +374,15 @@ std::tuple fused_qkvzba_split_re int64_t head_qk, int64_t head_v); +// fused_qkvzba_split_reshape_cat_cpu_contiguous +std::tuple 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 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( diff --git a/test/srt/cpu/test_mamba.py b/test/srt/cpu/test_mamba.py index f64c3713a..a6ab8e517 100644 --- a/test/srt/cpu/test_mamba.py +++ b/test/srt/cpu/test_mamba.py @@ -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__": diff --git a/test/srt/cpu/test_qwen3.py b/test/srt/cpu/test_qwen3.py index 0ea95be52..ab501451f 100644 --- a/test/srt/cpu/test_qwen3.py +++ b/test/srt/cpu/test_qwen3.py @@ -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()