ReshapeAndCacheOmniOperation

Applicable Products

Description

In the Decode phase, OmniAttention performs attention sparse computation on a specific head of a specific layer of the LLM to save the KV cache graphics memory, reduce the amount of data transferred from the on-chip memory to the cache, and reduce the attention computation amount.

Operator Context

Definition

struct ReshapeAndCacheOmniParam{
    uint8_t rsv[8] = {0};
}

Parameters

Omni Head Compression

• Function
  • Omni Head: Compression is performed. The KV cache only needs to store (First Sink Size + Recent Sink Size) tokens, and does not need to store seqlen tokens.
  • Full Head: no compression.

  In the model, this operator is used to compress specific headers using omni compression. Other headers are not compressed and the original ReshapeAndCache operator is used.

• Input

  Parameter	Dimension	Data Type	Format	Description
  key	[num_tokens, num_head, head_size]	float16/bf16	ND	Key to be cached. head_size must be a multiple of 32.
  value	[num_tokens, num_head, head_size]	float16/bf16	ND	Value to be cached.
  key_cache	[num_blocks, block_size, 1, head_size]	float16/bf16	ND	Compressed and cached keys.
  value_cache	[num_blocks, block_size, 1, head_size]	float16/bf16	ND	Compressed and cached values.
  slot_mapping	[batch * num_head]	int32	ND	Storage offset of the token key or value of each head in each batch in the cache, that is, (block_id × block_size + offset_in_block). The value range is [0, num_blocks × block_size), and the value must be unique.
  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.
  seqLen	[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[floor(i/num_head)]-wins[i]]. If the value is -1, compression is not performed. There is no requirement on the value of wins.

• Output

  Parameter	Dimension	Data Type	Format	Description
  key_cache	[num_blocks, block_size, 1, head_size]	float16/bf16	ND	key_cache of the key to be flashed.
  value_cache	[num_blocks, block_size, 1, head_size]	float16/bf16	ND	value_cache of the value to be flashed.

• Restrictions
  • 
  • batch * num_head<500
  • The value of head_size must be a multiple of 32 and cannot exceed 16 × 1024.
  • The data types of key, value, keyCache, valueCache, keyCacheOut, and valueCacheOut must be the same.