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
@@ -28,7 +28,7 @@ class GenerationBenchmark:
    def _initialize_kv_cache(self, batch_size: int) -> list:
        """初始化KV缓存"""
        config = self.config
-        shape = (batch_size, config.n_layer, config.m_len, config.n_kvhead, config.n_dim // config.n_head)
+        shape = (batch_size, config.m_len, config.n_layer, config.n_kvhead, config.n_dim // config.n_head)
        k_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
        v_cache = torch.zeros(shape, device=self.device, dtype=self.dtype)
        return (k_cache, v_cache)