From 177c048c68b38930a675fac1e0b860b4c668b95f Mon Sep 17 00:00:00 2001 From: jianan-gu Date: Thu, 9 Jul 2026 10:26:22 +0800 Subject: [PATCH] [Diffusion][CPU] Adding AMX optimizations for CPU platform (#28527) Co-authored-by: Ma Mingfei --- .../layers/attention/backends/amx_attn.py | 70 +++++++++++++++++++ .../multimodal_gen/runtime/layers/linear.py | 25 +++++++ .../component_loaders/text_encoder_loader.py | 12 +++- .../loader/component_loaders/vae_loader.py | 6 +- .../runtime/loader/fsdp_load.py | 13 ++++ .../runtime/models/vaes/wanvae.py | 4 +- .../multimodal_gen/runtime/platforms/cpu.py | 18 ++++- .../runtime/platforms/interface.py | 1 + python/sglang/srt/layers/attention/vision.py | 2 + sgl-kernel/csrc/cpu/flash_attn.cpp | 9 +-- sgl-kernel/csrc/cpu/torch_extension_cpu.cpp | 5 +- test/registered/cpu/test_flash_attn.py | 14 +++- 12 files changed, 166 insertions(+), 13 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/layers/attention/backends/amx_attn.py 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 new file mode 100644 index 000000000..0e1a61468 --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/layers/linear.py b/python/sglang/multimodal_gen/runtime/layers/linear.py index 94affa99b..2e280fa4d 100644 --- a/python/sglang/multimodal_gen/runtime/layers/linear.py +++ b/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 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 296910b79..9a466827b 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,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(): 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 b97ed08f5..0992a9a2f 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,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) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index afbfc5437..d178ebcd6 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py b/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py index 1bbad1241..1fadab5ad 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py +++ b/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() ) diff --git a/python/sglang/multimodal_gen/runtime/platforms/cpu.py b/python/sglang/multimodal_gen/runtime/platforms/cpu.py index e6f46b115..e393f07bd 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cpu.py +++ b/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" diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index ca136d39c..d6b66e7c3 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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() diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 9ffd181ab..12a5199d9 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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 diff --git a/sgl-kernel/csrc/cpu/flash_attn.cpp b/sgl-kernel/csrc/cpu/flash_attn.cpp index 4b3da3314..95c692ad8 100644 --- a/sgl-kernel/csrc/cpu/flash_attn.cpp +++ b/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& 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(head_size)); + double _sm_scale = sm_scale.has_value() ? sm_scale.value() : 1.0 / std::sqrt(static_cast(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); } diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 6aedaf536..80a910e7f 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/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& sm_scale); // linear attention std::tuple 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 diff --git a/test/registered/cpu/test_flash_attn.py b/test/registered/cpu/test_flash_attn.py index 7d79adefb..26227a492 100644 --- a/test/registered/cpu/test_flash_attn.py +++ b/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]