ai-edge-torch-nightly 0.2.0.dev20240610__py3-none-any.whl → 0.2.0.dev20240617__py3-none-any.whl

ai_edge_torch/generative/layers/unet/model_config.py

@@ -22,16 +22,28 @@ from typing import List, Optional
 import ai_edge_torch.generative.layers.model_config as layers_cfg
 
 
-@dataclass
+@enum.unique
 class SamplingType(enum.Enum):
   NEAREST = enum.auto()
   BILINEAR = enum.auto()
+  AVERAGE = enum.auto()
+  CONVOLUTION = enum.auto()
 
 
 @dataclass
-class SamplingConfig:
+class UpSamplingConfig:
+  mode: SamplingType
   scale_factor: float
+
+
+@dataclass
+class DownSamplingConfig:
   mode: SamplingType
+  in_channels: int
+  kernel_size: int
+  stride: int
+  padding: int
+  out_channels: Optional[int] = None
 
 
 @dataclass
@@ -46,9 +58,38 @@ class ResidualBlock2DConfig:
 
 @dataclass
 class AttentionBlock2DConfig:
-  dims: int
+  dim: int
+  normalization_config: layers_cfg.NormalizationConfig
+  attention_config: layers_cfg.AttentionConfig
+  enable_hlfb: bool = True
+  attention_batch_size: int = 1
+
+
+@dataclass
+class CrossAttentionBlock2DConfig:
+  query_dim: int
+  cross_dim: int
   normalization_config: layers_cfg.NormalizationConfig
   attention_config: layers_cfg.AttentionConfig
+  enable_hlfb: bool = True
+  attention_batch_size: int = 1
+
+
+@dataclass
+class FeedForwardBlock2DConfig:
+  dim: int
+  hidden_dim: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  use_bias: bool
+
+
+@dataclass
+class TransformerBlock2Dconfig:
+  pre_conv_normalization_config: layers_cfg.NormalizationConfig
+  attention_block_config: AttentionBlock2DConfig
+  cross_attention_block_config: CrossAttentionBlock2DConfig
+  feed_forward_block_config: FeedForwardBlock2DConfig
 
 
 @dataclass
@@ -58,14 +99,62 @@ class UpDecoderBlock2DConfig:
   normalization_config: layers_cfg.NormalizationConfig
   activation_config: layers_cfg.ActivationConfig
   num_layers: int
-  # Optional time embedding channels if the residual blocks take a time embedding context as input
+  # Optional time embedding channels if the residual blocks take a time embedding as input
   time_embedding_channels: Optional[int] = None
   # Whether to add upsample operation after residual blocks
   add_upsample: bool = True
   # Whether to add a conv2d layer after upsample
   upsample_conv: bool = True
   # Optional sampling config if add_upsample is True.
-  sampling_config: Optional[SamplingConfig] = None
+  sampling_config: Optional[UpSamplingConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2Dconfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
+
+
+@dataclass
+class SkipUpDecoderBlock2DConfig:
+  in_channels: int
+  out_channels: int
+  # The dimension of output channels of previous connected block
+  prev_out_channels: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  num_layers: int
+  # Optional time embedding channels if the residual blocks take a time embedding as input
+  time_embedding_channels: Optional[int] = None
+  # Whether to add upsample operation after residual blocks
+  add_upsample: bool = True
+  # Whether to add a conv2d layer after upsample
+  upsample_conv: bool = True
+  # Optional sampling config if add_upsample is True.
+  sampling_config: Optional[UpSamplingConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2Dconfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
+
+
+@dataclass
+class DownEncoderBlock2DConfig:
+  in_channels: int
+  out_channels: int
+  normalization_config: layers_cfg.NormalizationConfig
+  activation_config: layers_cfg.ActivationConfig
+  num_layers: int
+  # Padding for the downsampling convolution.
+  padding: int = 1
+  # Optional time embedding channels if the residual blocks take a time embedding as input
+  time_embedding_channels: Optional[int] = None
+  # Whether to add downsample operation after residual blocks
+  add_downsample: bool = True
+  # Optional sampling config if add_upsample is True.
+  sampling_config: Optional[DownSamplingConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2Dconfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
 
 
 @dataclass
@@ -78,6 +167,10 @@ class MidBlock2DConfig:
   time_embedding_channels: Optional[int] = None
   # Optional config of attention blocks interleaved with residual blocks
   attention_block_config: Optional[AttentionBlock2DConfig] = None
+  # Optional config of transformer blocks interleaved with residual blocks
+  transformer_block_config: Optional[TransformerBlock2Dconfig] = None
+  # Optional dimension of context tensor if context tensor is given as input.
+  context_dim: Optional[int] = None
 
 
 @dataclass
@@ -115,3 +208,62 @@ class AutoEncoderConfig:
 
   # The configuration of middle blocks, that is, after the last block of encoder and before the first block of decoder.
   mid_block_config: MidBlock2DConfig
+
+
+@dataclass
+class DiffusionModelConfig:
+  """Configurations of Diffusion model."""
+
+  # Number of channels in the input tensor.
+  in_channels: int
+
+  # Number of channels in the output tensor.
+  out_channels: int
+
+  # The output channels of each block.
+  block_out_channels: List[int]
+
+  # The layesr number of each block.
+  layers_per_block: int
+
+  # The padding to use for the downsampling.
+  downsample_padding: int
+
+  # Normalization config used in residual blocks.
+  residual_norm_config: layers_cfg.NormalizationConfig
+
+  # Activation config used in residual blocks
+  residual_activation_type: layers_cfg.ActivationType
+
+  # The batch size used in transformer blocks, for attention layers.
+  transformer_batch_size: int
+
+  # The number of attention heads used in transformer blocks.
+  transformer_num_attention_heads: int
+
+  # The dimension of cross attention used in transformer blocks.
+  transformer_cross_attention_dim: int
+
+  # Normalization config used in prev conv layer of transformer blocks.
+  transformer_pre_conv_norm_config: layers_cfg.NormalizationConfig
+
+  # Normalization config used in transformer blocks.
+  transformer_norm_config: layers_cfg.NormalizationConfig
+
+  # Activation type of feed forward used in transformer blocks.
+  transformer_ff_activation_type: layers_cfg.ActivationType
+
+  # Number of layers in mid block.
+  mid_block_layers: int
+
+  # Dimension of time embedding.
+  time_embedding_dim: int
+
+  # Time embedding dimensions for blocks.
+  time_embedding_blocks_dim: int
+
+  # Normalization config used for final layer
+  final_norm_config: layers_cfg.NormalizationConfig
+
+  # Activation type used in final layer
+  final_activation_type: layers_cfg.ActivationType

ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/translate_recipe.py

@@ -0,0 +1,164 @@
+# 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.
+# ==============================================================================
+
+import json
+
+from ai_edge_quantizer import quantizer
+
+from ai_edge_torch.generative.quantize import quant_attrs
+from ai_edge_torch.generative.quantize import quant_recipe
+
+_OpExecutionMode = quantizer.qtyping.OpExecutionMode
+_OpName = quantizer.qtyping.TFLOperationName
+_TensorQuantConfig = quantizer.qtyping.TensorQuantizationConfig
+_OpQuantConfig = quantizer.qtyping.OpQuantizationConfig
+
+_DEFAULT_REGEX_STR = '.*'
+_ATTENTION_IDX_REGEX_STR = (
+    'transformer_blocks\[{}\]/ai_edge_torch.generative.layers.attention'
+)
+_FEEDFORWARD_IDX_REGEX_STR = (
+    'transformer_blocks\[{}\]/ai_edge_torch.generative.layers.feed_forward'
+)
+_EMBEDDING_REGEX_STR = 'Embedding_tok_embedding'
+_ANY_TWO_DIGITS_REGEX_STR = '\d{1,2}'
+
+
+def _get_nbits_from_dtype(dtype: quant_attrs.Dtype) -> int:
+  if dtype == quant_attrs.Dtype.FP32:
+    return 32
+  elif dtype == quant_attrs.Dtype.FP16:
+    return 16
+  elif dtype == quant_attrs.Dtype.INT8:
+    return 8
+  raise ValueError('Unimplemented number of bits')
+
+
+def _get_dtype_from_dtype(dtype: quant_attrs.Dtype) -> quantizer.qtyping.TensorDataType:
+  if dtype == quant_attrs.Dtype.FP32 or dtype == quant_attrs.Dtype.FP16:
+    return quantizer.qtyping.TensorDataType.FLOAT
+  else:
+    return quantizer.qtyping.TensorDataType.INT
+
+
+def _get_execution_mode_from_mode(mode: quant_attrs.Mode) -> _OpExecutionMode:
+  if mode == quant_attrs.Mode.DYNAMIC_RANGE:
+    return _OpExecutionMode.DRQ
+  elif mode == quant_attrs.Mode.WEIGHT_ONLY:
+    return _OpExecutionMode.WEIGHT_ONLY
+  raise ValueError('Unimplemented execution mode')
+
+
+def _get_channelwise_from_granularity(granularity: quant_attrs.Granularity) -> bool:
+  if granularity == quant_attrs.Granularity.CHANNELWISE:
+    return True
+  elif granularity == quant_attrs.Granularity.NONE:
+    return False
+  raise ValueError('Unimplemented granularity')
+
+
+def _get_algorithm_key_from_algorithm(algo: quant_attrs.Algorithm) -> str:
+  if algo == quant_attrs.Algorithm.MIN_MAX:
+    return quantizer.algorithm_manager.AlgorithmName.MIN_MAX_UNIFORM_QUANT
+  elif algo == quant_attrs.Algorithm.FLOAT_CAST:
+    return quantizer.algorithm_manager.AlgorithmName.FLOAT_CASTING
+  raise ValueError('Unimplemented algorithm')
+
+
+def _set_quant_config(
+    rm: quantizer.recipe_manager.RecipeManager,
+    layer_recipe: quant_recipe.LayerQuantRecipe,
+    regex: str,
+):
+  support_op_list = [_OpName.FULLY_CONNECTED, _OpName.CONV_2D]
+  if layer_recipe.algorithm == quant_attrs.Algorithm.MIN_MAX:
+    support_op_list += [_OpName.BATCH_MATMUL, _OpName.EMBEDDING_LOOKUP]
+  for op_name in support_op_list:
+    rm.add_quantization_config(
+        regex=regex,
+        operation_name=op_name,
+        op_config=_OpQuantConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=_get_nbits_from_dtype(layer_recipe.weight_dtype),
+                symmetric=True,
+                channel_wise=_get_channelwise_from_granularity(
+                    layer_recipe.granularity
+                ),
+                dtype=_get_dtype_from_dtype(layer_recipe.weight_dtype),
+            ),
+            execution_mode=_get_execution_mode_from_mode(layer_recipe.mode),
+        ),
+        algorithm_key=_get_algorithm_key_from_algorithm(layer_recipe.algorithm),
+        override_algorithm=True,
+    )
+
+
+def translate_to_ai_edge_recipe(
+    recipe: quant_recipe.GenerativeQuantRecipe,
+) -> quantizer.recipe_manager.ModelQuantizationRecipe:
+  rm = quantizer.recipe_manager.RecipeManager()
+
+  if recipe.default is not None:
+    _set_quant_config(rm, recipe.default, _DEFAULT_REGEX_STR)
+
+  if recipe.embedding is not None:
+    _set_quant_config(rm, recipe.embedding, _EMBEDDING_REGEX_STR)
+
+  if recipe.attention is not None:
+    if isinstance(recipe.attention, dict):
+      for idx, layer in recipe.attention.items():
+        _set_quant_config(rm, layer, _ATTENTION_IDX_REGEX_STR.format(idx))
+    else:
+      _set_quant_config(
+          rm,
+          recipe.attention,
+          _ATTENTION_IDX_REGEX_STR.format(_ANY_TWO_DIGITS_REGEX_STR),
+      )
+
+  if recipe.feedforward is not None:
+    if isinstance(recipe.feedforward, dict):
+      for idx, layer in recipe.feedforward.items():
+        _set_quant_config(rm, layer, _FEEDFORWARD_IDX_REGEX_STR.format(idx))
+    else:
+      _set_quant_config(
+          rm,
+          recipe.feedforward,
+          _FEEDFORWARD_IDX_REGEX_STR.format(_ANY_TWO_DIGITS_REGEX_STR),
+      )
+
+  return rm.get_quantization_recipe()
+
+
+def quantize_model(
+    model: bytearray, recipe: quantizer.recipe_manager.ModelQuantizationRecipe
+) -> bytearray:
+  # TODO(b/336599483): Remove tempfile and use bytearray instead
+  tmp_model_path = '/tmp/tmp.tflite'
+  tmp_recipe_path = '/tmp/recipe.json'
+  with open(tmp_model_path, 'wb') as fp:
+    fp.write(model)
+  with open(tmp_recipe_path, 'w') as rp:
+    rp.write(json.dumps(recipe))
+
+  qt = quantizer.Quantizer(tmp_model_path, tmp_recipe_path)
+  result = qt.quantize()
+
+  # TODO(b/336599483): Remove tempfile and use bytearray instead
+  import os
+
+  os.remove(tmp_model_path)
+  os.remove(tmp_recipe_path)
+
+  return result.quantized_model

ai_edge_torch/generative/quantize/quant_attrs.py

@@ -32,9 +32,11 @@ class Algorithm(enum.Enum):
   Attributes:
     MIN_MAX: Maps the min/max of floating point space to the min/max of
       quantized space and quantize uniformly.
+    FLOAT_CAST: Casts a float to another float of a different type.
   """
 
   MIN_MAX = enum.auto()
+  FLOAT_CAST = enum.auto()
 
 
 @enum.unique

ai_edge_torch/generative/quantize/quant_recipe.py

@@ -14,8 +14,7 @@
 # ==============================================================================
 
 from dataclasses import dataclass
-import enum
-from typing import Optional
+from typing import Optional, Union
 
 from ai_edge_torch.generative.quantize import quant_attrs
 from ai_edge_torch.generative.quantize import supported_schemes
@@ -80,18 +79,50 @@ class LayerQuantRecipe:
 
 
 @dataclass
-class TransformerQuantRecipe:
+class GenerativeQuantRecipe:
   """Quantization recipe for a model composed of the Edge Generative API layers.
 
+  Some layers can be specified with different `LayerQuantRecipe` for each block by
+  providing a dictionary keyed by the TransformerBlock index, e.g. attention
+  and feedforward. For example,
+
+  ```
+  default = LayerQuantRecipeA
+  attention = { 2: LayerQuantRecipeB }
+  feedforward = { 3: LayerQuantRecipeC }
+  ```
+
+  will apply LayerQuantRecipeA to the entire model, overriden by
+  LayerQuantRecipeB for the TransformerBlock[2].attention layer and
+  LayerQuantRecipeC for the TransformerBlock[3].feedforward layer. Any config
+  with invalid indices will be ignored.
+
   Attributes:
     default: The quantization recipe for global scope of the model.
+    embedding: Recipe for the embedding table.
+    attention: Recipe for the attention blocks. This could be specified with
+      different LayerQuantRecipe for each block by providing a dictionary
+      keyed by the TransformerBlock index.
+    feedforward: Recipe for the feedforward layers. This could be specified with
+      different LayerQuantRecipe for each block by providing a dictionary
+      keyed by the TransformerBlock index.
   """
 
   default: Optional[LayerQuantRecipe] = None
+  embedding: Optional[LayerQuantRecipe] = None
+  attention: Union[
+      Optional[LayerQuantRecipe], Optional[dict[int, LayerQuantRecipe]]
+  ] = None
+  feedforward: Union[
+      Optional[LayerQuantRecipe], Optional[dict[int, LayerQuantRecipe]]
+  ] = None
 
   def __str__(self):
-    return f"""TransformerQuantRecipe(
+    return f"""GenerativeQuantRecipe(
   Default: {self.default}
+  Embedding: {self.embedding}
+  Attention: {self.attention}
+  Feedforward: {self.feedforward}
 )"""
 
   __repr__ = __str__
@@ -104,3 +135,17 @@ class TransformerQuantRecipe:
     """
     if self.default is not None:
       self.default.verify()
+    if self.embedding is not None:
+      self.embedding.verify()
+    if self.attention is not None:
+      if isinstance(self.attention, dict):
+        for recipe in self.attention.values():
+          recipe.verify()
+      else:
+        self.attention.verify()
+    if self.feedforward is not None:
+      if isinstance(self.feedforward, dict):
+        for recipe in self.feedforward.values():
+          recipe.verify()
+      else:
+        self.feedforward.verify()

ai_edge_torch/generative/quantize/quant_recipe_utils.py

@@ -22,7 +22,7 @@ Typical usage example:
 
 1. Applying a single layer recipe to the entire model
 
-  quant_recipe.TransformerQuantRecipe(
+  quant_recipe.GenerativeQuantRecipe(
     default=quant_recipe_utils.create_layer_quant_int8_dynamic()
   )
 """
@@ -46,6 +46,6 @@ def create_layer_quant_fp16() -> quant_recipe.LayerQuantRecipe:
       activation_dtype=quant_attrs.Dtype.FP32,
       weight_dtype=quant_attrs.Dtype.FP16,
       mode=quant_attrs.Mode.WEIGHT_ONLY,
-      algorithm=quant_attrs.Algorithm.MIN_MAX,
+      algorithm=quant_attrs.Algorithm.FLOAT_CAST,
       granularity=quant_attrs.Granularity.NONE,
   )

ai_edge_torch/generative/quantize/quant_recipes.py

@@ -34,15 +34,15 @@ from ai_edge_torch.quantize import quant_config
 
 def full_linear_int8_dynamic_recipe() -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
-      transformer_recipe=quant_recipe.TransformerQuantRecipe(
-          default=quant_recipe_utils.create_layer_quant_int8_dynamic()
+      generative_recipe=quant_recipe.GenerativeQuantRecipe(
+          default=quant_recipe_utils.create_layer_quant_int8_dynamic(),
       )
   )
 
 
 def full_fp16_recipe() -> quant_config.QuantConfig:
   return quant_config.QuantConfig(
-      transformer_recipe=quant_recipe.TransformerQuantRecipe(
+      generative_recipe=quant_recipe.GenerativeQuantRecipe(
           default=quant_recipe_utils.create_layer_quant_fp16()
       )
   )

ai_edge_torch/generative/quantize/supported_schemes.py

@@ -27,5 +27,6 @@ def get_supported_layer_schemes():
 
   return [
       (_t.FP32, _t.INT8, _m.DYNAMIC_RANGE, _a.MIN_MAX, _g.CHANNELWISE),
-      (_t.FP32, _t.FP16, _m.WEIGHT_ONLY, _a.MIN_MAX, _g.NONE),
+      (_t.FP32, _t.INT8, _m.WEIGHT_ONLY, _a.MIN_MAX, _g.CHANNELWISE),
+      (_t.FP32, _t.FP16, _m.WEIGHT_ONLY, _a.FLOAT_CAST, _g.NONE),
   ]

ai_edge_torch/generative/test/test_model_conversion.py

@@ -55,6 +55,30 @@ class TestModelConversion(unittest.TestCase):
           )
       )
 
+  def test_toy_model_with_multi_batches(self):
+    config = toy_model_with_kv_cache.get_model_config()
+    config.batch_size = 2
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
+    idx, input_pos = torch.tensor([[1], [2]], dtype=torch.long), torch.tensor(
+        [10], dtype=torch.int64
+    )
+
+    edge_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+
+    # TODO(b/338288901): re-enable test to check output tensors.
+    skip_output_check = True
+    if skip_output_check is False:
+      self.assertTrue(
+          model_coverage.compare_tflite_torch(
+              edge_model,
+              pytorch_model,
+              (idx, input_pos),
+              num_valid_inputs=1,
+              atol=1e-5,
+              rtol=1e-5,
+          )
+      )
+
   def test_toy_model_with_kv_cache_with_hlfb(self):
     config = toy_model_with_kv_cache.get_model_config()
     config.enable_hlfb = True

ai_edge_torch/generative/test/test_quantize.py

@@ -21,11 +21,13 @@ import torch
 import ai_edge_torch
 from ai_edge_torch.generative.examples.test_models import toy_model_with_kv_cache  # NOQA
 from ai_edge_torch.generative.quantize import quant_recipe
+from ai_edge_torch.generative.quantize import quant_recipe_utils
 from ai_edge_torch.generative.quantize import quant_recipes
 from ai_edge_torch.generative.quantize.quant_attrs import Algorithm
 from ai_edge_torch.generative.quantize.quant_attrs import Dtype
 from ai_edge_torch.generative.quantize.quant_attrs import Granularity
 from ai_edge_torch.generative.quantize.quant_attrs import Mode
+from ai_edge_torch.quantize import quant_config
 from ai_edge_torch.testing import model_coverage
 
 
@@ -34,34 +36,47 @@ class TestVerifyRecipes(unittest.TestCase):
 
   @parameterized.expand(
       [
-          (Dtype.FP32, Dtype.FP32, Mode.DYNAMIC_RANGE),
-          (Dtype.INT8, Dtype.INT8, Mode.DYNAMIC_RANGE),
-          (Dtype.INT8, Dtype.FP16, Mode.DYNAMIC_RANGE),
-          (Dtype.FP16, Dtype.INT8, Mode.DYNAMIC_RANGE),
-          (Dtype.FP32, Dtype.FP32, Mode.WEIGHT_ONLY),
-          (Dtype.INT8, Dtype.INT8, Mode.WEIGHT_ONLY),
-          (Dtype.FP16, Dtype.INT8, Mode.WEIGHT_ONLY),
-          (Dtype.INT8, Dtype.FP16, Mode.WEIGHT_ONLY),
-          (Dtype.FP16, Dtype.FP16, Mode.WEIGHT_ONLY),
+          (Dtype.FP32, Dtype.FP32),
+          (Dtype.INT8, Dtype.INT8),
+          (Dtype.INT8, Dtype.FP16),
+          (Dtype.FP16, Dtype.INT8),
+          (Dtype.FP16, Dtype.FP16),
       ]
   )
   def test_verify_invalid_recipes(
       self,
       activation,
       weight,
-      mode,
-      algo=Algorithm.MIN_MAX,
-      granularity=Granularity.CHANNELWISE,
   ):
-    with self.assertRaises(ValueError):
-      quant_recipe.LayerQuantRecipe(
-          activation, weight, mode, algo, granularity
-      ).verify()
+    for m in Mode:
+      for a in Algorithm:
+        for g in Granularity:
+          with self.assertRaises(ValueError):
+            quant_recipe.LayerQuantRecipe(activation, weight, m, a, g).verify()
 
   @parameterized.expand(
       [
-          (Dtype.FP32, Dtype.INT8, Mode.DYNAMIC_RANGE, Granularity.CHANNELWISE),
-          (Dtype.FP32, Dtype.FP16, Mode.WEIGHT_ONLY, Granularity.NONE),
+          (
+              Dtype.FP32,
+              Dtype.INT8,
+              Mode.DYNAMIC_RANGE,
+              Algorithm.MIN_MAX,
+              Granularity.CHANNELWISE,
+          ),
+          (
+              Dtype.FP32,
+              Dtype.INT8,
+              Mode.WEIGHT_ONLY,
+              Algorithm.MIN_MAX,
+              Granularity.CHANNELWISE,
+          ),
+          (
+              Dtype.FP32,
+              Dtype.FP16,
+              Mode.WEIGHT_ONLY,
+              Algorithm.FLOAT_CAST,
+              Granularity.NONE,
+          ),
       ]
   )
   def test_verify_valid_recipes(
@@ -69,8 +84,8 @@ class TestVerifyRecipes(unittest.TestCase):
       activation,
       weight,
       mode,
+      algo,
       granularity,
-      algo=Algorithm.MIN_MAX,
   ):
     quant_recipe.LayerQuantRecipe(activation, weight, mode, algo, granularity).verify()
 
@@ -78,7 +93,47 @@ class TestVerifyRecipes(unittest.TestCase):
 class TestQuantizeConvert(unittest.TestCase):
   """Test conversion with quantization."""
 
-  def test_quantize_convert_toy(self):
+  def _attention_1_int8_dynamic_recipe() -> quant_config.QuantConfig:
+    return quant_config.QuantConfig(
+        generative_recipe=quant_recipe.GenerativeQuantRecipe(
+            attention={1: quant_recipe_utils.create_layer_quant_int8_dynamic()},
+        )
+    )
+
+  def _feedforward_0_int8_dynamic_recipe() -> quant_config.QuantConfig:
+    return quant_config.QuantConfig(
+        generative_recipe=quant_recipe.GenerativeQuantRecipe(
+            feedforward={0: quant_recipe_utils.create_layer_quant_int8_dynamic()},
+        )
+    )
+
+  @parameterized.expand(
+      [
+          (quant_recipes.full_fp16_recipe(), 0.75),
+          (quant_recipes.full_linear_int8_dynamic_recipe(), 0.64),
+          (_attention_1_int8_dynamic_recipe(), 0.95),
+          (_feedforward_0_int8_dynamic_recipe(), 0.87),
+      ]
+  )
+  def test_quantize_convert_toy_sizes(self, quant_config, expected_compression):
+    self.skipTest("b/346896669")
+    config = toy_model_with_kv_cache.get_model_config()
+    pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)
+    idx, input_pos = torch.tensor([[1]], dtype=torch.long), torch.tensor(
+        [10], dtype=torch.int64
+    )
+
+    quantized_model = ai_edge_torch.convert(
+        pytorch_model, (idx, input_pos), quant_config=quant_config
+    )
+    float_model = ai_edge_torch.convert(pytorch_model, (idx, input_pos))
+    self.assertAlmostEqual(
+        len(quantized_model._tflite_model) / len(float_model._tflite_model),
+        expected_compression,
+        delta=0.01,
+    )
+
+  def test_quantize_convert_compare_toy(self):
     self.skipTest("b/338288901")
     config = toy_model_with_kv_cache.get_model_config()
     pytorch_model = toy_model_with_kv_cache.ToyModelWithKV(config)