diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 78f2ffabe..cce4be83a 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -314,8 +314,8 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--hicache-size` | The size of host KV cache memory pool in gigabytes, which will override the hicache_ratio if set. | `0` | Type: int | | `--hicache-write-policy` | The write policy of hierarchical cache. | `write_through` | `write_back`, `write_through`, `write_through_selective` | | `--radix-eviction-policy` | The eviction policy of radix trees. 'lru' stands for Least Recently Used, 'lfu' stands for Least Frequently Used. | `lru` | `lru`, `lfu` | -| `--hicache-io-backend` | The IO backend for KV cache transfer between CPU and GPU | `kernel` | `direct`, `kernel` | -| `--hicache-mem-layout` | The layout of host memory pool for hierarchical cache. | `layer_first` | `layer_first`, `page_first`, `page_first_direct` | +| `--hicache-io-backend` | The IO backend for KV cache transfer between CPU and GPU | `kernel` | `direct`, `kernel`, `kernel_ascend` | +| `--hicache-mem-layout` | The layout of host memory pool for hierarchical cache. | `layer_first` | `layer_first`, `page_first`, `page_first_direct`, `page_first_kv_split` | | `--hicache-storage-backend` | The storage backend for hierarchical KV cache. Built-in backends: file, mooncake, hf3fs, nixl, aibrix. For dynamic backend, use --hicache-storage-backend-extra-config to specify: backend_name (custom name), module_path (Python module path), class_name (backend class name). | `None` | `file`, `mooncake`, `hf3fs`, `nixl`, `aibrix`, `dynamic`, `eic` | | `--hicache-storage-prefetch-policy` | Control when prefetching from the storage backend should stop. | `best_effort` | `best_effort`, `wait_complete`, `timeout` | | `--hicache-storage-backend-extra-config` | A dictionary in JSON string format containing extra configuration for the storage backend. | `None` | Type: str | diff --git a/python/sglang/srt/layers/attention/ascend_backend.py b/python/sglang/srt/layers/attention/ascend_backend.py index 82526f0e8..9eb141fe1 100644 --- a/python/sglang/srt/layers/attention/ascend_backend.py +++ b/python/sglang/srt/layers/attention/ascend_backend.py @@ -403,6 +403,11 @@ class AscendAttnBackend(AttentionBackend): assert ( layer.qk_head_dim != layer.v_head_dim ), "FIA only supports qk_head_dim != v_head_dim" + + # Wait for the KV transfer to complete before performing attention computation. + forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id) + num_token_padding = q.shape[0] q, k, v = [ data[: forward_batch.num_token_non_padded_cpu] for data in [q, k, v] diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index b5c4aa172..901ddc873 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -41,22 +41,25 @@ from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, ) from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool +from sglang.srt.utils import get_device_module logger = logging.getLogger(__name__) +device_module = get_device_module() + class LayerLoadingEvent: def __init__(self, num_layers: int): self._num_layers = num_layers - self.load_events = [torch.cuda.Event() for _ in range(num_layers)] - self.start_event = torch.cuda.Event() # start event on controller stream + self.load_events = [device_module.Event() for _ in range(num_layers)] + self.start_event = device_module.Event() # start event on controller stream def complete(self, layer_index: int): assert 0 <= layer_index < self._num_layers self.load_events[layer_index].record() def wait(self, layer_index: int): - torch.cuda.current_stream().wait_event(self.load_events[layer_index]) + device_module.current_stream().wait_event(self.load_events[layer_index]) @property def finish_event(self): @@ -136,8 +139,8 @@ class CacheOperation: class HiCacheAck(NamedTuple): - start_event: torch.cuda.Event - finish_event: torch.cuda.Event + start_event: device_module.Event + finish_event: device_module.Event node_ids: List[int] @@ -343,8 +346,8 @@ class HiCacheController: self.stop_event, buffer_count=10, max_buffer_size=100 ) - self.write_stream = torch.cuda.Stream() - self.load_stream = torch.cuda.Stream() + self.write_stream = device_module.Stream() + self.load_stream = device_module.Stream() if self.enable_storage: self.prefetch_thread = threading.Thread( @@ -445,11 +448,11 @@ class HiCacheController: host_indices, device_indices = self.move_indices(op) self.write_queue.clear() - start_event = torch.cuda.Event() - finish_event = torch.cuda.Event() + start_event = device_module.Event() + finish_event = device_module.Event() start_event.record() - with torch.cuda.stream(self.write_stream): + with device_module.stream(self.write_stream): start_event.wait(self.write_stream) self.mem_pool_host.backup_from_device_all_layer( self.mem_pool_device, host_indices, device_indices, self.io_backend @@ -496,6 +499,8 @@ class HiCacheController: return host_indices, device_indices.index_select(0, idx) elif self.mem_pool_host.layout == "page_first_direct": return host_indices, device_indices.cpu() + elif self.io_backend == "kernel_ascend": + return host_indices, device_indices else: raise ValueError(f"Unsupported io backend") @@ -510,7 +515,7 @@ class HiCacheController: producer_event = self.layer_done_counter.events[producer_id] producer_event.start_event.record() - with torch.cuda.stream(self.load_stream): + with device_module.stream(self.load_stream): producer_event.start_event.wait(self.load_stream) for i in range(self.layer_num): self.mem_pool_host.load_to_device_per_layer( diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index d3e9d97ea..8729e21db 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -26,9 +26,15 @@ if not (_is_npu or _is_xpu): transfer_kv_per_layer_mla_pf_lf, transfer_kv_per_layer_pf_lf, ) +if _is_npu: + from sgl_kernel_npu.kvcacheio import TransferDirection, transfer_kv_dim_exchange logger = logging.getLogger(__name__) +SUPPORT_PIN_MEMORY = not _is_npu +if SUPPORT_PIN_MEMORY: + logger.warning("Current platform not support pin_memory") + def synchronized(func): @wraps(func) @@ -54,7 +60,7 @@ class HostKVCache(abc.ABC): self.device_pool = device_pool self.page_size = page_size self.layout = layout - self.pin_memory = pin_memory + self.pin_memory = pin_memory and SUPPORT_PIN_MEMORY self.device = device self.dtype = device_pool.store_dtype @@ -63,9 +69,9 @@ class HostKVCache(abc.ABC): self.size = int(host_size * 1e9 // self.size_per_token) else: self.size = int(device_pool.size * host_to_device_ratio) - # Align the host memory pool size to the page size - self.size = self.size - (self.size % self.page_size) - self.page_num = self.size // self.page_size + # Align up the host memory pool size to the page size + self.page_num = self.size // self.page_size + 1 + self.size = self.page_num * self.page_size self.start_layer = device_pool.start_layer self.end_layer = device_pool.end_layer @@ -316,6 +322,22 @@ class MHATokenToKVPoolHost(HostKVCache): ) else: raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "kernel_ascend": + if self.layout == "page_first_direct": + # Ascend-specific: transfer KV data for all layers when layer_id == 0 + if 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, + page_size=self.page_size, + direction=TransferDirection.H2D, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") else: raise ValueError(f"Unsupported IO backend: {io_backend}") @@ -367,6 +389,20 @@ class MHATokenToKVPoolHost(HostKVCache): ) else: raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "kernel_ascend": + if self.layout == "page_first_direct": + 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, + page_size=self.page_size, + direction=TransferDirection.D2H, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") else: raise ValueError(f"Unsupported IO backend: {io_backend}") @@ -549,6 +585,28 @@ class MLATokenToKVPoolHost(HostKVCache): 1, self.kv_lora_rank + self.qk_rope_head_dim, ) + # Ascend-specific: Aligns with AscendMLAPagedTokenToKVPool 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, + ) + self.k_buffer = torch.empty( + (*base_dims, self.kv_lora_rank), + dtype=self.dtype, + device=self.device, + ) + self.v_buffer = torch.empty( + (*base_dims, self.qk_rope_head_dim), + dtype=self.dtype, + device=self.device, + ) + # Return k_buffer to preserve original kv_buffer and data_refs init logic, + # though Ascend doesn't use these parameters. + return self.k_buffer else: raise ValueError(f"Unsupported layout: {self.layout}") self.token_stride_size = ( @@ -610,6 +668,24 @@ class MLATokenToKVPoolHost(HostKVCache): ) else: 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 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, + page_size=self.page_size, + direction=TransferDirection.H2D, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend @@ -655,6 +731,20 @@ class MLATokenToKVPoolHost(HostKVCache): ) else: raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "kernel_ascend": + if self.layout == "page_first_kv_split": + 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, + page_size=self.page_size, + direction=TransferDirection.D2H, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") else: raise ValueError(f"Unsupported IO backend: {io_backend}") diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9709f9ef7..0686ad1ad 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1506,6 +1506,20 @@ class ServerArgs: "Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes." ) + # Below are the only parameters currently supported on Ascend + if self.enable_hierarchical_cache and is_npu(): + # FIXME(iforgetmyname) fix decode_attention_backend on ascend + self.decode_attention_backend = "ascend" + self.hicache_io_backend = "kernel_ascend" + if self.use_mla_backend(): + self.hicache_mem_layout = "page_first_kv_split" + else: + self.hicache_mem_layout = "page_first_direct" + logger.warning( + f"Ascend NPU Platform detected, change `hicache_io_backend` to `kernel_ascend` and " + f"`hicache_mem_layout` to `{self.hicache_mem_layout}`" + ) + def _handle_speculative_decoding(self): if self.speculative_algorithm == "NEXTN": self.speculative_algorithm = "EAGLE" @@ -2991,14 +3005,19 @@ class ServerArgs: parser.add_argument( "--hicache-io-backend", type=str, - choices=["direct", "kernel"], + choices=["direct", "kernel", "kernel_ascend"], default=ServerArgs.hicache_io_backend, help="The IO backend for KV cache transfer between CPU and GPU", ) parser.add_argument( "--hicache-mem-layout", type=str, - choices=["layer_first", "page_first", "page_first_direct"], + choices=[ + "layer_first", + "page_first", + "page_first_direct", + "page_first_kv_split", + ], default=ServerArgs.hicache_mem_layout, help="The layout of host memory pool for hierarchical cache.", ) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 3a0c02c2a..f1bee1b14 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -653,6 +653,11 @@ def wait_cmo_stream(): cur_stream.wait_stream(get_cmo_stream()) +@lru_cache(maxsize=1) +def get_device_module(): + return torch.get_device_module() + + def set_random_seed(seed: int) -> None: """Set the random seed for all libraries.""" random.seed(seed) diff --git a/scripts/ci/npu_ci_install_dependency.sh b/scripts/ci/npu_ci_install_dependency.sh index 9dc355151..71f2042b3 100755 --- a/scripts/ci/npu_ci_install_dependency.sh +++ b/scripts/ci/npu_ci_install_dependency.sh @@ -59,7 +59,7 @@ wget -O "${BISHENG_NAME}" "${BISHENG_URL}" && chmod a+x "${BISHENG_NAME}" && "./ ### Install sgl-kernel-npu -SGL_KERNEL_NPU_TAG="20251030" +SGL_KERNEL_NPU_TAG="20251106" git clone --depth 1 https://github.com/sgl-project/sgl-kernel-npu.git --branch ${SGL_KERNEL_NPU_TAG} # pin wheel to 0.45.1 ref: https://github.com/pypa/wheel/issues/662 pip install wheel==0.45.1 diff --git a/test/srt/ascend/test_ascend_hicache_mha.py b/test/srt/ascend/test_ascend_hicache_mha.py new file mode 100644 index 000000000..d6f1aa9c2 --- /dev/null +++ b/test/srt/ascend/test_ascend_hicache_mha.py @@ -0,0 +1,98 @@ +import unittest +from types import SimpleNamespace +from urllib.parse import urlparse + +from sglang.srt.utils import kill_process_tree +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, + run_bench_offline_throughput, +) + +TEST_MODEL_MATRIX = { + "Qwen/Qwen2.5-7B-Instruct": { + "accuracy": 0.85, + "latency": 150, + "output_throughput": 30, + }, +} + + +class TestAscendMhaHicache(CustomTestCase): + + @classmethod + def setUpClass(cls): + cls.models = TEST_MODEL_MATRIX.keys() + cls.base_url = DEFAULT_URL_FOR_TEST + cls.url = urlparse(DEFAULT_URL_FOR_TEST) + cls.common_args = [ + "--trust-remote-code", + "--mem-fraction-static", + 0.8, + "--attention-backend", + "ascend", + "--enable-hierarchical-cache", + "--hicache-ratio", + 1.2, + ] + + def test_a_gsm8k(self): + for model in self.models: + with self.subTest(model=model): + print(f"##=== Testing accuracy: {model} ===##") + + process = popen_launch_server( + model, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + *self.common_args, + ], + ) + + try: + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=1319, + max_new_tokens=512, + parallel=128, + host=f"http://{self.url.hostname}", + port=int(self.url.port), + ) + + metrics = run_eval_few_shot_gsm8k(args) + self.assertGreaterEqual( + metrics["accuracy"], + TEST_MODEL_MATRIX[model]["accuracy"], + ) + finally: + kill_process_tree(process.pid) + + def test_b_throughput(self): + for model in self.models: + with self.subTest(model=model): + print(f"##=== Testing throughput: {model} ===##") + + output_throughput = run_bench_offline_throughput( + model, + [ + *self.common_args, + ], + ) + + print(f"##=== {model} throughput: {output_throughput} ===##") + + if is_in_ci(): + self.assertGreater( + output_throughput, + TEST_MODEL_MATRIX[model]["output_throughput"], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index e8a44da80..f77701844 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -521,6 +521,7 @@ suite_ascend = { "per-commit-1-ascend-npu": [ TestFile("ascend/test_ascend_graph_tp1_bf16.py", 400), TestFile("ascend/test_ascend_tp1_bf16.py", 400), + TestFile("ascend/test_ascend_hicache_mha.py", 400), ], "per-commit-2-ascend-npu": [ TestFile("ascend/test_ascend_graph_tp2_bf16.py", 400),