[Diffusion] Revert CPU AMX optimizations

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

@@ -38,15 +38,7 @@
 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 @@ def create_weights(
         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.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 @@ def load_model(
                     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.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.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 @@
 
 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 @@
     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 @@ def get_attn_backend_cls_str(
         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 @@ def forward(
         cu_seqlens: torch.Tensor | SingletonCache | None,
         bsz: int,
         seq_len: int,
-        softmax_scale: Optional[float] = None,
         **kwargs,
     ) -> torch.Tensor:
         r"""
@@ -806,7 +805,6 @@ def forward(
             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,
-    const std::optional<double>& sm_scale) {
+    bool causal) {
   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,
-    const std::optional<double>& sm_scale);
+    bool causal);
 
 // 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