diff --git a/gemma/kv_cache.h b/gemma/kv_cache.h
index 2b122760..f66cd0a2 100644
--- a/gemma/kv_cache.h
+++ b/gemma/kv_cache.h
@@ -184,6 +184,7 @@ struct KVCache {
         .kv_cache = kv_cache,
         .k_cache = k_cache,
         .v_cache = v_cache,
+        .cache = this,
     };
   }