From 1cfbbc42d82ce5ef985c7d6e021526a665d16d1d Mon Sep 17 00:00:00 2001 From: Johnsonms Date: Tue, 4 Nov 2025 13:57:32 -0800 Subject: [PATCH] [Bug] Fix NSA Backend KV-Buffer Shape Mismatch in DeepSeek-V3.2 (#12645) --- python/sglang/srt/mem_cache/memory_pool.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index f93ebd47f..9d33a4b8b 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1568,6 +1568,7 @@ class MLATokenToKVPool(KVCache): class NSATokenToKVPool(MLATokenToKVPool): quant_block_size = 128 index_k_with_scale_buffer_dtype = torch.uint8 + rope_storage_dtype = torch.bfloat16 # rope is always stored in bf16 def __init__( self, @@ -1589,10 +1590,11 @@ class NSATokenToKVPool(MLATokenToKVPool): # Calculate override_kv_cache_dim for FP8 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 = ( kv_lora_rank + 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__(