From 81e1c3f18339ee7101c4528203a374215b62e0f4 Mon Sep 17 00:00:00 2001 From: ZhouXiaoya25 Date: Thu, 15 Oct 2026 12:47:37 +0800 Subject: [PATCH] [XPU][MLA] Integrate XPU MLA support into the existing FLASH_ATTN_MLA backend (#58297) Signed-off-by: baodii Signed-off-by: Zhouxiaoya25 Signed-off-by: ZhouXiaoya25 Co-authored-by: baodii Co-authored-by: Kunshang Ji --- tests/v1/attention/test_mla_backends.py | 39 ++--- vllm/_xpu_ops.py | 6 +- vllm/platforms/xpu.py | 32 +++- vllm/v1/attention/backends/fa_utils.py | 2 + .../attention/backends/mla/flashattn_mla.py | 137 +++++++++++++++++- 5 files changed, 190 insertions(+), 26 deletions(-) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 91bd9008f0..b696e35b77 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -358,6 +358,8 @@ if current_platform.is_cuda(): ) elif current_platform.is_rocm(): PREFILL_BACKENDS_TO_TEST.append(MLAPrefillBackendEnum.ROCM_AITER_FA) +elif current_platform.is_xpu(): + PREFILL_BACKENDS_TO_TEST.append(MLAPrefillBackendEnum.FLASH_ATTN) MLA_DIMENSIONS_TO_TEST = [ ("deepseek", 128, 128), @@ -370,25 +372,20 @@ def _prefill_backend_dimension_params(): params = [] for prefill_backend in PREFILL_BACKENDS_TO_TEST: for dimensions_id, qk_nope_head_dim, v_head_dim in MLA_DIMENSIONS_TO_TEST: - if device_capability is None: - invalid_reasons = ["device capability unavailable"] - else: - try: - invalid_reasons = ( - prefill_backend.get_class().validate_configuration( - device_capability, - MLAPrefillSelectorConfig( - dtype=torch.bfloat16, - mla_dimensions=MLADimensions( - qk_nope_head_dim=qk_nope_head_dim, - qk_rope_head_dim=64, - v_head_dim=v_head_dim, - ), - ), - ) - ) - except ImportError: - invalid_reasons = ["ImportError"] + try: + invalid_reasons = prefill_backend.get_class().validate_configuration( + device_capability, + MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=MLADimensions( + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=64, + v_head_dim=v_head_dim, + ), + ), + ) + except ImportError: + invalid_reasons = ["ImportError"] marks = [] if invalid_reasons: @@ -1668,6 +1665,10 @@ def _run_backend_correctness( block_size=default_block_size, hf_config_override=hf_config_override, ) + if hf_config_override and current_platform.is_xpu(): + vllm_config.model_config.model_arch_config.total_num_attention_heads = ( + hf_config_override["num_attention_heads"] + ) vllm_config.cache_config.cache_dtype = kv_cache_dtype # For spec decode tests, add a speculative_config to set the reorder_batch_threshold diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index a384ac7215..166534e5b1 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -1064,7 +1064,11 @@ class xpu_ops: "when block_table is disabled, cu_seqlens_k is needed" ) if out is None: - out = torch.empty(q.shape, dtype=q.dtype, device=q.device) + out = torch.empty( + q.shape[:-1] + (v.shape[-1],), + dtype=q.dtype, + device=q.device, + ) real_window_size: tuple[int, int] if window_size is None: real_window_size = (-1, -1) diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 992bfae053..71da7631ea 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -158,8 +158,36 @@ class XPUPlatform(Platform): logger.info_once("Using XPU MLA Sparse backend.") return AttentionBackendEnum.XPU_MLA_SPARSE.get_path() if attn_selector_config.use_mla: - logger.info_once("Using Triton MLA backend on V1 engine.") - return AttentionBackendEnum.TRITON_MLA.get_path() + if selected_backend == AttentionBackendEnum.TRITON_MLA: + logger.info_once("Using Triton MLA backend on V1 engine.") + return AttentionBackendEnum.TRITON_MLA.get_path() + if selected_backend not in (None, AttentionBackendEnum.FLASH_ATTN_MLA): + raise ValueError( + f"Invalid attention backend for {cls.device_name}, " + f"with use_mla: {attn_selector_config.use_mla}" + ) + from vllm.v1.attention.backends.mla.flashattn_mla import ( + XPUFlashAttnMLABackend, + ) + + invalid_reasons = XPUFlashAttnMLABackend.validate_configuration( + device_capability=cls.get_device_capability(), # type: ignore[arg-type] + **attn_selector_config._asdict(), + ) + if invalid_reasons: + if selected_backend == AttentionBackendEnum.FLASH_ATTN_MLA: + raise ValueError( + f"FlashAttnMLA on XPU is not valid for this " + f"configuration: {invalid_reasons}" + ) + logger.info_once( + "FlashAttnMLA on XPU is not valid for this configuration " + "(%s); falling back to Triton MLA backend.", + tuple(invalid_reasons), + ) + return AttentionBackendEnum.TRITON_MLA.get_path() + logger.info_once("Using Flash Attention MLA backend on XPU.") + return "vllm.v1.attention.backends.mla.flashattn_mla.XPUFlashAttnMLABackend" if selected_backend == AttentionBackendEnum.TRITON_ATTN: logger.info_once("Using Triton backend.") return AttentionBackendEnum.TRITON_ATTN.get_path() diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 83348bb94a..a20d825b9e 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -340,6 +340,8 @@ def flash_attn_supports_mla(): except (ImportError, AssertionError): pass + elif current_platform.is_xpu(): + return True return False diff --git a/vllm/v1/attention/backends/mla/flashattn_mla.py b/vllm/v1/attention/backends/mla/flashattn_mla.py index b685be102e..f2ac83851d 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla.py @@ -29,13 +29,11 @@ from vllm.v1.attention.backend import ( ) from vllm.v1.attention.backends.fa_utils import ( flash_attn_supports_mla, - get_flash_attn_version, -) -from vllm.v1.kv_cache_interface import AttentionSpec -from vllm.vllm_flash_attn import ( # type: ignore[attr-defined] flash_attn_varlen_func, + get_flash_attn_version, get_scheduler_metadata, ) +from vllm.v1.kv_cache_interface import AttentionSpec logger = init_logger(__name__) @@ -90,6 +88,28 @@ class FlashAttnMLABackend(MLACommonBackend): return None +class XPUFlashAttnMLABackend(FlashAttnMLABackend): + @staticmethod + def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: + return [MultipleOf(64)] + + @staticmethod + def get_builder_cls() -> type["XPUFlashAttnMLAMetadataBuilder"]: + return XPUFlashAttnMLAMetadataBuilder + + @staticmethod + def get_impl_cls() -> type["XPUFlashAttnMLAImpl"]: + return XPUFlashAttnMLAImpl + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + return True + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [576] + + @dataclass class FlashAttnMLADecodeMetadata(MLACommonDecodeMetadata): query_start_loc: torch.Tensor @@ -256,6 +276,27 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata] return metadata +class XPUFlashAttnMLAMetadataBuilder(FlashAttnMLAMetadataBuilder): + _cudagraph_support: ClassVar[AttentionCGSupport] = ( + AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE + ) + query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.SINGLE_ONLY + reorder_batch_threshold: int = 1 + + def _schedule_decode( + self, + num_reqs, + cu_query_lens, + max_query_len, + seqlens, + max_seq_len, + causal, + max_num_splits, + ): + # XPU uses the FA2 varlen interface, which has no AOT scheduler metadata. + return None + + class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]): can_return_lse_for_decode: bool = True supports_dcp: bool = True @@ -367,3 +408,91 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]): else: o = attn_out return o, None + + +class XPUFlashAttnMLAImpl(FlashAttnMLAImpl): + supports_dcp: bool = False + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None, + attn_type: str, + kv_sharing_target_layer_name: str | None, + # MLA Specific Arguments + **mla_args, + ) -> None: + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + **mla_args, + ) + + if self.dcp_world_size > 1: + raise NotImplementedError( + "FlashAttnMLA on XPU does not support decode context parallelism" + ) + + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: FlashAttnMLAMetadata, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + assert kv_c_and_k_pe_cache.numel() > 0 + assert attn_metadata.decode is not None + num_decodes = attn_metadata.num_decodes + decode_cu_seqlens_q = attn_metadata.query_start_loc[: num_decodes + 1] + + cache = kv_c_and_k_pe_cache + if cache.dim() == 3: + cache = cache.unsqueeze(-2) + assert cache.dim() == 4 and cache.size(-2) == 1, ( + "kv_c_and_k_pe_cache must be [num_blocks, block_size, " + "(num_heads_kv=1)?, kv_lora_rank+qk_rope_head_dim]" + ) + + if type(q) is tuple: # noqa: SIM108 + q_in = torch.cat(q, dim=-1) + else: + q_in = q + if not q_in.is_contiguous(): + q_in = q_in.contiguous() + + attn_out = flash_attn_varlen_func( + q_in, + cache, + cache.narrow(-1, 0, self.kv_lora_rank), + max_seqlen_q=1, + cu_seqlens_q=decode_cu_seqlens_q, + max_seqlen_k=attn_metadata.decode.max_seq_len, + seqused_k=attn_metadata.decode.seq_lens, + block_table=attn_metadata.decode.block_table, + softmax_scale=self.scale, + # Single-token paged decode has no future KV to mask; seqused_k + # already restricts each request to its valid KV prefix. + causal=False, + fa_version=2, + return_softmax_lse=self.need_to_return_lse_for_decode, + ) + + if self.need_to_return_lse_for_decode: + o, lse = attn_out + # FA returns LSE in shape [ H, B ] but DCP wants [ B, H ] + return o, lse.transpose(0, 1) # [ H, B ] -> [ B, H ] + return attn_out, None