diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/amx_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/amx_attn.py deleted file mode 100644 index 0e1a61468..000000000 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/amx_attn.py +++ /dev/null @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/layers/linear.py b/python/sglang/multimodal_gen/runtime/layers/linear.py index 2e280fa4d..94affa99b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/linear.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index 9a466827b..296910b79 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/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(): diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 0992a9a2f..b97ed08f5 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index d178ebcd6..afbfc5437 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py index 1fadab5ad..1bbad1241 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py +++ b/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() ) diff --git a/python/sglang/multimodal_gen/runtime/platforms/cpu.py b/python/sglang/multimodal_gen/runtime/platforms/cpu.py index e393f07bd..e6f46b115 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cpu.py +++ b/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" diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index d6b66e7c3..ca136d39c 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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() diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 87405bb25..9b063b015 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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 diff --git a/sgl-kernel/csrc/cpu/flash_attn.cpp b/sgl-kernel/csrc/cpu/flash_attn.cpp index 95c692ad8..4b3da3314 100644 --- a/sgl-kernel/csrc/cpu/flash_attn.cpp +++ b/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& 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(head_size)); + double sm_scale = 1.0 / std::sqrt(static_cast(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); } diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index eea0fe06a..a147d1fc2 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/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& sm_scale); + bool causal); // linear attention std::tuple 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 diff --git a/test/registered/cpu/test_flash_attn.py b/test/registered/cpu/test_flash_attn.py index 26227a492..7d79adefb 100644 --- a/test/registered/cpu/test_flash_attn.py +++ b/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]