Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
This commit is contained in:
Woosuk Kwon
2026-04-10 04:07:00 +00:00
parent 571e7d3cac
commit 24151eb438
2 changed files with 3 additions and 7 deletions
@@ -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