[Diffusion][CPU] Adding AMX optimizations for CPU platform (#28527)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

python/sglang/multimodal_gen/runtime/layers/attention/backends/amx_attn.py

@@ -0,0 +1,70 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (  # FlashAttentionMetadata,
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+flash_attn_varlen_func = torch.ops.sgl_kernel.flash_attn_varlen_func
+
+
+class AMXAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_enum() -> AttentionBackendEnum:
+        return AttentionBackendEnum.AMX_ATTN
+
+    @staticmethod
+    def get_impl_cls() -> type["AMXATTNImpl"]:
+        return AMXATTNImpl
+
+
+class AMXATTNImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        max_seqlen_q = query.shape[1]
+        max_seqlen_k = key.shape[1]
+        return flash_attn_varlen_func(
+            query[0],
+            key[0],
+            value[0],
+            torch.tensor([0, max_seqlen_q]).to(torch.int),
+            torch.tensor([0, max_seqlen_k]).to(torch.int),
+            max_seqlen_q,
+            max_seqlen_k,
+            self.causal,
+            self.softmax_scale,
+        ).unsqueeze(0)

python/sglang/multimodal_gen/runtime/layers/linear.py

@@ -38,7 +38,15 @@ from sglang.multimodal_gen.runtime.models.parameter import (
 from sglang.multimodal_gen.runtime.models.utils import set_weight_attrs
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
+from sglang.srt.utils import (
+    cpu_has_amx_support,
+    is_cpu,
+    use_intel_amx_backend,
+)
 
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_cpu = is_cpu()
 logger = init_logger(__name__)
 
 IS_AMP_SUPPORTED = current_platform.is_amp_supported()
@@ -152,9 +160,26 @@ class UnquantizedLinearMethod(LinearMethodBase):
         layer.register_parameter("weight", weight)
         set_weight_attrs(weight, extra_weight_attrs)
 
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        if _is_cpu and _is_cpu_amx_available:
+            _amx_process_weight_after_loading(layer, ["weight"])
+
     def apply(
         self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None
     ) -> torch.Tensor:
+        if use_intel_amx_backend(layer):
+            x_shapes = x.shape
+            if len(x_shapes) == 3:
+                x = x.view(-1, x.shape[-1])
+            output = torch.ops.sgl_kernel.weight_packed_linear(
+                x.to(layer.weight.dtype),
+                layer.weight,
+                bias,
+                True,  # is_vnni
+            )
+            if len(x_shapes) == 3:
+                output = output.view(x_shapes[0], x_shapes[1], -1)
+            return output
         output = (
             F.linear(x, layer.weight, bias)
             if IS_AMP_SUPPORTED or bias is None

python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py

@@ -53,6 +53,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.runtime.utils.precision import precision_to_dtype
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 from sglang.srt.environ import envs
+from sglang.srt.model_loader.loader import device_loading_context
 
 logger = init_logger(__name__)
 
@@ -424,7 +425,16 @@ class TextEncoderLoader(ComponentLoader):
                     to_cpu=should_offload,
                 )
             )
-
+            for _, module in model.named_modules():
+                quant_method = getattr(module, "quant_method", None)
+                if quant_method is not None:
+                    # When quant methods need to process weights after loading
+                    # (for repacking, quantizing, etc), they expect parameters
+                    # to be on the global target device. This scope is for the
+                    # case where cpu offloading is used, where we will move the
+                    # parameters onto device for processing and back off after.
+                    with device_loading_context(module, local_torch_device):
+                        quant_method.process_weights_after_loading(module)
             if should_offload:
                 # Disable FSDP for MPS as it's not compatible
                 if current_platform.is_mps():

python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py

@@ -72,7 +72,11 @@ def _should_use_channels_last_3d(
     if component_name not in (
         "vae",
         "video_vae",
-    ) or not (current_platform.is_cuda() or current_platform.is_rocm()):
+    ) or not (
+        current_platform.is_cuda()
+        or current_platform.is_rocm()
+        or current_platform.is_cpu()
+    ):
         return False
 
     override = os.getenv(VAE_CHANNELS_LAST_3D_ENV)

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -24,6 +24,7 @@ from torch.distributed.fsdp import (
 from torch.nn.modules.module import _IncompatibleKeys
 
 from sglang.multimodal_gen.configs.models.fsdp import is_module_list_entry_in
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod
 from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
     attach_bitsandbytes_4bit_quant_states,
@@ -42,6 +43,7 @@ from sglang.multimodal_gen.runtime.loader.weight_utils import (
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import set_mixed_precision_policy
+from sglang.srt.model_loader.loader import device_loading_context
 from sglang.srt.utils import is_npu
 
 _is_npu = is_npu()
@@ -336,6 +338,17 @@ def maybe_load_fsdp_model(
         # Avoid unintended computation graph accumulation during inference
         if isinstance(p, torch.nn.Parameter):
             p.requires_grad = False
+    local_torch_device = get_local_torch_device()
+    for _, module in model.named_modules():
+        quant_method = getattr(module, "quant_method", None)
+        if quant_method is not None:
+            # When quant methods need to process weights after loading
+            # (for repacking, quantizing, etc), they expect parameters
+            # to be on the global target device. This scope is for the
+            # case where cpu offloading is used, where we will move the
+            # parameters onto device for processing and back off after.
+            with device_loading_context(module, local_torch_device):
+                quant_method.process_weights_after_loading(module)
 
     # 4. deferred cpu offload
     if defer_cpu_offload:

python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py

@@ -64,7 +64,9 @@ first_chunk = contextvars.ContextVar("first_chunk", default=None)
 
 def _channels_last_3d_supported_by_platform() -> bool:
     return hasattr(torch, "channels_last_3d") and (
-        current_platform.is_cuda() or current_platform.is_rocm()
+        current_platform.is_cuda()
+        or current_platform.is_rocm()
+        or current_platform.is_cpu()
     )
 
 

python/sglang/multimodal_gen/runtime/platforms/cpu.py

@@ -17,7 +17,13 @@ from sglang.multimodal_gen.runtime.platforms.interface import (
     PlatformEnum,
 )
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.srt.utils import (
+    cpu_has_amx_support,
+    is_cpu,
+)
 
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_cpu = is_cpu()
 logger = init_logger(__name__)
 
 
@@ -102,12 +108,18 @@ class CpuPlatform(Platform):
         head_size: int,
         dtype: torch.dtype,
     ) -> str:
-        if selected_backend not in (None, AttentionBackendEnum.TORCH_SDPA):
+        if selected_backend not in (
+            None,
+            AttentionBackendEnum.TORCH_SDPA,
+            AttentionBackendEnum.AMX_ATTN,
+        ):
             logger.warning(
-                "%s is not supported on CPU; falling back to Torch SDPA.",
+                "%s is not supported on CPU; falling back to auto selection SDPA or AMX_ATTN",
                 selected_backend,
             )
-
+        if _is_cpu and _is_cpu_amx_available:
+            logger.info("Using AMX Attention backend for CPU.")
+            return "sglang.multimodal_gen.runtime.layers.attention.backends.amx_attn.AMXAttentionBackend"
         logger.info("Using Torch SDPA backend for CPU.")
         return (
             "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -42,6 +42,7 @@ class AttentionBackendEnum(enum.Enum):
     BLOCK_SPARSE_ATTN = enum.auto()
     RAIN_FUSION_ATTN = enum.auto()
     NO_ATTENTION = enum.auto()
+    AMX_ATTN = enum.auto()
 
     def __str__(self):
         return self.name.lower()

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

@@ -775,6 +775,7 @@ class VisionAMXAttention(nn.Module):
         cu_seqlens: torch.Tensor | SingletonCache | None,
         bsz: int,
         seq_len: int,
+        softmax_scale: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
         r"""
@@ -805,6 +806,7 @@ class VisionAMXAttention(nn.Module):
             max_seqlen_q=max_seqlen,
             max_seqlen_k=max_seqlen,
             causal=False,
+            sm_scale=softmax_scale,
         )
 
         return output

sgl-kernel/csrc/cpu/flash_attn.cpp

@@ -443,7 +443,8 @@ at::Tensor flash_attn_varlen_func(
     const at::Tensor& cu_seqlens_k,
     int64_t max_seqlen_q,
     int64_t max_seqlen_k,
-    bool causal) {
+    bool causal,
+    const std::optional<double>& sm_scale) {
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(q);
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(k);
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(v);
@@ -480,7 +481,7 @@ at::Tensor flash_attn_varlen_func(
   TORCH_CHECK(head_size_v % 2 == 0, "invalid head_size_v ", head_size_v);
 
   // softmax scale
-  double sm_scale = 1.0 / std::sqrt(static_cast<double>(head_size));
+  double _sm_scale = sm_scale.has_value() ? sm_scale.value() : 1.0 / std::sqrt(static_cast<double>(head_size));
 
   // check whether the batch has variant lengths
   const bool is_varlen =
@@ -522,7 +523,7 @@ at::Tensor flash_attn_varlen_func(
           k_strideH,
           v_strideN,
           v_strideH,
-          sm_scale,
+          _sm_scale,
           sz,
           causal);
     } else {
@@ -545,7 +546,7 @@ at::Tensor flash_attn_varlen_func(
           k_strideH,
           v_strideN,
           v_strideH,
-          sm_scale,
+          _sm_scale,
           sz,
           causal);
     }

sgl-kernel/csrc/cpu/torch_extension_cpu.cpp

@@ -153,7 +153,8 @@ at::Tensor flash_attn_varlen_func(
     const at::Tensor& cu_seqlens_k,
     int64_t max_seqlen_q,
     int64_t max_seqlen_k,
-    bool causal);
+    bool causal,
+    const std::optional<double>& sm_scale);
 
 // linear attention
 std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
@@ -533,7 +534,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   // flash attn
   m.def(
       "flash_attn_varlen_func(Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q, Tensor cu_seqlens_k, "
-      "int max_seqlen_q, int max_seqlen_k, bool causal) -> Tensor");
+      "int max_seqlen_q, int max_seqlen_k, bool causal, float? sm_scale) -> Tensor");
   m.impl("flash_attn_varlen_func", torch::kCPU, &flash_attn_varlen_func);
 
   // linear attn

test/registered/cpu/test_flash_attn.py

@@ -23,6 +23,7 @@ def flash_attn_varlen_ref(
     cu_seqlens_k,
     is_causal,
     enable_gqa,
+    softmax_scale=None,
 ):
     cu_q = cu_seqlens_q.tolist()
     cu_k = cu_seqlens_k.tolist()
@@ -43,6 +44,7 @@ def flash_attn_varlen_ref(
             v[:, :, start_k:end_k, :],
             is_causal=is_causal,
             enable_gqa=enable_gqa,
+            scale=softmax_scale,
         )
 
     # [1, H, T, D] -> [T, H, D]
@@ -58,6 +60,7 @@ def flash_attn_non_varlen_ref(
     cu_seqlens_k,
     is_causal,
     enable_gqa,
+    softmax_scale=None,
 ):
     cu_q = cu_seqlens_q.tolist()
     cu_k = cu_seqlens_k.tolist()
@@ -75,6 +78,7 @@ def flash_attn_non_varlen_ref(
         v,
         is_causal=is_causal,
         enable_gqa=enable_gqa,
+        scale=softmax_scale,
     )
     # [B, H, T, D] -> [B * T, H, D]
     return out.transpose(1, 2).reshape(batch * T, H, D)
@@ -91,6 +95,7 @@ class TestFlashAttn(CustomTestCase):
         head_dim=[32, 48],  # test when D is not 32x
         head_dim_v=[32],
         is_causal=[True, False],
+        softmax_scale=[None, 0.2],
     )
     def test_flash_attn_varlen(
         self,
@@ -102,6 +107,7 @@ class TestFlashAttn(CustomTestCase):
         head_dim,
         head_dim_v,
         is_causal,
+        softmax_scale,
     ):
         dtype = torch.bfloat16
 
@@ -127,6 +133,7 @@ class TestFlashAttn(CustomTestCase):
             cu_seqlens_k,
             is_causal=is_causal,
             enable_gqa=num_heads != num_heads_kv,
+            softmax_scale=softmax_scale,
         )
 
         out = flash_attn_varlen_func(
@@ -138,6 +145,7 @@ class TestFlashAttn(CustomTestCase):
             seqlens_q.max().item(),
             seqlens_k.max().item(),
             is_causal,
+            softmax_scale,
         )
 
         atol = rtol = precision[dtype]
@@ -153,6 +161,7 @@ class TestFlashAttn(CustomTestCase):
         head_dim=[32],
         head_dim_v=[32],
         is_causal=[False],
+        softmax_scale=[None, 0.2],
     )
     def test_flash_attn_large_size(
         self,
@@ -164,6 +173,7 @@ class TestFlashAttn(CustomTestCase):
         head_dim,
         head_dim_v,
         is_causal,
+        softmax_scale,
     ):
         dtype = torch.bfloat16
 
@@ -190,6 +200,7 @@ class TestFlashAttn(CustomTestCase):
             cu_seqlens_k,
             is_causal=is_causal,
             enable_gqa=num_heads != num_heads_kv,
+            softmax_scale=softmax_scale,
         )
 
         out = flash_attn_varlen_func(
@@ -201,6 +212,7 @@ class TestFlashAttn(CustomTestCase):
             seqlens_q.max().item(),
             seqlens_k.max().item(),
             is_causal,
+            softmax_scale,
         )
 
         atol = rtol = precision[dtype]
@@ -226,7 +238,7 @@ class TestFlashAttn(CustomTestCase):
             q, k, v, cu_seqlens, cu_seqlens, is_causal=True, enable_gqa=True
         )
         out = flash_attn_varlen_func(
-            q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, True
+            q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, True, None
         )
 
         atol = rtol = precision[dtype]