ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl

ai_edge_torch/generative/examples/t5/t5_attention.py

@@ -0,0 +1,231 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Attention modules for the T5 encoder-decoder model family.
+
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from ai_edge_torch.generative.layers.attention import CrossAttention
+import ai_edge_torch.generative.layers.builder as builder
+from ai_edge_torch.generative.layers.kv_cache import KVCache
+import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention  # NOQA
+from ai_edge_torch.generative.layers.scaled_dot_product_attention import scaled_dot_product_attention_with_hlfb  # NOQA
+
+BATCH_SIZE = 1
+
+
+class EncoderDecoderBlock(nn.Module):
+
+  def __init__(
+      self, config: cfg.ModelConfig, has_relative_attention_bias: bool = False
+  ) -> None:
+    """Initialize an instance of the EncoderDecoderBlock.
+
+    Args:
+      config (cfg.ModelConfig): the configuration object
+        for this transformer block.
+      has_relative_attention_bias (bool): whether the
+        self attention block has relative bias.
+    """
+
+    super().__init__()
+    self.atten_func = T5Attention(
+        BATCH_SIZE,
+        config.embedding_dim,
+        config.attn_config,
+        config.pre_attention_norm_config,
+        config.kv_cache_max,
+        config.enable_hlfb,
+        has_relative_attention_bias=has_relative_attention_bias,
+    )
+    # For a decoder, we add a cross attention.
+    if config.is_decoder:
+      self.cross_atten_func = T5Attention(
+          BATCH_SIZE,
+          config.embedding_dim,
+          config.attn_config,
+          config.pre_attention_norm_config,
+          config.kv_cache_max,
+          config.enable_hlfb,
+          # Cross Attention does not have relative attention bias.
+          has_relative_attention_bias=False,
+      )
+    else:
+      self.cross_atten_func = None
+
+    self.pre_ff_norm = builder.build_norm(
+        config.embedding_dim, config.pre_ff_norm_config
+    )
+    self.ff = builder.build_ff(config.embedding_dim, config.ff_config)
+    self.config = config
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      input_pos: Optional[torch.Tensor] = None,
+      mask: Optional[torch.Tensor] = None,
+      relative_position: Optional[torch.Tensor] = None,
+      position_bias: Optional[torch.Tensor] = None,
+      encoder_hidden_states: Optional[torch.Tensor] = None,
+      encoder_attention_mask: Optional[torch.Tensor] = None,
+      encoder_decoder_position_bias: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
+    """Forward function of the EncoderDecoderBlock.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+
+    Returns:
+      output activation from this transformer block.
+    """
+
+    hidden_states, position_bias = self.atten_func(
+        x,
+        input_pos=input_pos,
+        mask=mask,
+        relative_position=relative_position,
+        position_bias=position_bias,
+    )
+
+    attn_out = hidden_states + x
+
+    if self.cross_atten_func:
+      hidden_states, encoder_decoder_position_bias = self.cross_atten_func(
+          attn_out,
+          input_pos=input_pos,
+          key_value_states=encoder_hidden_states,
+          mask=encoder_attention_mask,
+          relative_position=relative_position,
+          position_bias=encoder_decoder_position_bias,
+      )
+      attn_out = hidden_states + attn_out
+
+    forwarded = self.pre_ff_norm(attn_out)
+    forwarded = self.ff(forwarded)
+    hidden_states = attn_out + forwarded
+
+    # encoder_deocder_position_bias is from CrossAttention
+    return hidden_states, position_bias, encoder_decoder_position_bias
+
+
+class T5Attention(CrossAttention):
+
+  def __init__(
+      self,
+      batch: int,
+      dim: int,
+      config: cfg.AttentionConfig,
+      norm_config: cfg.NormalizationConfig,
+      kv_cache_max: int,
+      enable_hlfb: bool,
+      has_relative_attention_bias=False,
+  ) -> None:
+    """Initialize an instance of T5Attention.
+
+    Args:
+      dim (int): causal attention's input/output dimmension.
+      config (cfg.AttentionConfig): attention specific configurations.
+      norm_config (cfg.NormalizationConfig): normalization configure before attention.
+      kv_cache_max (int): determines the size of the KV Cache buffer, if enabled.
+      enable_hlfb (bool): whether hlfb is enabled or not.
+      has_relative_attention_bias (bool): whether we compute relative bias.
+    """
+    super().__init__(batch, dim, dim, config, kv_cache_max, enable_hlfb)
+    self.pre_atten_norm = builder.build_norm(dim, norm_config)
+
+    self.has_relative_attention_bias = has_relative_attention_bias
+    self.relative_attention_num_buckets = config.relative_attention_num_buckets
+    if self.has_relative_attention_bias:
+      self.relative_attention_bias = nn.Embedding(
+          self.relative_attention_num_buckets, self.n_heads
+      )
+
+  def forward(
+      self,
+      x: torch.Tensor,
+      input_pos: Optional[torch.Tensor] = None,
+      key_value_states: Optional[torch.Tensor] = None,
+      mask: Optional[torch.Tensor] = None,
+      relative_position: Optional[torch.Tensor] = None,
+      position_bias: Optional[torch.Tensor] = None,
+  ) -> torch.Tensor:
+    """Forward function of the T5Attention layer.
+
+    Args:
+      x (torch.Tensor): the input tensor.
+      rope (Tuple[torch.Tensor, torch.Tensor]): the input rope tensor.
+      mask (torch.Tensor): the optional mask tensor.
+      input_pos (torch.Tensor): the optional input position tensor.
+
+    Returns:
+      output activation from this self attention layer.
+    """
+
+    x = self.pre_atten_norm(x)
+    B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)
+    query_states = self.q_projection(x)
+    query_states = query_states.reshape(B, T, -1, self.head_dim)  # (B, T, nh_q, hs)
+
+    if key_value_states is not None:
+      (
+          kvB,
+          kvT,
+          kvC,
+      ) = (
+          key_value_states.size()
+      )  # batch size, sequence length, embedding dimensionality (n_embd)
+      key_states = self.k_projection(key_value_states)
+      value_states = self.v_projection(key_value_states)
+      key_states = key_states.reshape(kvB, kvT, -1, self.head_dim)
+      value_states = value_states.reshape(kvB, kvT, -1, self.head_dim)
+    else:
+      key_states = self.k_projection(x)
+      value_states = self.v_projection(x)
+      key_states = key_states.reshape(B, T, -1, self.head_dim)
+      value_states = value_states.reshape(B, T, -1, self.head_dim)
+
+    if key_value_states is None and self.kv_cache is not None:
+      key_states, value_states = self.kv_cache.update_cache(
+          input_pos, key_states, value_states
+      )
+
+    if position_bias is None:
+      # handle the encoder case first
+      if self.has_relative_attention_bias:
+        position_bias = self.relative_attention_bias(
+            relative_position
+        )  # shape (query_length, key_length, num_heads)
+        position_bias = position_bias.permute([0, 1, 4, 2, 3]).squeeze(
+            0
+        )  # shape (1, num_heads, query_length, key_length)
+      else:
+        # position_bias = torch.zeros(B, self.n_heads, T, self.head_dim, dtype=torch.float32)
+        position_bias = torch.zeros_like(mask, dtype=torch.float32)
+
+    mask = mask + position_bias
+    y = self.sdpa_func(
+        query_states, key_states, value_states, self.head_dim, mask=mask, scale=1.0
+    )
+    y = y.reshape(B, T, C)  # re-assemble all head outputs side by side
+    # output projection
+    y = self.output_projection(y)
+    return y, position_bias

ai_edge_torch/generative/examples/test_models/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

ai_edge_torch/generative/examples/test_models/toy_model.py

@@ -0,0 +1,122 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# A toy example which has a single-layer transformer block.
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+import ai_edge_torch
+from ai_edge_torch.generative.layers.attention import TransformerBlock
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.builder as builder
+import ai_edge_torch.generative.layers.model_config as cfg
+
+RoPECache = Tuple[torch.Tensor, torch.Tensor]
+KV_CACHE_MAX_LEN = 100
+
+
+class ToySingleLayerModel(torch.nn.Module):
+
+  def __init__(self, config: cfg.ModelConfig) -> None:
+    super().__init__()
+    self.lm_head = nn.Linear(
+        config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+    )
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.transformer_block = TransformerBlock(config)
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.max_seq_len,
+        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        base=10_000,
+        condense_ratio=1,
+        dtype=torch.float32,
+        device=torch.device('cpu'),
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.max_seq_len, dtype=torch.float32, device=torch.device('cpu')
+    )
+    self.config = config
+
+  @torch.inference_mode
+  def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
+    x = self.tok_embedding(idx)
+    cos, sin = self.rope_cache
+
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+    mask = self.mask_cache.index_select(2, input_pos)
+    mask = mask[:, :, :, : self.config.max_seq_len]
+
+    x = self.transformer_block(x, (cos, sin), mask, input_pos)
+    x = self.final_norm(x)
+    return self.lm_head(x)
+
+
+def get_model_config() -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=32, num_query_groups=4, rotary_percentage=1.0, enable_kv_cache=False
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+      intermediate_size=256,
+  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  config = cfg.ModelConfig(
+      vocab_size=400,
+      num_layers=1,
+      max_seq_len=KV_CACHE_MAX_LEN,
+      embedding_dim=128,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      pre_ff_norm_config=norm_config,
+      final_norm_config=norm_config,
+  )
+  return config
+
+
+def define_and_run() -> None:
+  model = ToySingleLayerModel(get_model_config())
+  idx = torch.unsqueeze(torch.arange(0, KV_CACHE_MAX_LEN), 0)
+  input_pos = torch.arange(0, KV_CACHE_MAX_LEN)
+  print('running an inference')
+  print(
+      model.forward(
+          idx,
+          input_pos,
+      )
+  )
+
+  # Convert model to tflite.
+  print('converting model to tflite')
+  edge_model = ai_edge_torch.convert(
+      model,
+      (
+          idx,
+          input_pos,
+      ),
+  )
+  edge_model.export('/tmp/toy_model.tflite')
+
+
+if __name__ == '__main__':
+  define_and_run()

ai_edge_torch/generative/examples/test_models/toy_model_with_external_kv_cache.py

@@ -0,0 +1,161 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# A toy example which has basic transformer block (w/ externalized KV-Cache).
+
+from typing import List, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch_xla
+
+import ai_edge_torch
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.builder as builder
+from ai_edge_torch.generative.layers.experimental.attention import TransformerBlock  # NOQA
+import ai_edge_torch.generative.layers.model_config as cfg
+
+RoPECache = Tuple[torch.Tensor, torch.Tensor]
+
+
+class ToyModelWithExternalKV(torch.nn.Module):
+
+  def __init__(self, config: cfg.ModelConfig) -> None:
+    super().__init__()
+    self.lm_head = nn.Linear(
+        config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+    )
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.transformer_blocks = nn.ModuleList(
+        TransformerBlock(config) for _ in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.max_seq_len,
+        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        base=10_000,
+        condense_ratio=1,
+        dtype=torch.float32,
+        device=torch.device('cpu'),
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.max_seq_len, dtype=torch.float32, device=torch.device('cpu')
+    )
+    self.config = config
+
+  def forward(
+      self,
+      idx: torch.Tensor,
+      input_pos: torch.Tensor,
+      k_caches: torch.Tensor,
+      v_caches: torch.Tensor,
+  ) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+    x = self.tok_embedding(idx)
+    cos, sin = self.rope_cache
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+    mask = self.mask_cache.index_select(2, input_pos)
+    mask = mask[:, :, :, : self.config.max_seq_len]
+
+    for i, block in enumerate(self.transformer_blocks):
+      input_k, input_v = k_caches[i], v_caches[i]
+      x, (updated_k, updated_v) = block(
+          x, (cos, sin), mask, input_pos, (input_k, input_v)
+      )
+      k_caches[i], v_caches[i] = updated_k, updated_v
+
+    x = self.final_norm(x)
+    return self.lm_head(x), k_caches, v_caches
+
+
+def _export_stablehlo_mlir(model, args):
+  ep = torch.export.export(model, args)
+  stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
+  return stablehlo_gm.get_stablehlo_text()
+
+
+def get_model_config() -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=32, num_query_groups=4, rotary_percentage=1.0
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+      intermediate_size=256,
+  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  config = cfg.ModelConfig(
+      vocab_size=150,
+      num_layers=2,
+      max_seq_len=100,
+      embedding_dim=128,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      pre_ff_norm_config=norm_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_sample_prefill_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+  idx = torch.unsqueeze(torch.arange(0, 100), 0)
+  input_pos = torch.arange(0, 100)
+  return idx, input_pos
+
+
+def get_sample_decode_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+  idx = torch.tensor([[1]], dtype=torch.long)
+  input_pos = torch.tensor([10])
+  return idx, input_pos
+
+
+def define_and_run() -> None:
+  dump_mlir = False
+
+  config = get_model_config()
+  model = ToyModelWithExternalKV(config)
+  print('running an inference')
+  k_caches = torch.zeros((2, 1, 100, 4, 4), dtype=torch.float32)
+  v_caches = torch.zeros((2, 1, 100, 4, 4), dtype=torch.float32)
+
+  idx, input_pos = get_sample_prefill_inputs()
+  decode_idx, decode_input_pos = get_sample_decode_inputs()
+  print(model.forward(idx, input_pos, k_caches, v_caches))
+
+  if dump_mlir:
+    mlir_text = _export_stablehlo_mlir(model, (idx, input_pos, k_caches, v_caches))
+    with open('/tmp/toy_model_with_external_kv.stablehlo.mlir', 'w') as f:
+      f.write(mlir_text)
+
+  # Convert model to tflite with 2 signatures (prefill + decode).
+  # TODO(b/344014416): currently conversion will fail, because we generate int64 index
+  # in dynamic update slice op.
+  print('converting toy model to tflite with 2 signatures (prefill + decode)')
+  edge_model = (
+      ai_edge_torch.signature('prefill', model, (idx, input_pos, k_caches, v_caches))
+      .signature('decode', model, (decode_idx, decode_input_pos, k_caches, v_caches))
+      .convert()
+  )
+  edge_model.export('/tmp/toy_external_kv_cache.tflite')
+
+
+if __name__ == '__main__':
+  with torch.inference_mode():
+    define_and_run()

ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py

@@ -0,0 +1,143 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# A toy example which has basic transformer block (w/ KV-Cache).
+from typing import List, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch_xla
+
+import ai_edge_torch
+from ai_edge_torch.generative.layers.attention import TransformerBlock
+import ai_edge_torch.generative.layers.attention_utils as attn_utils
+import ai_edge_torch.generative.layers.builder as builder
+import ai_edge_torch.generative.layers.model_config as cfg
+
+RoPECache = Tuple[torch.Tensor, torch.Tensor]
+
+
+class ToyModelWithKV(torch.nn.Module):
+
+  def __init__(self, config: cfg.ModelConfig) -> None:
+    super().__init__()
+    self.lm_head = nn.Linear(
+        config.embedding_dim, config.vocab_size, bias=config.lm_head_use_bias
+    )
+    self.tok_embedding = nn.Embedding(config.vocab_size, config.embedding_dim)
+    self.transformer_blocks = nn.ModuleList(
+        TransformerBlock(config) for _ in range(config.num_layers)
+    )
+    self.final_norm = builder.build_norm(
+        config.embedding_dim,
+        config.final_norm_config,
+    )
+    self.rope_cache = attn_utils.build_rope_cache(
+        size=config.max_seq_len,
+        dim=int(config.attn_config.rotary_percentage * config.head_dim),
+        base=10_000,
+        condense_ratio=1,
+        dtype=torch.float32,
+        device=torch.device('cpu'),
+    )
+    self.mask_cache = attn_utils.build_causal_mask_cache(
+        size=config.max_seq_len, dtype=torch.float32, device=torch.device('cpu')
+    )
+    self.config = config
+
+  @torch.inference_mode
+  def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
+    x = self.tok_embedding(idx)
+    cos, sin = self.rope_cache
+    cos = cos.index_select(0, input_pos)
+    sin = sin.index_select(0, input_pos)
+    mask = self.mask_cache.index_select(2, input_pos)
+    mask = mask[:, :, :, : self.config.max_seq_len]
+    for i, block in enumerate(self.transformer_blocks):
+      x = block(x, (cos, sin), mask, input_pos)
+    x = self.final_norm(x)
+    return self.lm_head(x)
+
+
+def _export_stablehlo_mlir(model, args):
+  ep = torch.export.export(model, args)
+  stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
+  return stablehlo_gm.get_stablehlo_text()
+
+
+def get_model_config() -> cfg.ModelConfig:
+  attn_config = cfg.AttentionConfig(
+      num_heads=32, num_query_groups=4, rotary_percentage=1.0
+  )
+  ff_config = cfg.FeedForwardConfig(
+      type=cfg.FeedForwardType.GATED,
+      activation=cfg.ActivationConfig(cfg.ActivationType.SILU),
+      intermediate_size=256,
+  )
+  norm_config = cfg.NormalizationConfig(type=cfg.NormalizationType.RMS_NORM)
+  config = cfg.ModelConfig(
+      vocab_size=150,
+      num_layers=2,
+      max_seq_len=500,
+      embedding_dim=128,
+      attn_config=attn_config,
+      ff_config=ff_config,
+      pre_attention_norm_config=norm_config,
+      pre_ff_norm_config=norm_config,
+      final_norm_config=norm_config,
+      enable_hlfb=True,
+  )
+  return config
+
+
+def get_sample_prefill_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+  idx = torch.unsqueeze(torch.arange(0, 100), 0)
+  input_pos = torch.arange(0, 100)
+  return idx, input_pos
+
+
+def get_sample_decode_inputs() -> Tuple[torch.Tensor, torch.Tensor]:
+  idx = torch.tensor([[1]], dtype=torch.long)
+  input_pos = torch.tensor([10], dtype=torch.int64)
+  return idx, input_pos
+
+
+def define_and_run() -> None:
+  dump_mlir = False
+
+  config = get_model_config()
+  model = ToyModelWithKV(config)
+  print('running an inference')
+  idx, input_pos = get_sample_prefill_inputs()
+  decode_idx, decode_input_pos = get_sample_decode_inputs()
+  print(model.forward(idx, input_pos))
+
+  if dump_mlir:
+    mlir_text = _export_stablehlo_mlir(model, (idx, input_pos))
+    with open('/tmp/toy_model_with_kv.stablehlo.mlir', 'w') as f:
+      f.write(mlir_text)
+
+  # Convert model to tflite with 2 signatures (prefill + decode).
+  print('converting toy model to tflite with 2 signatures (prefill + decode)')
+  edge_model = (
+      ai_edge_torch.signature('prefill', model, (idx, input_pos))
+      .signature('decode', model, (decode_idx, decode_input_pos))
+      .convert()
+  )
+  edge_model.export('/tmp/toy_kv_cache.tflite')
+
+
+if __name__ == '__main__':
+  define_and_run()