ttnn.experimental.indexed_fused_update_cache

ttnn.experimental.indexed_fused_update_cache(cache_tensor1: ttnn.Tensor, input_tensor1: ttnn.Tensor, cache_tensor2: ttnn.Tensor, input_tensor2: ttnn.Tensor, physical_update_idxs_tensor: ttnn.Tensor) → Tuple[ttnn.Tensor, ttnn.Tensor]

Writes packed rows from two input tensors into two cache tensors in parallel.

physical_update_idxs_tensor is a row-major INT32 tensor with shape [1, num_indices]. Entry i gives the physical cache row for input row i; negative and out-of-range entries are skipped. Physical rows flatten cache dimensions 0 and 2: page = index // cache.shape[2] and row_in_page = index % cache.shape[2]. This operation does not perform logical page-table translation.

Both caches and inputs must be interleaved BF16 TILE tensors. Cache shape is [num_pages, num_heads, rows_per_page, head_dim] and packed input shape is [1, num_heads, packed_rows, head_dim], where 1 <= packed_rows <= 256 and packed_rows <= num_indices <= 256. The operation updates the cache tensors in place and serializes all rows owned by a head/width worker, so multiple rows may safely target the same physical page.

Replicated tensors are supported on single- or multi-device meshes. The four cache/input tensors may also use one identical tensor-parallel mesh topology sharded over num_heads (dimension 1), head_dim (dimension 3), or both. Physical update indices must use the same mesh shape and coordinates and remain replicated across those mesh axes. Cache page/row sharding, sharded physical indices, and per-device index remapping are not supported. Initial hardware support is Wormhole and Blackhole.

Parameters:

  • cache_tensor1 (ttnn.Tensor)First paged cache, updated in place.

  • input_tensor1 (ttnn.Tensor)Packed rows for cache_tensor1.

  • cache_tensor2 (ttnn.Tensor)Second paged cache, updated in place.

  • input_tensor2 (ttnn.Tensor)Packed rows for cache_tensor2.

  • physical_update_idxs_tensor (ttnn.Tensor)Physical destination rows.

Returns:

Tuple[ttnn.Tensor, ttnn.Tensor]The two in-place cache tensors.