Compute one grouped-query attention decode step and return the extended key-value (KV) caches. Let B be batch size, Hq the query-head count, Hkv the KV-head count, D the head width, and S the cached sequence length. Query has shape (B,Hq,D), each new key or value has shape (B,Hkv,D), and each cache has shape (B,Hkv,S,D). Each KV head serves a contiguous group of query heads of size:
G=HkvHq.The new key and value occupy position S after all cached positions. Let b index a batch item, h a KV head, g a query within its group, and s a cache position. Query qb,h,g is original query head hG+g; kb,h,s and vb,h,s are the extended key and value vectors. Attention weights a and output vectors y are:
ab,h,g,s=softmaxs(Dqb,h,gTkb,h,s). yb,h,g=s=0∑Sab,h,g,svb,h,s.Softmax is over all cache positions, including the new token. Return a Python dictionary with output, a PyTorch tensor of shape (B, Hq, D), and new_key_cache and new_value_cache, PyTorch tensors of shape (B, Hkv, S + 1, D). All three preserve the input dtype and device.
Input: query = [[[0.1,0.2]]], new_key = [[[0.2,0.4]]], new_value = [[[0.3,0.6]]], key_cache = [[[[0.2,0.4]]]], value_cache = [[[[0.1,0.2]]]], num_query_heads = 1, num_kv_heads = 1
Output: {"output":[[[0.2,0.4]]],"new_key_cache":[[[[0.2,0.4],[0.2,0.4]]]],"new_value_cache":[[[[0.1,0.2],[0.3,0.6]]]]}
Explanation: The two identical keys receive equal attention weights, so the output is the average of their values.
Input: query = [[[0]]], new_key = [[[2]]], new_value = [[[6]]], key_cache = [[[[1]]]], value_cache = [[[[2]]]], num_query_heads = 1, num_kv_heads = 1
Output: {"output":[[[4]]],"new_key_cache":[[[[1],[2]]]],"new_value_cache":[[[[2],[6]]]]}
Input: query = [[[0],[1]]], new_key = [[[1.09861229]]], new_value = [[[6]]], key_cache = [[[[0]]]], value_cache = [[[[2]]]], num_query_heads = 2, num_kv_heads = 1
Output: {"output":[[[4],[5]]],"new_key_cache":[[[[0],[1.098612]]]],"new_value_cache":[[[[2],[6]]]]}
torch.cat joins tensors along an existing dimension.
torch.matmul broadcasts leading batch dimensions.
torch.softmax accepts the axis to normalize.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Accepts: number
Compute one grouped-query attention decode step and return the extended key-value (KV) caches. Let B be batch size, Hq the query-head count, Hkv the KV-head count, D the head width, and S the cached sequence length. Query has shape (B,Hq,D), each new key or value has shape (B,Hkv,D), and each cache has shape (B,Hkv,S,D). Each KV head serves a contiguous group of query heads of size:
G=HkvHq.The new key and value occupy position S after all cached positions. Let b index a batch item, h a KV head, g a query within its group, and s a cache position. Query qb,h,g is original query head hG+g; kb,h,s and vb,h,s are the extended key and value vectors. Attention weights a and output vectors y are:
ab,h,g,s=softmaxs(Dqb,h,gTkb,h,s). yb,h,g=s=0∑Sab,h,g,svb,h,s.Softmax is over all cache positions, including the new token. Return a Python dictionary with output, a PyTorch tensor of shape (B, Hq, D), and new_key_cache and new_value_cache, PyTorch tensors of shape (B, Hkv, S + 1, D). All three preserve the input dtype and device.
Input: query = [[[0.1,0.2]]], new_key = [[[0.2,0.4]]], new_value = [[[0.3,0.6]]], key_cache = [[[[0.2,0.4]]]], value_cache = [[[[0.1,0.2]]]], num_query_heads = 1, num_kv_heads = 1
Output: {"output":[[[0.2,0.4]]],"new_key_cache":[[[[0.2,0.4],[0.2,0.4]]]],"new_value_cache":[[[[0.1,0.2],[0.3,0.6]]]]}
Explanation: The two identical keys receive equal attention weights, so the output is the average of their values.
Input: query = [[[0]]], new_key = [[[2]]], new_value = [[[6]]], key_cache = [[[[1]]]], value_cache = [[[[2]]]], num_query_heads = 1, num_kv_heads = 1
Output: {"output":[[[4]]],"new_key_cache":[[[[1],[2]]]],"new_value_cache":[[[[2],[6]]]]}
Input: query = [[[0],[1]]], new_key = [[[1.09861229]]], new_value = [[[6]]], key_cache = [[[[0]]]], value_cache = [[[[2]]]], num_query_heads = 2, num_kv_heads = 1
Output: {"output":[[[4],[5]]],"new_key_cache":[[[[0],[1.098612]]]],"new_value_cache":[[[[2],[6]]]]}
torch.cat joins tensors along an existing dimension.
torch.matmul broadcasts leading batch dimensions.
torch.softmax accepts the axis to normalize.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Accepts: number
