ttnn.kv_cache.update_cache_for_token_

ttnn.kv_cache.update_cache_for_token_(cache: ttnn.Tensor, token: ttnn.Tensor, update_index: int, batch_offset: int) → ttnn.Tensor

Updates the cache tensor in-place with values from input at update_index and batch_offset.

Parameters:

• cache (ttnn.Tensor) – the cache tensor to be written to.

• token (ttnn.Tensor) – the token tensor to be written to the cache.

• update_index (int) – the index into the cache tensor.

• batch_offset (int) – the batch_offset into the cache tensor.

Returns:

ttnn.Tensor – the output tensor.