Write the keys and values of L new tokens into a fixed-size key cache and value cache, starting at slot position along the sequence axis, and report the new number of filled slots.
key_cache and value_cache have shape (B, S, H, D), where B is the batch size, S is the number of cache slots, H is the number of heads and D is the head size. new_keys and new_values have shape (B, L, H, D). In every batch row and every head, slots position to position + L - 1 receive the new vectors in order, and every other slot keeps its old contents. The caches are stored in float32 or bfloat16, while the new vectors always arrive in float32.
Your function is called under jax.jit with position as a traced int32 scalar, so it has no concrete value while the function runs.
Return a dict with three keys. key_cache and value_cache hold the updated caches, each with the same shape and dtype as the cache passed in, and length holds position + L as an int32 JAX array of shape ().
Input: key_cache = [[[[0,0]],[[0,0]],[[0,0]],[[0,0]]]], value_cache = [[[[0,0]],[[0,0]],[[0,0]],[[0,0]]]], new_keys = [[[[1,2]]]], new_values = [[[[3,4]]]], position = 0
Output: {"key_cache":[[[[1.0,2.0]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]]],"value_cache":[[[[3.0,4.0]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]]],"length":1}
Explanation: The cache has four slots and position is 0, so the single new key [1, 2] and the single new value [3, 4] land in slot 0, the other three slots keep their zeros, and length becomes 0 + 1 = 1.
Input: key_cache = [[[[1,1]],[[2,2]],[[9,9]],[[9,9]]]], value_cache = [[[[-1,-1]],[[-2,-2]],[[9,9]],[[9,9]]]], new_keys = [[[[3,3]],[[4,4]]]], new_values = [[[[-3,-3]],[[-4,-4]]]], position = 2
Output: {"key_cache":[[[[1.0,1.0]],[[2.0,2.0]],[[3.0,3.0]],[[4.0,4.0]]]],"value_cache":[[[[-1.0,-1.0]],[[-2.0,-2.0]],[[-3.0,-3.0]],[[-4.0,-4.0]]]],"length":4}
Input: key_cache = [[[[1],[2]],[[0],[0]],[[0],[0]]],[[[5],[6]],[[0],[0]],[[0],[0]]]], value_cache = [[[[0.1],[0.2]],[[0],[0]],[[0],[0]]],[[[0.5],[0.6]],[[0],[0]],[[0],[0]]]], new_keys = [[[[3],[4]]],[[[7],[8]]]], new_values = [[[[0.3],[0.4]]],[[[0.7],[0.8]]]], position = 1
Output: {"key_cache":[[[[1.0],[2.0]],[[3.0],[4.0]],[[0.0],[0.0]]],[[[5.0],[6.0]],[[7.0],[8.0]],[[0.0],[0.0]]]],"value_cache":[[[[0.1],[0.2]],[[0.3],[0.4]],[[0.0],[0.0]]],[[[0.5],[0.6]],[[0.7],[0.8]],[[0.0],[0.0]]]],"length":2}
jax.lax.dynamic_update_slice(cache, block, (0, pos, 0, 0)) accepts a traced pos.
block.astype(cache.dtype) before writing.
Shapes are plain Python ints even under jax.jit.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: integer
Write the keys and values of L new tokens into a fixed-size key cache and value cache, starting at slot position along the sequence axis, and report the new number of filled slots.
key_cache and value_cache have shape (B, S, H, D), where B is the batch size, S is the number of cache slots, H is the number of heads and D is the head size. new_keys and new_values have shape (B, L, H, D). In every batch row and every head, slots position to position + L - 1 receive the new vectors in order, and every other slot keeps its old contents. The caches are stored in float32 or bfloat16, while the new vectors always arrive in float32.
Your function is called under jax.jit with position as a traced int32 scalar, so it has no concrete value while the function runs.
Return a dict with three keys. key_cache and value_cache hold the updated caches, each with the same shape and dtype as the cache passed in, and length holds position + L as an int32 JAX array of shape ().
Input: key_cache = [[[[0,0]],[[0,0]],[[0,0]],[[0,0]]]], value_cache = [[[[0,0]],[[0,0]],[[0,0]],[[0,0]]]], new_keys = [[[[1,2]]]], new_values = [[[[3,4]]]], position = 0
Output: {"key_cache":[[[[1.0,2.0]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]]],"value_cache":[[[[3.0,4.0]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]]],"length":1}
Explanation: The cache has four slots and position is 0, so the single new key [1, 2] and the single new value [3, 4] land in slot 0, the other three slots keep their zeros, and length becomes 0 + 1 = 1.
Input: key_cache = [[[[1,1]],[[2,2]],[[9,9]],[[9,9]]]], value_cache = [[[[-1,-1]],[[-2,-2]],[[9,9]],[[9,9]]]], new_keys = [[[[3,3]],[[4,4]]]], new_values = [[[[-3,-3]],[[-4,-4]]]], position = 2
Output: {"key_cache":[[[[1.0,1.0]],[[2.0,2.0]],[[3.0,3.0]],[[4.0,4.0]]]],"value_cache":[[[[-1.0,-1.0]],[[-2.0,-2.0]],[[-3.0,-3.0]],[[-4.0,-4.0]]]],"length":4}
Input: key_cache = [[[[1],[2]],[[0],[0]],[[0],[0]]],[[[5],[6]],[[0],[0]],[[0],[0]]]], value_cache = [[[[0.1],[0.2]],[[0],[0]],[[0],[0]]],[[[0.5],[0.6]],[[0],[0]],[[0],[0]]]], new_keys = [[[[3],[4]]],[[[7],[8]]]], new_values = [[[[0.3],[0.4]]],[[[0.7],[0.8]]]], position = 1
Output: {"key_cache":[[[[1.0],[2.0]],[[3.0],[4.0]],[[0.0],[0.0]]],[[[5.0],[6.0]],[[7.0],[8.0]],[[0.0],[0.0]]]],"value_cache":[[[[0.1],[0.2]],[[0.3],[0.4]],[[0.0],[0.0]]],[[[0.5],[0.6]],[[0.7],[0.8]],[[0.0],[0.0]]]],"length":2}
jax.lax.dynamic_update_slice(cache, block, (0, pos, 0, 0)) accepts a traced pos.
block.astype(cache.dtype) before writing.
Shapes are plain Python ints even under jax.jit.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: integer
