diff --git a/astrai/inference/core/cache.py b/astrai/inference/core/cache.py
index cb6d395..e73db69 100644
--- a/astrai/inference/core/cache.py
+++ b/astrai/inference/core/cache.py
@@ -435,24 +435,13 @@ class ContiguousCacheView(CacheView):
             pos = self._write_positions
             self._cache.k[layer_id, indices, pos] = k.squeeze(1)
             self._cache.v[layer_id, indices, pos] = v.squeeze(1)
-            for s, p in zip(indices.tolist(), pos.tolist()):
-                cur = self._cache._slot_len.get(s, 0)
-                if p + 1 > cur:
-                    self._cache._slot_len[s] = p + 1
         else:
             start_pos = self._total_len - seq_len
             self._cache.k[layer_id, indices, start_pos : start_pos + seq_len] = k
             self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v
-            new_len = start_pos + seq_len
-            for s in indices.tolist():
-                cur = self._cache._slot_len.get(s, 0)
-                if new_len > cur:
-                    self._cache._slot_len[s] = new_len
 
     def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
-        max_len = max(
-            self._cache._slot_len.get(int(s), 0) for s in self._batch_indices.tolist()
-        )
+        max_len = self._total_len
         indices = self._batch_indices
         k = self._cache.k[layer_id, indices, :max_len]
         v = self._cache.v[layer_id, indices, :max_len]
@@ -528,6 +517,9 @@ class ContiguousCache(KVCache):
     ) -> ContiguousCacheView:
         slots = [self._task_slot[tid] for tid in task_ids]
         batch_indices = torch.tensor(slots, dtype=torch.long, device=device)
+        for slot in slots:
+            if total_len > self._slot_len.get(slot, 0):
+                self._slot_len[slot] = total_len
         return ContiguousCacheView(
             self, batch_indices, total_len, write_positions=write_positions
         )
diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py
index 57c9c53..5c3d61a 100644
--- a/astrai/inference/core/executor.py
+++ b/astrai/inference/core/executor.py
@@ -104,28 +104,25 @@ class Executor:
         )
 
         history_lists = []
-        mask_lists = []
+        history_lens = []
         for t in tasks:
             window = t.rep_window
             prompt_part = t.prompt_ids[-window:]
             ids = prompt_part + t.output_ids
             history_lists.append(ids)
-            mask_lists.append([True] * len(ids))
+            history_lens.append(len(ids))
 
-        max_len = max(len(h) for h in history_lists)
+        max_len = max(history_lens) if history_lens else 0
         padded_ids = torch.zeros(
             len(tasks), max_len, dtype=torch.long, device=self.device
         )
         padded_mask = torch.zeros(
             len(tasks), max_len, dtype=torch.bool, device=self.device
         )
-        for i, (h, m) in enumerate(zip(history_lists, mask_lists)):
-            padded_ids[i, : len(h)] = torch.tensor(
-                h, dtype=torch.long, device=self.device
-            )
-            padded_mask[i, : len(m)] = torch.tensor(
-                m, dtype=torch.bool, device=self.device
-            )
+        for i, h in enumerate(history_lists):
+            L = history_lens[i]
+            padded_ids[i, :L] = torch.as_tensor(h, dtype=torch.long, device=self.device)
+            padded_mask[i, :L] = True
 
         with torch.inference_mode():
             outputs = self.model(