Update

[ghstack-poisoned]

Author: swolchok

examples/models/llama/source_transformation/quantized_kv_cache.py

@@ -205,8 +205,11 @@ def replace_kv_cache_with_quantized_kv_cache(module):
     # This is needed to ensure that custom ops are registered
     from executorch.extension.llm.custom_ops import custom_ops  # noqa: F401
 
+    import resource
+    maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+    nswap = resource.getrusage(resource.RUSAGE_SELF).ru_nswap
     logging.warning(
-        "Replacing KVCache with QuantizedKVCache. This modifies the model in place."
+        f"Replacing KVCache with QuantizedKVCache. This modifies the model in place. (HACK: maxrss: {maxrss} nswap: {nswap})"
     )
     for name, child in module.named_children():
         if isinstance(child, KVCache) or isinstance(child, CustomKVCache):
@@ -270,8 +273,11 @@ def replace_kv_cache_with_custom_kv_cache(module):
     This is because the custom op treats second dim as sequence dim.
     Future work: support [B, H, S, D]
     """
+    import resource
+    maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+    nswap = resource.getrusage(resource.RUSAGE_SELF).ru_nswap
     logging.warning(
-        "Replacing KVCache with CustomKVCache. This modifies the model in place."
+        f"Replacing KVCache with CustomKVCache. This modifies the model in place. (HACK: maxrss: {maxrss} nswap: {nswap})"
     )
     for name, child in module.named_children():
         if isinstance(child, KVCache):