diff --git a/kvcache/causal.go b/kvcache/causal.go
index d519cf602..9a79fa577 100644
--- a/kvcache/causal.go
+++ b/kvcache/causal.go
@@ -545,7 +545,7 @@ func (c *Causal) shift(seq int, beginIndex, offset int32) error {
 		}
 	}
 
-	kShift, err := ctx.FromIntSlice(offsets, len(offsets))
+	kShift, err := ctx.Input().FromIntSlice(offsets, len(offsets))
 	if err != nil {
 		return err
 	}