From 8787cf4566683c343fb5e437b8401b8623b16d3a Mon Sep 17 00:00:00 2001 From: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com> Date: Thu, 12 Mar 2026 10:33:11 +0800 Subject: [PATCH] Fix the scope of io_backend in NSATokenToKVPoolHost (#20327) --- python/sglang/srt/mem_cache/memory_pool_host.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index d71fd21cb..d98824420 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1209,7 +1209,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): ) else: raise ValueError(f"Unsupported layout: {self.layout}") - else: + elif io_backend == "direct": if self.layout == "layer_first": transfer_kv_direct( src_layers=[self.index_k_with_scale_buffer[layer_id]], @@ -1229,6 +1229,8 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): ) else: raise ValueError(f"Unsupported layout: {self.layout}") + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") def _backup_indexer_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend @@ -1259,7 +1261,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): ) else: raise ValueError(f"Unsupported layout: {self.layout}") - else: + elif io_backend == "direct": if self.layout == "layer_first": transfer_kv_direct( src_layers=device_pool.index_k_with_scale_buffer, @@ -1278,6 +1280,8 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): ) else: raise ValueError(f"Unsupported layout: {self.layout}") + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") def load_to_device_per_layer( self,