For bs=1 set kv_cache_type to paged.

Signed-off-by: MaheshRavishankar <[email protected]>
nod-ai · Jan 7, 2025 · 05358b1 · 05358b1
1 parent 922b1c2
commit 05358b1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sharktank/sharktank/examples/export_paged_llm_v1.py b/sharktank/sharktank/examples/export_paged_llm_v1.py
@@ -79,7 +79,7 @@ def main():
         hp,
         tensor_parallelism_size=tensor_parallelism_size,
         use_hf=False,
-        kv_cache_type="direct" if args.bs == [1] else "paged",
+        kv_cache_type="paged",
         attention_kernel=args.attention_kernel,
         block_seq_stride=args.block_seq_stride,
     )