From 87c3171aaa5589eaf945dc4b7dce61149edfd5ea Mon Sep 17 00:00:00 2001 From: blzheng Date: Wed, 27 May 2026 08:56:09 +0800 Subject: [PATCH] [CPU] Add support for Qwen3-vl and Qwen3-omni (#12662) Co-authored-by: Ma Mingfei Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> --- python/sglang/srt/layers/amx_utils.py | 59 +++++++++++-- python/sglang/srt/layers/attention/vision.py | 84 ++++++++++++++++--- python/sglang/srt/layers/conv.py | 19 +++++ python/sglang/srt/layers/logits_processor.py | 18 +++- python/sglang/srt/models/qwen3_omni_moe.py | 75 ++++++++++++++--- python/sglang/srt/models/qwen3_vl.py | 9 +- .../srt/multimodal/processors/qwen_vl.py | 18 ++++ python/sglang/srt/server_args.py | 1 + 8 files changed, 249 insertions(+), 34 deletions(-) diff --git a/python/sglang/srt/layers/amx_utils.py b/python/sglang/srt/layers/amx_utils.py index d69e71060..b2be65058 100755 --- a/python/sglang/srt/layers/amx_utils.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 2b67e5edb..2718ccfe4 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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 diff --git a/python/sglang/srt/layers/conv.py b/python/sglang/srt/layers/conv.py index d2885f0ef..6c72ae107 100644 --- a/python/sglang/srt/layers/conv.py +++ b/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) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index e14895e97..e7486b3ce 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/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( diff --git a/python/sglang/srt/models/qwen3_omni_moe.py b/python/sglang/srt/models/qwen3_omni_moe.py index badb6be08..fd488faae 100644 --- a/python/sglang/srt/models/qwen3_omni_moe.py +++ b/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, diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 44dddf1bf..3ffe4dde7 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/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 diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py index 99a4b12e6..b8774ebad 100644 --- a/python/sglang/srt/multimodal/processors/qwen_vl.py +++ b/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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c2d455a9f..660735bc8 100644 --- a/python/sglang/srt/server_args.py +++ b/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.",