Add memory diagnostics at each attention call to see if VRAM usage is growing unbounded or just spiking once

mvdiffusion/models_unclip/attn_processors.py

@@ -11,6 +11,22 @@ from einops import rearrange, repeat
 import math
 
 import torch.nn.functional as F
+
+
+_sdpa_call_count = [0]
+
+def _sdpa_diag(query, key, value, attention_mask):
+    _sdpa_call_count[0] += 1
+    allocated = torch.cuda.memory_allocated() / 1e9
+    reserved = torch.cuda.memory_reserved() / 1e9
+    free, total = torch.cuda.mem_get_info()
+    print(
+        f"MEM_DIAG call={_sdpa_call_count[0]} q={tuple(query.shape)} "
+        f"allocated={allocated:.2f}GB reserved={reserved:.2f}GB "
+        f"free={free/1e9:.2f}GB total={total/1e9:.2f}GB",
+        flush=True,
+    )
+    return F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
 if is_xformers_available():
     import xformers
     import xformers.ops
@@ -117,7 +133,7 @@ class XFormersMVAttnProcessor:
         key = attn.head_to_batch_dim(key)
         value = attn.head_to_batch_dim(value)
 
-        hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
+        hidden_states = _sdpa_diag(query, key, value, attention_mask)
         hidden_states = attn.batch_to_head_dim(hidden_states)
 
         # linear proj
@@ -245,7 +261,7 @@ class XFormersIPCDAttnProcessor:
         key = attn.head_to_batch_dim(key).contiguous()
         value = attn.head_to_batch_dim(value).contiguous()
 
-        hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
+        hidden_states = _sdpa_diag(query, key, value, attention_mask)
         hidden_states = attn.batch_to_head_dim(hidden_states)
 
         # linear proj
@@ -376,7 +392,7 @@ class XFormersIPCrossAttnProcessor:
         key = attn.head_to_batch_dim(key).contiguous()
         value = attn.head_to_batch_dim(value).contiguous()
         
-        hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
+        hidden_states = _sdpa_diag(query, key, value, attention_mask)
         hidden_states = attn.batch_to_head_dim(hidden_states)
         
         # ip attn
@@ -610,7 +626,7 @@ class XFormersCDAttnProcessor:
         key = attn.head_to_batch_dim(key).contiguous()
         value = attn.head_to_batch_dim(value).contiguous()
 
-        hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
+        hidden_states = _sdpa_diag(query, key, value, attention_mask)
         hidden_states = attn.batch_to_head_dim(hidden_states)
         
         # linear proj