Operator Context

Figure 1 ReshapeAndCacheOperation operator context

The preceding figure shows the structure of the PA decoder layer in the typical rope scenario. ReshapeAndCache is used together with paged attention to store the historical key and value of each token at a layer. Different from the ATB KVCache operator, the ReshapeAndCache operator maintains only the KV cache of a certain layer instead of all KV caches of a service. Therefore, if the model has n layers, n k_cache tensors and n v_cache tensors are required to manage historical KVs when ReshapeAndCache is used.