Kimish comment

Author: jackzhxng

torchtune/modules/attention.py

@@ -274,7 +274,7 @@ def forward(
 
             # Update key-value cache
             if self.kv_cache is not None:
-                SDPA.kv_cache_update(input_pos, k, v)
+                self._sdpa.kv_cache_update(input_pos, k, v)
 
-        output = SDPA.sdpa(q, k, v, b, s_x)
+        output = self._sdpa.sdpa(q, k, v, b, s_x)
         return self.output_proj(output)