forked from Karylab-cklius/vllm
@@ -65,7 +65,7 @@ def dsa(
|
||||
output.zero_()
|
||||
return output
|
||||
|
||||
num_actual_toks = mla_attn_metadata.num_actual_tokens # noqa
|
||||
num_actual_toks = mla_attn_metadata.num_actual_tokens # type: ignore[attr-defined]
|
||||
if num_actual_toks == 0:
|
||||
output.zero_()
|
||||
return output
|
||||
@@ -77,7 +77,7 @@ def dsa(
|
||||
mla_k_scale = None
|
||||
idx_meta = attn_metadata.get(attn.indexer_k_cache.prefix)
|
||||
if idx_meta is not None:
|
||||
slot_mapping = idx_meta.slot_mapping # noqa
|
||||
slot_mapping = idx_meta.slot_mapping # type: ignore[attr-defined]
|
||||
indexer_k_cache = attn.indexer_k_cache.kv_cache
|
||||
mla_kv_cache = attn.mla_attn.kv_cache
|
||||
mla_k_scale = attn.mla_attn._k_scale
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
import torch
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.forward_context import get_forward_context
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
@@ -15,11 +16,6 @@ from vllm.v1.attention.backends.mla.indexer import (
|
||||
from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
|
||||
from vllm.v1.worker.workspace import current_workspace_manager
|
||||
|
||||
if current_platform.is_cuda_alike():
|
||||
from vllm import _custom_ops as ops
|
||||
elif current_platform.is_xpu():
|
||||
from vllm._xpu_ops import xpu_ops as ops
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
|
||||
|
||||
Reference in New Issue
Block a user