diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d3eb3470d..c20d8602b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -166,6 +166,7 @@ class Envs: SGLANG_TORCH_PROFILER_DIR = EnvStr("/tmp") SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS = EnvInt(500) SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE = EnvInt(64) + SGLANG_NATIVE_MOVE_KV_CACHE = EnvBool(False) # Scheduler: memory leak test SGLANG_TEST_RETRACT = EnvBool(False) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 99550e64a..271a2eb34 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -17,8 +17,10 @@ from __future__ import annotations import dataclasses from dataclasses import dataclass +from typing import List from sglang.srt.configs.mamba_utils import BaseLinearStateParams +from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa import index_buf_accessor from sglang.srt.layers.attention.nsa.quant_k_cache import quantize_k_cache from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter @@ -797,6 +799,10 @@ class MHATokenToKVPool(KVCache): self.v_buffer[layer_id - self.start_layer][loc] = cache_v def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + if envs.SGLANG_NATIVE_MOVE_KV_CACHE.get(): + move_kv_cache_native(self.k_buffer, self.v_buffer, tgt_loc, src_loc) + return + N = tgt_loc.numel() if N == 0: return @@ -1893,6 +1899,22 @@ class DoubleSparseTokenToKVPool(KVCache): self.label_buffer[layer_id - self.start_layer][loc] = cache_label +def move_kv_cache_native( + k_buffer: List[torch.Tensor], + v_buffer: List[torch.Tensor], + tgt_loc: torch.Tensor, + src_loc: torch.Tensor, +): + if tgt_loc.numel() == 0: + return + + tgt_loc_flat = tgt_loc.view(-1).long() + src_loc_flat = src_loc.view(-1).long() + for k_cache, v_cache in zip(k_buffer, v_buffer): + k_cache[tgt_loc_flat] = k_cache[src_loc_flat] + v_cache[tgt_loc_flat] = v_cache[src_loc_flat] + + @triton.jit def copy_all_layer_kv_cache_tiled( data_ptrs,