ttnn.fill_cache

ttnn.fill_cache() → None

Fills the cache tensor in place with the values from input at the specified batch_idx, optionally offset along the sequence dimension by update_idx (tile-aligned).

:param * cache_tensor: The cache tensor to be written to. :type * cache_tensor: ttnn.Tensor :param * input_tensor: The token tensor to be written to the cache. :type * input_tensor: ttnn.Tensor :param * batch_idx: The index into the cache tensor. :type * batch_idx: int

:keyword * update_idx: seq-dim offset within the user slot, must be a multiple of TILE_HEIGHT. Default = 0. :kwtype * update_idx: int

Example

>>> cache_tensor = ttnn.from_torch(torch.zeros((4, 1, 128, 64), dtype=torch.bfloat16), device=device)
>>> input_tensor = ttnn.from_torch(torch.randn((1, 1, 32, 64), dtype=torch.bfloat16), device=device)
>>> output = ttnn.fill_cache(cache_tensor, input_tensor, batch_idx=0, update_idx=32)