ai-edge-torch-nightly 0.5.0.dev20250423__py3-none-any.whl → 0.5.0.dev20250425__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -35,13 +35,11 @@ def _run_convert_passes(
   )
 
   passes = [
-      fx_passes.CastInputsBf16ToF32Pass(),
-      fx_passes.BuildInterpolateCompositePass(),
-      fx_passes.CanonicalizePass(),
       fx_passes.OptimizeLayoutTransposesPass(),
       fx_passes.CanonicalizePass(),
       fx_passes.BuildAtenCompositePass(),
       fx_passes.RemoveNonUserOutputsPass(),
+      fx_passes.CastInputsBf16ToF32Pass(),
   ]
 
   # Debuginfo is not injected automatically by odml_torch. Only inject

ai_edge_torch/_convert/fx_passes/__init__.py

@@ -16,7 +16,6 @@
 from typing import Sequence, Union
 
 from ai_edge_torch._convert.fx_passes.build_aten_composite_pass import BuildAtenCompositePass
-from ai_edge_torch._convert.fx_passes.build_interpolate_composite_pass import BuildInterpolateCompositePass
 from ai_edge_torch._convert.fx_passes.cast_inputs_bf16_to_f32_pass import CastInputsBf16ToF32Pass
 from ai_edge_torch._convert.fx_passes.inject_mlir_debuginfo_pass import InjectMlirDebuginfoPass
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import OptimizeLayoutTransposesPass

ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py

@@ -20,7 +20,8 @@ import torch
 import torch.utils._pytree as pytree
 
 _composite_builders: dict[
-    Callable, Callable[[torch.fx.GraphModule, torch.fx.Node], None]
+    Callable[[Any, ...], Any],
+    Callable[[torch.fx.GraphModule, torch.fx.Node], None],
 ] = {}
 
 
@@ -272,13 +273,73 @@ def _aten_embedding(gm: torch.fx.GraphModule, node: torch.fx.Node):
     output = op(**full_kwargs)
     output = builder.mark_outputs(output)
 
-    # Explicitly reshape back to the original shape. This places the ReshapeOp outside of the HLFB.
+    # Explicitly reshape back to the original shape. This places the ReshapeOp
+    # outside of the HLFB.
     output = torch.reshape(output, (*(original_idx_shape), embedding_dim))
     return output
 
   node.target = embedding
 
 
+@_register_composite_builder(torch.ops.aten.upsample_bilinear2d.vec)
+def _aten_upsample_bilinear2d_vec(_, node: torch.fx.Node):
+  """Build a composite for aten.upsample_bilinear2d.vec."""
+  op = node.target
+  args_mapper = TorchOpArgumentsMapper(op)
+  # Assumes later FX passes does not change the args/kwargs of the op.
+  # Which is a valid assumption for, given that composite/mark_tensor wrapper
+  # should semantically prevents any future mutations on the op.
+  output_h, output_w = node.meta["val"].shape[-2:]
+
+  def upsample_bilinear2d_vec(*args, **kwargs):
+    nonlocal op, args_mapper
+    full_kwargs = args_mapper.get_full_kwargs(args, kwargs)
+
+    builder = lowertools.StableHLOCompositeBuilder(
+        name="odml.upsample_bilinear2d",
+        attr={
+            "size": (int(output_h), int(output_w)),
+            "align_corners": full_kwargs["align_corners"],
+            "is_nchw_op": True,
+        },
+    )
+    full_kwargs["input"] = builder.mark_inputs(full_kwargs["input"])
+    output = op(**full_kwargs)
+    output = builder.mark_outputs(output)
+    return output
+
+  node.target = upsample_bilinear2d_vec
+
+
+@_register_composite_builder(torch.ops.aten.upsample_nearest2d.vec)
+def _aten_upsample_nearest2d_vec(_, node: torch.fx.Node):
+  """Build a composite for aten.upsample_nearest2d.vec."""
+  op = node.target
+  args_mapper = TorchOpArgumentsMapper(op)
+  # Assumes later FX passes does not change the args/kwargs of the op.
+  # Which is a valid assumption for, given that composite/mark_tensor wrapper
+  # should semantically prevents any future mutations on the op.
+  output_h, output_w = node.meta["val"].shape[-2:]
+
+  def upsample_nearest2d_vec(*args, **kwargs):
+    nonlocal op, args_mapper
+    full_kwargs = args_mapper.get_full_kwargs(args, kwargs)
+
+    builder = lowertools.StableHLOCompositeBuilder(
+        name="tfl.resize_nearest_neighbor",
+        attr={
+            "size": (int(output_h), int(output_w)),
+            "is_nchw_op": True,
+        },
+    )
+    full_kwargs["input"] = builder.mark_inputs(full_kwargs["input"])
+    output = op(**full_kwargs)
+    output = builder.mark_outputs(output)
+    return output
+
+  node.target = upsample_nearest2d_vec
+
+
 class BuildAtenCompositePass(fx_infra.PassBase):
 
   def call(self, graph_module: torch.fx.GraphModule):

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py

@@ -17,6 +17,7 @@
 import operator
 
 import ai_edge_torch
+from ai_edge_torch import lowertools
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import op_func_registry
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import utils
@@ -24,7 +25,7 @@ import torch
 import torch.utils._pytree as pytree
 
 aten = torch.ops.aten
-StableHLOCompositeBuilder = ai_edge_torch.hlfb.StableHLOCompositeBuilder
+StableHLOCompositeBuilder = lowertools.StableHLOCompositeBuilder
 
 __all__ = ["rewrite_nhwc_node", "has_nhwc_rewriter"]
 

ai_edge_torch/generative/examples/deepseek/convert_to_tflite.py

@@ -17,11 +17,43 @@
 
 from absl import app
 from ai_edge_torch.generative.examples.deepseek import deepseek
+from ai_edge_torch.generative.layers import kv_cache
 from ai_edge_torch.generative.utilities import converter
-from ai_edge_torch.generative.utilities import export_config
+from ai_edge_torch.generative.utilities.model_builder import export_cfg
+import torch
+
+flags = converter.define_conversion_flags('deepseek')
+ExportConfig = export_cfg.ExportConfig
+
+
+def _create_mask(mask_len, kv_cache_max_len):
+  mask = torch.full(
+      (mask_len, kv_cache_max_len), float('-inf'), dtype=torch.float32
+  )
+  mask = torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
+  return mask
+
+
+def _create_export_config(
+    prefill_seq_lens: list[int], kv_cache_max_len: int
+) -> ExportConfig:
+  """Creates the export config for the model."""
+  export_config = ExportConfig()
+  if isinstance(prefill_seq_lens, list):
+    prefill_mask = [_create_mask(i, kv_cache_max_len) for i in prefill_seq_lens]
+  else:
+    prefill_mask = _create_mask(prefill_seq_lens, kv_cache_max_len)
+
+  export_config.prefill_mask = prefill_mask
+
+  decode_mask = torch.full(
+      (1, kv_cache_max_len), float('-inf'), dtype=torch.float32
+  )
+  decode_mask = torch.triu(decode_mask, diagonal=1).unsqueeze(0).unsqueeze(0)
+  export_config.decode_mask = decode_mask
+  export_config.kvcache_layout = kv_cache.KV_LAYOUT_TRANSPOSED
+  return export_config
 
-flags = converter.define_conversion_flags("deepseek")
-ExportConfig = export_config.ExportConfig
 
 def main(_):
   pytorch_model = deepseek.build_model(
@@ -34,7 +66,9 @@ def main(_):
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=ExportConfig(),
+      export_config=_create_export_config(
+          flags.FLAGS.prefill_seq_lens, flags.FLAGS.kv_cache_max_len
+      ),
   )
 
 

ai_edge_torch/generative/examples/deepseek/deepseek.py

@@ -53,6 +53,7 @@ def get_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   norm_config = cfg.NormalizationConfig(
       type=cfg.NormalizationType.RMS_NORM,
       epsilon=1e-06,
+      enable_hlfb=True,
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,

ai_edge_torch/generative/examples/gemma3/decoder.py

@@ -17,10 +17,10 @@
 
 from typing import List, Optional, Tuple
 
+from ai_edge_torch.generative.layers import attention
 from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
-from ai_edge_torch.generative.layers.experimental import attention
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 from ai_edge_torch.generative.utilities import export_config as export_cfg

ai_edge_torch/generative/examples/qwen/convert_to_tflite.py

@@ -17,13 +17,14 @@
 
 from absl import app
 from ai_edge_torch.generative.examples.qwen import qwen
+from ai_edge_torch.generative.layers import kv_cache
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
+import torch
 
 flags = converter.define_conversion_flags('qwen')
 ExportConfig = export_config.ExportConfig
 
-
 _MODEL_SIZE = flags.DEFINE_enum(
     'model_size',
     '3b',
@@ -37,6 +38,36 @@ _BUILDER = {
     '3b': qwen.build_3b_model,
 }
 
+
+def _create_mask(mask_len, kv_cache_max_len):
+  mask = torch.full(
+      (mask_len, kv_cache_max_len), float('-inf'), dtype=torch.float32
+  )
+  mask = torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0)
+  return mask
+
+
+def _create_export_config(
+    prefill_seq_lens: list[int], kv_cache_max_len: int
+) -> ExportConfig:
+  """Creates the export config for the model."""
+  export_config = ExportConfig()
+  if isinstance(prefill_seq_lens, list):
+    prefill_mask = [_create_mask(i, kv_cache_max_len) for i in prefill_seq_lens]
+  else:
+    prefill_mask = _create_mask(prefill_seq_lens, kv_cache_max_len)
+
+  export_config.prefill_mask = prefill_mask
+
+  decode_mask = torch.full(
+      (1, kv_cache_max_len), float('-inf'), dtype=torch.float32
+  )
+  decode_mask = torch.triu(decode_mask, diagonal=1).unsqueeze(0).unsqueeze(0)
+  export_config.decode_mask = decode_mask
+  export_config.kvcache_layout = kv_cache.KV_LAYOUT_TRANSPOSED
+  return export_config
+
+
 def main(_):
   pytorch_model = _BUILDER[_MODEL_SIZE.value](
       flags.FLAGS.checkpoint_path, kv_cache_max_len=flags.FLAGS.kv_cache_max_len
@@ -48,7 +79,11 @@ def main(_):
       prefill_seq_len=flags.FLAGS.prefill_seq_lens,
       quantize=flags.FLAGS.quantize,
       lora_ranks=flags.FLAGS.lora_ranks,
-      export_config=ExportConfig(),
+      export_config=_create_export_config(
+          flags.FLAGS.prefill_seq_lens, flags.FLAGS.kv_cache_max_len
+      )
+      if flags.FLAGS.transpose_kv_cache
+      else ExportConfig(),
   )
 
 

ai_edge_torch/generative/examples/qwen/qwen.py

@@ -53,6 +53,7 @@ def get_3b_model_config(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
   norm_config = cfg.NormalizationConfig(
       type=cfg.NormalizationType.RMS_NORM,
       epsilon=1e-06,
+      enable_hlfb=True,
   )
   block_config = cfg.TransformerBlockConfig(
       attn_config=attn_config,

ai_edge_torch/generative/layers/attention.py

@@ -21,6 +21,7 @@ from ai_edge_torch.generative.layers import builder
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.layers import lora as lora_utils
 from ai_edge_torch.generative.layers import scaled_dot_product_attention as sdpa
+from ai_edge_torch.generative.layers import sdpa_with_kv_update
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 import torch
@@ -142,11 +143,6 @@ class CausalSelfAttention(nn.Module):
     self.key_norm = builder.build_norm(config.head_dim, config.key_norm_config)
     self.config = config
     self.enable_hlfb = enable_hlfb
-    self.sdpa_func = (
-        sdpa.scaled_dot_product_attention_with_hlfb
-        if enable_hlfb
-        else sdpa.scaled_dot_product_attention
-    )
 
   def forward(
       self,
@@ -174,7 +170,7 @@ class CausalSelfAttention(nn.Module):
         KV Cach Entry (if passed in).
     """
     # Batch size, sequence length, embedding dimensionality.
-    B, T, E = x.size()
+    B, T, _ = x.size()
     qkv = self.qkv_projection(x)
 
     # Assemble into a number of query groups to support MHA, MQA and GQA.
@@ -218,19 +214,9 @@ class CausalSelfAttention(nn.Module):
       cos, sin = rope
       q, k = rotary_pos_emb.apply_rope_inline(q, k, cos, sin)
 
-    if kv_cache is not None:
-      kv_cache = kv_utils.update(kv_cache, input_pos, k, v)
-      k, v = kv_cache.k_cache, kv_cache.v_cache
-
-    sdpa_out = self.sdpa_func(
-        q,
-        k,
-        v,
-        self.config.head_dim,
-        mask=mask,
-        softcap=self.config.logit_softcap,
+    sdpa_out, kv_cache = sdpa_with_kv_update.sdpa_with_kv_update(
+        q, k, v, kv_cache, input_pos, mask, self.config, self.enable_hlfb
     )
-    sdpa_out = sdpa_out.reshape(B, T, -1)
 
     # Compute the output projection.
     y = self.output_projection(sdpa_out)

ai_edge_torch/generative/layers/scaled_dot_product_attention.py

@@ -17,6 +17,8 @@
 import math
 from typing import Optional
 
+from ai_edge_torch.generative.custom_ops import bmm_4d as bmm_lib
+from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 import torch
 import torch.nn.functional as F
@@ -142,3 +144,52 @@ def scaled_dot_product_attention_with_hlfb(
   result = y.transpose(1, 2)
   result = builder.mark_outputs(result)
   return result
+
+
+def scaled_dot_product_attention_transposed(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    head_size: int,
+    mask: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+    softcap: Optional[float] = None,
+):
+  """Scaled dot product attention with transposed key and value.
+
+  Args:
+    query: Query tensor, with shape [B, T, N, H].
+    key: Key tensor, with shape [B, T, KV_LEN, H].
+    value: Value tensor, with shape [B, T, KV_LEN, H].
+    head_size (int): head dimension.
+    mask (torch.Tensor): the optional mask tensor.
+    scale (float): the optional scale factor.
+    softcap (float): the optional softcap for the logits.
+
+  Returns:
+    The output tensor of scaled_dot_product_attention_transposed.
+  """
+
+  if scale is None:
+    scale = 1.0 / math.sqrt(head_size)
+
+  query = query * scale
+
+  assert mask is not None, "Mask should not be None!"
+  t = mask.shape[2]
+
+  logits = bmm_lib.bmm_4d(query, key)
+
+  _, bk, gt, s = logits.shape
+  g = gt // t
+  logits = logits.reshape((bk, g, t, s))
+  if softcap is not None:
+    logits = torch.tanh(logits / softcap)
+    logits = logits * softcap
+
+  padded_logits = logits + mask
+  padded_logits = padded_logits.reshape(1, bk, gt, s)
+  probs = F.softmax(padded_logits, dim=-1).type_as(key)
+  encoded = bmm_lib.bmm_4d(probs, value)
+
+  return encoded  # 1, bk, gt, h

ai_edge_torch/generative/layers/sdpa_with_kv_update.py

@@ -12,16 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Common utility functions for data loading etc.
-from dataclasses import dataclass
+
+"""Common utility functions for data loading etc."""
+
 from typing import Tuple
+
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
-from ai_edge_torch.generative.layers import scaled_dot_product_attention as sdpa_default
+from ai_edge_torch.generative.layers import scaled_dot_product_attention as sdpa
 from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils_experimental
-from ai_edge_torch.generative.layers.experimental import scaled_dot_product_attention as sdpa
 import ai_edge_torch.generative.layers.model_config as cfg
-from ai_edge_torch.generative.utilities import types
-from multipledispatch import dispatch
 import torch
 
 
@@ -33,32 +32,27 @@ def sdpa_with_kv_update(
     input_pos: torch.Tensor,
     mask: torch.Tensor,
     config: cfg.AttentionConfig,
+    enable_hlfb: bool,
 ) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
-  return sdpa_with_kv_update_impl(
-      kv.kv_layout[0](),  # key layout
-      kv.kv_layout[1](),  # value layout
-      query=query,
-      key=key,
-      value=value,
-      kv=kv,
-      input_pos=input_pos,
-      mask=mask,
-      config=config,
+  """Wrapper function for scaled dot product attention with KV cache update."""
+  if kv is not None and kv.kv_layout == kv_utils.KV_LAYOUT_TRANSPOSED:
+    return _sdpa_with_kv_update_transposed(
+        query, key, value, kv, input_pos, mask, config
+    )
+  return _sdpa_with_kv_update_default(
+      query, key, value, kv, input_pos, mask, config, enable_hlfb
   )
 
 
-@dispatch(types.BNTH, types.BNHT)
-def sdpa_with_kv_update_impl(
-    k_type, v_type, *args, **kwargs
+def _sdpa_with_kv_update_transposed(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    kv: kv_utils.KVCacheEntry,
+    input_pos: torch.Tensor,
+    mask: torch.Tensor,
+    config: cfg.AttentionConfig,
 ) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
-  query = kwargs["query"]
-  key = kwargs["key"]
-  value = kwargs["value"]
-  kv = kwargs["kv"]
-  input_pos = kwargs["input_pos"]
-  mask = kwargs["mask"]
-  config = kwargs["config"]
-
   # Transpose k/v to specific layout for GPU implementation.
   b, seq_len, n, h = query.shape
   g = n // config.num_query_groups
@@ -74,12 +68,10 @@ def sdpa_with_kv_update_impl(
       1, -1, config.head_dim, seq_len
   )  # 1, bk, h, s
 
-  if kv is not None:
-    kv = kv_utils_experimental.update(kv, input_pos, key, value)
-    key, value = kv.k_cache, kv.v_cache
+  kv = kv_utils_experimental.update(kv, input_pos, key, value)
+  key, value = kv.k_cache, kv.v_cache
 
-  sdpa_out = sdpa.scaled_dot_product_attention(
-      kv,
+  sdpa_out = sdpa.scaled_dot_product_attention_transposed(
       query,
       key,
       value,
@@ -95,24 +87,26 @@ def sdpa_with_kv_update_impl(
   return sdpa_out, kv
 
 
-@dispatch(object, object)
-def sdpa_with_kv_update_impl(
-    k_type, v_type, *args, **kwargs
+def _sdpa_with_kv_update_default(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    kv: kv_utils.KVCacheEntry,
+    input_pos: torch.Tensor,
+    mask: torch.Tensor,
+    config: cfg.AttentionConfig,
+    enable_hlfb: bool,
 ) -> Tuple[torch.Tensor, kv_utils.KVCacheEntry]:
-  query = kwargs["query"]
-  key = kwargs["key"]
-  value = kwargs["value"]
-  kv = kwargs["kv"]
-  input_pos = kwargs["input_pos"]
-  mask = kwargs["mask"]
-  config = kwargs["config"]
-
   b, seq_len, _, _ = query.shape
   if kv is not None:
     kv = kv_utils.update(kv, input_pos, key, value)
     key, value = kv.k_cache, kv.v_cache
 
-  sdpa_out = sdpa_default.scaled_dot_product_attention(
+  if enable_hlfb:
+    sdpa_func = sdpa.scaled_dot_product_attention_with_hlfb
+  else:
+    sdpa_func = sdpa.scaled_dot_product_attention
+  sdpa_out = sdpa_func(
       query,
       key,
       value,

ai_edge_torch/generative/test/test_model_conversion.py

@@ -32,16 +32,14 @@ class TestModelConversion(googletest.TestCase):
 
   def setUp(self):
     super().setUp()
-    # Builder function for an Interpreter that supports custom ops.
     self._interpreter_builder = (
-        lambda tflite_model: lambda: interpreter.InterpreterWithCustomOps(
-            custom_op_registerers=["GenAIOpsRegisterer"],
+        lambda tflite_model: lambda: interpreter.Interpreter(
             model_content=tflite_model,
             experimental_default_delegate_latest_features=True,
         )
     )
 
-  def _get_params(self, enable_hlfb: bool):
+  def _get_params(self, enable_hlfb: bool, kv_layout: kv_cache.KVLayout):
     """Returns a model, edge model and the kwargs to use for testing."""
     config = toy_model_with_kv_cache.get_model_config()
     config.enable_hlfb = enable_hlfb
@@ -49,7 +47,7 @@ class TestModelConversion(googletest.TestCase):
     tokens, input_pos = torch.tensor([[1]], dtype=torch.int), torch.tensor(
         [10], dtype=torch.int
     )
-    kv = kv_cache.KVCache.from_model_config(config)
+    kv = kv_cache.KVCache.from_model_config(config, kv_layout=kv_layout)
     kwargs = {
         "tokens": tokens,
         "input_pos": input_pos,
@@ -65,8 +63,12 @@ class TestModelConversion(googletest.TestCase):
     )
     return pytorch_model, edge_model, kwargs
 
-  def _test_model_with_kv_cache(self, enable_hlfb: bool):
-    pytorch_model, edge_model, kwargs = self._get_params(enable_hlfb)
+  def _test_model_with_kv_cache(
+      self,
+      enable_hlfb: bool = False,
+      kv_layout: kv_cache.KVLayout = kv_cache.KV_LAYOUT_DEFAULT,
+  ):
+    pytorch_model, edge_model, kwargs = self._get_params(enable_hlfb, kv_layout)
 
     self.assertTrue(
         test_utils.compare_tflite_torch(
@@ -81,38 +83,34 @@ class TestModelConversion(googletest.TestCase):
         )
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_toy_model_with_kv_cache(self):
     self._test_model_with_kv_cache(enable_hlfb=False)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_toy_model_with_kv_cache_with_hlfb(self):
     self._test_model_with_kv_cache(enable_hlfb=True)
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
+  def test_toy_model_with_kv_cache_transposed(self):
+    self._test_model_with_kv_cache(kv_layout=kv_cache.KV_LAYOUT_TRANSPOSED)
+
   def test_toy_model_has_dus_op(self):
     """Tests that the model has the dynamic update slice op."""
-    _, edge_model, _ = self._get_params(enable_hlfb=True)
-    interpreter_ = interpreter.InterpreterWithCustomOps(
-        custom_op_registerers=["GenAIOpsRegisterer"],
-        model_content=edge_model.tflite_model(),
-        experimental_default_delegate_latest_features=True,
+    _, edge_model, _ = self._get_params(
+        enable_hlfb=True, kv_layout=kv_cache.KV_LAYOUT_DEFAULT
     )
+    interpreter = self._interpreter_builder(edge_model.tflite_model())()
 
     # pylint: disable=protected-access
-    op_names = [op["op_name"] for op in interpreter_._get_ops_details()]
+    op_names = [op["op_name"] for op in interpreter._get_ops_details()]
     self.assertIn("DYNAMIC_UPDATE_SLICE", op_names)
 
-  def _test_multisig_model(self, config, pytorch_model, atol, rtol):
+  def _test_multisig_model(
+      self,
+      config,
+      pytorch_model,
+      atol,
+      rtol,
+      kv_layout=kv_cache.KV_LAYOUT_DEFAULT,
+  ):
     # prefill
     seq_len = 10
     prefill_tokens = torch.zeros((1, seq_len), dtype=torch.int, device="cpu")
@@ -124,7 +122,7 @@ class TestModelConversion(googletest.TestCase):
     decode_token = torch.tensor([[1]], dtype=torch.int)
     decode_input_pos = torch.tensor([5], dtype=torch.int)
 
-    kv = kv_cache.KVCache.from_model_config(config)
+    kv = kv_cache.KVCache.from_model_config(config, kv_layout=kv_layout)
 
     edge_model = (
         ai_edge_torch.signature(
@@ -160,7 +158,7 @@ class TestModelConversion(googletest.TestCase):
             kv,
             signature_name="prefill",
             atol=atol,
-            rtol=atol,
+            rtol=rtol,
         )
     )
 
@@ -173,19 +171,26 @@ class TestModelConversion(googletest.TestCase):
             kv,
             signature_name="decode",
             atol=atol,
-            rtol=atol,
+            rtol=rtol,
         )
     )
 
-  @googletest.skipIf(
-      ai_edge_torch.config.in_oss,
-      reason="tests with custom ops are not supported in oss",
-  )
   def test_tiny_llama_multisig(self):
     config = tiny_llama.get_fake_model_config()
     pytorch_model = tiny_llama.TinyLlama(config).eval()
     self._test_multisig_model(config, pytorch_model, atol=1e-5, rtol=1e-5)
 
+  def test_tiny_llama_multisig_kv_layout_transposed(self):
+    config = tiny_llama.get_fake_model_config()
+    pytorch_model = tiny_llama.TinyLlama(config).eval()
+    self._test_multisig_model(
+        config,
+        pytorch_model,
+        atol=1e-5,
+        rtol=1e-5,
+        kv_layout=kv_cache.KV_LAYOUT_TRANSPOSED,
+    )
+
 
 if __name__ == "__main__":
   googletest.main()