diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index c0a2144dc..07ab21e2d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -923,6 +923,7 @@ class Envs: SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False) SGLANG_NPU_USE_MULTI_STREAM = EnvBool(False) SGLANG_NPU_USE_MLAPO = EnvBool(False) + SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD = EnvBool(False) # Forward native implementation for activation gelu tanh for model Skywork-Reward-Gemma-2-27B-v0.2 SGLANG_NPU_FORWARD_NATIVE_GELUTANH = EnvBool(False) # Forward native implementation for gemma rms norm for model Skywork-Reward-Gemma-2-27B-v0.2 diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index e8d5e1313..b316333ef 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -19,6 +19,10 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import ( is_fia_nz, is_mla_preprocess_enabled, ) +from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config import ( + get_sparsity_driven_kv_offload_sparse_context_len, + is_sparsity_driven_kv_offload_enabled, +) from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.radix_attention import AttentionType @@ -341,6 +345,29 @@ class AscendAttnBackend(AttentionBackend): self.req_to_token = model_runner.req_to_token_pool.req_to_token self.graph_mode = False self.use_fa = get_bool_env_var("ASCEND_USE_FA", "False") + self.enable_sparsity_driven_kv_offload = is_sparsity_driven_kv_offload_enabled( + model_config=model_runner.model_config, + use_mla_backend=model_runner.use_mla_backend, + ) + self.sparse_kv_manager = None + if self.enable_sparsity_driven_kv_offload: + from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.manager import ( + SparseKVCacheManager, + register_sparse_kv_manager, + ) + + self.sparse_kv_manager = SparseKVCacheManager( + model_runner.req_to_token_pool, + model_runner.token_to_kv_pool_allocator, + sparse_context_len=get_sparsity_driven_kv_offload_sparse_context_len( + model_config=model_runner.model_config + ), + ) + register_sparse_kv_manager(self.sparse_kv_manager) + logger.info( + "Sparsity-driven KV offload is enabled with manager %s.", + self.sparse_kv_manager, + ) self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") self.enable_torch_compile = get_flags().capture.enable_torch_compile self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens @@ -635,6 +662,14 @@ class AscendAttnBackend(AttentionBackend): ) -> ForwardMetadata: """Create and store the per-bs ForwardMetadata for CUDA graph capture.""" metadata = ForwardMetadata() + if self.enable_sparsity_driven_kv_offload: + from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.host_callback import ( + register_npu_host_callback_stream, + ) + + register_npu_host_callback_stream( + torch.npu.current_stream(self.device), self.device + ) metadata.block_tables = self.graph_metadata["block_tables"][:bs, :] if self.is_hybrid_swa: metadata.block_tables_swa = self.graph_metadata["block_tables_swa"][:bs, :] @@ -1198,6 +1233,23 @@ class AscendAttnBackend(AttentionBackend): k_rope=k_rope, ) if topk_indices is not None: + if self.enable_sparsity_driven_kv_offload: + from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.attention import ( + forward_sparsity_driven_kv_offload, + ) + + return forward_sparsity_driven_kv_offload( + self, + q, + k, + v, + layer, + forward_batch, + save_kv_cache, + q_rope, + k_rope, + topk_indices, + ) return self.forward_sparse( q, k, @@ -2551,6 +2603,23 @@ class AscendAttnBackend(AttentionBackend): # MLAPO does saving kv_cache save_kv_cache = False if topk_indices is not None: + if self.enable_sparsity_driven_kv_offload: + from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.attention import ( + forward_sparsity_driven_kv_offload, + ) + + return forward_sparsity_driven_kv_offload( + self, + q, + k, + v, + layer, + forward_batch, + save_kv_cache, + q_rope, + k_rope, + topk_indices, + ) return self.forward_sparse( q, k, 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 d893fba92..60fff6e85 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -552,33 +552,42 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): self.kv_lora_rank = kv_lora_rank self.qk_rope_head_dim = qk_rope_head_dim self.index_head_dim = index_head_dim + self.enable_sparsity_driven_kv_offload = ( + envs.SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD.get() + ) + if self.enable_sparsity_driven_kv_offload and self.index_head_dim is None: + raise ValueError("Sparsity-driven KV offload requires an index KV cache.") self.custom_mem_pool = None with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): # The padded slot 0 is used for writing dummy outputs from padded tokens. - self.k_buffer = torch.zeros( - ( - layer_num, - self.size // self.page_size + 1, - self.page_size, - 1, - self.kv_lora_rank, - ), - dtype=self.store_dtype, - device=self.device, - ) - self.v_buffer = torch.zeros( - ( - layer_num, - self.size // self.page_size + 1, - self.page_size, - 1, - self.qk_rope_head_dim, - ), - dtype=self.store_dtype, - device=self.device, - ) + if self.enable_sparsity_driven_kv_offload: + self.k_buffer = None + self.v_buffer = None + else: + self.k_buffer = torch.zeros( + ( + layer_num, + self.size // self.page_size + 1, + self.page_size, + 1, + self.kv_lora_rank, + ), + dtype=self.store_dtype, + device=self.device, + ) + self.v_buffer = torch.zeros( + ( + layer_num, + self.size // self.page_size + 1, + self.page_size, + 1, + self.qk_rope_head_dim, + ), + dtype=self.store_dtype, + device=self.device, + ) self.index_k_buffer = None if self.index_head_dim is not None: self.index_k_buffer = torch.zeros( @@ -596,22 +605,34 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): self._finalize_allocation_log(size) def get_kv_size_bytes(self): - assert hasattr(self, "k_buffer") - assert hasattr(self, "v_buffer") kv_size_bytes = 0 - for k_cache in self.k_buffer: - kv_size_bytes += get_tensor_size_bytes(k_cache) - for v_cache in self.v_buffer: - kv_size_bytes += get_tensor_size_bytes(v_cache) - if self.index_head_dim is not None: - assert hasattr(self, "index_k_buffer") + + if getattr(self, "k_buffer", None) is not None: + for k_cache in self.k_buffer: + kv_size_bytes += get_tensor_size_bytes(k_cache) + if getattr(self, "v_buffer", None) is not None: + for v_cache in self.v_buffer: + kv_size_bytes += get_tensor_size_bytes(v_cache) + if getattr(self, "index_k_buffer", None) is not None: for index_k_cache in self.index_k_buffer: kv_size_bytes += get_tensor_size_bytes(index_k_cache) return kv_size_bytes + def _raise_if_native_kv_cache_disabled(self): + if ( + getattr(self, "k_buffer", None) is None + or getattr(self, "v_buffer", None) is None + ): + raise RuntimeError( + "Native NPU MLA device KV cache is disabled; " + "k_buffer/v_buffer are not available. Use the sparse KV manager " + "path, or re-enable native device KV cache for this code path." + ) + def get_kv_buffer(self, layer_id: int): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + self._raise_if_native_kv_cache_disabled() return ( self.k_buffer[layer_id - self.start_layer], self.v_buffer[layer_id - self.start_layer], @@ -628,6 +649,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): def get_key_buffer(self, layer_id: int): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + self._raise_if_native_kv_cache_disabled() if self.store_dtype != self.dtype: return self.k_buffer[layer_id - self.start_layer].view(self.dtype) @@ -636,6 +658,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): def get_value_buffer(self, layer_id: int): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + self._raise_if_native_kv_cache_disabled() if self.store_dtype != self.dtype: return self.v_buffer[layer_id - self.start_layer].view(self.dtype) @@ -644,6 +667,8 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): def get_index_k_buffer(self, layer_id: int): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + if getattr(self, "index_k_buffer", None) is None: + raise RuntimeError("NPU MLA index KV cache is not allocated.") if self.store_dtype != self.dtype: return self.index_k_buffer[layer_id - self.start_layer].view(self.dtype) @@ -651,6 +676,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): # for disagg 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) @@ -681,6 +707,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): cache_v: torch.Tensor, ): loc, _, _ = unwrap_write_loc(loc_info) + self._raise_if_native_kv_cache_disabled() layer_id = layer.layer_id if cache_k.dtype != self.dtype: cache_k = cache_k.to(self.dtype) diff --git a/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/attention.py b/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/attention.py new file mode 100644 index 000000000..4c3f63cfc --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/attention.py @@ -0,0 +1,318 @@ +"""Ascend attention path backed by sparsity-driven KV offload.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch +import torch_npu + +from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.manager import ( + _wait_stream_event, + normalize_batch_topk_indices, +) +from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp + +if TYPE_CHECKING: + from sglang.srt.hardware_backend.npu.attention.ascend_backend import ( + AscendAttnBackend, + ) + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +def _get_sparse_kv_manager(backend: AscendAttnBackend): + if backend.sparse_kv_manager is None: + raise RuntimeError( + "Sparsity-driven KV offload is disabled or was not initialized." + ) + return backend.sparse_kv_manager + + +def _expand_dsa_sparse_indices(topk_indices: torch.Tensor) -> torch.Tensor: + """Expand [T, K] to [T, 1, K] for NPU sparse attention.""" + if topk_indices.dim() == 2: + return topk_indices.unsqueeze(-2) + return topk_indices + + +def forward_sparsity_driven_kv_offload( + backend: AscendAttnBackend, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, +): + """Run sparse attention using host-offloaded compact MLA KV.""" + del v + if q_rope is None or k_rope is None or topk_indices is None: + raise ValueError( + "Sparsity-driven KV offload requires q_rope, k_rope, and topk_indices." + ) + + is_prefill = forward_batch.forward_mode.is_extend_without_speculative() + + q_nope, q_pe = q, q_rope + k_nope = k.view(-1, layer.tp_k_head_num, backend.kv_lora_rank).contiguous() + k_pe = k_rope.view(-1, layer.tp_k_head_num, backend.qk_rope_head_dim).contiguous() + sparse_kv_manager = _get_sparse_kv_manager(backend) + stream = torch.npu.current_stream(backend.device) + + if save_kv_cache: + sparse_kv_manager.offload_v2(k_nope, k_pe, layer, forward_batch, stream) + + if is_prefill: + if backend.forward_metadata.actual_seq_lengths_q is not None: + actual_seq_qlen = backend.forward_metadata.actual_seq_lengths_q + else: + actual_seq_qlen = torch.cumsum(forward_batch.extend_seq_lens, dim=0) + elif backend.forward_metadata.actual_seq_lengths_q is None: + if ( + forward_batch.forward_mode.is_draft_extend_v2() + or forward_batch.forward_mode.is_target_verify() + ): + actual_seq_qlen = ( + torch.arange( + backend.speculative_num_draft_tokens, + backend.speculative_num_draft_tokens + q.shape[0], + backend.speculative_num_draft_tokens, + dtype=torch.int32, + ) + .to(q.device) + .to(torch.int32) + ) + else: + actual_seq_qlen = ( + torch.arange(1, q.shape[0] + 1).to(q.device).to(torch.int32) + ) + else: + actual_seq_qlen = backend.forward_metadata.actual_seq_lengths_q + + if backend.forward_metadata.actual_seq_lengths_kv is not None: + actual_seq_lengths_kv = backend.forward_metadata.actual_seq_lengths_kv + elif backend.forward_metadata.seq_lens_cpu_int is not None: + actual_seq_lengths_kv = backend.forward_metadata.seq_lens_cpu_int + else: + actual_seq_lengths_kv = backend.forward_metadata.seq_lens + + if ( + is_prefill + and is_dsa_enable_prefill_cp() + and forward_batch.attn_cp_metadata is not None + ): + attn_out = backend.do_cp_balance_attn( + q_nope, + k_nope, + q_pe, + k_pe, + topk_indices, + layer, + actual_seq_qlen, + actual_seq_lengths_kv, + ) + elif forward_batch.forward_mode.is_decode(): + batch_size = forward_batch.batch_size + num_kv_heads = layer.tp_k_head_num + num_query_heads = layer.tp_q_head_num + nope_head_dim = backend.kv_lora_rank + rope_head_dim = backend.qk_rope_head_dim + + topk_2d_input = normalize_batch_topk_indices(topk_indices) + effective_topk_length = topk_2d_input.shape[1] + selected_kv_length = sparse_kv_manager.sparse_context_len + if effective_topk_length <= 0: + raise RuntimeError("SFA BSND compact path expects a positive top-k length.") + if effective_topk_length > selected_kv_length: + raise RuntimeError( + "DSA top-k length exceeds sparse KV device cache capacity: " + f"topk_len={effective_topk_length}, " + f"sparse_context_len={selected_kv_length}." + ) + if effective_topk_length == selected_kv_length: + topk_2d = topk_2d_input.contiguous() + else: + topk_2d = torch.full( + (batch_size, selected_kv_length), + -1, + dtype=topk_2d_input.dtype, + device=topk_2d_input.device, + ) + topk_2d[:, :effective_topk_length] = topk_2d_input + topk_2d = topk_2d.contiguous() + + assert num_kv_heads == 1, ( + "FIA_v2 MLA selected KV path expects KV_N == 1, " + f"got num_kv_heads={num_kv_heads}" + ) + + padded_query_heads = q_nope.numel() // (batch_size * nope_head_dim) + assert padded_query_heads >= num_query_heads, ( + "query head count mismatch: " + f"padded_query_heads={padded_query_heads}, " + f"num_query_heads={num_query_heads}" + ) + + # Materialize the compact top-k KV for the current attention step: + # device-cache hits and host misses are copied into this buffer, the + # device cache is refilled from it, and sparse attention consumes it + # directly. + selected_kv_buffer = torch.zeros( + ( + batch_size, + selected_kv_length, + num_kv_heads, + nope_head_dim + rope_head_dim, + ), + dtype=k.dtype, + device=backend.device, + ) + sparse_kv_manager.materialize_selected_kv( + layer, forward_batch, topk_indices, selected_kv_buffer, stream + ) + + _wait_stream_event(stream, sparse_kv_manager.hit_done) + _wait_stream_event(stream, sparse_kv_manager.miss_done) + + selected_k_nope, selected_k_rope = selected_kv_buffer.split( + [nope_head_dim, rope_head_dim], dim=-1 + ) + + topk_valid = topk_2d >= 0 + if forward_batch.seq_lens is not None: + valid_rows = (forward_batch.seq_lens[:batch_size] > 0).view(batch_size, 1) + topk_valid = topk_valid & valid_rows + + actual_seq_lengths_kv = ( + topk_valid.sum(dim=1) + .clamp(min=1, max=selected_kv_length) + .to(device=q_nope.device, dtype=torch.int32) + .contiguous() + ) + actual_seq_lengths_query = torch.ones( + batch_size, dtype=torch.int32, device=q_nope.device + ).contiguous() + + compact_indices = ( + torch.arange(selected_kv_length, device=q_nope.device, dtype=torch.int32) + .view(1, 1, 1, selected_kv_length) + .expand(batch_size, 1, num_kv_heads, selected_kv_length) + .clone() + ) + compact_valid = topk_valid.view(batch_size, 1, 1, selected_kv_length).expand( + batch_size, 1, num_kv_heads, selected_kv_length + ) + sparse_indices = torch.where( + compact_valid, + compact_indices, + torch.full_like(compact_indices, -1), + ).contiguous() + + empty_rows = (topk_valid.sum(dim=1) == 0).view(batch_size, 1, 1) + sparse_indices[:, :, :, 0] = torch.where( + empty_rows.expand(batch_size, 1, num_kv_heads), + torch.zeros( + (batch_size, 1, num_kv_heads), + dtype=torch.int32, + device=q_nope.device, + ), + sparse_indices[:, :, :, 0], + ) + + q_nope_sfa = q_nope.view( + batch_size, 1, padded_query_heads, nope_head_dim + ).contiguous() + q_rope_sfa = q_pe.view( + batch_size, 1, padded_query_heads, rope_head_dim + ).contiguous() + k_nope_sfa = selected_k_nope.contiguous() + k_rope_sfa = selected_k_rope.contiguous() + + assert q_nope_sfa.shape == ( + batch_size, + 1, + padded_query_heads, + nope_head_dim, + ) + assert q_rope_sfa.shape == ( + batch_size, + 1, + padded_query_heads, + rope_head_dim, + ) + assert k_nope_sfa.shape == ( + batch_size, + selected_kv_length, + num_kv_heads, + nope_head_dim, + ) + assert k_rope_sfa.shape == ( + batch_size, + selected_kv_length, + num_kv_heads, + rope_head_dim, + ) + + ret = torch_npu.npu_sparse_flash_attention( + q_nope_sfa, + k_nope_sfa, + k_nope_sfa, + sparse_indices, + layer.scaling, + actual_seq_lengths_query=actual_seq_lengths_query, + actual_seq_lengths_kv=actual_seq_lengths_kv, + query_rope=q_rope_sfa, + key_rope=k_rope_sfa, + sparse_block_size=1, + layout_query="BSND", + layout_kv="BSND", + sparse_mode=0, + attention_mode=2, + return_softmax_lse=False, + ) + + _wait_stream_event(stream, sparse_kv_manager.refill_done) + _wait_stream_event(stream, sparse_kv_manager.slot_map_done) + + attn_out = ret[0] if isinstance(ret, tuple) else ret + attn_out = attn_out[:, :, :num_query_heads, :].reshape( + batch_size, num_query_heads * nope_head_dim + ) + else: + if is_prefill: + k_nope_sfa, k_pe_sfa = sparse_kv_manager.get_forward_kv( + layer, forward_batch, stream + ) + forward_actual_seq_lengths_kv = torch.cumsum(forward_batch.seq_lens, dim=0) + else: + k_nope_sfa, k_pe_sfa = k_nope, k_pe + forward_actual_seq_lengths_kv = actual_seq_lengths_kv + + topk_indices = _expand_dsa_sparse_indices(topk_indices) + attn_out, _, _ = torch_npu.npu_sparse_flash_attention( + query=q_nope, + key=k_nope_sfa, + value=k_nope_sfa, + query_rope=q_pe, + key_rope=k_pe_sfa, + sparse_indices=topk_indices, + scale_value=layer.scaling, + actual_seq_lengths_query=actual_seq_qlen.to( + device=q_nope.device, dtype=torch.int32 + ), + actual_seq_lengths_kv=forward_actual_seq_lengths_kv.to( + device=q_nope.device, dtype=torch.int32 + ), + sparse_block_size=1, + layout_query="TND", + layout_kv="TND", + sparse_mode=3, + attention_mode=2, + return_softmax_lse=False, + ) + + return attn_out diff --git a/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/config.py b/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/config.py new file mode 100644 index 000000000..31389c7d5 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/config.py @@ -0,0 +1,96 @@ +"""Configuration and validation for sparsity-driven KV offload.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from sglang.srt.configs.model_config import ( + get_dsa_index_head_dim, + get_dsa_index_topk, + is_deepseek_dsa, +) +from sglang.srt.environ import envs +from sglang.srt.runtime_context import attention_backends, get_schedule +from sglang.srt.utils.common import is_npu + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + + +def is_sparsity_driven_kv_offload_enabled( + *, + model_config: ModelConfig, + use_mla_backend: bool, +) -> bool: + if not envs.SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD.get(): + return False + + prefill_attention_backend, decode_attention_backend = attention_backends() + if not ( + is_npu() + and prefill_attention_backend == "ascend" + and decode_attention_backend == "ascend" + and use_mla_backend + and is_deepseek_dsa(model_config.hf_config) + ): + raise ValueError( + "SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD requires an NPU " + "DSA-family MLA model " + "(for example DeepSeek V3.2 or GLM-5.x) using the Ascend MLA " + "attention backend." + ) + if get_schedule().max_running_requests is None: + raise ValueError( + "SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD requires max_running_requests " + "to be set to bound the per-process host KV allocation." + ) + return True + + +def get_sparsity_driven_kv_offload_sparse_context_len( + *, + model_config: ModelConfig, +) -> int: + """Return the per-request on-device sparse KV window size.""" + sparse_context_len = int(get_dsa_index_topk(model_config.hf_config)) + if sparse_context_len <= 0: + raise ValueError( + "Sparsity-driven KV offload requires a positive DSA index_topk, " + f"got {sparse_context_len}." + ) + return sparse_context_len + + +def get_sparsity_driven_kv_offload_index_head_dim( + *, + model_config: ModelConfig, +) -> int: + index_head_dim = getattr(model_config, "index_head_dim", None) + if index_head_dim is None: + index_head_dim = get_dsa_index_head_dim(model_config.hf_config) + index_head_dim = int(index_head_dim) + if index_head_dim <= 0: + raise ValueError( + "Sparsity-driven KV offload requires a positive DSA index_head_dim, " + f"got {index_head_dim}." + ) + return index_head_dim + + +def get_sparsity_driven_kv_offload_cell_size( + *, + model_config: ModelConfig, + use_mla_backend: bool, + num_layers: int, + element_size: int, +) -> Optional[int]: + if not is_sparsity_driven_kv_offload_enabled( + model_config=model_config, + use_mla_backend=use_mla_backend, + ): + return None + + index_head_dim = get_sparsity_driven_kv_offload_index_head_dim( + model_config=model_config + ) + return index_head_dim * num_layers * element_size diff --git a/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/host_callback.py b/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/host_callback.py new file mode 100644 index 000000000..511ae1d60 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/host_callback.py @@ -0,0 +1,80 @@ +"""ACL host-callback support for sparsity-driven KV offload.""" + +import atexit +import threading + +import acl +import torch + +_REPORTERS = {} + + +def _get_stream_ptr(stream) -> int: + for attr in ("npu_stream", "stream_ptr", "cuda_stream"): + if hasattr(stream, attr): + value = getattr(stream, attr) + value = value() if callable(value) else value + return int(value) + raise RuntimeError("cannot get raw NPU stream ptr") + + +class _AclReportThread: + def __init__(self, device_index: int): + self.device_index = device_index + self.ready = threading.Event() + self.stop = threading.Event() + self.thread_id = None + self.streams = set() + self.thread = threading.Thread(target=self._loop, daemon=True) + self.thread.start() + self.ready.wait() + atexit.register(self.close) + + def _loop(self): + torch.npu.set_device(self.device_index) + try: + acl.rt.set_device(self.device_index) + except Exception: + pass + + self.thread_id = threading.current_thread().ident + self.ready.set() + + while not self.stop.is_set(): + acl.rt.process_report(100) + + def subscribe(self, stream): + stream_ptr = _get_stream_ptr(stream) + if stream_ptr in self.streams: + return + + ret = acl.rt.subscribe_report(self.thread_id, stream_ptr) + if ret != 0: + raise RuntimeError( + f"acl.rt.subscribe_report failed, ret={ret}, " + f"thread_id={self.thread_id}, stream_ptr={stream_ptr}" + ) + self.streams.add(stream_ptr) + + def close(self): + for stream_ptr in list(self.streams): + try: + acl.rt.unsubscribe_report(self.thread_id, stream_ptr) + except Exception: + pass + self.streams.clear() + self.stop.set() + self.thread.join(timeout=1) + + +def register_npu_host_callback_stream(stream, device): + device_index = device.index + if device_index is None: + device_index = torch.npu.current_device() + + reporter = _REPORTERS.get(device_index) + if reporter is None: + reporter = _AclReportThread(device_index) + _REPORTERS[device_index] = reporter + + reporter.subscribe(stream) diff --git a/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/manager.py b/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/manager.py new file mode 100644 index 000000000..527679a8b --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/sparsity_driven_kv_offload/manager.py @@ -0,0 +1,980 @@ +"""Sparsity-driven KV offload manager for the Ascend backend.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, List, Optional, Union + +import torch +from sgl_kernel_npu.sparsity_driven_kv_offload import ( + create_shm_tensor, + slot_map_lookup, + unidex_copy_inplace, +) + +from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE +from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator +from sglang.srt.mem_cache.memory_pool import ( + MLATokenToKVPool, + ReqToTokenPool, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter + +if TYPE_CHECKING: + import torch.npu + + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.managers.schedule_batch import Req + +logger = logging.getLogger(__name__) + + +def _record_stream_event(stream, event) -> None: + if hasattr(stream, "record_event"): + stream.record_event(event) + else: + event.record(stream) + + +def _wait_stream_event(stream, event) -> None: + if hasattr(stream, "wait_event"): + stream.wait_event(event) + else: + event.wait(stream) + + +def normalize_batch_topk_indices(topk_indices: torch.Tensor) -> torch.Tensor: + """Normalize DSA top-k indices to [batch, topk] for compact KV copies.""" + if topk_indices.dim() == 2: + return topk_indices + if topk_indices.dim() == 3 and topk_indices.shape[1] == 1: + return topk_indices[:, 0, :] + if ( + topk_indices.dim() == 4 + and topk_indices.shape[1] == 1 + and topk_indices.shape[2] == 1 + ): + return topk_indices[:, 0, 0, :] + raise RuntimeError( + "Sparsity-driven KV offload expects DSA top-k indices with shape " + f"[batch, topk], [batch, 1, topk], or [batch, 1, 1, topk], got " + f"{tuple(topk_indices.shape)}." + ) + + +class SparseKVCacheManager: + copy_stream = None + miss_shm_cpu_tensor: list = [] + miss_shm_dev_ptr: Optional[int] = None + miss_shm_shape: list = [] + miss_shm_dtype: list = [] + + def __init__( + self, + req_to_token_pool: ReqToTokenPool, + token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, + sparse_context_len: int, + ) -> None: + enable_memory_saver = False + memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=enable_memory_saver + ) + + # Include the padding row because real request IDs can equal the + # configured capacity when row 0 is reserved for graph padding. + self.size = int(req_to_token_pool.req_to_token.shape[0]) + self.max_context_len = req_to_token_pool.max_context_len + self.sparse_context_len = int(sparse_context_len) + if self.sparse_context_len <= 0: + raise ValueError( + "SparseKVCacheManager requires a positive sparse_context_len, " + f"got {self.sparse_context_len}." + ) + self.device = req_to_token_pool.device + paged_kv_cache = token_to_kv_pool_allocator.get_kvcache() + if not isinstance(paged_kv_cache, MLATokenToKVPool): + raise TypeError( + "SparseKVCacheManager requires an MLATokenToKVPool, " + f"got {type(paged_kv_cache).__name__}" + ) + self.paged_kv_cache = paged_kv_cache + self.start_layer = paged_kv_cache.start_layer + # MLA params + self.head_num = 1 + self.kv_lora_rank = self.paged_kv_cache.kv_lora_rank + self.qk_rope_head_dim = self.paged_kv_cache.qk_rope_head_dim + # kv_cache_dim = kv_lora_rank + qk_rope_head_dim + self.head_dim = ( + self.paged_kv_cache.kv_lora_rank + self.paged_kv_cache.qk_rope_head_dim + ) + self.store_dtype = self.paged_kv_cache.store_dtype + self.layer_num = self.paged_kv_cache.layer_num + self._materialize_d2d_hit_stream = torch.npu.Stream() + self._materialize_h2d_miss_stream = torch.npu.Stream() + self._materialize_refill_stream = torch.npu.Stream() + self._materialize_slot_map_stream = torch.npu.Stream() + + self.hit_done = torch.npu.Event() + self.miss_done = torch.npu.Event() + self.refill_done = torch.npu.Event() + self.slot_map_done = torch.npu.Event() + + # device KV buffer + try: + with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + # [bs, ctx_len, head_num, head_dim] for each layer + # The padded slot 0 is used for writing dummy outputs from padded tokens. + self.device_kv_buffer: list[torch.Tensor] = [ + torch.empty( + ( + self.size, + self.sparse_context_len, + self.head_num, + self.head_dim, + ), + dtype=self.store_dtype, + device=self.device, + ) + for _ in range(self.layer_num) + ] + except Exception as e: + self._raise_buffer_allocation_error("device_kv_buffer", e) + + try: + with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + # Reserve the last row for padded requests and ensure token index + # `max_context_len` is a valid sentinel column for masked writes. + # The row width is also aligned to eight int32 values (32 bytes). + self.device_slot_map: list[torch.Tensor] = [ + torch.full( + ( + self.size + 1, + (self.max_context_len // 8 + 1) * 8, + ), + -1, + dtype=torch.int32, + device=self.device, + ) + for _ in range(self.layer_num) + ] + # Pre-filled all-(-1) template used to reset the slot map with a + # fast device-to-device copy instead of an indexed fill. + self._device_slot_map_minus_one = torch.full_like( + self.device_slot_map[0], -1 + ) + except Exception as e: + self._raise_buffer_allocation_error("device_slot_map", e) + + # Host KV buffer + # [bs, ctx_len, head_num, head_dim] for each layer + # The padded slot 0 is used for writing dummy outputs from padded tokens. + self.host_kv_buffer: list[torch.Tensor] = [] + self.host_ptr_list: list[int] = [] + self.dev_ptr_list: list[int] = [] + + host_kv_shape = ( + self.size, + self.max_context_len, + self.head_num, + self.head_dim, + ) + logger.info("Sparse KV host buffer shape: %s", host_kv_shape) + device_id = torch.npu.current_device() + + try: + for layer_idx in range(self.layer_num): + shm_cpu_tensor, host_ptr, dev_ptr = create_shm_tensor( + shape=host_kv_shape, + dtype=self.store_dtype, + device_id=device_id, + name=f"host_kv_layer_{layer_idx}_rank_{device_id}", + ) + self.host_kv_buffer.append(shm_cpu_tensor) + self.host_ptr_list.append(host_ptr) + self.dev_ptr_list.append(dev_ptr) + except Exception as e: + self._raise_buffer_allocation_error("host_kv_buffer", e) + self.host_kv_ctx_len = torch.zeros( + (self.size, self.max_context_len), dtype=torch.int32, device="cpu" + ) + self.topk_indices_cpu = None + self.token_on_device_cpu = None + self.device_token_pos_cpu = None + self.current_req_indices_cpu = None + + self._device_cache_slot_ids = torch.arange( + self.sparse_context_len, dtype=torch.long, device=self.device + ) + self._slot_map_width = (self.max_context_len // 8 + 1) * 8 + + self._install_req_alloc_hook(req_to_token_pool) + + def _raise_buffer_allocation_error( + self, + buffer_name: str, + exc: Exception, + ) -> None: + raise RuntimeError( + "Failed to allocate sparse KV buffer " + f"{buffer_name}: req_capacity={self.size}, " + f"max_context_len={self.max_context_len}, " + f"sparse_context_len={self.sparse_context_len}. " + "The sparse KV request capacity may be too large; set a smaller " + "--max-running-requests for sparse KV offload." + ) from exc + + def init_req(self, req: Req) -> None: + if req.is_chunked > 0: + return + rid = req.req_pool_idx + if rid is None: + raise RuntimeError( + "Cannot initialize sparse KV state before allocating a request pool slot" + ) + current_len = len(req.origin_input_ids) + self.host_kv_ctx_len[rid] = current_len + self.reset_requests([rid]) + + def reset_requests(self, req_ids: List[int]) -> None: + if not req_ids: + return + + req_ids_tensor = torch.tensor( + req_ids, dtype=torch.long, device=self.device + ).contiguous() + for layer_idx in range(self.layer_num): + self.device_slot_map[layer_idx].index_fill_(0, req_ids_tensor, -1) + + def _install_req_alloc_hook(self, req_to_token_pool: ReqToTokenPool) -> None: + original_alloc = getattr( + req_to_token_pool, "_sparse_kv_original_alloc", req_to_token_pool.alloc + ) + setattr(req_to_token_pool, "_sparse_kv_original_alloc", original_alloc) + + def alloc_with_sparse_reset(reqs: list[Req]) -> Optional[List[int]]: + newly_allocated = [req.req_pool_idx is None for req in reqs] + req_pool_indices = original_alloc(reqs) + if req_pool_indices is not None: + self.reset_requests( + [ + req_pool_indices[i] + for i, is_new in enumerate(newly_allocated) + if is_new + ] + ) + return req_pool_indices + + setattr(req_to_token_pool, "alloc", alloc_with_sparse_reset) + + def offload( + self, + k: torch.Tensor, + k_rope: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + stream: torch.npu.Stream, + ): + layer_idx = layer.layer_id - self.start_layer + device = k.device + + # k: [total_token_slots, nhead, dim] + # k_rope: [total_token_slots, nhead, dim] + # kv_device: [total_token_slots, nhead, 2*dim] + kv_device = torch.cat([k, k_rope], dim=-1) + + # Source row indices into kv_device. + # In graph mode this tensor is expected to have a static shape. + src_token_indices = forward_batch.out_cache_loc.to(torch.long).contiguous() + static_token_slots = int(src_token_indices.numel()) + + if forward_batch.forward_mode.is_decode(): + # decode graph mode: + # one token slot per request, padded requests are masked out + req_ids = forward_batch.req_pool_indices.to(torch.long) + token_pos = (forward_batch.seq_lens - 1).to(torch.long) + + dst_token_indices = ( + req_ids * self.max_context_len + token_pos + ).contiguous() + + # Existing graph decode convention: + # padded decode requests carry seq_len == 1. + valid_mask = ( + (forward_batch.seq_lens != 1) + & (src_token_indices >= 0) + & (req_ids >= 0) + ).contiguous() + + else: + # prefill graph mode: + # assume out_cache_loc is laid out as [B, TOKENS_PER_REQ] flattened row-major + if ( + forward_batch.extend_seq_lens is None + or forward_batch.extend_prefix_lens is None + ): + raise RuntimeError( + "Sparse graph prefill offload requires extend_seq_lens and " + "extend_prefix_lens in ForwardBatch." + ) + + batch_size = int(forward_batch.req_pool_indices.shape[0]) + if batch_size <= 0: + return + + if static_token_slots % batch_size != 0: + raise RuntimeError( + f"out_cache_loc length {static_token_slots} is not divisible by " + f"batch size {batch_size}. Cannot infer graph static token layout." + ) + + tokens_per_req = static_token_slots // batch_size + + req_ids = forward_batch.req_pool_indices.to(torch.long) + extend_seq_lens = forward_batch.extend_seq_lens.to(torch.long) + extend_prefix_lens = forward_batch.extend_prefix_lens.to(torch.long) + + local_offsets = ( + torch.arange( + tokens_per_req, + device=device, + dtype=torch.long, + ) + .unsqueeze(0) + .expand(batch_size, tokens_per_req) + ) + + req_ids_2d = req_ids.unsqueeze(1).expand(batch_size, tokens_per_req) + dst_pos_2d = extend_prefix_lens.unsqueeze(1) + local_offsets + + dst_token_indices = ( + (req_ids_2d * self.max_context_len + dst_pos_2d) + .reshape(-1) + .contiguous() + ) + + valid_mask = ( + (local_offsets < extend_seq_lens.unsqueeze(1)).reshape(-1) + & (src_token_indices >= 0) + & (req_ids_2d.reshape(-1) >= 0) + ).contiguous() + + # Layout check: + # kv_device rows: [token_slot] + # host_kv_buffer[layer] rows: [req_id, seq_pos] + # block dims must match on [nhead, 2*dim] + assert kv_device.shape[1:] == self.host_kv_buffer[layer_idx].shape[2:] + # torch.npu.synchronize() + actual_stream = stream if stream is not None else torch.npu.current_stream() + with torch.npu.stream(actual_stream): + unidex_copy_inplace( + kv_device, + self.host_kv_buffer[layer_idx], + src_token_indices, + dst_token_indices, + valid_mask, + 1, # kv_device: [token_slot, nhead, 2*dim] + 2, # host_kv_buffer: [num_req, max_context_len, nhead, 2*dim] + block_dim=48, + dst_ptr=self.dev_ptr_list[layer_idx], + ) + # torch.npu.synchronize() + + def offload_v2( + self, + k: torch.Tensor, + k_rope: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + stream: torch.npu.Stream, + ): + """Offload compact per-forward KV rows into the sparse host KV buffer. + + v1 expects k/k_rope to be full native KV-cache views and therefore uses + forward_batch.out_cache_loc as source rows. v2 expects k/k_rope to be + compact rows produced by the current forward pass, so source rows are + simply [0, num_new_tokens). The native cache slot is kept only as + validity metadata. + + Keep src_tensor/dst_tensor/src_index/dst_index/valid_mask explicit so + the final copy can be swapped to a custom kernel without changing the + graph-friendly index construction. + """ + layer_idx = layer.layer_id - self.start_layer + device = k.device + + # k: [num_new_tokens, nhead, kv_lora_rank] + # k_rope: [num_new_tokens, nhead, qk_rope_head_dim] + # kv_device: [num_new_tokens, nhead, kv_lora_rank + qk_rope_head_dim] + src_tensor = torch.cat([k, k_rope], dim=-1).contiguous() + dst_tensor = self.host_kv_buffer[layer_idx] + src_index = torch.arange(src_tensor.shape[0], device=device, dtype=torch.long) + num_src_rows = int(src_tensor.shape[0]) + + if forward_batch.forward_mode.is_decode(): + req_ids = forward_batch.req_pool_indices.to(torch.long) + token_pos = (forward_batch.seq_lens - 1).to(torch.long) + cache_loc = forward_batch.out_cache_loc.to(torch.long) + + if int(req_ids.shape[0]) != num_src_rows: + raise RuntimeError( + "Sparse v2 decode offload expects compact KV rows to match " + f"batch size, got {num_src_rows} and {int(req_ids.shape[0])}." + ) + if int(cache_loc.shape[0]) != num_src_rows: + raise RuntimeError( + "Sparse v2 decode offload expects out_cache_loc rows to match " + f"compact KV rows, got {int(cache_loc.shape[0])} and " + f"{num_src_rows}." + ) + + dst_index = (req_ids * self.max_context_len + token_pos).contiguous() + valid_mask = ( + (forward_batch.seq_lens != 1) + & (cache_loc >= 0) + & (req_ids >= 0) + & (token_pos >= 0) + & (token_pos < self.max_context_len) + ).contiguous() + else: + if ( + forward_batch.extend_seq_lens is None + or forward_batch.extend_prefix_lens is None + ): + raise RuntimeError( + "Sparse v2 prefill offload requires extend_seq_lens and " + "extend_prefix_lens in ForwardBatch." + ) + + req_ids = forward_batch.req_pool_indices.to(torch.long) + extend_seq_lens = forward_batch.extend_seq_lens.to(torch.long) + extend_prefix_lens = forward_batch.extend_prefix_lens.to(torch.long) + + batch_size = int(req_ids.shape[0]) + if batch_size <= 0: + return + + if int(extend_seq_lens.shape[0]) != batch_size: + raise RuntimeError( + "Sparse v2 prefill offload expects extend_seq_lens to be " + f"padded to batch size {batch_size}, got " + f"{int(extend_seq_lens.shape[0])}." + ) + + prefix_len_size = int(extend_prefix_lens.shape[0]) + if prefix_len_size < batch_size: + extend_prefix_lens = torch.cat( + [ + extend_prefix_lens, + torch.zeros( + batch_size - prefix_len_size, + device=device, + dtype=torch.long, + ), + ], + dim=0, + ) + elif prefix_len_size > batch_size: + raise RuntimeError( + "Sparse v2 prefill offload expects extend_prefix_lens length " + f"<= batch size {batch_size}, got {prefix_len_size}." + ) + + if forward_batch.extend_seq_lens_cpu is not None: + extend_seq_lens_sum = int( + sum(forward_batch.extend_seq_lens_cpu[:batch_size]) + ) + else: + extend_seq_lens_sum = int(extend_seq_lens.sum().item()) + if extend_seq_lens_sum == num_src_rows: + # Chunk prefill emits compact rows as [req0 tokens][req1 tokens]... + # instead of graph-captured padded [B, tokens_per_req] rows. + seq_starts = torch.cumsum(extend_seq_lens, dim=0) - extend_seq_lens + flat_req_ids = torch.repeat_interleave( + req_ids, extend_seq_lens, output_size=num_src_rows + ) + flat_seq_starts = torch.repeat_interleave( + seq_starts, extend_seq_lens, output_size=num_src_rows + ) + flat_prefix_lens = torch.repeat_interleave( + extend_prefix_lens, extend_seq_lens, output_size=num_src_rows + ) + token_pos = ( + flat_prefix_lens + + torch.arange(num_src_rows, device=device, dtype=torch.long) + - flat_seq_starts + ) + dst_index = ( + flat_req_ids * self.max_context_len + token_pos + ).contiguous() + valid_mask = ( + (flat_req_ids >= 0) + & (token_pos >= 0) + & (token_pos < self.max_context_len) + ) + else: + if num_src_rows % batch_size != 0: + raise RuntimeError( + "Sparse v2 prefill offload expects either compact ragged " + "layout with rows=sum(extend_seq_lens) or graph-style " + f"row-major layout [B, tokens_per_req], got rows={num_src_rows}, " + f"batch={batch_size}, extend_seq_lens_sum={extend_seq_lens_sum}." + ) + + # Graph-friendly static layout: + # compact rows are interpreted as [batch_size, tokens_per_req]. + # Invalid padded columns are masked by local_offsets < extend_seq_lens. + tokens_per_req = num_src_rows // batch_size + + local_offsets = ( + torch.arange(tokens_per_req, device=device, dtype=torch.long) + .unsqueeze(0) + .expand(batch_size, tokens_per_req) + ) + req_ids_2d = req_ids.unsqueeze(1).expand(batch_size, tokens_per_req) + token_pos_2d = extend_prefix_lens.unsqueeze(1) + local_offsets + + dst_index = ( + (req_ids_2d * self.max_context_len + token_pos_2d) + .reshape(-1) + .contiguous() + ) + valid_mask = ( + (local_offsets < extend_seq_lens.unsqueeze(1)) + & (req_ids_2d >= 0) + & (token_pos_2d >= 0) + & (token_pos_2d < self.max_context_len) + ).reshape(-1) + + if ( + forward_batch.out_cache_loc is not None + and int(forward_batch.out_cache_loc.numel()) == num_src_rows + ): + valid_mask = valid_mask & ( + forward_batch.out_cache_loc.to(torch.long) >= 0 + ) + valid_mask = valid_mask.contiguous() + + assert src_tensor.shape[1:] == dst_tensor.shape[2:] + assert src_index.shape == dst_index.shape == valid_mask.shape + + actual_stream = stream if stream is not None else torch.npu.current_stream() + with torch.npu.stream(actual_stream): + unidex_copy_inplace( + src_tensor, + dst_tensor, + src_index, + dst_index, + valid_mask, + 1, # src_tensor: [num_new_tokens, nhead, head_dim] + 2, # dst_tensor: [num_req, max_context_len, nhead, head_dim] + block_dim=48, + dst_ptr=self.dev_ptr_list[layer_idx], + ) + + def get_forward_kv( + self, + layer: Union[RadixAttention, int], + forward_batch: ForwardBatch, + stream: Optional[torch.npu.Stream] = None, + ): + """Gather full request KV from sparse storage as compact TND tensors. + + This helper is intended for the prefill/extend sparse path. It returns + KV in the same request order as forward_batch.req_pool_indices: + [req0 tokens][req1 tokens]... . The SFA caller should pair this with + actual_seq_lengths_kv = cumsum(forward_batch.seq_lens). + + TODO: Add a prefill-resident device KV cache shaped like + [max_prefill_parallel_reqs, max_prefill_len, nhead, head_dim]. Reuse a + slot while the same req_id continues chunked prefill, and fully + overwrite it when a new req_id takes that slot. This mirrors the decode + device cache idea and avoids repeatedly copying the full prefix KV from + host for every chunk. + """ + layer_id = layer.layer_id if hasattr(layer, "layer_id") else int(layer) + layer_idx = layer_id - self.start_layer + if layer_idx < 0 or layer_idx >= self.layer_num: + raise RuntimeError( + f"Invalid sparse KV layer id {layer_id}; start_layer=" + f"{self.start_layer}, layer_num={self.layer_num}." + ) + + if forward_batch.req_pool_indices is None or forward_batch.seq_lens is None: + raise RuntimeError( + "get_forward_kv requires req_pool_indices and seq_lens in ForwardBatch." + ) + + device = ( + forward_batch.req_pool_indices.device + if forward_batch.req_pool_indices.device.type == "npu" + else self.device + ) + req_ids = forward_batch.req_pool_indices.to(device=device, dtype=torch.long) + seq_lens = forward_batch.seq_lens.to(device=device, dtype=torch.long) + + if int(req_ids.numel()) != int(seq_lens.numel()): + raise RuntimeError( + "get_forward_kv expects req_pool_indices and seq_lens to have the " + f"same length, got {int(req_ids.numel())} and " + f"{int(seq_lens.numel())}." + ) + + valid_reqs = (seq_lens > 0) & (req_ids >= 0) + req_ids = req_ids[valid_reqs].contiguous() + seq_lens = seq_lens[valid_reqs].contiguous() + + if int(seq_lens.numel()) == 0: + empty_nope = torch.empty( + (0, self.head_num, self.kv_lora_rank), + dtype=self.store_dtype, + device=device, + ) + empty_pe = torch.empty( + (0, self.head_num, self.qk_rope_head_dim), + dtype=self.store_dtype, + device=device, + ) + return empty_nope, empty_pe + + if bool((req_ids >= self.size).any().item()): + raise RuntimeError( + f"get_forward_kv got req id outside sparse pool size {self.size}." + ) + if bool((seq_lens > self.max_context_len).any().item()): + raise RuntimeError( + "get_forward_kv got seq_len larger than max_context_len " + f"{self.max_context_len}." + ) + + total_tokens = int(seq_lens.sum().item()) + kv_cat = torch.empty( + (total_tokens, self.head_num, self.head_dim), + dtype=self.store_dtype, + device=device, + ) + + seq_starts = torch.cumsum(seq_lens, dim=0) - seq_lens + src_req_ids = torch.repeat_interleave( + req_ids, seq_lens, output_size=total_tokens + ) + src_seq_starts = torch.repeat_interleave( + seq_starts, seq_lens, output_size=total_tokens + ) + token_pos = ( + torch.arange(total_tokens, device=device, dtype=torch.long) - src_seq_starts + ) + src_index = (src_req_ids * self.max_context_len + token_pos).contiguous() + dst_index = torch.arange(total_tokens, device=device, dtype=torch.long) + valid_mask = torch.ones(total_tokens, device=device, dtype=torch.bool) + + actual_stream = stream if stream is not None else torch.npu.current_stream() + with torch.npu.stream(actual_stream): + unidex_copy_inplace( + self.host_kv_buffer[layer_idx], + kv_cat, + src_index, + dst_index, + valid_mask, + 2, # host_kv_buffer: [num_req, max_context_len, nhead, head_dim] + 1, # kv_cat: [total_tokens, nhead, head_dim] + block_dim=48, + src_ptr=self.dev_ptr_list[layer_idx], + ) + + k_nope, k_pe = kv_cat.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + return k_nope.contiguous(), k_pe.contiguous() + + def materialize_selected_kv( + self, + layer: RadixAttention, + forward_batch: ForwardBatch, + topk_indices: torch.Tensor, + selected_kv_buffer: torch.Tensor, + stream: torch.npu.Stream, + ) -> None: + """Materialize top-k KV entries and refresh the device cache metadata.""" + layer_idx = layer.layer_id - self.start_layer + stream = stream if stream is not None else torch.npu.current_stream() + + with torch.npu.stream(stream): + # Route invalid requests to sentinel rows without changing graph shape. + # slot_map_row_indices: invalid -> self.size (reserved slot-map row) + # device_cache_row_indices: invalid -> 0 (masked by valid_topk_mask) + req_pool_indices = forward_batch.req_pool_indices + req_pool_indices = req_pool_indices.to(torch.long).contiguous() + valid_req_mask = (req_pool_indices >= 0) & (req_pool_indices < self.size) + slot_map_row_indices = torch.where( + valid_req_mask, + req_pool_indices, + torch.full_like(req_pool_indices, self.size), + ) + device_cache_row_indices = torch.where( + valid_req_mask, + req_pool_indices, + torch.zeros_like(req_pool_indices), + ) + + # Normalize top-k indices and mask invalid requests and token IDs. + topk_indices = normalize_batch_topk_indices(topk_indices) + batch_size, topk_len = topk_indices.shape + if topk_len > self.sparse_context_len: + raise RuntimeError( + "DSA top-k length exceeds sparse KV device cache capacity: " + f"topk_len={topk_len}, sparse_context_len={self.sparse_context_len}." + ) + valid_topk_mask = ( + (topk_indices >= 0) + & (topk_indices < self.max_context_len) + & valid_req_mask.unsqueeze(1) + ) + + # Query the slot map for device-cache hits and their slot positions. + slot_lookup_req_indices = slot_map_row_indices.to( + dtype=torch.int32 + ).contiguous() + slot_lookup_topk_indices = topk_indices.to(dtype=torch.int32).contiguous() + token_on_device, device_token_pos = slot_map_lookup( + self.device_slot_map[layer_idx], + slot_lookup_req_indices, + slot_lookup_topk_indices, + ) + token_on_device = token_on_device.to(torch.bool) & valid_topk_mask + + # Build copy indices on the main stream, then protect their use on + # the hit and miss streams with copy_ready. + hit_src_index, hit_dst_index, hit_valid_mask = _build_hit_src_dst_index( + token_on_device, + device_token_pos, + device_cache_row_indices, + self.sparse_context_len, + ) + + host_miss_mask = (~token_on_device) & valid_topk_mask + miss_src_index, miss_dst_index, miss_valid_mask = _build_miss_src_dst_index( + host_miss_mask, + topk_indices, + device_cache_row_indices, + self.max_context_len, + ) + + cache_slot_ids = self._device_cache_slot_ids[:topk_len] + request_cache_offsets = ( + device_cache_row_indices.unsqueeze(1) * self.sparse_context_len + ) + refill_src_index = torch.arange( + batch_size * topk_len, + dtype=torch.long, + device=topk_indices.device, + ) + refill_dst_index = ( + (request_cache_offsets + cache_slot_ids).reshape(-1).contiguous() + ) + refill_valid_mask = valid_topk_mask.reshape(-1).contiguous() + + copy_ready = torch.npu.Event() + _record_stream_event(stream, copy_ready) + + # Copy device-cache hits into the selected KV buffer. + with torch.npu.stream(self._materialize_d2d_hit_stream): + _wait_stream_event(self._materialize_d2d_hit_stream, copy_ready) + unidex_copy_inplace( + self.device_kv_buffer[layer_idx], + selected_kv_buffer, + hit_src_index, + hit_dst_index, + hit_valid_mask, + 2, + 2, # + block_dim=24, + ) + _record_stream_event(self._materialize_d2d_hit_stream, self.hit_done) + + # Copy host shared-memory misses into the selected KV buffer. + with torch.npu.stream(self._materialize_h2d_miss_stream): + _wait_stream_event(self._materialize_h2d_miss_stream, copy_ready) + unidex_copy_inplace( + self.host_kv_buffer[layer_idx], + selected_kv_buffer, + miss_src_index, + miss_dst_index, + miss_valid_mask, + 2, + 2, + block_dim=24, + src_ptr=self.dev_ptr_list[layer_idx], + ) + _record_stream_event(self._materialize_h2d_miss_stream, self.miss_done) + + # Refill the device cache with the current top-k after hit and miss + # copies complete, so the next step can reuse these entries. + with torch.npu.stream(self._materialize_refill_stream): + _wait_stream_event(self._materialize_refill_stream, self.hit_done) + _wait_stream_event(self._materialize_refill_stream, self.miss_done) + unidex_copy_inplace( + selected_kv_buffer, + self.device_kv_buffer[layer_idx], + refill_src_index, + refill_dst_index, + refill_valid_mask, + 2, + 2, + block_dim=24, + ) + _record_stream_event(self._materialize_refill_stream, self.refill_done) + + # Replace the slot-map row with the current top-k mapping. Invalid + # entries use the max_context_len sentinel column to preserve shape. + with torch.npu.stream(self._materialize_slot_map_stream): + _wait_stream_event(self._materialize_slot_map_stream, copy_ready) + self.device_slot_map[layer_idx].copy_(self._device_slot_map_minus_one) + + slot_map_token_indices = torch.where( + valid_topk_mask, + topk_indices.to(torch.long), + torch.full_like(topk_indices, self.max_context_len, dtype=torch.long), + ) + slot_map_slot_values = torch.where( + valid_topk_mask, + cache_slot_ids.to(torch.int32), + torch.full_like(cache_slot_ids, -1, dtype=torch.int32), + ) + slot_map_flat_indices = ( + slot_map_row_indices.unsqueeze(1) * self._slot_map_width + + slot_map_token_indices + ).reshape(-1) + unidex_copy_inplace( + slot_map_slot_values.reshape(-1, 1).contiguous(), + self.device_slot_map[layer_idx].view(-1, 1), + torch.arange( + slot_map_slot_values.numel(), + dtype=torch.long, + device=topk_indices.device, + ), + slot_map_flat_indices, + valid_topk_mask.reshape(-1), + 1, + 1, + block_dim=48, + ) + _record_stream_event(self._materialize_slot_map_stream, self.slot_map_done) + + +_global_sparse_kv_manager: Optional[SparseKVCacheManager] = None + + +def register_sparse_kv_manager(manager: SparseKVCacheManager) -> None: + global _global_sparse_kv_manager + _global_sparse_kv_manager = manager + + +def get_sparse_kv_manager() -> Optional[SparseKVCacheManager]: + return _global_sparse_kv_manager + + +def _build_hit_src_dst_index( + token_on_device: torch.Tensor, + device_token_pos: torch.Tensor, + current_req_indices: torch.Tensor, + sparse_context_len: int, +): + """ + token_on_device: [bs, topk], bool + device_token_pos: [bs, topk], int64 or int32 + current_req_indices: [bs], int64 + + Return: + src_index_full: [bs * topk], int64 + dst_index_full: [bs * topk], int64 + valid_mask: [bs * topk], bool + + Flattening rule: + src row = req_id * sparse_context_len + device_token_pos + dst row = batch_id * topk + topk_pos + """ + if token_on_device.dim() != 2 or device_token_pos.dim() != 2: + raise RuntimeError( + f"token_on_device and device_token_pos must be 2-D, got " + f"{token_on_device.dim()} and {device_token_pos.dim()}" + ) + if token_on_device.shape != device_token_pos.shape: + raise RuntimeError( + f"token_on_device and device_token_pos must have the same shape, got " + f"{tuple(token_on_device.shape)} and {tuple(device_token_pos.shape)}" + ) + if current_req_indices.dim() != 1: + raise RuntimeError( + f"current_req_indices must be 1-D, got {current_req_indices.dim()}" + ) + + bs, topk = token_on_device.shape + if current_req_indices.numel() != bs: + raise RuntimeError( + f"current_req_indices length mismatch: " + f"{current_req_indices.numel()} vs batch {bs}" + ) + if sparse_context_len <= 0: + raise RuntimeError( + f"sparse_context_len must be positive, got {sparse_context_len}" + ) + + device = token_on_device.device + + valid_mask = token_on_device.reshape(-1).contiguous() + + flat_dst_index_all = torch.arange( + bs * topk, + device=device, + dtype=torch.int64, + ) + + req_offsets = current_req_indices.to(torch.int64).unsqueeze(1) * sparse_context_len + src_index_2d = req_offsets + device_token_pos.to(torch.int64) + flat_src_index_all = src_index_2d.reshape(-1).contiguous() + + return flat_src_index_all, flat_dst_index_all, valid_mask + + +def _build_miss_src_dst_index( + token_from_host: torch.Tensor, + topk_indices: torch.Tensor, + current_req_indices: torch.Tensor, + max_context_len: int, +): + if token_from_host.dim() != 2 or topk_indices.dim() != 2: + raise RuntimeError( + f"token_from_host and topk_indices must be 2-D, got " + f"{token_from_host.dim()} and {topk_indices.dim()}" + ) + if token_from_host.shape != topk_indices.shape: + raise RuntimeError( + f"token_from_host and topk_indices must have the same shape, got " + f"{tuple(token_from_host.shape)} and {tuple(topk_indices.shape)}" + ) + if current_req_indices.dim() != 1: + raise RuntimeError( + f"current_req_indices must be 1-D, got {current_req_indices.dim()}" + ) + if current_req_indices.numel() != token_from_host.shape[0]: + raise RuntimeError( + f"current_req_indices length mismatch: " + f"{current_req_indices.numel()} vs batch {token_from_host.shape[0]}" + ) + + bs, topk = token_from_host.shape + device = token_from_host.device + + valid_2d = token_from_host & (topk_indices >= 0) & (topk_indices < max_context_len) + valid_mask = valid_2d.reshape(-1).contiguous() + + flat_dst_index_all = torch.arange( + bs * topk, + device=device, + dtype=torch.int64, + ) + + req_offsets = current_req_indices.to(torch.int64).unsqueeze(1) * max_context_len + src_index_2d = req_offsets + topk_indices.to(torch.int64) + flat_src_index_all = src_index_2d.reshape(-1).contiguous() + + return flat_src_index_all, flat_dst_index_all, valid_mask diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 4e70ad77f..8b5ae00e9 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -310,6 +310,21 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): dcp_size = get_parallel().attn_dcp_size if kvc.use_mla_backend: + if envs.SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD.get(): + # NPU sparse KV offload uses an index-only device pool. + from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config import ( + get_sparsity_driven_kv_offload_cell_size, + ) + + offload_cell_size = get_sparsity_driven_kv_offload_cell_size( + model_config=model_config, + use_mla_backend=kvc.use_mla_backend, + num_layers=num_layers, + element_size=kv_size, + ) + if offload_cell_size is not None: + return offload_cell_size + from sglang.srt.mem_cache.kv_cache_configurator import ( calculate_mla_kv_cache_dim, ) diff --git a/test/registered/unit/npu/test_sparsity_driven_kv_offload_config.py b/test/registered/unit/npu/test_sparsity_driven_kv_offload_config.py new file mode 100644 index 000000000..8a136121f --- /dev/null +++ b/test/registered/unit/npu/test_sparsity_driven_kv_offload_config.py @@ -0,0 +1,121 @@ +import os +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config import ( + get_sparsity_driven_kv_offload_cell_size, + get_sparsity_driven_kv_offload_sparse_context_len, + is_sparsity_driven_kv_offload_enabled, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +def _make_glm51_model_config(): + hf_config = SimpleNamespace( + architectures=["GlmMoeDsaForCausalLM"], + index_head_dim=128, + index_topk=1536, + ) + hf_config.get_text_config = lambda: hf_config + return SimpleNamespace( + hf_config=hf_config, + index_head_dim=128, + ) + + +class TestSparsityDrivenKVOffloadConfig(unittest.TestCase): + def test_glm_dsa_model_enables_sparse_kv_offload(self): + with ( + patch.dict( + os.environ, + {"SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD": "1"}, + ), + patch( + "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.is_npu", + return_value=True, + ), + patch( + "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.attention_backends", + return_value=("ascend", "ascend"), + ), + patch( + "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.get_schedule", + return_value=SimpleNamespace(max_running_requests=8), + ), + ): + model_config = _make_glm51_model_config() + + self.assertTrue( + is_sparsity_driven_kv_offload_enabled( + model_config=model_config, + use_mla_backend=True, + ) + ) + self.assertEqual( + get_sparsity_driven_kv_offload_sparse_context_len( + model_config=model_config + ), + 1536, + ) + self.assertEqual( + get_sparsity_driven_kv_offload_cell_size( + model_config=model_config, + use_mla_backend=True, + num_layers=2, + element_size=2, + ), + 512, + ) + + def test_split_attention_backend_rejects_sparse_kv_offload(self): + with ( + patch.dict( + os.environ, + {"SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD": "1"}, + ), + patch( + "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.is_npu", + return_value=True, + ), + patch( + "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.attention_backends", + return_value=("ascend", "torch_native"), + ), + ): + with self.assertRaisesRegex(ValueError, "Ascend MLA attention backend"): + is_sparsity_driven_kv_offload_enabled( + model_config=_make_glm51_model_config(), + use_mla_backend=True, + ) + + def test_missing_request_capacity_rejects_sparse_kv_offload(self): + with ( + patch.dict( + os.environ, + {"SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD": "1"}, + ), + patch( + "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.is_npu", + return_value=True, + ), + patch( + "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.attention_backends", + return_value=("ascend", "ascend"), + ), + patch( + "sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config.get_schedule", + return_value=SimpleNamespace(max_running_requests=None), + ), + ): + with self.assertRaisesRegex(ValueError, "max_running_requests"): + is_sparsity_driven_kv_offload_enabled( + model_config=_make_glm51_model_config(), + use_mla_backend=True, + ) + + +if __name__ == "__main__": + unittest.main()