File tree Expand file tree Collapse file tree 1 file changed +4
-2
lines changed
Expand file tree Collapse file tree 1 file changed +4
-2
lines changed Original file line number Diff line number Diff line change @@ -2848,7 +2848,7 @@ def _allocate_kv_cache_tensors(
28482848 v_tensor = self ._align_memory (
28492849 v_tensor , alignment )[:v_tensor_size ]
28502850 #### k cache: for deepseek sparse attention
2851- if dsa_k_cache_factor is not None :
2851+ if dsa_k_cache_factor is not None and dsa_k_cache_size is not None :
28522852 k_cache_tensor = torch .zeros (dsa_k_cache_size +
28532853 alignment ,
28542854 dtype = torch .int8 ,
@@ -2873,6 +2873,8 @@ def _allocate_kv_cache_tensors(
28732873 assert layer_names == set (kv_cache_raw_tensors .keys (
28742874 )), "Some layers are not correctly initialized"
28752875
2876+ return kv_cache_raw_tensors
2877+
28762878 def _reshape_kv_cache_tensors (
28772879 self ,
28782880 kv_cache_config : KVCacheConfig ,
@@ -2951,7 +2953,7 @@ def _reshape_kv_cache_tensors(
29512953 k_cache = self ._convert_torch_format (k_cache )
29522954 v_cache = raw_v_tensor .view (dtype ).view (kv_cache_shape [1 :])
29532955 v_cache = self ._convert_torch_format (v_cache )
2954- if self .use_sparse :
2956+ if self .use_sparse and raw_dsa_k_cache is not None :
29552957 dsa_k_cache_shape = (num_blocks , block_size , 1 , 128 )
29562958 dsa_k_cache = raw_dsa_k_cache .view (dtype ).view (
29572959 dsa_k_cache_shape )
You can’t perform that action at this time.
0 commit comments