KV Cache

[manupak]

Example

// prompt length = 32
// head_size = 8
// Prefill phase

gemm0  = Q x K : [32 x 8] x [8 x 32] -> [32 x 32]
QK         = softmax( Q x K ) : [32 x 32] -> [32 x 32]
gemm1  = QK x V : [32 x 32] x [ 32 x 8 ] -> [32 x 8]

// Decoding phase
// You get one token

// the sequence length has increased to 33 after the first inference
// we would give a buffer of max_seq_len (say 4096)
gemm0 = Q x K : [1 x 8] x [8 x (32 + 1)] -> [1 x (32 + 1)] 
QK = softmax( Q x K ) : [1 x (32 + 1)] -> [1 x (32 + 1)]
gemm1  = QK x V : [1 x (32 + 1)] x [ (32 + 1) x 8 ] -> [1 x 8]

Short summary of notes

• The N dimensions of gemm1 is growing.
• We can assume the N dimension of gemm1 to be capped at max_seq_len (say 4096)
• This means we need to "mask" out cleverly the N dimension of otherwise would be 4096.
• This would also mean the K dimension of gemm2 is growing which is also maxed at 4096.

• I have just realized N dimension of gemm1 is also a temporal axis in attention as there is a reduction across it
• This means we just need to do an early exit on the gemm0NperBlock loop.
• We need to mask out from current n_block_idx * gemm0NperBlock to current_seq_len

For a short term solution :

• There would be a runtime argument provided as current_seq_len.
• For gemm0N loop (outermost) :
  • We just add early termination of gemm0NPerBlock * n_block_idx > current_seq_len
  • We add early termination for nRepeats loops
  • The rest we mask gemm0OutputViews to be -inf (for softmax).
• For gemm1K loop (inner) :
  • We just add early termination of gemm1KPerBlock (aka gemm0NPerBlock) * k_block_idx > current_seq_len
  • We might need to mask within kPack. i.e. some values in KPack should be set to zero on some threads for k reduction MFMAs.

NOTE 1 : I think above needed to be understood in transposed manner as we do : (Vt x ( Kt x Qt ))t
NOTE 2 : This is how do this w/o touching coordinate transforms.

[manupak]

@krzysz00 Ive written this. I think there would be no early workgroup exits.
The one thing missed in the conversation before was, all the early exists are loop induction variables because in attention gemm0N dimension is iterated (blockwise) as it is a reduction dimension.