Replication of index was causing issues for kv cache writes (#715)

Signed-off-by: Rob Suderman <[email protected]>

Author: rsuderman

sharktank/sharktank/examples/export_paged_llm_v1.py

@@ -334,7 +334,8 @@ def _(
 
     bsizes = []
     for bs in args.bs:
-        generate_batch_prefill(bs)
+        if not args.skip_prefill:
+            generate_batch_prefill(bs)
         if not args.skip_decode:
             generate_batch_decode(bs)
         bsizes.append(bs)

sharktank/sharktank/layers/kv_cache.py

@@ -456,12 +456,25 @@ def write_timestep(
             page_offset = (seq_positions % self.block_seq_stride).unsqueeze(1)
 
             # [1, 1]
-            partitions = torch.tensor(idx).unsqueeze(0)
+            if isinstance(seq_positions, ReplicatedTensor):
+                partitions = [
+                    torch.tensor(idx).unsqueeze(0)
+                    for _ in range(seq_positions.shard_count)
+                ]
+
+                transformer_block = [
+                    torch.full((bs, 1), transformer_block_index, device=device)
+                    for _ in range(seq_positions.shard_count)
+                ]
+
+                partitions = ReplicatedTensor(ts=partitions)
+                transformer_block = ReplicatedTensor(ts=transformer_block)
+            else:
+                partitions = torch.tensor(idx).unsqueeze(0)
+                transformer_block = torch.full(
+                    (bs, 1), transformer_block_index, device=device
+                )
 
-            # [bs, 1]
-            transformer_block = torch.full(
-                (bs, 1), transformer_block_index, device=device
-            )
             partitions = partitions.repeat(bs, 1)
 
             indices = (page_id, transformer_block, partitions, page_offset)

sharktank/sharktank/utils/cli.py

@@ -69,9 +69,14 @@ def add_model_options(parser: argparse.ArgumentParser):
         default="torch",
         choices=["decomposed", "torch"],
     )
+    parser.add_argument(
+        "--skip-prefill",
+        help="Skips exporting prefill",
+        action="store_true",
+    )
     parser.add_argument(
         "--skip-decode",
-        help="Enables prefill only, skips decode",
+        help="Skips export decode",
         action="store_true",
     )