From 209654c42090a7a33896a2849b33ab3a50af32c5 Mon Sep 17 00:00:00 2001 From: James <445169590@qq.com> Date: Fri, 11 Sep 2026 00:23:36 +0800 Subject: [PATCH] [NPU][Hicache] Optimize HiCache L2 IO with Memfabric acc_offload (#38826) --- .../hardware_backend/npu/memory_pool_npu.py | 27 +- .../sglang/srt/hardware_backend/npu/utils.py | 7 +- .../sglang/srt/managers/cache_controller.py | 18 + python/sglang/srt/mem_cache/pool_host/mha.py | 18 +- python/sglang/srt/mem_cache/pool_host/mla.py | 402 ++++++++++++++++-- .../srt/mem_cache/pool_host/npu_memfabric.py | 159 +++++++ 6 files changed, 589 insertions(+), 42 deletions(-) create mode 100644 python/sglang/srt/mem_cache/pool_host/npu_memfabric.py diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index a562c7ccd..59b296707 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -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 diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index 32e4fe577..a1f782edf 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/python/sglang/srt/hardware_backend/npu/utils.py @@ -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", diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index d42980ca9..70892adf4 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -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") diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 8af8e8089..2bb1d3776 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/python/sglang/srt/mem_cache/pool_host/mha.py @@ -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, diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py index f3762c064..1d394aa95 100644 --- a/python/sglang/srt/mem_cache/pool_host/mla.py +++ b/python/sglang/srt/mem_cache/pool_host/mla.py @@ -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, ) diff --git a/python/sglang/srt/mem_cache/pool_host/npu_memfabric.py b/python/sglang/srt/mem_cache/pool_host/npu_memfabric.py new file mode 100644 index 000000000..84605fb7d --- /dev/null +++ b/python/sglang/srt/mem_cache/pool_host/npu_memfabric.py @@ -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