49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
# PR vllm-project/vllm#32623 introduced a new MLAPrefillBackend abstraction.
|
|
# When MLAAttention.__init__ calls get_mla_prefill_backend(), the upstream
|
|
# selector sees that Ascend NPU returns None for get_device_capability() and
|
|
# falls back to FlashAttnPrefillBackend, which asserts flash_attn_varlen_func
|
|
# is available — crashing on Ascend.
|
|
#
|
|
# Ascend's AscendSFAImpl/AscendMLAImpl handles the full forward pass (including
|
|
# prefill) via impl.forward(), so prefill_backend.run_prefill_* is never called.
|
|
# We register a no-op AscendMLAPrefillBackend and patch get_mla_prefill_backend
|
|
# so that MLAAttention.__init__ completes without error.
|
|
|
|
import torch
|
|
import vllm.model_executor.layers.attention.mla_attention
|
|
from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
|
|
|
|
|
|
class AscendMLAPrefillBackend(MLAPrefillBackend):
|
|
@staticmethod
|
|
def get_name() -> str:
|
|
return "ASCEND"
|
|
|
|
@classmethod
|
|
def is_available(cls) -> bool:
|
|
return True
|
|
|
|
def run_prefill_new_tokens(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
return_softmax_lse: bool,
|
|
) -> torch.Tensor:
|
|
raise NotImplementedError("Ascend MLA prefill is handled by AscendSFAImpl/AscendMLAImpl")
|
|
|
|
def run_prefill_context_chunk(
|
|
self,
|
|
chunk_idx: int,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
raise NotImplementedError("Ascend MLA prefill is handled by AscendSFAImpl/AscendMLAImpl")
|
|
|
|
|
|
vllm.model_executor.layers.attention.mla_attention.get_mla_prefill_backend = lambda vllm_config: AscendMLAPrefillBackend
|