ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(model): 调整 KV Cache 的维度顺序以匹配新的索引逻辑

Browse Source
This commit is contained in:
ViperEkura
2025-11-19 18:26:15 +08:00
parent e12ed0a72b
commit d9ff662e3a
3 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
khaosz/model/module.py
+4 -4
View File
@@ -182,12 +182,12 @@ class GQA(nn.Module):
k_cache, v_cache = kv_cache
# copy to cache
k_cache[:bsz, self.layer_id, start_pos:start_pos + seq_len] = k
v_cache[:bsz, self.layer_id, start_pos:start_pos + seq_len] = v
k_cache[:bsz, start_pos:start_pos + seq_len, self.layer_id] = k
v_cache[:bsz, start_pos:start_pos + seq_len, self.layer_id] = v
# get cache
k = k_cache[:bsz, self.layer_id, :start_pos + seq_len]
v = v_cache[:bsz, self.layer_id, :start_pos + seq_len]
k = k_cache[:bsz, :start_pos + seq_len, self.layer_id]
v = v_cache[:bsz, :start_pos + seq_len, self.layer_id]
k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)