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

2025-11-19 18:26:15 +08:00
parent e12ed0a72b
commit d9ff662e3a
3 changed files with 7 additions and 7 deletions
@@ -206,11 +206,11 @@ class KVCacheManager:

    def _initialize(self):
        k_cache = torch.zeros(
-            (self.batch_size, self.num_layers, self.max_len, self.num_heads, self.head_dim),
+            (self.batch_size, self.max_len, self.num_layers, self.num_heads, self.head_dim),
            device=self.device, dtype=self.dtype
        )
        v_cache = torch.zeros(
-            (self.batch_size, self.num_layers, self.max_len, self.num_heads, self.head_dim),
+            (self.batch_size, self.max_len, self.num_layers, self.num_heads, self.head_dim),
            device=self.device, dtype=self.dtype
        )
        self._kv_cache = (k_cache, v_cache)