[Diffusion] Revert CPU AMX optimizations (#30716)

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

@@ -1,70 +0,0 @@
-# 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,15 +38,7 @@ 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()
@@ -160,26 +152,9 @@ 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,7 +53,6 @@ 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__)
 
@@ -425,16 +424,7 @@ 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,11 +72,7 @@ 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 current_platform.is_cpu()
-    ):
+    ) or not (current_platform.is_cuda() or current_platform.is_rocm()):
         return False
 
     override = os.getenv(VAE_CHANNELS_LAST_3D_ENV)

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

@@ -24,7 +24,6 @@ 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,
@@ -43,7 +42,6 @@ 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()
@@ -338,17 +336,6 @@ 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,9 +64,7 @@ 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()
-        or current_platform.is_cpu()
+        current_platform.is_cuda() or current_platform.is_rocm()
     )
 
 

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

@@ -17,13 +17,7 @@ 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__)
 
 
@@ -108,18 +102,12 @@ class CpuPlatform(Platform):
         head_size: int,
         dtype: torch.dtype,
     ) -> str:
-        if selected_backend not in (
-            None,
-            AttentionBackendEnum.TORCH_SDPA,
-            AttentionBackendEnum.AMX_ATTN,
-        ):
+        if selected_backend not in (None, AttentionBackendEnum.TORCH_SDPA):
             logger.warning(
-                "%s is not supported on CPU; falling back to auto selection SDPA or AMX_ATTN",
+                "%s is not supported on CPU; falling back to Torch SDPA.",
                 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,7 +42,6 @@ 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,7 +775,6 @@ class VisionAMXAttention(nn.Module):
         cu_seqlens: torch.Tensor | SingletonCache | None,
         bsz: int,
         seq_len: int,
-        softmax_scale: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
         r"""
@@ -806,7 +805,6 @@ class VisionAMXAttention(nn.Module):
             max_seqlen_q=max_seqlen,
             max_seqlen_k=max_seqlen,
             causal=False,
-            sm_scale=softmax_scale,
         )
 
         return output