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

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>

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

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,
+            )

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)

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
         )

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)

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

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

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
 

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)

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(

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