add get_text_config

Author: xadupre

onnx_diagnostic/helpers/cache_helper.py

@@ -361,6 +361,9 @@ def __init__(self):
             self.num_hidden_layers = len(key_value_pairs)
             self.dtype = dtype
 
+        def get_text_config(self):
+            return self
+
     cache = MambaCache(
         _config(),
         max_batch_size=key_value_pairs[0][0].shape[0],
@@ -401,6 +404,9 @@ def __init__(self):
             self.num_hidden_layers = len(key_value_pairs)
             self.sliding_window = key_value_pairs[0][0].shape[2]
 
+        def get_text_config(self):
+            return self
+
     cache = transformers.cache_utils.SlidingWindowCache(
         config=_config(),
         max_batch_size=key_value_pairs[0][0].shape[0],
@@ -566,6 +572,9 @@ class _config:
         sliding_window = _sliding_window
         num_key_value_heads = key_value_pairs[0][1].shape[1]  # transformers 4.48.3
 
+        def get_text_config(self):
+            return self
+
     if layer_types:
         _config.layer_types = layer_types  # type: ignore[attr-defined]
 

onnx_diagnostic/helpers/helper.py

@@ -1479,8 +1479,12 @@ def max_diff(
     # backup function in case pytorch does not know how to serialize.
     if expected.__class__.__name__ == "DynamicCache":
         if got.__class__.__name__ == "DynamicCache":
+            from .cache_helper import CacheKeyValue
+
             if verbose >= 6:
                 print(f"[max_diff] DynamicCache: {string_type(expected)} ? {string_type(got)}")
+            expected = CacheKeyValue(expected)
+            got = CacheKeyValue(got)
             return max_diff(
                 [expected.key_cache, expected.value_cache],
                 [got.key_cache, got.value_cache],