[CPU] Add support for Qwen3-vl and Qwen3-omni (#12662)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>

python/sglang/srt/layers/amx_utils.py

@@ -1,6 +1,7 @@
 import logging
 
 import torch
+import transformers
 
 from sglang.srt.utils import cpu_has_amx_support
 
@@ -21,17 +22,60 @@ class CPUQuantAlgo(IntEnum):
     GPTQ = 1
 
 
+def fast_preprocess_cpu(
+    self,
+    images: list["torch.Tensor"],
+    do_resize: bool,
+    size,
+    interpolation,
+    do_rescale: bool,
+    rescale_factor: float,
+    do_normalize: bool,
+    image_mean,
+    image_std,
+    patch_size: int,
+    temporal_patch_size: int,
+    merge_size: int,
+    disable_grouping,
+    return_tensors,
+    **kwargs,
+):
+    pixel_values, image_grid_thw = torch.ops.sgl_kernel.image_preprocess_cpu(
+        images,
+        True,
+        do_resize,
+        size["shortest_edge"],
+        size["longest_edge"],
+        "bicubic",
+        do_rescale,
+        rescale_factor,
+        do_normalize,
+        image_mean,
+        image_std,
+        patch_size,
+        temporal_patch_size,
+        merge_size,
+        True,
+        torch.bfloat16,
+    )
+    return transformers.image_processing_base.BatchFeature(
+        data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw},
+        tensor_type=return_tensors,
+    )
+
+
 def amx_process_weight_after_loading(weight, is_conv=False):
     if weight.device != torch.device("cpu"):
         return weight
     if not cpu_has_amx_support():
         return weight
     if is_conv:
+        if weight.dim() == 5:
+            return torch.ops.sgl_kernel.conv3d_embed_weight_pack(weight)
         return torch.ops.sgl_kernel.causal_conv1d_weight_pack(
             weight.view(-1, weight.size(-1))
         )
-    else:
-        return torch.ops.sgl_kernel.convert_weight_packed(weight)
+    return torch.ops.sgl_kernel.convert_weight_packed(weight)
 
 
 # TODO: currently gemm kernel has the below requirements:
@@ -58,7 +102,7 @@ def dtype_is_supported(weight):
 
 
 def is_dim_conv_weight(weight):
-    return weight.dim() == 3 and weight.size(1) == 1
+    return (weight.dim() == 3 and weight.size(1) == 1) or weight.dim() == 5
 
 
 def _init_amx_conv_state(conv_state):
@@ -120,7 +164,7 @@ def _amx_process_weight_after_loading(
             )
             packed_weight.__dict__ = weight_tensor.__dict__
             setattr(module, weight_name, packed_weight)
-            if is_conv_weight:
+            if is_conv_weight and weight_tensor.dim() != 5:
                 # need to use inplace copy for conv weight amx packing,
                 # as its usage in radix_linear_attention will use the original conv weight.
                 weight_tensor = weight_tensor.view(-1, weight_tensor.size(-1))
@@ -159,7 +203,12 @@ def _amx_process_weight_after_loading(
         and hasattr(module, "bias")
         and module.bias is not None
     ):
-        module.bias = torch.nn.Parameter(module.bias.data.float(), requires_grad=False)
+        if is_conv_weight and module.weight.data.dim() == 5:
+            module.bias = torch.nn.Parameter(module.bias.data, requires_grad=False)
+        else:
+            module.bias = torch.nn.Parameter(
+                module.bias.data.float(), requires_grad=False
+            )
 
 
 class PackWeightMethod:

python/sglang/srt/layers/attention/vision.py

@@ -17,9 +17,11 @@ from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size
 from sglang.srt.models.utils import apply_qk_norm
 from sglang.srt.utils import (
+    cpu_has_amx_support,
     get_bool_env_var,
     get_device_capability,
     is_blackwell_supported,
+    is_cpu,
     is_cuda,
     is_hip,
     is_musa,
@@ -32,10 +34,12 @@ from sglang.srt.utils.multi_stream_utils import (
     with_multi_stream,
 )
 
+_is_cpu = is_cpu()
 _is_cuda = is_cuda()
 _is_musa = is_musa()
 _is_npu = is_npu()
 _is_hip = is_hip()
+_is_cpu_amx_available = cpu_has_amx_support()
 _is_xpu = is_xpu()
 
 if _is_cuda:
@@ -45,6 +49,9 @@ if _is_cuda:
         flash_attn_varlen_func,
     )
 
+if _is_cpu and _is_cpu_amx_available:
+    flash_attn_varlen_func = torch.ops.sgl_kernel.flash_attn_varlen_func
+
 if _is_musa:
     from flash_attn_interface import flash_attn_varlen_func
 
@@ -730,6 +737,60 @@ class VisionAscendAttention(nn.Module):
         return output
 
 
+class VisionAMXAttention(nn.Module):
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        if not _is_cpu or not _is_cpu_amx_available:
+            raise Exception(
+                "VisionAMXAttention is only available for cpu with amx support"
+            )
+        super().__init__()
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens: torch.Tensor | SingletonCache | None,
+        bsz: int,
+        seq_len: int,
+        **kwargs,
+    ) -> torch.Tensor:
+        r"""
+        Args:
+            cu_seqlens: [b]
+        Returns:
+             [b * s, h, head_size]
+        """
+        if cu_seqlens is None:
+            cu_seqlens = _get_cu_seqlens_for_shape(bsz, seq_len, device=q.device)
+        elif isinstance(cu_seqlens, SingletonCache):
+            if cu_seqlens.empty():
+                cu_seqlens.set_data(
+                    _get_cu_seqlens_for_shape(bsz, seq_len, device=q.device)
+                )
+            cu_seqlens = cu_seqlens.get_data()
+
+        cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
+        seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+        max_seqlen = seq_lens.max().item()
+
+        output = flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=max_seqlen,
+            max_seqlen_k=max_seqlen,
+            causal=False,
+        )
+
+        return output
+
+
 QKV_BACKEND_IMPL = {
     "triton_attn": VisionTritonAttention,
     "sdpa": VisionSdpaAttention,
@@ -738,6 +799,7 @@ QKV_BACKEND_IMPL = {
     "flashinfer_cudnn": VisionFlashInferAttention,
     "ascend_attn": VisionAscendAttention,
     "aiter_attn": VisionAiterAttention,
+    "amx_attn": VisionAMXAttention,
 }
 
 
@@ -965,6 +1027,8 @@ class VisionAttention(nn.Module):
                 backend = "aiter_attn"
             else:
                 backend = "triton_attn"
+        elif _is_cpu and _is_cpu_amx_available:
+            backend = "amx_attn"
         elif _is_xpu:
             backend = "triton_attn"
         else:
@@ -1068,11 +1132,9 @@ class VisionAttention(nn.Module):
             q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
 
             # [b, s, embed_dim] --> [b * s, head, head_size]
-            q = q.reshape(bsz * s, head, -1).contiguous()
-            k = k.reshape(bsz * s, kv_head, -1).contiguous()
-            v = v.reshape(bsz * s, kv_head, -1).contiguous()
-            if self.qk_normalization_by_head_size:
-                q, k = self._apply_qk_norm_head_size(q, k)
+            q = q.reshape(bsz * s, head, -1)
+            k = k.reshape(bsz * s, kv_head, -1)
+            v = v.reshape(bsz * s, kv_head, -1)
         else:
             # [b, s, embed_dim] --> [s, b, embed_dim]
             x = rearrange(x, "b s ... -> s b ...")
@@ -1090,12 +1152,14 @@ class VisionAttention(nn.Module):
             q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
 
             # [s, b, head, head_size] --> [b, s, head, head_size]
-            q, k, v = [
-                rearrange(x, "s b ... -> b s ...").contiguous() for x in (q, k, v)
-            ]
+            q, k, v = [rearrange(x, "s b ... -> b s ...") for x in (q, k, v)]
 
-            if self.qk_normalization_by_head_size:
-                q, k = self._apply_qk_norm_head_size(q, k)
+        if not (_is_cpu and _is_cpu_amx_available):
+            q = q.contiguous()
+            k = k.contiguous()
+            v = v.contiguous()
+        if self.qk_normalization_by_head_size:
+            q, k = self._apply_qk_norm_head_size(q, k)
 
         cos = None
         sin = None

python/sglang/srt/layers/conv.py

@@ -14,7 +14,14 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from sglang.srt.layers.amx_utils import PackWeightMethod
 from sglang.srt.layers.utils.multi_platform import MultiPlatformOp
+from sglang.srt.utils import cpu_has_amx_support, is_cpu, use_intel_amx_backend
+
+_is_cpu = is_cpu()
+_is_cpu_amx_available = cpu_has_amx_support()
+if _is_cpu and _is_cpu_amx_available:
+    conv3d_embed = torch.ops.sgl_kernel.conv3d_embed_cpu
 
 _VALID_PADDING_STRINGS = {"same", "valid"}
 _VALID_PADDING_MODES = {"zeros", "reflect", "replicate", "circular"}
@@ -251,6 +258,8 @@ class Conv3dLayer(MultiPlatformOp):
         else:
             self.register_parameter("bias", None)
 
+        if _is_cpu and _is_cpu_amx_available and self.bias is not None:
+            self.quant_method = PackWeightMethod(weight_names=["weight"])
         self._reset_parameters()
 
     def _reset_parameters(self):
@@ -289,6 +298,16 @@ class Conv3dLayer(MultiPlatformOp):
             self.groups,
         )
 
+    def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
+        if use_intel_amx_backend(self):
+            return conv3d_embed(
+                x,
+                self.weight,
+                self.bias,
+                is_vnni=True,
+            )
+        return self.forward_native(x)
+
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         if self.enable_linear:
             return self._forward_mulmat(x)

python/sglang/srt/layers/logits_processor.py

@@ -56,11 +56,17 @@ from sglang.srt.model_executor.forward_batch_info import (
     ForwardMode,
 )
 from sglang.srt.server_args import get_global_server_args
-from sglang.srt.utils.common import is_npu, use_intel_amx_backend
+from sglang.srt.utils.common import (
+    is_cpu,
+    is_npu,
+    is_pin_memory_available,
+    use_intel_amx_backend,
+)
 
 logger = logging.getLogger(__name__)
 
 _is_npu = is_npu()
+_is_cpu = is_cpu()
 
 
 @dataclasses.dataclass
@@ -539,10 +545,14 @@ class LogitsProcessor(nn.Module):
             # Build the index tensors via pinned host memory + non-blocking H2D
             # so the small copy doesn't drain the stream.
             sample_indices = torch.tensor(
-                sample_indices, dtype=torch.int64, pin_memory=True
+                sample_indices,
+                dtype=torch.int64,
+                pin_memory=is_pin_memory_available(),
             ).to(pruned_states.device, non_blocking=True)
             input_logprob_indices = torch.tensor(
-                input_logprob_indices, dtype=torch.int64, pin_memory=True
+                input_logprob_indices,
+                dtype=torch.int64,
+                pin_memory=is_pin_memory_available(),
             ).to(pruned_states.device, non_blocking=True)
 
         return (
@@ -613,7 +623,7 @@ class LogitsProcessor(nn.Module):
         pruned_lens = torch.tensor(
             logits_metadata.extend_logprob_pruned_lens_cpu,
             dtype=torch.int64,
-            pin_memory=True,
+            pin_memory=is_pin_memory_available(),
         ).to(device, non_blocking=True)
         if logits_metadata.temp_scaled_logprobs:
             logits_metadata.temperature = torch.repeat_interleave(

python/sglang/srt/models/qwen3_omni_moe.py

@@ -31,8 +31,15 @@ from sglang.srt.configs.qwen3_omni import (
     Qwen3OmniMoeVisionEncoderConfig,
 )
 from sglang.srt.configs.qwen3_vl import Qwen3VLMoeConfig
+from sglang.srt.distributed import (
+    get_tensor_model_parallel_world_size,
+)
 from sglang.srt.layers.attention.vision import VisionAttention
-from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
+from sglang.srt.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.managers.schedule_batch import MultimodalDataItem
@@ -43,7 +50,26 @@ from sglang.srt.models.qwen3_vl_moe import (
     Qwen3VLMoeForConditionalGeneration,
     load_fused_expert_weights,
 )
-from sglang.srt.utils import add_prefix, is_npu, logger
+from sglang.srt.utils import add_prefix, is_cpu, is_npu, logger
+
+_is_cpu = is_cpu()
+
+
+def get_head_dim_and_projection_size(
+    embed_dim: int,
+    num_heads: int,
+    original_num_heads: Optional[int] = None,
+) -> Tuple[Optional[int], int]:
+    if (not _is_cpu) or original_num_heads is None:
+        return None, embed_dim
+
+    # On CPU, TP may pad num_heads (e.g. for tp=3/6). In that case we keep the
+    # original per-head width (from original_num_heads) and recompute projection_size
+    # with padded num_heads, so attention tensor shapes stay TP-friendly while
+    # preserving checkpoint semantics.
+    head_dim = embed_dim // original_num_heads
+    projection_size = num_heads * head_dim
+    return head_dim, projection_size
 
 
 class Qwen3OmniMoeAudioEncoderLayer(nn.Module):
@@ -56,10 +82,18 @@ class Qwen3OmniMoeAudioEncoderLayer(nn.Module):
         super().__init__()
         embed_dim = config.d_model
         self.embed_dim = config.d_model
+        head_dim, projection_size = get_head_dim_and_projection_size(
+            embed_dim=embed_dim,
+            num_heads=config.encoder_attention_heads,
+            original_num_heads=getattr(
+                config, "original_encoder_attention_heads", None
+            ),
+        )
         self.self_attn = VisionAttention(
             embed_dim=embed_dim,
             num_heads=config.encoder_attention_heads,
-            projection_size=embed_dim,
+            head_dim=head_dim,
+            projection_size=projection_size,
             use_qkv_parallel=True,
             proj_bias=True,
             flatten_batch=True,
@@ -70,15 +104,21 @@ class Qwen3OmniMoeAudioEncoderLayer(nn.Module):
         self.dropout = config.dropout
         self.activation_fn = ACT2FN[config.activation_function]
         self.activation_dropout = config.activation_dropout
-        self.fc1 = ColumnParallelLinear(
+        tp_size = get_tensor_model_parallel_world_size()
+        use_replicated = config.encoder_ffn_dim % tp_size != 0
+        fc1_cls = ReplicatedLinear if use_replicated else ColumnParallelLinear
+        fc2_cls = ReplicatedLinear if use_replicated else RowParallelLinear
+        self.fc1 = fc1_cls(
             self.embed_dim,
             config.encoder_ffn_dim,
+            quant_config=quant_config,
             bias=True,
             prefix=f"{prefix}.fc1",
         )
-        self.fc2 = RowParallelLinear(
+        self.fc2 = fc2_cls(
             config.encoder_ffn_dim,
             self.embed_dim,
+            quant_config=quant_config,
             bias=True,
             prefix=f"{prefix}.fc2",
         )
@@ -162,7 +202,7 @@ def _get_feat_extract_output_lengths(input_lengths):
 class Qwen3OmniMoeAudioEncoder(PreTrainedModel):
     config: Qwen3OmniMoeAudioEncoderConfig
 
-    def __init__(self, config: Qwen3OmniMoeAudioEncoderConfig):
+    def __init__(self, config: Qwen3OmniMoeAudioEncoderConfig, quant_config=None):
         super().__init__(config)
         self.dropout = config.dropout
 
@@ -200,10 +240,19 @@ class Qwen3OmniMoeAudioEncoder(PreTrainedModel):
         conv_out_dim = config.downsample_hidden_size * (
             (((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2
         )
-        self.conv_out = nn.Linear(conv_out_dim, config.d_model, bias=False)
-        self.proj1 = nn.Linear(config.d_model, config.d_model)
+        self.conv_out = ReplicatedLinear(
+            conv_out_dim,
+            config.d_model,
+            bias=False,
+            quant_config=quant_config,
+        )
+        self.proj1 = ReplicatedLinear(
+            config.d_model, config.d_model, quant_config=quant_config
+        )
         self.act = ACT2FN[config.activation_function]
-        self.proj2 = nn.Linear(config.d_model, config.output_dim)
+        self.proj2 = ReplicatedLinear(
+            config.d_model, config.output_dim, quant_config=quant_config
+        )
         self.n_window_infer = self.config.n_window_infer
         self.conv_chunksize = self.config.conv_chunksize
 
@@ -277,7 +326,7 @@ class Qwen3OmniMoeAudioEncoder(PreTrainedModel):
         b, c, f, t = padded_embed.size()
         padded_embed = self.conv_out(
             padded_embed.permute(0, 3, 1, 2).contiguous().view(b, t, c * f)
-        )
+        )[0]
 
         positional_embedding = (
             self.positional_embedding.positional_embedding[: padded_embed.shape[1], :]
@@ -313,9 +362,9 @@ class Qwen3OmniMoeAudioEncoder(PreTrainedModel):
             hidden_states = layer_outputs[0]
 
         hidden_states = self.ln_post(hidden_states)
-        hidden_states = self.proj1(hidden_states)
+        hidden_states = self.proj1(hidden_states)[0]
         hidden_states = self.act(hidden_states)
-        hidden_states = self.proj2(hidden_states)
+        hidden_states = self.proj2(hidden_states)[0]
         return BaseModelOutput(last_hidden_state=hidden_states)
 
     # Ignore copy
@@ -447,7 +496,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(Qwen3VLMoeForConditionalGenera
         super().__init__(
             config, quant_config, prefix, language_model_cls=Qwen3MoeLLMModel
         )
-        self.audio_tower = Qwen3OmniMoeAudioEncoder(config.audio_config)
+        self.audio_tower = Qwen3OmniMoeAudioEncoder(config.audio_config, quant_config)
         self.visual = Qwen3OmniMoeVisionEncoder(
             config.vision_config,
             quant_config=quant_config,

python/sglang/srt/models/qwen3_vl.py

@@ -348,8 +348,13 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
         else:
             self.pos_embed = PPMissingLayer()
 
-        norm_layer = partial(nn.LayerNorm, eps=norm_eps)
-        if is_cpu() and hasattr(vision_config, "original_num_heads"):
+        if _is_cpu and _is_cpu_amx_available:
+            from sglang.srt.layers.layernorm import LayerNorm
+
+            norm_layer = partial(LayerNorm, eps=norm_eps, dtype=self.dtype)
+        else:
+            norm_layer = partial(nn.LayerNorm, eps=norm_eps)
+        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

python/sglang/srt/multimodal/processors/qwen_vl.py

@@ -34,6 +34,7 @@ from sglang.srt.multimodal.processors.base_processor import (
 from sglang.srt.multimodal.processors.base_processor import (
     MultimodalSpecialTokens,
 )
+from sglang.srt.utils import cpu_has_amx_support, is_cpu
 from sglang.srt.utils.video_decoder import VideoDecoderWrapper
 from sglang.utils import logger
 
@@ -59,6 +60,23 @@ FPS_MIN_FRAMES = 4
 FPS_MAX_FRAMES = 768
 
 
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_cpu = is_cpu()
+if _is_cpu and _is_cpu_amx_available:
+    try:
+        import transformers
+
+        from sglang.srt.layers.amx_utils import fast_preprocess_cpu
+
+        transformers.models.qwen2_vl.image_processing_qwen2_vl_fast.Qwen2VLImageProcessorFast._preprocess = (
+            fast_preprocess_cpu
+        )
+    except Exception as e:
+        logger.warning(
+            f"Failed to hack Qwen2VLImageProcessorFast with AMX optimization: {e}"
+        )
+
+
 def smart_resize(
     height: int,
     width: int,

python/sglang/srt/server_args.py

@@ -5442,6 +5442,7 @@ class ServerArgs:
                 "ascend_attn",
                 "aiter_attn",
                 "flashinfer_cudnn",
+                "amx_attn",
             ],
             default=ServerArgs.mm_attention_backend,
             help="Set multimodal attention backend.",