iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[XPU][MLA] Integrate XPU MLA support into the existing FLASH_ATTN_MLA backend (#58297)

Browse Source 
Signed-off-by: baodii <di.bao@intel.com>
Signed-off-by: Zhouxiaoya25 <xiaoya.zhou@intel.com>
Signed-off-by: ZhouXiaoya25 <xiaoya.zhou@intel.com>
Co-authored-by: baodii <di.bao@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
ZhouXiaoya25
2026-10-15 12:47:37 +08:00
committed by GitHub
co-authored by baodii Kunshang Ji
parent 96b5b0ee4b
commit 81e1c3f183
5 changed files with 190 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_mla_backends.py
+20 -19
View File
@@ -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
vllm/_xpu_ops.py
+5 -1
View File
@@ -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)
vllm/platforms/xpu.py
+30 -2
View File
@@ -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()
vllm/v1/attention/backends/fa_utils.py
+2
View File
@@ -340,6 +340,8 @@ def flash_attn_supports_mla():
except (ImportError, AssertionError):
pass
elif current_platform.is_xpu():
return True
return False
vllm/v1/attention/backends/mla/flashattn_mla.py
+133 -4
View File
@@ -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