Enabling the Multi-Head Compression Function in the Rope Scenario

key

[num_tokens, num_head, head_size]

float16/bf16

ND

Key of multiple tokens in the current step. The last dimension must be a multiple of 32.

value

[num_tokens, num_head, head_size]

float16/bf16

ND

Value of multiple tokens in the current step. The last dimension must be a multiple of 32.

keyCache

[num_blocks, block_size, 1, head_size]

float16/bf16

ND

All key caches.

valueCache

[num_blocks, block_size, 1, head_size]

float16/bf16

ND

All value caches.

slotMapping

[batch*num_head]

int32

ND

• Storage offset of the key or value in the cache for each batch of each head, that is, (block_id × block_size + offset_in_block).
• The value range is [0, num_blocks × kv_heads × block_size), and the value must be unique.

  That is, num_tokens <= num_blocks × block_size.

wins

[batch*num_head]

int32

ND

Compression. wins[i] < seqLens[floor(i/num_head)]. The value in wins must be greater than or equal to 0. If the value is 0, compression is not performed.

seqLens

[batch]

int32

ND

Actual seqLen of each batch. The internal value is greater than 0.

offsetIndex

[batch*num_head]

int32

ND

Compression start point of each head of each batch. The value range is [-1, seqLens[i]-wins[i]]. If the value is -1, compression is not performed. There is no requirement on the value of wins.

keyCacheOut

[num_blocks, block_size, 1, head_size]

float16/bf16

ND

All key caches. The address is the same as that of the input keyCache.

valueCacheOut

[num_blocks, block_size, 1, head_size]

float16/bf16

ND

All value caches. The address is the same as that of valueCache.