Skip to content

Commit 1fe3431

Browse files
committed
fix lint
Signed-off-by: MengqingCao <[email protected]>
1 parent 7b4b053 commit 1fe3431

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

vllm_ascend/worker/model_runner_v1.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff 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)

0 commit comments

Comments
 (0)