Compute the attention output of one new token for every query head, reading keys and values from a fixed-size cache whose heads are shared by groups of query heads.
query has shape (B, Hq, D) for B sequences, Hq query heads and head size D. key_cache and value_cache have shape (B, S, Hkv, D), where S is the number of cache slots and Hkv is the number of cache heads. With the group size G = Hq / Hkv, query head h reads cache head g=⌊h/G⌋, so query heads 0 to G - 1 share cache head 0, the next G query heads share cache head 1, and so on.
Only slots j<lengths[b] of sequence b hold real tokens, and the new token's own key and value are already stored among them. For query head h of sequence b, with query vector q and with Kj and Vj the key and value at slot j of cache head g,
sj=Dq⋅Kj,wj=∑m<lengths[b]esmesj,o=j<lengths[b]∑wjVj.Slots at or beyond the length hold stale numbers from earlier requests and must receive zero weight.
Your function is called under jax.jit with lengths traced, so a length cannot set the size of any array.
Return a float32 JAX array of shape (B, Hq, D) whose entry for sequence b and head h is the output o.
Input: query = [[[1,0],[0,1]]], key_cache = [[[[1,0]],[[0,1]],[[5,5]]]], value_cache = [[[[1,2]],[[3,4]],[[50,50]]]], lengths = [2]
Output: [[[1.660477,2.660477],[2.339523,3.339523]]]
Explanation: Both query heads read the single cache head, and only slots 0 and 1 are valid, so the stale value [50, 50] in slot 2 gets no weight. Head 0 scores the two keys as 1 / sqrt(2) = 0.707107 and 0, the softmax gives weights 0.669762 and 0.330238, and the output is 0.669762 * [1, 2] + 0.330238 * [3, 4] = [1.660477, 2.660477].
Input: query = [[[1,0],[0,1],[1,1],[-1,0]],[[0.5,0.5],[1,-1],[0,2],[2,0]]], key_cache = [[[[1,0],[0,1]],[[0,1],[1,0]],[[1,1],[-1,1]]],[[[1,0],[0,1]],[[9,9],[9,9]],[[9,9],[9,9]]]], value_cache = [[[[1,0],[0,1]],[[0,1],[1,0]],[[2,2],[-2,2]]],[[[1,2],[3,4]],[[90,90],[90,90]],[[90,90],[90,90]]]], lengths = [3,1]
Output: [[[1.203336,1.0],[1.0,1.203336],[0.00556,0.796664],[-1.011921,1.435946]],[[1.0,2.0],[1.0,2.0],[3.0,4.0],[3.0,4.0]]]
Input: query = [[[1,0,1,0],[0,2,0,-1]]], key_cache = [[[[1,1,0,0],[0,1,0,1]],[[0,0,1,1],[1,0,1,0]],[[1,0,0,1],[0,0,2,0]],[[0,1,1,0],[1,1,1,1]]]], value_cache = [[[[1,0,0,0],[0,1,0,0]],[[0,1,0,0],[0,0,1,0]],[[0,0,1,0],[0,0,0,1]],[[0,0,0,1],[1,0,0,0]]]], lengths = [4]
Output: [[[0.25,0.25,0.25,0.25],[0.31123,0.31123,0.18877,0.18877]]]
q.reshape(B, num_kv_heads, group, d) groups query heads by cache head.
jnp.einsum("bkgd,bskd->bkgs", q, keys) contracts D against a (B, S, Hkv, D) cache.
jnp.arange(T)[None, :] < lengths[:, None] builds a fixed-shape validity mask.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Compute the attention output of one new token for every query head, reading keys and values from a fixed-size cache whose heads are shared by groups of query heads.
query has shape (B, Hq, D) for B sequences, Hq query heads and head size D. key_cache and value_cache have shape (B, S, Hkv, D), where S is the number of cache slots and Hkv is the number of cache heads. With the group size G = Hq / Hkv, query head h reads cache head g=⌊h/G⌋, so query heads 0 to G - 1 share cache head 0, the next G query heads share cache head 1, and so on.
Only slots j<lengths[b] of sequence b hold real tokens, and the new token's own key and value are already stored among them. For query head h of sequence b, with query vector q and with Kj and Vj the key and value at slot j of cache head g,
sj=Dq⋅Kj,wj=∑m<lengths[b]esmesj,o=j<lengths[b]∑wjVj.Slots at or beyond the length hold stale numbers from earlier requests and must receive zero weight.
Your function is called under jax.jit with lengths traced, so a length cannot set the size of any array.
Return a float32 JAX array of shape (B, Hq, D) whose entry for sequence b and head h is the output o.
Input: query = [[[1,0],[0,1]]], key_cache = [[[[1,0]],[[0,1]],[[5,5]]]], value_cache = [[[[1,2]],[[3,4]],[[50,50]]]], lengths = [2]
Output: [[[1.660477,2.660477],[2.339523,3.339523]]]
Explanation: Both query heads read the single cache head, and only slots 0 and 1 are valid, so the stale value [50, 50] in slot 2 gets no weight. Head 0 scores the two keys as 1 / sqrt(2) = 0.707107 and 0, the softmax gives weights 0.669762 and 0.330238, and the output is 0.669762 * [1, 2] + 0.330238 * [3, 4] = [1.660477, 2.660477].
Input: query = [[[1,0],[0,1],[1,1],[-1,0]],[[0.5,0.5],[1,-1],[0,2],[2,0]]], key_cache = [[[[1,0],[0,1]],[[0,1],[1,0]],[[1,1],[-1,1]]],[[[1,0],[0,1]],[[9,9],[9,9]],[[9,9],[9,9]]]], value_cache = [[[[1,0],[0,1]],[[0,1],[1,0]],[[2,2],[-2,2]]],[[[1,2],[3,4]],[[90,90],[90,90]],[[90,90],[90,90]]]], lengths = [3,1]
Output: [[[1.203336,1.0],[1.0,1.203336],[0.00556,0.796664],[-1.011921,1.435946]],[[1.0,2.0],[1.0,2.0],[3.0,4.0],[3.0,4.0]]]
Input: query = [[[1,0,1,0],[0,2,0,-1]]], key_cache = [[[[1,1,0,0],[0,1,0,1]],[[0,0,1,1],[1,0,1,0]],[[1,0,0,1],[0,0,2,0]],[[0,1,1,0],[1,1,1,1]]]], value_cache = [[[[1,0,0,0],[0,1,0,0]],[[0,1,0,0],[0,0,1,0]],[[0,0,1,0],[0,0,0,1]],[[0,0,0,1],[1,0,0,0]]]], lengths = [4]
Output: [[[0.25,0.25,0.25,0.25],[0.31123,0.31123,0.18877,0.18877]]]
q.reshape(B, num_kv_heads, group, d) groups query heads by cache head.
jnp.einsum("bkgd,bskd->bkgs", q, keys) contracts D against a (B, S, Hkv, D) cache.
jnp.arange(T)[None, :] < lengths[:, None] builds a fixed-shape validity mask.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
