Compatibility change

kwen2501 · kwen2501 · commit 3ba19ea3c7da · 2024-09-21T00:10:06.000-07:00
diff --git a/torchchat/export.py b/torchchat/export.py
@@ -152,9 +152,9 @@ def __init__(self, attention: Attention):
             self.wo = attention.wo
 
             max_batch_size, n_heads, max_seq_length, head_dim = (
-                attention.kv_cache.k_cache.shape
+                attention.kv_cache[0].k_cache.shape
             )
-            cache_dtype = attention.kv_cache.k_cache.dtype
+            cache_dtype = attention.kv_cache[0].k_cache.dtype
             self.kv_cache = CustomKVCache(
                 max_batch_size, max_seq_length, n_heads, head_dim, cache_dtype
             )

Original file line number	Diff line number	Diff line change
`@@ -152,9 +152,9 @@ def __init__(self, attention: Attention):`
`152`	`152`	`self.wo = attention.wo`
`153`	`153`
`154`	`154`	`max_batch_size, n_heads, max_seq_length, head_dim = (`
`155`		`- attention.kv_cache.k_cache.shape`
	`155`	`+ attention.kv_cache[0].k_cache.shape`
`156`	`156`	`)`
`157`		`- cache_dtype = attention.kv_cache.k_cache.dtype`
	`157`	`+ cache_dtype = attention.kv_cache[0].k_cache.dtype`
`158`	`158`	`self.kv_cache = CustomKVCache(`
`159`	`159`	`max_batch_size, max_seq_length, n_heads, head_dim, cache_dtype`
`160`	`160`	`)`