[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 (
+    ) or not (current_platform.is_cuda() or current_platform.is_rocm()):
-        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,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()
+        current_platform.is_cuda() or current_platform.is_rocm()
-        or current_platform.is_rocm()
-        or current_platform.is_cpu()
     )
 
 

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 (
+        if selected_backend not in (None, AttentionBackendEnum.TORCH_SDPA):
-            None,
-            AttentionBackendEnum.TORCH_SDPA,
-            AttentionBackendEnum.AMX_ATTN,
-        ):
             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

sgl-kernel/csrc/cpu/flash_attn.cpp

@@ -443,8 +443,7 @@ 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);
@@ -481,7 +480,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 = sm_scale.has_value() ? sm_scale.value() : 1.0 / std::sqrt(static_cast<double>(head_size));
+  double sm_scale = 1.0 / std::sqrt(static_cast<double>(head_size));
 
   // check whether the batch has variant lengths
   const bool is_varlen =
@@ -523,7 +522,7 @@ at::Tensor flash_attn_varlen_func(
           k_strideH,
           v_strideN,
           v_strideH,
-          _sm_scale,
+          sm_scale,
           sz,
           causal);
     } else {
@@ -546,7 +545,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

@@ -235,8 +235,7 @@ 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(
@@ -677,7 +676,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, float? sm_scale) -> Tensor");
+      "int max_seqlen_q, int max_seqlen_k, bool causal) -> Tensor");
   m.impl("flash_attn_varlen_func", torch::kCPU, &flash_attn_varlen_func);
 
   // linear attn

test/registered/cpu/test_flash_attn.py

@@ -23,7 +23,6 @@ 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()
@@ -44,7 +43,6 @@ 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]
@@ -60,7 +58,6 @@ 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()
@@ -78,7 +75,6 @@ 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)
@@ -95,7 +91,6 @@ 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,
@@ -107,7 +102,6 @@ class TestFlashAttn(CustomTestCase):
         head_dim,
         head_dim_v,
         is_causal,
-        softmax_scale,
     ):
         dtype = torch.bfloat16
 
@@ -133,7 +127,6 @@ 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(
@@ -145,7 +138,6 @@ class TestFlashAttn(CustomTestCase):
             seqlens_q.max().item(),
             seqlens_k.max().item(),
             is_causal,
-            softmax_scale,
         )
 
         atol = rtol = precision[dtype]
@@ -161,7 +153,6 @@ 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,
@@ -173,7 +164,6 @@ class TestFlashAttn(CustomTestCase):
         head_dim,
         head_dim_v,
         is_causal,
-        softmax_scale,
     ):
         dtype = torch.bfloat16
 
@@ -200,7 +190,6 @@ 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(
@@ -212,7 +201,6 @@ class TestFlashAttn(CustomTestCase):
             seqlens_q.max().item(),
             seqlens_k.max().item(),
             is_causal,
-            softmax_scale,
         )
 
         atol = rtol = precision[dtype]
@@ -238,7 +226,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, None
+            q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, True
         )
 
         atol = rtol = precision[dtype]