[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
+                update_config(
-                model_config.linear_num_value_heads_cpu = linear_num_value_heads_cpu
+                    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
+                update_config(
-                model_config.linear_num_value_heads = linear_num_value_heads_cpu
+                    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 = (
+                update_config(
-                    model_config.linear_num_key_heads
+                    model_config,
+                    "linear_num_key_heads_cpu",
+                    model_config.linear_num_key_heads,
                 )
-                model_config.linear_num_value_heads_cpu = (
+                update_config(
-                    model_config.linear_num_value_heads
+                    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 = (
+    for config in [model_config.hf_config, model_config.hf_text_config]:
-        model_config.num_attention_heads
+        update_config(
-    )
+            config,
-    model_config.hf_text_config.original_num_attention_heads = (
+            "original_num_attention_heads",
-        model_config.num_attention_heads
+            model_config.num_attention_heads,
-    )
+        )
-
+        update_config(
-    model_config.hf_config.original_total_num_kv_heads = (
+            config,
-        model_config.get_total_num_kv_heads()
+            "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()
-    )
 
     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 = (
+        head_dim = resolve_head_dim(
-            model_config.hf_config.qk_head_dim
+            model_config, model_config.num_attention_heads, True
-            if hasattr(model_config.hf_config, "qk_head_dim")
-            else model_config.hf_config.head_dim
         )
+
         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
+        for config in [
-        model_config.hf_config.num_attention_heads = num_attention_heads
+            model_config,
-        model_config.hf_text_config.num_attention_heads = num_attention_heads
+            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(
+    for moe_intermediate_attr in [
-        model_config, "moe_intermediate_size", intermediate_padding_size
+        "moe_intermediate_size",
-    )
+        "intermediate_size",
-    model_config = update_intermediate_size(
+        "intermediate_size_mlp",
-        model_config, "intermediate_size", intermediate_padding_size
+        "shared_expert_intermediate_size",
-    )
+    ]:
-    model_config = update_intermediate_size(
+        model_config = update_intermediate_size(
-        model_config, "intermediate_size_mlp", intermediate_padding_size
+            model_config, moe_intermediate_attr, 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
         )
-        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)
+    multimodal_config = [
-            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",
-        model_config.hf_config.vision_config = update_intermediate_size(
+            "siglip_vision_model",
-            model_config.hf_config.vision_config,
+            "num_attention_heads",
-            "intermediate_size",
+        ],
-            intermediate_padding_size,
+        [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(
+        if _use_cpu:
-            x,
+            assert (
-            g,
+                self.activation == "silu"
-            self.weight,
+            ), "CPU rmsnorm_gated currently only supports activation silu"
-            self.bias,
+            return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(
-            self.activation,
+                x, self.weight, g, self.eps
-            residual=residual,
+            )
-            eps=self.eps,
+        else:
-            prenorm=prenorm,
+            return rms_norm_gated(
-            residual_in_fp32=residual_in_fp32,
+                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_v_heads = (
-        self.num_k_heads = config.linear_num_key_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, (
+                if _is_cpu and _is_amx_available:
-                    f"Shape mismatch: param {param.data.shape} vs "
+                    slices = tuple(slice(0, s) for s in loaded_weight.shape)
-                    f"loaded {loaded_weight.shape}"
+                    param.data.zero_()
-                )
+                    param.data[slices].copy_(loaded_weight)
-                param.data.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:
+    except Exception as e:
-        logger.info(msg)
+        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

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

@@ -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>(
+  CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
-        q.data_ptr<scalar_t>(),
+      q.scalar_type(), A_log.scalar_type(), "fused_sigmoid_gating_delta_rule_update_kernel_impl", [&] {
-        k.data_ptr<scalar_t>(),
+        fused_sigmoid_gating_delta_rule_update_kernel_impl<scalar_t, param_t>(
-        v.data_ptr<scalar_t>(),
+            q.data_ptr<scalar_t>(),
-        A_log.data_ptr<float>(),
+            k.data_ptr<scalar_t>(),
-        a.data_ptr<scalar_t>(),
+            v.data_ptr<scalar_t>(),
-        dt_bias.data_ptr<scalar_t>(),
+            A_log.data_ptr<param_t>(),
-        b.data_ptr<scalar_t>(),
+            a.data_ptr<scalar_t>(),
-        initial_state_indices.data_ptr<int32_t>(),
+            dt_bias.data_ptr<scalar_t>(),
-        initial_state_source.data_ptr<float>(),
+            b.data_ptr<scalar_t>(),
-        core_attn_out.data_ptr<scalar_t>(),
+            initial_state_indices.data_ptr<int32_t>(),
-        qk_scale_buf.data_ptr<float>(),
+            initial_state_source.data_ptr<float>(),
-        seq_len,
+            core_attn_out.data_ptr<scalar_t>(),
-        batch_size,
+            qk_scale_buf.data_ptr<float>(),
-        num_heads,
+            seq_len,
-        head_dim,
+            batch_size,
-        v_num_heads,
+            num_heads,
-        v_head_dim,
+            head_dim,
-        q_strideB,
+            v_num_heads,
-        q_strideS,
+            v_head_dim,
-        q_strideH,
+            q_strideB,
-        k_strideB,
+            q_strideS,
-        k_strideS,
+            q_strideH,
-        k_strideH,
+            k_strideB,
-        v_strideB,
+            k_strideS,
-        v_strideS,
+            k_strideH,
-        v_strideH,
+            v_strideB,
-        use_qk_l2norm_in_kernel,
+            v_strideS,
-        softplus_threshold);
+            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

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

@@ -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) {
+      if (nb_offset == 0) {
-        int32_t index = A_ids[m] / topk;
+        // 1.a load A
-        copy_stub(A + m * K, input + index * K, K);
+        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

@@ -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) {
+      if (nb_offset == 0) {
-        int32_t index = A_ids[m] / topk;
+        // 1.a load A
-        copy_stub(A + m * K, input + index * K, K);
+        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

@@ -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) {
+      if (nb_offset == 0) {
-        int32_t index = A_ids[m] / topk;
+        // 1.a load A
-        copy_stub(A + m * K, Aq_tmp + index * K, K);
+        const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
-        As[m] = As_tmp[index];
+        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

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

@@ -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)
+            for A_log_dtype in [torch.float32, torch.bfloat16]:
-            a = torch.rand(1024, dim, dtype=torch.bfloat16)
+                A_log = torch.rand(dim, dtype=A_log_dtype)
-            b = torch.rand(1024, dim, dtype=torch.bfloat16)
+                a = torch.rand(1024, dim, dtype=torch.bfloat16)
-            dt_bias = torch.rand(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, beta = torch_gdn_gating(A_log, a, b, dt_bias)
-            g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(
+                g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(
-                A_log, a, b, dt_bias
+                    A_log, a, b, dt_bias
-            )
+                )
-            atol = rtol = precision[g.dtype]
+                atol = rtol = precision[g.dtype]
-            atol2 = rtol2 = precision[beta.dtype]
+                atol2 = rtol2 = precision[beta.dtype]
-            torch.testing.assert_close(g, g_sgl, atol=atol, rtol=rtol)
+                torch.testing.assert_close(g, g_sgl, atol=atol, rtol=rtol)
-            torch.testing.assert_close(beta, beta_sgl, atol=atol2, rtol=rtol2)
+                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(
+        for A_log_dtype in [torch.float32, torch.bfloat16]:
-            query_ref.transpose(0, 1),
+            A_log = A_log.to(A_log_dtype)
-            key_ref.transpose(0, 1),
+            core_attn_out_ref, last_recurrent_state_ref = (
-            value.transpose(0, 1),
+                sigmoid_gating_delta_rule_update(
-            A_log,
+                    query_ref.transpose(0, 1),
-            a,
+                    key_ref.transpose(0, 1),
-            dt_bias,
+                    value.transpose(0, 1),
-            b,
+                    A_log,
-            initial_state=ssm_states[cache_indices],
+                    a,
-            output_final_state=True,
+                    dt_bias,
-            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+                    b,
-        )
+                    initial_state=ssm_states[cache_indices],
-        core_attn_out = torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
+                    output_final_state=True,
-            A_log=A_log,
+                    use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
-            dt_bias=dt_bias,
+                )
-            q=query,
+            )
-            k=key,
+            core_attn_out = (
-            v=value,
+                torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
-            a=a,
+                    A_log=A_log,
-            b=b,
+                    dt_bias=dt_bias,
-            initial_state_source=ssm_states,
+                    q=query,
-            initial_state_indices=cache_indices,
+                    k=key,
-            cu_seqlens=query_start_loc,
+                    v=value,
-            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+                    a=a,
-            softplus_beta=1.0,
+                    b=b,
-            softplus_threshold=20.0,
+                    initial_state_source=ssm_states,
-        )
+                    initial_state_indices=cache_indices,
-        last_recurrent_state = ssm_states[cache_indices]
+                    cu_seqlens=query_start_loc,
-        atol = rtol = precision[core_attn_out.dtype]
+                    use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
-        torch.testing.assert_close(
+                    softplus_beta=1.0,
-            core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol
+                    softplus_threshold=20.0,
-        )
+                )
-        torch.testing.assert_close(
+            )
-            last_recurrent_state, last_recurrent_state_ref, atol=atol, rtol=rtol
+            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

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