tico 0.1.0.dev250714__py3-none-any.whl → 0.1.0.dev251102__py3-none-any.whl

tico/serialize/operators/op_repeat.py

@@ -21,7 +21,11 @@ import torch
 from circle_schema import circle
 
 from tico.serialize.circle_graph import CircleSubgraph
-from tico.serialize.circle_mapping import extract_circle_dtype, extract_shape
+from tico.serialize.circle_mapping import (
+    extract_circle_dtype,
+    extract_shape,
+    to_circle_shape,
+)
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
 from tico.serialize.operators.utils import create_builtin_operator, get_op_index
@@ -70,11 +74,16 @@ class RepeatVisitor(NodeVisitor):
             if r > 1:
                 # Except last created concat, a tensor should be created.
                 if repeat_dim_cnt > 1:
-                    repeated_shape = list(tensor_shape)
+                    repeated_shape: List[int | torch.SymInt] = list(tensor_shape)
                     repeated_shape[idx] = repeated_shape[idx] * r
+
+                    repeated_cshape, repeated_cshape_signature = to_circle_shape(
+                        repeated_shape
+                    )
                     concat_output = self.graph.add_tensor_from_scratch(
                         prefix=f"{node.name}_concat_{idx}",
-                        shape=repeated_shape,
+                        shape=repeated_cshape,
+                        shape_signature=repeated_cshape_signature,
                         dtype=tensor_dtype,
                         source_node=node,
                     )

tico/serialize/operators/op_reshape.py

@@ -66,7 +66,7 @@ class ReshapeVisitor(NodeVisitor):
             circle.BuiltinOptions.BuiltinOptions.ReshapeOptions
         )
         option = circle.ReshapeOptions.ReshapeOptionsT()
-        option.newShape = size_i32
+        option.newShape = size_i32.tolist()
 
         operator.builtinOptions = option
 

tico/serialize/operators/op_rmsnorm.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# 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.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import CircleRMSNormArgs
+
+
+@register_node_visitor
+class RMSNormVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.circle_custom.rms_norm.default,
+    ]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = CircleRMSNormArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        weight = args.weight
+        eps = args.eps
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.RMS_NORM, self._op_codes
+        )
+
+        inputs = [input, weight]
+        outputs = [node]
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.RmsNormOptions
+        )
+        option = circle.RmsNormOptions.RmsNormOptionsT()
+        option.epsilon = eps
+
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_softmax.py

@@ -24,25 +24,18 @@ from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
 from tico.serialize.operators.utils import create_builtin_operator, get_op_index
 from tico.utils.errors import NotYetSupportedError
-from tico.utils.utils import HAS_TORCH_OVER_25
 from tico.utils.validate_args_kwargs import SafeSoftmaxArgs, SoftmaxArgs
 
 
 @register_node_visitor
 class SoftMaxVisitor(NodeVisitor):
-    target: List[torch._ops.OpOverload] = (
-        [
-            torch.ops.aten._softmax.default,
-            # NOTE: Let's treat _safe_softmax as normal _softmax as its usage is for training.
-            # In order for optimization during inference, it can be replaced to softmax.
-            # ref: https://github.com/pytorch/pytorch/pull/133882
-            torch.ops.aten._safe_softmax.default,
-        ]
-        if HAS_TORCH_OVER_25
-        else [
-            torch.ops.aten._softmax.default,
-        ]
-    )
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten._softmax.default,
+        # NOTE: Let's treat _safe_softmax as normal _softmax as its usage is for training.
+        # In order for optimization during inference, it can be replaced to softmax.
+        # ref: https://github.com/pytorch/pytorch/pull/133882
+        torch.ops.aten._safe_softmax.default,
+    ]
 
     def __init__(self, op_codes: Dict[OpCode, int], graph):
         super().__init__(op_codes, graph)

tico/serialize/operators/op_split_with_sizes.py

@@ -58,12 +58,14 @@ class SplitWithSizesVisitor(NodeVisitor):
         inputs = [input, split_sizes_i32, axis_i32]
 
         """
-        `split_with_sizes` has multiple output tensors and they are represented as `getitem`.
-        Therefore, unlike other ops, node itself doesn't become a circle tensor. Instead, each `getitem` will be
+        `split_with_sizes` has multiple output tensors along with `getitem`.
+        Unlike other ops, node itself doesn't become a circle tensor. Instead, each `getitem` will be
         a circle tensor.
-        Further, torch module having `split_with_sizes` may somtimes return selected outputs. At that time, `getitem`
-        nodes are generated only for the ouptut selected. Since one-compiler assumes that `CircleSplitV` always has 
-        all the outputs, let's add unused output tensors to compensate this restriction.
+        
+        torch module having `split_with_sizes` may return selected outputs by using `getitem`.
+        However, one-compiler assumes that `CircleSplitV` always have all outputs.
+        
+        So, let's add unused output tensors to compensate this restriction.
         """
         outputs: List[Union[circle.Tensor.TensorT, torch.fx.node.Node]] = []
         sorted_users = sorted(node.users.keys(), key=lambda x: x.args[1])  # type: ignore[arg-type, return-value]
@@ -80,11 +82,17 @@ class SplitWithSizesVisitor(NodeVisitor):
                 fake_tensor = node_val[idx]
                 assert isinstance(fake_tensor, FakeTensor)
                 shape = list(fake_tensor.size())
+
+                if any(isinstance(s, torch.SymInt) for s in shape):
+                    # TODO: support dynamic shape
+                    raise NotImplementedError("Dynamic shape is not supported yet.")
+
                 dtype = to_circle_dtype(fake_tensor.dtype)
                 tensor = self.graph.add_tensor_from_scratch(
-                    f"{node.name}_unused_{idx}",
-                    shape,
-                    dtype,
+                    prefix=f"{node.name}_unused_{idx}",
+                    shape=shape,
+                    shape_signature=None,  # TODO: support dynamic shape
+                    dtype=dtype,
                     source_node=node,
                 )
                 outputs.append(tensor)

tico/serialize/operators/op_transpose_conv.py

@@ -23,7 +23,8 @@ from circle_schema import circle
 from tico.serialize.circle_mapping import (
     circle_legalize_dtype_to,
     extract_circle_dtype,
-    extract_shape,
+    extract_circle_shape,
+    to_circle_shape,
 )
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
@@ -76,15 +77,13 @@ class TransposeConvVisitor(NodeVisitor):
         bias = args.bias
         stride = args.stride
         padding = args.padding
-        output_padding = args.output_padding
         groups = args.groups
-        dilation = args.dilation
 
         assert groups == 1, "Only support group 1"
 
-        input_shape = list(extract_shape(input_))
-        output_shape = list(extract_shape(node))
-        weight_shape = list(extract_shape(weight))
+        input_shape, input_shape_signature = extract_circle_shape(input_)
+        output_shape, _ = extract_circle_shape(node)
+        weight_shape, _ = extract_circle_shape(weight)
         assert len(input_shape) == 4, len(input_shape)
         assert len(output_shape) == 4, len(output_shape)
         assert len(weight_shape) == 4, len(weight_shape)
@@ -103,17 +102,21 @@ class TransposeConvVisitor(NodeVisitor):
                 ],
                 dtype=torch.int32,
             )
-            pad_output_shape = [
+            pad_output_shape: List[int | torch.SymInt] = [
                 input_shape[0],
                 input_shape[1] + pad_h * 2,
                 input_shape[2] + pad_w * 2,
                 input_shape[3],
             ]
+            pad_output_cshape, pad_output_cshape_signature = to_circle_shape(
+                pad_output_shape
+            )
             # create padded output tensor
             input_qparam: Optional[QuantParam] = input_.meta.get(QPARAM_KEY)
             pad_output = self.graph.add_tensor_from_scratch(
                 prefix=f"{node.name}_input_pad_output",
-                shape=pad_output_shape,
+                shape=pad_output_cshape,
+                shape_signature=pad_output_cshape_signature,
                 dtype=extract_circle_dtype(input_),
                 qparam=input_qparam,
                 source_node=node,

tico/serialize/operators/op_view.py

@@ -56,6 +56,7 @@ class ViewVisitor(NodeVisitor):
         if isinstance(size, int):
             raise Exception("scalar size conversion is not supported yet.")
 
+        # TODO: support dynamic shape
         size_i32 = circle_legalize_dtype_to(size, dtype=torch.int32)
         inputs = [input, size_i32]
         outputs = [node]
@@ -67,7 +68,7 @@ class ViewVisitor(NodeVisitor):
             circle.BuiltinOptions.BuiltinOptions.ReshapeOptions
         )
         option = circle.ReshapeOptions.ReshapeOptionsT()
-        option.newShape = size_i32
+        option.newShape = size_i32.tolist()
 
         operator.builtinOptions = option
 

tico/serialize/quant_param.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+
 """
 This is a key for torch.fx.Node's meta dict to save QuantParam
 
@@ -19,11 +24,6 @@ QuantParam can be retrieved as node.meta[QPARAM_KEY]
 """
 QPARAM_KEY = "_quantization_parameters_"
 
-from dataclasses import dataclass
-from typing import List, Optional
-
-import torch
-
 
 @dataclass
 class QuantParam:

tico/utils/convert.py

@@ -20,25 +20,14 @@ import torch
 from torch.export import export, ExportedProgram
 
 from tico.config import CompileConfigBase, get_default_config
-from tico.experimental.quantization.passes.fold_quant_ops import FoldQuantOps
-from tico.experimental.quantization.passes.insert_quantize_on_dtype_mismatch import (
-    InsertQuantizeOnDtypeMismatch,
-)
-from tico.experimental.quantization.passes.propagate_qparam_backward import (
-    PropagateQParamBackward,
-)
-from tico.experimental.quantization.passes.propagate_qparam_forward import (
-    PropagateQParamForward,
-)
-from tico.experimental.quantization.passes.quantize_bias import QuantizeBias
-from tico.experimental.quantization.passes.remove_weight_dequant_op import (
-    RemoveWeightDequantOp,
-)
 from tico.passes.cast_aten_where_arg_type import CastATenWhereArgType
+from tico.passes.cast_clamp_mixed_type_args import CastClampMixedTypeArgs
 from tico.passes.cast_mixed_type_args import CastMixedTypeArgs
 from tico.passes.const_prop_pass import ConstPropPass
 from tico.passes.convert_conv1d_to_conv2d import ConvertConv1dToConv2d
+from tico.passes.convert_expand_to_slice_cat import ConvertExpandToSliceCat
 from tico.passes.convert_layout_op_to_reshape import ConvertLayoutOpToReshape
+from tico.passes.convert_matmul_to_linear import ConvertMatmulToLinear
 from tico.passes.convert_repeat_to_expand_copy import ConvertRepeatToExpandCopy
 from tico.passes.convert_to_relu6 import ConvertToReLU6
 from tico.passes.decompose_addmm import DecomposeAddmm
@@ -71,6 +60,14 @@ from tico.passes.remove_redundant_slice import RemoveRedundantSlice
 from tico.passes.remove_redundant_to_copy import RemoveRedundantToCopy
 from tico.passes.restore_linear import RestoreLinear
 from tico.passes.segment_index_select import SegmentIndexSelectConst
+from tico.quantization.passes.fold_quant_ops import FoldQuantOps
+from tico.quantization.passes.insert_quantize_on_dtype_mismatch import (
+    InsertQuantizeOnDtypeMismatch,
+)
+from tico.quantization.passes.propagate_qparam_backward import PropagateQParamBackward
+from tico.quantization.passes.propagate_qparam_forward import PropagateQParamForward
+from tico.quantization.passes.quantize_bias import QuantizeBias
+from tico.quantization.passes.remove_weight_dequant_op import RemoveWeightDequantOp
 from tico.serialize.circle_serializer import build_circle
 from tico.serialize.operators.node_visitor import get_support_targets
 from tico.utils import logging
@@ -105,6 +102,7 @@ def traced_run_decompositions(exported_program: ExportedProgram):
             torch.ops.aten._safe_softmax.default,
             torch.ops.aten.relu6.default,  # Do not decompose to hardtanh
             torch.ops.aten.linear.default,
+            torch.ops.aten.upsample_nearest2d.vec,
         )
         ep = ep.run_decompositions(_preserve_ops=_preserve_ops)
 
@@ -123,6 +121,7 @@ def traced_run_decompositions(exported_program: ExportedProgram):
             torch.ops.aten.relu6.default,  # Do not decompose to hardtanh
             torch.ops.aten.prelu.default,
             torch.ops.aten.linear.default,
+            torch.ops.aten.upsample_nearest2d.vec,
         )
         for op in _preserve_ops:
             if op in _decomp_table:
@@ -137,6 +136,8 @@ def traced_run_decompositions(exported_program: ExportedProgram):
         torch.__version__.startswith("2.6")
         or torch.__version__.startswith("2.7")
         or torch.__version__.startswith("2.8")
+        or torch.__version__.startswith("2.9")
+        or torch.__version__.startswith("2.10")
     ):
         return run_decompositions(exported_program)
     else:
@@ -153,7 +154,7 @@ def check_unsupported_target(exported_program: ExportedProgram):
     for n in exported_program.graph.nodes:
         if n.op != "call_function":
             continue
-        if not n.target in supported_target:
+        if n.target not in supported_target:
             unsupported.append(n)
 
     if unsupported:
@@ -245,12 +246,21 @@ def convert_exported_module_to_circle(
             ConstPropPass(),
             SegmentIndexSelectConst(),
             LegalizeCausalMaskValue(enabled=config.get("legalize_causal_mask_value")),
+            ConvertExpandToSliceCat(enabled=config.get("convert_expand_to_slice_cat")),
+            ConvertMatmulToLinear(
+                enable_lhs_const=config.get("convert_lhs_const_mm_to_fc"),
+                enable_rhs_const=config.get("convert_rhs_const_mm_to_fc"),
+                enable_single_batch_lhs_const_bmm=config.get(
+                    "convert_single_batch_lhs_const_bmm_to_fc"
+                ),
+            ),
             LowerToResizeNearestNeighbor(),
             LegalizePreDefinedLayoutOperators(),
             LowerPow2ToMul(),
             ConvertConv1dToConv2d(),
             *LowerToSlicePasses(),
             FuseLeadingUnsqueezeReshape(),
+            CastClampMixedTypeArgs(),
         ]
     )
     circle_legalize.run(exported_program)
@@ -282,7 +292,7 @@ def convert_exported_module_to_circle(
 
     check_unsupported_target(exported_program)
     check_training_ops(exported_program)
-    circle_program = build_circle(exported_program)
+    circle_program = build_circle(exported_program, config)
 
     return circle_program
 
@@ -291,6 +301,7 @@ def convert(
     mod: torch.nn.Module,
     args: Tuple[Any, ...],
     kwargs: Optional[Dict[str, Any]] = None,
+    dynamic_shapes: Optional[dict] = None,
     strict: bool = True,
     config: CompileConfigBase = get_default_config(),
 ) -> CircleModel:
@@ -301,7 +312,9 @@ def convert(
         )
 
     with torch.no_grad():
-        exported_program = export(mod, args, kwargs, strict=strict)
+        exported_program = export(
+            mod, args, kwargs, dynamic_shapes=dynamic_shapes, strict=strict
+        )
 
     circle_binary = convert_exported_module_to_circle(exported_program, config=config)
 

tico/utils/dtype.py

@@ -0,0 +1,42 @@
+import numpy as np
+import torch
+
+from circle_schema import circle
+
+NUMPY_TO_TORCH_DTYPE_DICT = {
+    np.dtype("float32"): torch.float32,
+    np.dtype("float64"): torch.float64,
+    np.dtype("float16"): torch.float16,
+    np.dtype("complex64"): torch.complex64,
+    np.dtype("complex128"): torch.complex128,
+    np.dtype("int64"): torch.int64,
+    np.dtype("int32"): torch.int32,
+    np.dtype("int16"): torch.int16,
+    np.dtype("int8"): torch.int8,
+    np.dtype("uint8"): torch.uint8,
+    np.dtype("bool"): torch.bool,
+}
+
+CIRCLE_TO_TORCH_DTYPE_DICT = {
+    circle.TensorType.TensorType.FLOAT32: torch.float32,
+    circle.TensorType.TensorType.UINT8: torch.uint8,
+    circle.TensorType.TensorType.INT8: torch.int8,
+    circle.TensorType.TensorType.INT16: torch.int16,
+    circle.TensorType.TensorType.INT32: torch.int32,
+    circle.TensorType.TensorType.INT64: torch.int64,
+    circle.TensorType.TensorType.BOOL: torch.bool,
+}
+
+
+def numpy_dtype_to_torch_dtype(np_dtype: np.dtype) -> torch.dtype:
+    return NUMPY_TO_TORCH_DTYPE_DICT[np_dtype]
+
+
+def circle_dtype_to_torch_dtype(circle_dtype: int) -> torch.dtype:
+    assert isinstance(circle_dtype, int)
+    if circle_dtype not in CIRCLE_TO_TORCH_DTYPE_DICT:
+        raise RuntimeError(f"Unsupported dtype {circle_dtype}")
+
+    torch_dtype = CIRCLE_TO_TORCH_DTYPE_DICT[circle_dtype]
+    assert torch_dtype is not None
+    return torch_dtype

tico/utils/graph.py

@@ -24,7 +24,7 @@ import torch
 from torch.export import ExportedProgram
 from torch.export.exported_program import InputKind, InputSpec, TensorArgument
 
-from tico.utils.utils import get_fake_mode, set_new_meta_val
+from tico.utils.utils import get_fake_mode
 
 
 def is_torch_param(node: torch.fx.Node, ep: ExportedProgram):

tico/utils/model.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from pathlib import Path
 from typing import Any
 
 from tico.interpreter import infer
@@ -32,6 +33,6 @@ class CircleModel:
             buf = bytes(f.read())
         return CircleModel(buf)
 
-    def save(self, circle_path: str) -> None:
+    def save(self, circle_path: str | Path) -> None:
         with open(circle_path, "wb") as f:
             f.write(self.circle_binary)

tico/utils/padding.py

@@ -39,8 +39,8 @@ class ConvPaddingInfo(NamedTuple):
 
 def identify_padding(
     padding: PaddingValue,
-    input_shape: Sequence[int],
-    output_shape: Sequence[int],
+    input_shape: Sequence[int | torch.SymInt] | torch.Size,
+    output_shape: Sequence[int | torch.SymInt] | torch.Size,
     stride: Sequence[int],
 ) -> ConvPaddingInfo:
     """

tico/utils/pytree_utils.py

@@ -0,0 +1,134 @@
+import threading
+
+import torch
+from packaging.version import Version
+
+from tico.utils import logging
+from tico.utils.installed_packages import is_transformers_installed
+
+__all__ = ["register_dynamic_cache"]
+
+
+def register_dynamic_cache():
+    PyTreeRegistryHelper().register_dynamic_cache()
+
+
+class PyTreeRegistryHelper:
+    """
+    Thread-safe singleton helper class for registering custom PyTree nodes.
+
+    This class provides functionality to register DynamicCache as a PyTree node
+    for torch.export compatibility. This registration is only needed for
+    transformers versions below 4.50.0.
+
+    Thread Safety:
+    - Uses a class-level threading.Lock() to ensure thread-safe singleton instantiation
+    - Uses the same lock to protect the registration process from concurrent calls
+    """
+
+    _instance = None  # Class variable to hold the singleton instance
+    _has_called = False  # Flag to track if registration has been performed
+    _lock = threading.Lock()  # Class-level lock for thread-safe operations
+
+    def __init__(self):
+        """Private constructor to prevent direct instantiation"""
+        pass
+
+    def __new__(cls, *args, **kwargs):
+        """
+        Thread-safe singleton instance creation using double-checked locking pattern.
+
+        Returns:
+            PyTreeRegistryHelper: The singleton instance of this class
+        """
+        if not cls._instance:
+            with cls._lock:  # Acquire lock for thread-safe instantiation
+                if not cls._instance:  # Double-check after acquiring lock
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def register_dynamic_cache(self):
+        """
+        Registers DynamicCache as a PyTree node for torch.export compatibility.
+
+        This method is thread-safe and idempotent - it will only perform the
+        registration once, even if called multiple times from different threads.
+
+        Note:
+            This registration is only needed for transformers versions below 4.50.0.
+
+        Raises:
+            ImportError: If transformers package is not installed
+        """
+        with self._lock:  # Acquire lock for thread-safe registration
+            if self.__class__._has_called:
+                logger = logging.getLogger(__name__)
+                logger.debug("register_dynamic_cache already called, skipping")
+                return
+
+            self.__class__._has_called = True
+            logger = logging.getLogger(__name__)
+            logger.info("Registering DynamicCache PyTree node")
+
+        if not is_transformers_installed:  # type: ignore[truthy-function]
+            raise ImportError("transformers package is not installed")
+
+        import transformers
+
+        HAS_TRANSFORMERS_LESS_4_50_0 = Version(transformers.__version__) < Version(
+            "4.50.0"
+        )
+        if not HAS_TRANSFORMERS_LESS_4_50_0:
+            return
+
+        from transformers.cache_utils import DynamicCache
+
+        def _flatten_dynamic_cache(dynamic_cache: DynamicCache):
+            if not isinstance(dynamic_cache, DynamicCache):
+                raise RuntimeError(
+                    "This pytree flattening function should only be applied to DynamicCache"
+                )
+            HAS_TORCH_2_6_0 = Version(torch.__version__) >= Version("2.6.0")
+            if not HAS_TORCH_2_6_0:
+                logger = logging.getLogger(__name__)
+                logger.warning_once(
+                    "DynamicCache + torch.export is tested on torch 2.6.0+ and may not work on earlier versions."
+                )
+            dictionary = {
+                "key_cache": getattr(dynamic_cache, "key_cache"),
+                "value_cache": getattr(dynamic_cache, "value_cache"),
+            }
+            return torch.utils._pytree._dict_flatten(dictionary)
+
+        def _flatten_with_keys_dynamic_cache(dynamic_cache: DynamicCache):
+            dictionary = {
+                "key_cache": getattr(dynamic_cache, "key_cache"),
+                "value_cache": getattr(dynamic_cache, "value_cache"),
+            }
+            return torch.utils._pytree._dict_flatten_with_keys(dictionary)
+
+        def _unflatten_dynamic_cache(values, context: torch.utils._pytree.Context):
+            dictionary = torch.utils._pytree._dict_unflatten(values, context)
+            cache = DynamicCache()
+            for k, v in dictionary.items():
+                setattr(cache, k, v)
+            return cache
+
+        def _flatten_dynamic_cache_for_fx(cache, spec):
+            dictionary = {
+                "key_cache": getattr(cache, "key_cache"),
+                "value_cache": getattr(cache, "value_cache"),
+            }
+            return torch.fx._pytree._dict_flatten_spec(dictionary, spec)
+
+        torch.utils._pytree.register_pytree_node(
+            DynamicCache,
+            _flatten_dynamic_cache,
+            _unflatten_dynamic_cache,
+            serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
+            flatten_with_keys_fn=_flatten_with_keys_dynamic_cache,
+        )
+        # TODO: This won't be needed in torch 2.7+.
+        torch.fx._pytree.register_pytree_flatten_spec(
+            DynamicCache, _flatten_dynamic_cache_for_fx
+        )