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

feat(hicache): support NPU Mamba states with FIA and async IO (#32500)

Browse Source
This commit is contained in:
qyb233
2026-09-17 10:12:26 +08:00
committed by GitHub
parent 4c85172f3a
commit 4fb9b5b5ba
6 changed files with 652 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/memory_pool_npu.py
+19 -6
View File
@@ -105,7 +105,11 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
# The padded slot 0 is used for writing dummy outputs from padded tokens.
# Continuous memory improves the efficiency of Ascend`s transmission backend,
# while other backends remain unchanged.
self.k_buffer = torch.zeros(
# FIA exposes the KV cache as per-layer Python views so graph
# capture does not retain the full multi-layer tensor. HiCache's
# NPU exchange operator still requires the original contiguous
# [layer, page, token, head, dim] allocation.
self._hicache_k_buffer = torch.zeros(
(
self.layer_num,
self.size // self.page_size + 1,
@@ -116,7 +120,7 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
dtype=self.store_dtype,
device=self.device,
)
self.v_buffer = torch.zeros(
self._hicache_v_buffer = torch.zeros(
(
self.layer_num,
self.size // self.page_size + 1,
@@ -130,8 +134,11 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
# Keep a reference to the contiguous tensor for HiCache
# D2H/H2D transfers (transfer_kv_dim_exchange expects a
# tensor, not the per-layer list used in FIA mode below).
self.k_buffer_tensor = self.k_buffer
self.v_buffer_tensor = self.v_buffer
self.k_buffer_tensor = self._hicache_k_buffer
self.v_buffer_tensor = self._hicache_v_buffer
self.k_buffer = self._hicache_k_buffer
self.v_buffer = self._hicache_v_buffer
if self.use_fia:
# Use per-layer Python lists to avoid torch.compile capturing
@@ -139,14 +146,20 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
# Each layer view: [P*ps, 1, H, D], sharing the contiguous
# storage allocated above.
self.k_buffer = [
self.k_buffer[i].view(-1, 1, self.head_num, self.head_dim)
self._hicache_k_buffer[i].view(-1, 1, self.head_num, self.head_dim)
for i in range(self.layer_num)
]
self.v_buffer = [
self.v_buffer[i].view(-1, 1, self.head_num, self.v_head_dim)
self._hicache_v_buffer[i].view(
-1, 1, self.head_num, self.v_head_dim
)
for i in range(self.layer_num)
]
def get_hicache_transfer_buffers(self):
"""Return contiguous all-layer KV tensors for NPU HiCache IO."""
return self._hicache_k_buffer, self._hicache_v_buffer
def _init_kv_copy_and_warmup(self):
# implementation relies on self.data_strides / self.data_ptrs, which the
# NPU paged buffer layout never builds.
python/sglang/srt/mem_cache/pool_host/mamba.py
+90 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import logging
import os
import threading
from typing import Optional
@@ -18,10 +19,13 @@ from sglang.srt.mem_cache.pool_host.common import (
ALLOC_MEMORY_FUNCS,
get_allocator_from_storage,
)
from sglang.srt.utils import is_cuda, is_hip
from sglang.srt.utils import is_cuda, is_hip, is_npu
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_npu = is_npu()
transfer_state_per_layer_direct_pf_lf = None
transfer_state_all_layer_direct_lf_pf = None
if _is_cuda or _is_hip:
from sgl_kernel.kvcacheio import (
transfer_kv_all_layer_direct_lf_pf,
@@ -34,10 +38,32 @@ if _is_cuda or _is_hip:
transfer_kv_mamba_lf_pf,
transfer_kv_mamba_pf_lf,
)
if _is_npu:
try:
from sgl_kernel_npu.kvcacheio import (
transfer_state_all_layer_direct_lf_pf,
transfer_state_per_layer_direct_pf_lf,
)
except ImportError:
pass
logger = logging.getLogger(__name__)
_NPU_HICACHE_MAMBA_IO_ENV = "SGLANG_NPU_HICACHE_MAMBA_IO"
_NPU_HICACHE_MAMBA_IO_MODES = {"sync", "async"}
def _npu_hicache_mamba_io_mode() -> str:
mode = os.getenv(_NPU_HICACHE_MAMBA_IO_ENV, "sync").strip().lower()
if mode not in _NPU_HICACHE_MAMBA_IO_MODES:
raise ValueError(
f"{_NPU_HICACHE_MAMBA_IO_ENV} must be one of "
f"{sorted(_NPU_HICACHE_MAMBA_IO_MODES)}, got {mode!r}."
)
return mode
class MambaPoolHost(HostKVCache):
def __init__(
self,
@@ -128,10 +154,36 @@ class MambaPoolHost(HostKVCache):
]
self.kv_buffer = self.init_kv_buffer()
self._configure_npu_mamba_io()
self._init_write_back_staging_buffers()
self.lock = threading.RLock()
self.clear()
def _configure_npu_mamba_io(self) -> None:
mode = _npu_hicache_mamba_io_mode()
if mode == "sync":
logger.info("NPU HiCache Mamba state transfer mode: sync torch fallback.")
return
required_ops = (
transfer_state_per_layer_direct_pf_lf,
transfer_state_all_layer_direct_lf_pf,
)
required_torch_ops = (
"transfer_state_per_layer_direct_pf_lf",
"transfer_state_all_layer_direct_lf_pf",
)
if any(op is None for op in required_ops) or any(
not hasattr(torch.ops.npu, op_name) for op_name in required_torch_ops
):
raise RuntimeError(
"NPU HiCache Mamba async state transfer requires "
"the per-layer PF->LF and all-layer LF->PF direct operators "
"from sgl-kernel-npu."
)
logger.info("NPU HiCache Mamba state transfer mode: native async.")
def init_kv_buffer(self):
_host_alloc = ALLOC_MEMORY_FUNCS[self.device_pool.device]
@@ -361,6 +413,22 @@ class MambaPoolHost(HostKVCache):
layer_id=layer_id,
page_size=1,
)
elif io_backend == "kernel_ascend":
if _npu_hicache_mamba_io_mode() == "async":
transfer_state_per_layer_direct_pf_lf(
src=src,
dst=dst,
src_indices=src_indices,
dst_indices=dst_indices,
layer_id=layer_id,
)
else:
host_indices = src_indices.to(dtype=torch.int64, device=src.device)
device_indices = dst_indices.to(dtype=torch.int64, device=dst.device)
values = (
src.select(1, layer_id).index_select(0, host_indices).select(1, 0)
)
dst.index_copy_(0, device_indices, values.to(device=dst.device))
else:
raise ValueError(f"Unsupported io_backend: {io_backend}")
@@ -404,6 +472,27 @@ class MambaPoolHost(HostKVCache):
dst_indices=dst_indices,
page_size=1,
)
elif io_backend == "kernel_ascend":
if _npu_hicache_mamba_io_mode() == "async":
transfer_state_all_layer_direct_lf_pf(
device_states=[src_layers],
host_states=[dst],
device_indices=src_indices,
host_indices=dst_indices,
)
else:
device_indices = src_indices.to(
dtype=torch.int64, device=src_layers.device
)
host_indices = dst_indices.to(dtype=torch.int64, device=dst.device)
values = (
src_layers.index_select(1, device_indices)
.movedim(0, 1)
.unsqueeze(2)
.contiguous()
.to(device=dst.device)
)
dst.index_copy_(0, host_indices, values)
else:
raise ValueError(f"Unsupported io_backend: {io_backend}")
python/sglang/srt/mem_cache/pool_host/mha.py
+43 -31
View File
@@ -131,7 +131,7 @@ class MHATokenToKVPoolHost(HostKVCache):
self.device_pool.device,
host_memory_registered=self.pin_memory,
)
if self.mtp_draft_device_pools:
if self.mtp_draft_device_pools and not _is_npu:
device_pools = (self.device_pool, *self.mtp_draft_device_pools)
if not _is_npu:
self.packed_device_k_data_ptrs = torch.cat(
@@ -364,22 +364,22 @@ class MHATokenToKVPoolHost(HostKVCache):
if self.layout == "page_first_direct":
# Ascend-specific: transfer KV data for all layers when layer_id == 0
if host_layer_id == 0:
device_k = getattr(
device_pool, "k_buffer_tensor", device_pool.k_buffer
)
device_v = getattr(
device_pool, "v_buffer_tensor", device_pool.v_buffer
)
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_k,
host_k=self.k_buffer,
device_v=device_v,
host_v=self.v_buffer,
page_size=self.page_size,
direction=TransferDirection.H2D,
)
for (
device_k,
device_v,
host_k,
host_v,
) in self._npu_transfer_buffers(device_pool):
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_k,
host_k=host_k,
device_v=device_v,
host_v=host_v,
page_size=self.page_size,
direction=TransferDirection.H2D,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
@@ -400,6 +400,19 @@ class MHATokenToKVPoolHost(HostKVCache):
device_pool.v_buffer,
)
def _npu_transfer_buffers(self, target_device_pool):
layer_start = 0
for pool in (target_device_pool, *self.mtp_draft_device_pools):
device_k, device_v = pool.get_hicache_transfer_buffers()
layer_end = layer_start + device_k.shape[0]
yield (
device_k,
device_v,
self.k_buffer[:, layer_start:layer_end],
self.v_buffer[:, layer_start:layer_end],
)
layer_start = layer_end
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend
):
@@ -502,20 +515,19 @@ class MHATokenToKVPoolHost(HostKVCache):
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_direct":
# In FIA mode, k_buffer/v_buffer are per-layer lists;
# use the 5-D contiguous view for transfer_kv_dim_exchange.
device_k = getattr(device_pool, "k_buffer_tensor", device_pool.k_buffer)
device_v = getattr(device_pool, "v_buffer_tensor", device_pool.v_buffer)
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_k,
host_k=self.k_buffer,
device_v=device_v,
host_v=self.v_buffer,
page_size=self.page_size,
direction=TransferDirection.D2H,
)
for device_k, device_v, host_k, host_v in self._npu_transfer_buffers(
device_pool
):
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_k,
host_k=host_k,
device_v=device_v,
host_v=host_v,
page_size=self.page_size,
direction=TransferDirection.D2H,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else: