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

[Bug] Fix NSA Backend KV-Buffer Shape Mismatch in DeepSeek-V3.2 (#12645)

Browse Source
This commit is contained in:
Johnsonms
2025-11-04 13:57:32 -08:00
committed by GitHub
parent 55dfb539cf
commit 1cfbbc42d8
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/memory_pool.py
+3 -1
View File
@@ -1568,6 +1568,7 @@ class MLATokenToKVPool(KVCache):
class NSATokenToKVPool(MLATokenToKVPool): class NSATokenToKVPool(MLATokenToKVPool):
quant_block_size = 128 quant_block_size = 128
index_k_with_scale_buffer_dtype = torch.uint8 index_k_with_scale_buffer_dtype = torch.uint8
rope_storage_dtype = torch.bfloat16 # rope is always stored in bf16
def __init__( def __init__(
self, self,
@@ -1589,10 +1590,11 @@ class NSATokenToKVPool(MLATokenToKVPool):
# Calculate override_kv_cache_dim for FP8 storage: # Calculate override_kv_cache_dim for FP8 storage:
# kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage # kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage
# Note: rope dimension is stored in original dtype (bf16), not quantized to fp8
override_dim = ( override_dim = (
kv_lora_rank kv_lora_rank
+ kv_lora_rank // self.quant_block_size * 4 + kv_lora_rank // self.quant_block_size * 4
+ qk_rope_head_dim * dtype.itemsize + qk_rope_head_dim * self.rope_storage_dtype.itemsize
) )
super().__init__( super().__init__(