chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix the scope of io_backend in NSATokenToKVPoolHost (#20327)

Browse Source
This commit is contained in:
huangtingwei
2026-03-12 10:33:11 +08:00
committed by GitHub
parent 7771fcd4b3
commit 8787cf4566
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -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,