chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Ascend][feature] support L1+ L2 radixcache on ascend (#12214)

Browse Source 
Co-authored-by: khalilzhk <zhangkaikai12@huawei.com>
Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
This commit is contained in:
khalilzhk
2025-11-12 20:45:24 +08:00
committed by GitHub
parent ebaf86d441
commit 2d531946db
9 changed files with 243 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/layers/attention/ascend_backend.py
View File

@@ -403,6 +403,11 @@ class AscendAttnBackend(AttentionBackend):
assert (
layer.qk_head_dim != layer.v_head_dim
), "FIA only supports qk_head_dim != v_head_dim"
# Wait for the KV transfer to complete before performing attention computation.
forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id)
num_token_padding = q.shape[0]
q, k, v = [
data[: forward_batch.num_token_non_padded_cpu] for data in [q, k, v]

27
python/sglang/srt/managers/cache_controller.py
View File

@@ -41,22 +41,25 @@ from sglang.srt.layers.dp_attention import (
is_dp_attention_enabled,
)
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
from sglang.srt.utils import get_device_module
logger = logging.getLogger(__name__)
device_module = get_device_module()
class LayerLoadingEvent:
def __init__(self, num_layers: int):
self._num_layers = num_layers
self.load_events = [torch.cuda.Event() for _ in range(num_layers)]
self.start_event = torch.cuda.Event() # start event on controller stream
self.load_events = [device_module.Event() for _ in range(num_layers)]
self.start_event = device_module.Event() # start event on controller stream
def complete(self, layer_index: int):
assert 0 <= layer_index < self._num_layers
self.load_events[layer_index].record()
def wait(self, layer_index: int):
torch.cuda.current_stream().wait_event(self.load_events[layer_index])
device_module.current_stream().wait_event(self.load_events[layer_index])
@property
def finish_event(self):
@@ -136,8 +139,8 @@ class CacheOperation:
class HiCacheAck(NamedTuple):
start_event: torch.cuda.Event
finish_event: torch.cuda.Event
start_event: device_module.Event
finish_event: device_module.Event
node_ids: List[int]
@@ -343,8 +346,8 @@ class HiCacheController:
self.stop_event, buffer_count=10, max_buffer_size=100
)
self.write_stream = torch.cuda.Stream()
self.load_stream = torch.cuda.Stream()
self.write_stream = device_module.Stream()
self.load_stream = device_module.Stream()
if self.enable_storage:
self.prefetch_thread = threading.Thread(
@@ -445,11 +448,11 @@ class HiCacheController:
host_indices, device_indices = self.move_indices(op)
self.write_queue.clear()
start_event = torch.cuda.Event()
finish_event = torch.cuda.Event()
start_event = device_module.Event()
finish_event = device_module.Event()
start_event.record()
with torch.cuda.stream(self.write_stream):
with device_module.stream(self.write_stream):
start_event.wait(self.write_stream)
self.mem_pool_host.backup_from_device_all_layer(
self.mem_pool_device, host_indices, device_indices, self.io_backend
@@ -496,6 +499,8 @@ class HiCacheController:
return host_indices, device_indices.index_select(0, idx)
elif self.mem_pool_host.layout == "page_first_direct":
return host_indices, device_indices.cpu()
elif self.io_backend == "kernel_ascend":
return host_indices, device_indices
else:
raise ValueError(f"Unsupported io backend")
@@ -510,7 +515,7 @@ class HiCacheController:
producer_event = self.layer_done_counter.events[producer_id]
producer_event.start_event.record()
with torch.cuda.stream(self.load_stream):
with device_module.stream(self.load_stream):
producer_event.start_event.wait(self.load_stream)
for i in range(self.layer_num):
self.mem_pool_host.load_to_device_per_layer(

98
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -26,9 +26,15 @@ if not (_is_npu or _is_xpu):
transfer_kv_per_layer_mla_pf_lf,
transfer_kv_per_layer_pf_lf,
)
if _is_npu:
from sgl_kernel_npu.kvcacheio import TransferDirection, transfer_kv_dim_exchange
logger = logging.getLogger(__name__)
SUPPORT_PIN_MEMORY = not _is_npu
if SUPPORT_PIN_MEMORY:
logger.warning("Current platform not support pin_memory")
def synchronized(func):
@wraps(func)
@@ -54,7 +60,7 @@ class HostKVCache(abc.ABC):
self.device_pool = device_pool
self.page_size = page_size
self.layout = layout
self.pin_memory = pin_memory
self.pin_memory = pin_memory and SUPPORT_PIN_MEMORY
self.device = device
self.dtype = device_pool.store_dtype
@@ -63,9 +69,9 @@ class HostKVCache(abc.ABC):
self.size = int(host_size * 1e9 // self.size_per_token)
else:
self.size = int(device_pool.size * host_to_device_ratio)
# Align the host memory pool size to the page size
self.size = self.size - (self.size % self.page_size)
self.page_num = self.size // self.page_size
# Align up the host memory pool size to the page size
self.page_num = self.size // self.page_size + 1
self.size = self.page_num * self.page_size
self.start_layer = device_pool.start_layer
self.end_layer = device_pool.end_layer
@@ -316,6 +322,22 @@ class MHATokenToKVPoolHost(HostKVCache):
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_direct":
# Ascend-specific: transfer KV data for all layers when layer_id == 0
if layer_id == 0:
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_pool.k_buffer,
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
host_v=self.v_buffer,
page_size=self.page_size,
direction=TransferDirection.H2D,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
raise ValueError(f"Unsupported IO backend: {io_backend}")
@@ -367,6 +389,20 @@ class MHATokenToKVPoolHost(HostKVCache):
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_direct":
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_pool.k_buffer,
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
host_v=self.v_buffer,
page_size=self.page_size,
direction=TransferDirection.D2H,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
raise ValueError(f"Unsupported IO backend: {io_backend}")
@@ -549,6 +585,28 @@ class MLATokenToKVPoolHost(HostKVCache):
1,
self.kv_lora_rank + self.qk_rope_head_dim,
)
# Ascend-specific: Aligns with AscendMLAPagedTokenToKVPool layout
# Separately allocate k_buffer and v_buffer for easier data transfer.
elif self.layout == "page_first_kv_split":
base_dims = (
self.page_num,
self.layer_num,
self.page_size,
1,
)
self.k_buffer = torch.empty(
(*base_dims, self.kv_lora_rank),
dtype=self.dtype,
device=self.device,
)
self.v_buffer = torch.empty(
(*base_dims, self.qk_rope_head_dim),
dtype=self.dtype,
device=self.device,
)
# Return k_buffer to preserve original kv_buffer and data_refs init logic,
# though Ascend doesn't use these parameters.
return self.k_buffer
else:
raise ValueError(f"Unsupported layout: {self.layout}")
self.token_stride_size = (
@@ -610,6 +668,24 @@ class MLATokenToKVPoolHost(HostKVCache):
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_kv_split":
# Ascend-specific: transfer KV data for all layers when layer_id == 0
if layer_id == 0:
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_pool.k_buffer,
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
host_v=self.v_buffer,
page_size=self.page_size,
direction=TransferDirection.H2D,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
raise ValueError(f"Unsupported IO backend: {io_backend}")
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend
@@ -655,6 +731,20 @@ class MLATokenToKVPoolHost(HostKVCache):
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_kv_split":
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_pool.k_buffer,
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
host_v=self.v_buffer,
page_size=self.page_size,
direction=TransferDirection.D2H,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
raise ValueError(f"Unsupported IO backend: {io_backend}")

23
python/sglang/srt/server_args.py
View File

@@ -1506,6 +1506,20 @@ class ServerArgs:
"Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes."
)
# Below are the only parameters currently supported on Ascend
if self.enable_hierarchical_cache and is_npu():
# FIXME(iforgetmyname) fix decode_attention_backend on ascend
self.decode_attention_backend = "ascend"
self.hicache_io_backend = "kernel_ascend"
if self.use_mla_backend():
self.hicache_mem_layout = "page_first_kv_split"
else:
self.hicache_mem_layout = "page_first_direct"
logger.warning(
f"Ascend NPU Platform detected, change `hicache_io_backend` to `kernel_ascend` and "
f"`hicache_mem_layout` to `{self.hicache_mem_layout}`"
)
def _handle_speculative_decoding(self):
if self.speculative_algorithm == "NEXTN":
self.speculative_algorithm = "EAGLE"
@@ -2991,14 +3005,19 @@ class ServerArgs:
parser.add_argument(
"--hicache-io-backend",
type=str,
choices=["direct", "kernel"],
choices=["direct", "kernel", "kernel_ascend"],
default=ServerArgs.hicache_io_backend,
help="The IO backend for KV cache transfer between CPU and GPU",
)
parser.add_argument(
"--hicache-mem-layout",
type=str,
choices=["layer_first", "page_first", "page_first_direct"],
choices=[
"layer_first",
"page_first",
"page_first_direct",
"page_first_kv_split",
],
default=ServerArgs.hicache_mem_layout,
help="The layout of host memory pool for hierarchical cache.",
)

5
python/sglang/srt/utils/common.py
View File

@@ -653,6 +653,11 @@ def wait_cmo_stream():
cur_stream.wait_stream(get_cmo_stream())
@lru_cache(maxsize=1)
def get_device_module():
return torch.get_device_module()
def set_random_seed(seed: int) -> None:
"""Set the random seed for all libraries."""
random.seed(seed)