Add block_seq_stride flag (#692)

Add `block_seq_stride` flag --------- Co-authored-by: Rob Suderman <[email protected]>
nod-ai · Dec 13, 2024 · ec1424e · ec1424e
1 parent 77ca02f
commit ec1424e
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/sharktank/sharktank/examples/export_paged_llm_v1.py b/sharktank/sharktank/examples/export_paged_llm_v1.py
@@ -45,6 +45,12 @@ def main():
         type=lambda arg: [int(bs) for bs in arg.split(",")],
         default="4",
     )
+    parser.add_argument(
+        "--block-seq-stride",
+        help="Block sequence stride for paged KV cache, must divide evenly into the context length",
+        type=int,
+        default="16",
+    )
     parser.add_argument(
         "--verbose",
         help="Include verbose logging",
@@ -76,6 +82,7 @@ def main():
         static_tables=False,  # Rely on the compiler for hoisting tables.
         kv_cache_type="direct" if args.bs == [1] else "paged",
         attention_kernel=args.attention_kernel,
+        block_seq_stride=args.block_seq_stride,
     )
     llama_config.fake_quant = args.fake_quant