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

[NPU][Hicache] Optimize HiCache L2 IO with Memfabric acc_offload (#38826)

Browse Source
This commit is contained in:
James
2026-09-11 00:23:36 +08:00
committed by GitHub
parent 19b83b4b9d
commit 209654c420
6 changed files with 589 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/memory_pool_npu.py
+18 -9
View File
@@ -127,6 +127,11 @@ class NPUMHATokenToKVPool(MHATokenToKVPool):
dtype=self.store_dtype,
device=self.device,
)
# 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
if self.use_fia:
# Use per-layer Python lists to avoid torch.compile capturing
@@ -436,6 +441,10 @@ class NPUMHATokenToKOnlyPool(MHATokenToKOnlyPool):
dtype=self.store_dtype,
device=self.device,
)
# 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
if self.use_fia:
self.k_buffer = [
self.k_buffer[i].view(-1, 1, self.head_num, self.head_dim)
@@ -755,15 +764,15 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
def get_contiguous_buf_infos(self):
self._raise_if_native_kv_cache_disabled()
# MLA has only one kv_buffer, so only the information of this buffer needs to be returned.
kv_data_ptrs = [self.k_buffer[i].data_ptr() for i in range(self.layer_num)] + [
self.v_buffer[i].data_ptr() for i in range(self.layer_num)
]
kv_data_lens = [self.k_buffer[i].nbytes for i in range(self.layer_num)] + [
self.v_buffer[i].nbytes for i in range(self.layer_num)
]
kv_item_lens = [self.k_buffer[i][0].nbytes for i in range(self.layer_num)] + [
self.v_buffer[i][0].nbytes for i in range(self.layer_num)
]
kv_data_ptrs = [self.k_buffer[i].data_ptr() for i in range(self.layer_num)]
kv_data_lens = [self.k_buffer[i].nbytes for i in range(self.layer_num)]
kv_item_lens = [self.k_buffer[i][0].nbytes for i in range(self.layer_num)]
# When DSA KV cache is packed into the FP8 k_buffer, the v_buffer is
# intentionally empty (kr_cache_dim == 0). Its data_ptr() is null
if not getattr(self, "dsa_kv_cache_store_fp8", False):
kv_data_ptrs += [self.v_buffer[i].data_ptr() for i in range(self.layer_num)]
kv_data_lens += [self.v_buffer[i].nbytes for i in range(self.layer_num)]
kv_item_lens += [self.v_buffer[i][0].nbytes for i in range(self.layer_num)]
if self.index_head_dim is not None:
ptrs, lens, item_lens = self.get_state_buf_infos()
kv_data_ptrs += ptrs
python/sglang/srt/hardware_backend/npu/utils.py
+5 -2
View File
@@ -170,8 +170,11 @@ def set_default_server_args(args: "ServerArgs"):
disable_custom_all_reduce=True,
)
# handles hierarchical cache configs
if cfg.enable_hierarchical_cache:
# handles hierarchical cache / decode-offload configs
if (
cfg.enable_hierarchical_cache
or cfg.disaggregation_decode_enable_offload_kvcache
):
declare_resolution(
args,
"set_default_server_args",
python/sglang/srt/managers/cache_controller.py
+18
View File
@@ -866,6 +866,24 @@ class HiCacheController:
f"Unsupported layout {self.mem_pool_host.layout!r} for io backend 'direct'"
)
elif self.io_backend == "kernel_ascend":
from sglang.srt.mem_cache.pool_host.npu_memfabric import (
ascendc_io_enabled,
to_device_no_sync,
)
if ascendc_io_enabled():
# The fused acc_offload kv_exchange kernel reads the token
# indices directly on the device; keeping them there avoids
# the D2H sync that would serialize the layer-group pipeline.
# (The legacy memcpy2d exchange op still wants CPU indices and
# converts them itself.)
# Upload through pinned memory: host_indices comes from the
# radix-tree match as a pageable CPU tensor, and a pageable
# .to(device) completes with a stream synchronize that drains
# all compute queued on the current (default) stream.
if host_indices.device != self.device:
host_indices = to_device_no_sync(host_indices, self.device)
return host_indices, device_indices
return host_indices, device_indices.cpu()
else:
raise ValueError(f"Unsupported io backend")
python/sglang/srt/mem_cache/pool_host/mha.py
+14 -4
View File
@@ -359,12 +359,18 @@ 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_pool.k_buffer,
device_k=device_k,
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
device_v=device_v,
host_v=self.v_buffer,
page_size=self.page_size,
direction=TransferDirection.H2D,
@@ -491,12 +497,16 @@ 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_pool.k_buffer,
device_k=device_k,
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
device_v=device_v,
host_v=self.v_buffer,
page_size=self.page_size,
direction=TransferDirection.D2H,
python/sglang/srt/mem_cache/pool_host/mla.py
+375 -27
View File
@@ -30,6 +30,14 @@ from sglang.srt.mem_cache.pool_host.common import (
get_allocator_from_storage,
)
from sglang.srt.mem_cache.pool_host.hisparse import HiSparseHostPoolMixin
from sglang.srt.mem_cache.pool_host.npu_memfabric import (
alloc_with_memfabric,
ascendc_io_enabled,
ensure_memfabric_capacity,
memfabric_host_memory_enabled,
to_device_no_sync,
track_pinned_staging,
)
from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
_is_cuda = is_cuda()
@@ -116,9 +124,10 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
element_size=self.kv_cache_dim * self.dtype.itemsize
)
if self.layout == "page_first":
if self.layout in ("page_first", "page_first_kv_split"):
# Transpose [page, layer, ...] -> [layer, page, ...] to get per-layer views
# This swaps strides without copying data
# This swaps strides without copying data. For page_first_kv_split,
# kv_buffer is the k_buffer with the same page-major dims.
transposed = self.kv_buffer.transpose(0, 1)
self.data_refs = [transposed[i] for i in range(self.layer_num)]
else:
@@ -206,19 +215,50 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
if self._is_dummy:
return [], [], []
data_ptrs = [int(self.data_ptrs[i].item()) for i in range(self.layer_num)]
data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)]
if self.layout == "page_first_kv_split":
# data_refs are per-layer views of the k_buffer (page-major), so
# take the per-layer slab size instead of kv_buffer[i] (a page slab).
data_lens = [x.nbytes for x in self.data_refs]
else:
data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)]
item_lens = [self.token_stride_size * self.page_size] * self.layer_num
return data_ptrs, data_lens, item_lens
def get_size_per_token(self):
self.kv_lora_rank = self.device_pool.kv_lora_rank
self.qk_rope_head_dim = self.device_pool.qk_rope_head_dim
# FP8 DSA packs K/V into the single device k_buffer (device v_buffer is
# empty and never transferred). Exposed for the L3 store to skip the
# dead v component when generating per-page keys/pointers.
self.dsa_kv_cache_store_fp8 = getattr(
self.device_pool, "dsa_kv_cache_store_fp8", False
)
self.target_layer_num = self._effective_host_layer_num()
self.layer_num = self.target_layer_num + len(self.mtp_draft_device_pools)
self.kv_cache_dim = self.override_kv_cache_dim or (
self.kv_lora_rank + self.qk_rope_head_dim
)
return self.kv_cache_dim * self.dtype.itemsize * self.layer_num
size_per_token = self.kv_cache_dim * self.dtype.itemsize * self.layer_num
if (
self.layout == "page_first_kv_split"
and self.device_pool.index_head_dim is not None
):
# Indexer buffers only exist for physical Indexer layers, which can
# be a subset of all layers (e.g. GLM 5.2: 21 of 78). Mirror the
# layer count used by init_kv_buffer so host capacity sizing stays
# consistent with the actually-allocated buffers.
num_indexer_layers = getattr(self.device_pool, "num_indexer_layers", None)
if num_indexer_layers is None:
num_indexer_layers = self.layer_num
size_per_token += (
self.device_pool.index_head_dim
* self.dtype.itemsize
* num_indexer_layers
)
if getattr(self.device_pool, "index_k_scale_buffer", None) is not None:
# FP32 quantization scale per token per indexer layer.
size_per_token += 4 * num_indexer_layers
return size_per_token
def get_ksize_per_token(self):
return self.get_size_per_token()
@@ -249,15 +289,52 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
# Ascend-specific: Aligns with NPUMLATokenToKVPool 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,
)
base_dims = (self.page_num, self.layer_num, self.page_size, 1)
# Indexer buffers only exist for physical Indexer layers, which can
# be a subset of all layers (e.g. GLM 5.2: 21 of 78). The device
# pool packs them as (num_indexer_layers, page, ...); mirror that
# layer count here so transfer_kv_dim_exchange's layer check
# (device dim0 == host dim1) holds.
num_indexer_layers = getattr(self.device_pool, "num_indexer_layers", None)
if num_indexer_layers is None:
num_indexer_layers = self.layer_num
indexer_dims = (self.page_num, num_indexer_layers, self.page_size, 1)
alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device]
if getattr(self.device_pool, "dsa_kv_cache_store_fp8", False):
# FP8 DSA packs latent+RoPE+scale into the device k_buffer;
# mirror the packed width so the 2D memcpy row width matches.
k_width = self.device_pool.kv_cache_dim
else:
k_width = self.kv_lora_rank
if _is_npu and memfabric_host_memory_enabled():
# Memfabric-mapped host DRAM (single switch): required by the
# AIV sparse-copy IO path and usable by the legacy memcpy2d
# path unchanged. Size the GB-aligned reserve with the
# combined bytes of all buffers allocated below.
total_bytes = (
self.page_num
* self.page_size
* self.layer_num
* (k_width + self.qk_rope_head_dim)
* self.dtype.itemsize
)
if self.device_pool.index_head_dim is not None:
total_bytes += (
self.page_num
* self.page_size
* num_indexer_layers
* self.device_pool.index_head_dim
* self.dtype.itemsize
)
if getattr(self.device_pool, "index_k_scale_buffer", None) is not None:
# FP32 scale mirror
total_bytes += (
self.page_num * self.page_size * num_indexer_layers * 4
)
ensure_memfabric_capacity(total_bytes, torch.npu.current_device())
alloc_func = alloc_with_memfabric
self.k_buffer = alloc_func(
(*base_dims, self.kv_lora_rank),
(*base_dims, k_width),
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
@@ -273,12 +350,24 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
self.index_k_buffer = None
if self.device_pool.index_head_dim is not None:
self.index_k_buffer = alloc_func(
(*base_dims, self.device_pool.index_head_dim),
(*indexer_dims, self.device_pool.index_head_dim),
dtype=self.dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
# Host-side mirror of the NPU quantized-Indexer FP32 scale cache
# (see NPUMLATokenToKVPool.index_k_scale_buffer). Only present when
# the device pool carries one (FP8 DSA + npu_quant_lightning_indexer).
self.index_k_scale_buffer = None
if getattr(self.device_pool, "index_k_scale_buffer", None) is not None:
self.index_k_scale_buffer = alloc_func(
(*indexer_dims, 1),
dtype=torch.float32,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
# Return k_buffer to preserve original kv_buffer and data_refs init logic,
# though Ascend doesn't use these parameters.
return self.k_buffer
@@ -333,6 +422,221 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
device=self.device_pool.device,
)
def _indexer_slot_range_for_layer(self, device_pool, device_layer_id):
"""Map a device layer onto the indexer slot space.
Returns ``(slot, 1)`` when the layer is an indexer layer, ``(0, 0)``
to skip the indexer components, or ``(0, -1)`` (all indexer layers)
when the mapping is unavailable. ``indexer_layer_ids`` holds
absolute (PP-global) layer ids while ``device_layer_id`` is
pool-local; convert before matching.
"""
indexer_layer_ids = getattr(device_pool, "indexer_layer_ids", None)
if not indexer_layer_ids or self.index_k_buffer is None:
return 0, -1
start_layer = getattr(device_pool, "start_layer", 0)
for slot, layer in enumerate(indexer_layer_ids):
if layer - start_layer == device_layer_id:
return slot, 1
return 0, 0
def _transfer_ascendc_sparse_copy(
self,
device_pool,
host_indices,
device_indices,
direction: TransferDirection,
layer_start: int = 0,
layer_num: int = -1,
index_k_layer_start: int = 0,
index_k_layer_num: int = -1,
) -> None:
"""One-shot KV transfer via the Memfabric acc_offload fused AIV kernel.
Sends a compact metadata array (per-component layout pitches and
layer ranges) plus the device-resident token indices to the acc_offload
``kv_exchange_copy`` kernel, which derives every (page, layer, split)
block address on the device: no (src, dst, len) entry table is built
on the host and the indices never round-trip through the CPU, so the
transfer launch does not synchronize the load stream.
The layer range arguments restrict the transfer to a single layer
(per-layer H2D pipelining); the defaults transfer everything (used by
the one-shot D2H backup path).
The kernel de-references host pool pointers directly, so the host
pool must be Memfabric-mapped.
"""
device = device_pool.k_buffer.device
# The kernel reads the token indices directly from device memory.
# Upload without a stream sync: a plain .to(device) from pageable
# memory synchronizes the stream and would serialize the pipeline.
if host_indices.device.type != "npu":
host_indices = to_device_no_sync(host_indices, device)
if device_indices.device.type != "npu":
device_indices = to_device_no_sync(device_indices, device)
# The kernel runs on the current (load) stream while the indices were
# allocated on another stream; keep them alive until the copy retires.
stream = torch.npu.current_stream()
host_indices.record_stream(stream)
device_indices.record_stream(stream)
comps = self._build_ascendc_component_meta(
device_pool,
layer_start,
layer_num,
index_k_layer_start,
index_k_layer_num,
)
num_pages = host_indices.numel() // self.page_size
direction_value = (
direction.value
if isinstance(direction, TransferDirection)
else int(direction)
)
vals = [
len(comps),
num_pages,
self.page_size,
direction_value,
device_indices.data_ptr(),
host_indices.data_ptr(),
0, # host_layout_mode: 0 = page-first
0,
]
for comp in comps:
vals.extend(comp)
self._launch_ascendc_kv_exchange(vals, device)
@staticmethod
def _ascendc_comp_meta(dev_t, host_t, lo, hi):
"""Build the 9-int metadata tuple for one KV component."""
itemsize = dev_t.dtype.itemsize
width = 1
for dim in dev_t.shape[2:]:
width *= dim
# host: [page, layer, page_size, 1, width]
host_page_stride = host_t.stride(0) * itemsize
host_layer_stride = host_t.stride(1) * itemsize
return (
dev_t.data_ptr(),
host_t.data_ptr(),
# device is always layer-first
dev_t.stride(0) * itemsize,
dev_t.stride(1) * itemsize,
host_page_stride,
host_layer_stride,
width * itemsize,
lo,
hi,
)
def _build_ascendc_component_meta(
self,
device_pool,
layer_start: int,
layer_num: int,
index_k_layer_start: int,
index_k_layer_num: int,
) -> list:
"""Assemble the per-component metadata list for kv_exchange_copy."""
k_lo = layer_start
k_hi = device_pool.k_buffer.shape[0] if layer_num < 0 else k_lo + layer_num
# Both pools must share the layer index space (same limitation as the
# legacy memcpy2d exchange op); catches e.g. MTP draft pools, whose
# host rows live past the main pool's layers.
host_layer_num = self.k_buffer.shape[1]
device_layer_num = device_pool.k_buffer.shape[0]
if k_hi > host_layer_num or k_hi > device_layer_num:
raise RuntimeError(
f"AscendC kv_exchange layer range [{k_lo}, {k_hi}) exceeds the "
f"pool layer space (device={device_layer_num}, "
f"host={host_layer_num})"
)
comp_meta = self._ascendc_comp_meta
comps = [comp_meta(device_pool.k_buffer, self.k_buffer, k_lo, k_hi)]
# FP8 DSA packs V into the device k_buffer; the device v_buffer is
# empty and must be skipped.
if device_pool.v_buffer.numel() > 0 and self.v_buffer.numel() > 0:
comps.append(comp_meta(device_pool.v_buffer, self.v_buffer, k_lo, k_hi))
device_index_k = getattr(device_pool, "index_k_buffer", None)
if self.index_k_buffer is not None and device_index_k is not None:
if index_k_layer_num < 0:
ik_lo, ik_hi = 0, self.index_k_buffer.shape[1]
else:
ik_lo = index_k_layer_start
ik_hi = index_k_layer_start + index_k_layer_num
if ik_hi > ik_lo:
comps.append(
comp_meta(device_index_k, self.index_k_buffer, ik_lo, ik_hi)
)
device_scale = getattr(device_pool, "index_k_scale_buffer", None)
if self.index_k_scale_buffer is not None and device_scale is not None:
comps.append(
comp_meta(
device_scale,
self.index_k_scale_buffer,
ik_lo,
ik_hi,
)
)
if len(comps) > 4:
raise RuntimeError(
f"AscendC kv_exchange supports at most 4 components, got {len(comps)}"
)
return comps
@staticmethod
def _rewrite_host_base_to_dva(vals: list) -> list:
"""Convert host VAs to device VAs for AIV de-referencing."""
from memfabric_hybrid import offload
_KV_EXCHANGE_META_HEADER = 8
_KV_EXCHANGE_META_STRIDE = 9
_KV_EXCHANGE_MAX_COMPONENTS = 4
_KV_EXCHANGE_HOST_BASE_OFFSET = 1
num_components = int(vals[0])
if num_components < 0 or num_components > _KV_EXCHANGE_MAX_COMPONENTS:
raise ValueError(
f"kv_exchange: invalid num_components {num_components} in meta"
)
for c in range(num_components):
idx = (
_KV_EXCHANGE_META_HEADER
+ _KV_EXCHANGE_META_STRIDE * c
+ _KV_EXCHANGE_HOST_BASE_OFFSET
)
host_base = int(vals[idx])
if host_base == 0:
continue
dva = offload.get_dva(host_base)
if dva == 0:
raise ValueError(
f"kv_exchange: get_dva failed for host_base 0x{host_base:x}"
)
if dva != host_base:
vals[idx] = dva
return vals
def _launch_ascendc_kv_exchange(self, vals: list, device) -> None:
"""Rewrite host bases to DVAs and launch the AIV sparse-copy kernel."""
from memfabric_hybrid import offload
vals = self._rewrite_host_base_to_dva(vals)
pinned_meta = torch.tensor(vals, dtype=torch.int64, pin_memory=True)
meta = torch.empty(pinned_meta.shape, dtype=torch.int64, device=device)
meta.copy_(pinned_meta, non_blocking=True)
track_pinned_staging(pinned_meta)
ret = offload.kv_exchange_copy(meta, device)
if ret != 0:
raise RuntimeError(f"offload.kv_exchange_copy failed with code {ret}")
def load_to_device_per_layer(
self,
device_pool,
@@ -415,20 +719,48 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
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 device_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,
device_index_k=device_pool.index_k_buffer,
host_index_k=self.index_k_buffer,
page_size=self.page_size,
direction=TransferDirection.H2D,
if _is_npu and ascendc_io_enabled():
# The per-layer complete(i) event recorded by the caller lets
# later layers' DMA overlap the current layer's compute.
ik_start, ik_num = self._indexer_slot_range_for_layer(
device_pool, device_layer_id
)
self._transfer_ascendc_sparse_copy(
device_pool,
host_indices,
device_indices,
TransferDirection.H2D,
layer_start=device_layer_id,
layer_num=1,
index_k_layer_start=ik_start,
index_k_layer_num=ik_num,
)
return
# transfer_kv_dim_exchange transfers all layers in one call;
# only invoke it on the first owned layer to avoid duplicate
# work on subsequent per-layer iterations.
if device_layer_id != 0:
return
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=getattr(
device_pool, "k_buffer_tensor", device_pool.k_buffer
),
host_k=self.k_buffer,
device_v=getattr(
device_pool, "v_buffer_tensor", device_pool.v_buffer
),
host_v=self.v_buffer,
device_index_k=device_pool.index_k_buffer,
host_index_k=self.index_k_buffer,
device_index_k_scale=getattr(
device_pool, "index_k_scale_buffer", None
),
host_index_k_scale=self.index_k_scale_buffer,
page_size=self.page_size,
direction=TransferDirection.H2D,
)
else:
raise ValueError(f"Unsupported layout: {self.layout}")
else:
@@ -611,15 +943,31 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
raise ValueError(f"Unsupported layout: {self.layout}")
elif io_backend == "kernel_ascend":
if self.layout == "page_first_kv_split":
if _is_npu and ascendc_io_enabled():
self._transfer_ascendc_sparse_copy(
device_pool,
host_indices,
device_indices,
TransferDirection.D2H,
)
return
transfer_kv_dim_exchange(
device_indices=device_indices,
host_indices=host_indices,
device_k=device_pool.k_buffer,
device_k=getattr(
device_pool, "k_buffer_tensor", device_pool.k_buffer
),
host_k=self.k_buffer,
device_v=device_pool.v_buffer,
device_v=getattr(
device_pool, "v_buffer_tensor", device_pool.v_buffer
),
host_v=self.v_buffer,
device_index_k=device_pool.index_k_buffer,
host_index_k=self.index_k_buffer,
device_index_k_scale=getattr(
device_pool, "index_k_scale_buffer", None
),
host_index_k_scale=self.index_k_scale_buffer,
page_size=self.page_size,
direction=TransferDirection.D2H,
)
python/sglang/srt/mem_cache/pool_host/npu_memfabric.py
+159
View File
@@ -0,0 +1,159 @@
"""Memfabric-mapped host DRAM and AscendC sparse-copy IO path (NPU only).
torch pin_memory buffers are only reachable by the SDMA engine; AIV kernels
(e.g. offload.sparse_copy) can only de-reference host VAs that were mapped
into the device VA space via the Memfabric offload entity (DRAM_MAP_HOST_VA,
see acc_offload_local_dram_entry.cpp). SGLANG_HICACHE_HOST_MEM_BACKEND=memfabric
is the single switch for the whole feature: the HiCache host pool is
allocated through memfabric_hybrid.offload.empty AND the L2<->L1 IO uses
the AIV sparse-copy kernel (see ascendc_io_enabled).
"""
from __future__ import annotations
import logging
import os
import torch
logger = logging.getLogger(__name__)
_MEMFABRIC_GB = 1024**3
_memfabric_state = {
"offload": None,
"initialized": False,
"reserved_bytes": 0,
"allocated_bytes": 0,
"device_id": None,
}
def memfabric_host_memory_enabled() -> bool:
"""Single switch for the Memfabric host pool + AscendC IO path."""
return os.environ.get("SGLANG_HICACHE_HOST_MEM_BACKEND", "").lower() == (
"memfabric"
)
def ascendc_io_enabled() -> bool:
"""Use the acc_offload AIV sparse-copy kernel for HiCache L2<->L1 IO.
Rides on the single memfabric switch: SGLANG_HICACHE_HOST_MEM_BACKEND=
memfabric enables both the host pool allocation and this IO path (the
AIV kernel de-references host pool pointers, which requires
Memfabric-mapped memory).
"""
return memfabric_host_memory_enabled()
def _get_memfabric_offload():
if _memfabric_state["offload"] is None:
try:
from memfabric_hybrid import offload
except ImportError as exc:
raise ImportError(
"SGLANG_HICACHE_HOST_MEM_BACKEND=memfabric requires "
"the memfabric_hybrid package (provides the acc_offload host "
"memory allocator). Install it or unset the env var."
) from exc
_memfabric_state["offload"] = offload
return _memfabric_state["offload"]
def ensure_memfabric_capacity(total_bytes: int, device_id: int) -> None:
"""Lazily initialize the Memfabric offload entity, sized by total_bytes.
total_bytes is the combined size of all buffers the calling host pool is
about to allocate (ultimately derived from --hicache-size /
--hicache-ratio). The entity is sized by the first declaration; later
host pools in the same process must fit into what is left.
The C++ side aligns the reservation up to whole GBs, so the physical
reservation may be up to ~1GB larger than the value passed here.
"""
offload = _get_memfabric_offload()
if not _memfabric_state["initialized"]:
config = offload.OffloadConfig()
config.device_id = device_id
config.reserve_size = total_bytes
config.alloc_size = total_bytes
config.flags = offload.OFFLOAD_FLAG_URMA_POOL
config.scene = offload.Scene.LOCAL
assert offload.initialize(config) == 0, "offload.initialize failed"
_memfabric_state.update(
initialized=True, reserved_bytes=total_bytes, device_id=device_id
)
logger.info(
"[HiCache] memfabric host memory initialized: reserve=%.2fGB device=%d "
"(physically reserved up to %dGB after C++-side GB alignment)",
total_bytes / _MEMFABRIC_GB,
device_id,
(total_bytes + _MEMFABRIC_GB - 1) // _MEMFABRIC_GB,
)
remaining = _memfabric_state["reserved_bytes"] - _memfabric_state["allocated_bytes"]
if total_bytes > remaining:
raise RuntimeError(
f"memfabric host memory exhausted: need "
f"{total_bytes / _MEMFABRIC_GB:.2f}GB, "
f"only {remaining / _MEMFABRIC_GB:.2f}GB left of the "
f"{_memfabric_state['reserved_bytes'] / _MEMFABRIC_GB:.2f}GB reserve "
"(sized automatically from the first L2 host pool, i.e. from "
"--hicache-size / --hicache-ratio). All host pools of the "
"process share this reserve."
)
def alloc_with_memfabric(
dims: tuple,
dtype: torch.dtype,
device: str,
pin_memory: bool,
allocator: None,
) -> torch.Tensor:
"""Allocate host tensor backed by Memfabric-mapped DRAM (AIV-de-referencable)."""
offload = _get_memfabric_offload()
numel = 1
for d in dims:
numel *= d
tensor = offload.empty(list(dims), dtype=dtype)
_memfabric_state["allocated_bytes"] += numel * dtype.itemsize
return tensor
# ---------------------------------------------------------------------------
# Sync-free H2D upload (NPU)
#
# A pageable .to(device) / torch.tensor(..., device=npu) completes with an
# aclrtStreamSynchronize that drains EVERYTHING queued on the current stream
# (profiler-confirmed on the HiCache load path). When called on the default
# stream before entering the load/write stream, that sync stalls all queued
# compute; on the load stream it serializes the layer-group pipeline. Stage
# through pinned memory instead: pinned + non_blocking=True is a genuinely
# async enqueue. The pinned staging tensors are kept alive until their
# consumer copy retires, tracked via events (Event.query() is host-side and
# never synchronizes).
# ---------------------------------------------------------------------------
_pinned_inflight: list = []
def track_pinned_staging(pinned: torch.Tensor) -> None:
"""Keep a pinned staging tensor alive until its async consumer retires.
Completed entries are dropped on each call so the list stays small.
"""
done = torch.npu.Event()
done.record()
_pinned_inflight.append((pinned, done))
_pinned_inflight[:] = [e for e in _pinned_inflight if not e[1].query()]
def to_device_no_sync(cpu_tensor: torch.Tensor, device) -> torch.Tensor:
"""Upload a CPU tensor to the NPU without synchronizing the stream.
NPU-only helper (torch.npu.Event); callers are on the AscendC IO path.
"""
pinned = cpu_tensor.pin_memory()
out = torch.empty(pinned.shape, dtype=pinned.dtype, device=device)
out.copy_(pinned, non_blocking=True)
track_pinned_staging(pinned)
return out