Ascend attention backend(PA&MLA) (#7722)
Co-authored-by: Maksim <makcum888e@mail.ru> Co-authored-by: VDV1985 <vladdv85@mail.ru>
This commit is contained in:
ronnie_zheng
@@ -72,12 +72,15 @@ from sglang.srt.managers.schedule_batch import (
|
||||
global_server_args_dict,
|
||||
)
|
||||
from sglang.srt.mem_cache.allocator import (
|
||||
AscendPagedTokenToKVPoolAllocator,
|
||||
BaseTokenToKVPoolAllocator,
|
||||
PagedTokenToKVPoolAllocator,
|
||||
SWATokenToKVPoolAllocator,
|
||||
TokenToKVPoolAllocator,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool import (
|
||||
AscendMLAPagedTokenToKVPool,
|
||||
AscendTokenToKVPool,
|
||||
DoubleSparseTokenToKVPool,
|
||||
MHATokenToKVPool,
|
||||
MLATokenToKVPool,
|
||||
@@ -110,6 +113,7 @@ from sglang.srt.utils import (
|
||||
is_hip,
|
||||
is_hopper_with_cuda_12_3,
|
||||
is_no_spec_infer_or_topk_one,
|
||||
is_npu,
|
||||
monkey_patch_p2p_access_check,
|
||||
monkey_patch_vllm_gguf_config,
|
||||
set_cpu_offload_max_bytes,
|
||||
@@ -117,6 +121,7 @@ from sglang.srt.utils import (
|
||||
)
|
||||
|
||||
_is_hip = is_hip()
|
||||
_is_npu = is_npu()
|
||||
_is_cpu_amx_available = cpu_has_amx_support()
|
||||
|
||||
# Use a small KV cache pool size for tests in CI
|
||||
@@ -308,6 +313,7 @@ class ModelRunner:
|
||||
self.init_cuda_graphs()
|
||||
else:
|
||||
self.cuda_graph_runner = None
|
||||
self.cuda_graph_mem_usage = 0
|
||||
self.init_attention_backend()
|
||||
|
||||
# auxiliary hidden capture mode. TODO: expose this to server args?
|
||||
@@ -369,6 +375,8 @@ class ModelRunner:
|
||||
server_args.attention_backend = "fa3"
|
||||
elif _is_hip:
|
||||
server_args.attention_backend = "aiter"
|
||||
elif _is_npu:
|
||||
server_args.attention_backend = "ascend"
|
||||
else:
|
||||
server_args.attention_backend = (
|
||||
"flashinfer" if is_flashinfer_available() else "triton"
|
||||
@@ -388,6 +396,8 @@ class ModelRunner:
|
||||
server_args.attention_backend = "aiter"
|
||||
else:
|
||||
server_args.attention_backend = "triton"
|
||||
elif _is_npu:
|
||||
server_args.attention_backend = "ascend"
|
||||
else:
|
||||
server_args.attention_backend = "triton"
|
||||
logger.info(
|
||||
@@ -402,6 +412,7 @@ class ModelRunner:
|
||||
"triton",
|
||||
"flashmla",
|
||||
"cutlass_mla",
|
||||
"ascend",
|
||||
]:
|
||||
logger.info(
|
||||
f"MLA optimization is turned on. Use {server_args.attention_backend} backend."
|
||||
@@ -1096,7 +1107,35 @@ class ModelRunner:
|
||||
# Draft worker shares req_to_token_pool with the target worker.
|
||||
assert self.is_draft_worker
|
||||
|
||||
if self.use_mla_backend:
|
||||
if self.server_args.attention_backend == "ascend" and not self.use_mla_backend:
|
||||
self.token_to_kv_pool = AscendTokenToKVPool(
|
||||
self.max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
head_num=self.model_config.get_num_kv_heads(get_attention_tp_size()),
|
||||
head_dim=self.model_config.head_dim,
|
||||
layer_num=self.model_config.num_hidden_layers,
|
||||
device=self.device,
|
||||
enable_memory_saver=self.server_args.enable_memory_saver,
|
||||
)
|
||||
elif self.server_args.attention_backend == "ascend" and self.use_mla_backend:
|
||||
self.token_to_kv_pool = AscendMLAPagedTokenToKVPool(
|
||||
self.max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
kv_lora_rank=self.model_config.kv_lora_rank,
|
||||
qk_rope_head_dim=self.model_config.qk_rope_head_dim,
|
||||
layer_num=(
|
||||
self.model_config.num_hidden_layers
|
||||
if not self.is_draft_worker
|
||||
else self.model_config.hf_config.num_nextn_predict_layers
|
||||
), # PP is not compatible with mla backend
|
||||
device=self.device,
|
||||
enable_memory_saver=self.server_args.enable_memory_saver,
|
||||
start_layer=self.start_layer,
|
||||
end_layer=self.end_layer,
|
||||
)
|
||||
elif self.use_mla_backend:
|
||||
self.token_to_kv_pool = MLATokenToKVPool(
|
||||
self.max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
@@ -1176,13 +1215,22 @@ class ModelRunner:
|
||||
kvcache=self.token_to_kv_pool,
|
||||
)
|
||||
else:
|
||||
self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
|
||||
self.max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
device=self.device,
|
||||
kvcache=self.token_to_kv_pool,
|
||||
)
|
||||
if _is_npu:
|
||||
self.token_to_kv_pool_allocator = AscendPagedTokenToKVPoolAllocator(
|
||||
self.max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
device=self.device,
|
||||
kvcache=self.token_to_kv_pool,
|
||||
)
|
||||
else:
|
||||
self.token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator(
|
||||
self.max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
device=self.device,
|
||||
kvcache=self.token_to_kv_pool,
|
||||
)
|
||||
else:
|
||||
assert self.is_draft_worker
|
||||
|
||||
@@ -1229,6 +1277,10 @@ class ModelRunner:
|
||||
from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
|
||||
|
||||
return AiterAttnBackend(self)
|
||||
elif self.server_args.attention_backend == "ascend":
|
||||
from sglang.srt.layers.attention.ascend_backend import AscendAttnBackend
|
||||
|
||||
return AscendAttnBackend(self)
|
||||
elif self.server_args.attention_backend == "triton":
|
||||
assert not self.model_config.is_encoder_decoder, (
|
||||
"Cross attention is not supported in the triton attention backend. "
|
||||
|
||||
Reference in New Issue
Block a user