fix kv cache testing

apache · Jul 18, 2024 · ec72bc4 · ec72bc4
1 parent 32ee354
commit ec72bc4
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
@@ -344,6 +344,9 @@ def set_global_func():
 
 def create_kv_cache(rope_mode):
     support_sliding_window = 0
+    num_storage = head_dim
+    kv_storage_dtype = dtype
+
     cache = fcreate(
         tvm.runtime.ShapeTuple(
             [
@@ -361,7 +364,9 @@ def create_kv_cache(rope_mode):
         rope_mode,
         rope_scale,
         rope_theta,
+        num_storage,
         tvm.nd.empty((), dtype, device=device),
+        tvm.nd.empty((), kv_storage_dtype, device=device),
         ftranspose_append,
         fattention_prefill,
         fattention_decode,

diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
@@ -142,6 +142,9 @@ def set_global_func(head_dim, dtype):
 
 
 def create_kv_cache(head_dim, dtype, rope_mode, support_sliding_window):
+    num_storage = head_dim
+    kv_storage_dtype = dtype
+
     fcreate = tvm.get_global_func("vm.builtin.paged_attention_kv_cache_create_reduced")
     cache = fcreate(
         tvm.runtime.ShapeTuple(
@@ -160,7 +163,9 @@ def create_kv_cache(head_dim, dtype, rope_mode, support_sliding_window):
         rope_mode,
         rope_scale,
         rope_theta,
+        num_storage,
         tvm.nd.empty((), dtype, device=device),
+        tvm.nd.empty((), kv_storage_dtype, device=device),
         ftranspose_append,
         fattn_prefill,
         fattn_decode,