Constraints

ntokens = sum(seqlen[i]), i=0, 1...batch-1.
There are the following restrictions on the Atlas inference products:
- The compression function is not supported.
- If the data format is marked with NZ, the data format must be NZ.
- The dimensions of keyCache and valueCache must be [num_blocks, num_head*head_size/16, block_size, 16]. The last dimension must be 16, and block_size must be 16-aligned.
- When kvCacheCfg is set to K_CACHE_V_CACHE or K_CACHE_V_CACHE_NZ, the data format is NZ.
- head_size must be 32-byte aligned.
If the last dimensions of the key and value are different, only Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products are supported. In this scenario, keyCache, valueCache, keyCacheOut, and valueCacheOut support only the ND format.
The data types of key, value, keyCache, valueCache, keyCacheOut, and valueCacheOut must be the same.
The data formats of keyCache, valueCache, keyCacheOut, and valueCacheOut must be the same.
The single-input single-output function is supported only by the Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products and does not support compression.
In the compression scenario, the cache size must be greater than the size after compression.
If multi-head compression is disabled, the NZ format supports only the float16 data type on the Atlas 200I/500 A2 inference products.
When the function of converting the format of the input to NZ is enabled, the following restrictions apply:
- Only Atlas A2 training products/Atlas A2 inference products and Atlas A3 inference products/Atlas A3 training products are supported.
- The length of key can be different from that of value. That is, the last dimension of key can be different from that of value.
- The data type of key can be different from that of value.
When the multi-head compression function is enabled in the rope scenario, the value of head_size cannot exceed 16 × 1024.

Parent topic: ReshapeAndCacheOperation