Support proper KVCache

[dayo09]

What?

Previously, KVCache export was not supported for mainly TWO reasons.

(1) torch.export doesn't allow input mutation (e.g. in torch 2.7)

• This is resolved in torch 2.10

(2) ONE CIRCLE doesn't allow in-memory buffer update.

• TBD (Will they provide support?)

If above limitations are lifted, we could directly convert KV Cache into corresponding Circle.

Importance

• ASR/Llama/VLM models includes self-attention/cross-attention and they applies implies multiple kv cache handling.
  • • speculative decoding

KV Cache's aten origins

ERROR:tico.utils.convert:NOT SUPPORTED OPERATOR
        (op) index_put.default
        (trace)   File "/home/dayo/miniconda3/envs/py310-tvm/lib/python3.10/site-packages/torch/_dynamo/functional_export.py", line 216, in forward
    res = self._export_root(*args, **kwargs)
  File "/home/dayo/miniconda3/envs/py310-tvm/lib/python3.10/site-packages/transformers/models/whisper/modeling_whisper.py", line 337, in forward
    key_states, value_states = past_key_values.update(
  File "/home/dayo/miniconda3/envs/py310-tvm/lib/python3.10/site-packages/transformers/cache_utils.py", line 783, in update
    keys, values = self.layers[layer_idx].update(key_states, value_states, cache_kwargs)
  File "/home/dayo/miniconda3/envs/py310-tvm/lib/python3.10/site-packages/transformers/cache_utils.py", line 340, in update
    self.values.index_copy_(2, cache_position, value_states)

[dayo09]

Target Examples

class SimpleIndexCopy(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, keys, key_states):
        cache_position = torch.arange(key_states.shape[-2])
        return keys.index_copy_(2, cache_position, key_states)

    def get_example_inputs(self):
        batch, num_heads, max_cache_len, head_dim = (1, 20, 448, 64)
        keys = torch.randn(batch, num_heads, max_cache_len, head_dim)
        key_states = torch.randn(batch, num_heads, 1, head_dim)
        return (keys, key_states), {}


class SimpleIndexCopyUpdate(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, keys, key_states):
        cache_position = torch.arange(key_states.shape[-2])
        keys[:, :, cache_position] = key_states
        return keys

    def get_example_inputs(self):
        batch, num_heads, max_cache_len, head_dim = (1, 20, 448, 64)
        keys = torch.randn(batch, num_heads, max_cache_len, head_dim)
        key_states = torch.randn(batch, num_heads, 1, head_dim)
        return (keys, key_states), {}

[dayo09]

Current KVCache export

FYI, This example is exported like below.

class KVCacheSlice(TestModuleBase):
    def __init__(self):
        super().__init__()
        self.n_rep = 1
        self.register_buffer(
            "cache_k",
            torch.randn(
                [max_batch_size, n_local_kv_heads, max_seq_len, head_dim],
            ),
        )  # 8, 8, 64, 128
        self.register_buffer(
            "cache_v",
            torch.randn(
                [max_batch_size, n_local_kv_heads, max_seq_len, head_dim],
            ),
        )

    def forward(self, start_pos, seq_len):
        keys = self.cache_k[:bsz, :, : start_pos + seq_len, :]  # type: ignore[index]
        values = self.cache_v[:bsz, :, : start_pos + seq_len, :]  # type: ignore[index]

        keys = repeat_kv(keys, self.n_rep)
        values = repeat_kv(values, self.n_rep)
        return (keys, values, self.cache_k, self.cache_v)

    def get_example_inputs(self):
        start_pos = 8
        return (start_pos, seq_len), {}

[dayo09]

Circle mapping

• op: DynamicUpdateSlice
• link: https://github.com/Samsung/ONE/blob/master/runtime/compute/cker/include/cker/operation/DynamicUpdateSlice.h

This is the correspondant circle operation. It isn't sure if it performs 'in-place' update though. From guess, it only produces the result output as another tensor where in-place update for circle operator is supported in another runtime feature.