Gather the cached keys and values of every sequence in a batch from a shared pool of fixed-size pages, the way paged attention stores its KV cache.
The pool k_pages has shape (num_pages, page_size, num_heads, head_dim). Page p holds page_size consecutive tokens, and its slot s holds the key vectors of one token for every head. v_pages has the same shape and layout for values. The batch holds B sequences, one per row of block_tables. Sequence b splits its tokens into blocks of S = page_size tokens, and row b of block_tables lists the page that stores each block. Token position j of sequence b therefore lives in block ⌊j/S⌋, at slot jmodS of that block's page.
keys[b,j]={k_pages[block_tables[b,⌊j/S⌋], jmodS]0if j<seq_lens[b]otherwiseHere seq_lens[b] is the number of tokens sequence b has cached, and j runs from 0 to L−1 with L = max_blocks times page_size, where max_blocks is the width of block_tables. values follows the same rule with v_pages. Entries of block_tables past the blocks a sequence uses are -1, and slots past the end of a sequence inside its last page still hold old data from earlier requests, so none of them may reach the output.
Return a dict with three keys. keys and values are JAX arrays of shape (B, L, num_heads, head_dim) with the dtype of k_pages and v_pages, and mask is a bool JAX array of shape (B, L) that is True exactly at the positions j< seq_lens[b].
Input: k_pages = [[[[-1.6,-1.5]],[[1.3,-0.7]]],[[[-1.0,-2.0]],[[-0.1,0.5]]],[[[0.7,-1.2]],[[-1.5,-1.5]]],[[[-0.8,0.3]],[[-1.8,0.9]]]], v_pages = [[[[0.6,0.8]],[[-0.9,-0.4]]],[[[1.4,-0.5]],[[-0.9,0.8]]],[[[1.2,1.1]],[[-1.4,0.2]]],[[[-1.8,-0.3]],[[1.9,-0.9]]]], block_tables = [[2,0]], seq_lens = [3]
Output: {"keys":[[[[0.7,-1.2]],[[-1.5,-1.5]],[[-1.6,-1.5]],[[0.0,0.0]]]],"values":[[[[1.2,1.1]],[[-1.4,0.2]],[[0.6,0.8]],[[0.0,0.0]]]],"mask":[[true,true,true,false]]}
Explanation: With page_size 2, positions 0 and 1 form block 0, which lives in page 2, so they copy slots 0 and 1 of page 2, the keys [0.7, -1.2] and [-1.5, -1.5]. Position 2 is slot 0 of block 1, which lives in page 0, so it copies [-1.6, -1.5]. Position 3 is past seq_lens = 3, so it is zero even though slot 1 of page 0 holds [1.3, -0.7].
Input: k_pages = [[[[-1.9,0.5]],[[-0.8,0.2]]],[[[-1.7,0.9]],[[-0.8,2.0]]],[[[-0.9,-1.5]],[[-1.5,-0.7]]],[[[1.4,-0.5]],[[-1.1,-0.3]]],[[[0.1,-1.1]],[[0.7,-2.0]]],[[[0.5,0.6]],[[1.5,1.1]]]], v_pages = [[[[1.4,-0.8]],[[1.1,-1.8]]],[[[-0.8,-1.1]],[[1.4,-1.8]]],[[[-1.6,-0.5]],[[-0.2,-0.4]]],[[[-1.6,-0.2]],[[-0.4,0.5]]],[[[-1.6,-0.6]],[[1.7,-1.4]]],[[[0.1,1.9]],[[-0.5,-1.8]]]], block_tables = [[4,5,1],[3,-1,-1]], seq_lens = [5,2]
Output: {"keys":[[[[0.1,-1.1]],[[0.7,-2.0]],[[0.5,0.6]],[[1.5,1.1]],[[-1.7,0.9]],[[0.0,0.0]]],[[[1.4,-0.5]],[[-1.1,-0.3]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]]],"values":[[[[-1.6,-0.6]],[[1.7,-1.4]],[[0.1,1.9]],[[-0.5,-1.8]],[[-0.8,-1.1]],[[0.0,0.0]]],[[[-1.6,-0.2]],[[-0.4,0.5]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]]],"mask":[[true,true,true,true,true,false],[true,true,false,false,false,false]]}
Input: k_pages = [[[[1.9,1.6]],[[-0.9,-1.3]],[[-2.0,-1.0]]],[[[-1.8,-1.5]],[[-0.1,0.6]],[[-1.6,0.2]]],[[[-0.7,1.7]],[[0.5,-0.7]],[[1.2,-1.9]]]], v_pages = [[[[0.9,-1.6]],[[1.4,1.4]],[[-1.7,-1.3]]],[[[-0.8,-1.5]],[[-1.1,1.6]],[[-0.5,1.2]]],[[[1.5,-1.1]],[[-0.2,0.3]],[[-1.7,0.9]]]], block_tables = [[-1],[1]], seq_lens = [0,3]
Output: {"keys":[[[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]],[[[-1.8,-1.5]],[[-0.1,0.6]],[[-1.6,0.2]]]],"values":[[[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]],[[[-0.8,-1.5]],[[-1.1,1.6]],[[-0.5,1.2]]]],"mask":[[false,false,false],[true,true,true]]}
pos = jnp.arange(L); pos // page_size picks the block_tables column and pos % page_size the slot.
k_pages[pages, slots] with pages of shape (B, L) and slots of shape (L,) gathers every position at once.
jnp.where(mask[:, :, None, None], gathered, 0) zeroes padding and keeps the dtype.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Gather the cached keys and values of every sequence in a batch from a shared pool of fixed-size pages, the way paged attention stores its KV cache.
The pool k_pages has shape (num_pages, page_size, num_heads, head_dim). Page p holds page_size consecutive tokens, and its slot s holds the key vectors of one token for every head. v_pages has the same shape and layout for values. The batch holds B sequences, one per row of block_tables. Sequence b splits its tokens into blocks of S = page_size tokens, and row b of block_tables lists the page that stores each block. Token position j of sequence b therefore lives in block ⌊j/S⌋, at slot jmodS of that block's page.
keys[b,j]={k_pages[block_tables[b,⌊j/S⌋], jmodS]0if j<seq_lens[b]otherwiseHere seq_lens[b] is the number of tokens sequence b has cached, and j runs from 0 to L−1 with L = max_blocks times page_size, where max_blocks is the width of block_tables. values follows the same rule with v_pages. Entries of block_tables past the blocks a sequence uses are -1, and slots past the end of a sequence inside its last page still hold old data from earlier requests, so none of them may reach the output.
Return a dict with three keys. keys and values are JAX arrays of shape (B, L, num_heads, head_dim) with the dtype of k_pages and v_pages, and mask is a bool JAX array of shape (B, L) that is True exactly at the positions j< seq_lens[b].
Input: k_pages = [[[[-1.6,-1.5]],[[1.3,-0.7]]],[[[-1.0,-2.0]],[[-0.1,0.5]]],[[[0.7,-1.2]],[[-1.5,-1.5]]],[[[-0.8,0.3]],[[-1.8,0.9]]]], v_pages = [[[[0.6,0.8]],[[-0.9,-0.4]]],[[[1.4,-0.5]],[[-0.9,0.8]]],[[[1.2,1.1]],[[-1.4,0.2]]],[[[-1.8,-0.3]],[[1.9,-0.9]]]], block_tables = [[2,0]], seq_lens = [3]
Output: {"keys":[[[[0.7,-1.2]],[[-1.5,-1.5]],[[-1.6,-1.5]],[[0.0,0.0]]]],"values":[[[[1.2,1.1]],[[-1.4,0.2]],[[0.6,0.8]],[[0.0,0.0]]]],"mask":[[true,true,true,false]]}
Explanation: With page_size 2, positions 0 and 1 form block 0, which lives in page 2, so they copy slots 0 and 1 of page 2, the keys [0.7, -1.2] and [-1.5, -1.5]. Position 2 is slot 0 of block 1, which lives in page 0, so it copies [-1.6, -1.5]. Position 3 is past seq_lens = 3, so it is zero even though slot 1 of page 0 holds [1.3, -0.7].
Input: k_pages = [[[[-1.9,0.5]],[[-0.8,0.2]]],[[[-1.7,0.9]],[[-0.8,2.0]]],[[[-0.9,-1.5]],[[-1.5,-0.7]]],[[[1.4,-0.5]],[[-1.1,-0.3]]],[[[0.1,-1.1]],[[0.7,-2.0]]],[[[0.5,0.6]],[[1.5,1.1]]]], v_pages = [[[[1.4,-0.8]],[[1.1,-1.8]]],[[[-0.8,-1.1]],[[1.4,-1.8]]],[[[-1.6,-0.5]],[[-0.2,-0.4]]],[[[-1.6,-0.2]],[[-0.4,0.5]]],[[[-1.6,-0.6]],[[1.7,-1.4]]],[[[0.1,1.9]],[[-0.5,-1.8]]]], block_tables = [[4,5,1],[3,-1,-1]], seq_lens = [5,2]
Output: {"keys":[[[[0.1,-1.1]],[[0.7,-2.0]],[[0.5,0.6]],[[1.5,1.1]],[[-1.7,0.9]],[[0.0,0.0]]],[[[1.4,-0.5]],[[-1.1,-0.3]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]]],"values":[[[[-1.6,-0.6]],[[1.7,-1.4]],[[0.1,1.9]],[[-0.5,-1.8]],[[-0.8,-1.1]],[[0.0,0.0]]],[[[-1.6,-0.2]],[[-0.4,0.5]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]]],"mask":[[true,true,true,true,true,false],[true,true,false,false,false,false]]}
Input: k_pages = [[[[1.9,1.6]],[[-0.9,-1.3]],[[-2.0,-1.0]]],[[[-1.8,-1.5]],[[-0.1,0.6]],[[-1.6,0.2]]],[[[-0.7,1.7]],[[0.5,-0.7]],[[1.2,-1.9]]]], v_pages = [[[[0.9,-1.6]],[[1.4,1.4]],[[-1.7,-1.3]]],[[[-0.8,-1.5]],[[-1.1,1.6]],[[-0.5,1.2]]],[[[1.5,-1.1]],[[-0.2,0.3]],[[-1.7,0.9]]]], block_tables = [[-1],[1]], seq_lens = [0,3]
Output: {"keys":[[[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]],[[[-1.8,-1.5]],[[-0.1,0.6]],[[-1.6,0.2]]]],"values":[[[[0.0,0.0]],[[0.0,0.0]],[[0.0,0.0]]],[[[-0.8,-1.5]],[[-1.1,1.6]],[[-0.5,1.2]]]],"mask":[[false,false,false],[true,true,true]]}
pos = jnp.arange(L); pos // page_size picks the block_tables column and pos % page_size the slot.
k_pages[pages, slots] with pages of shape (B, L) and slots of shape (L,) gathers every position at once.
jnp.where(mask[:, :, None, None], gathered, 0) zeroes padding and keeps the dtype.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
