Address attention context review feedback

comfy/ldm/cosmos/predict2.py

@@ -41,7 +41,7 @@ def torch_attention_op(
     q_B_S_H_D: torch.Tensor,
     k_B_S_H_D: torch.Tensor,
     v_B_S_H_D: torch.Tensor,
-    transformer_options: Optional[dict] = {},
+    transformer_options: dict = {},
     is_self_attention: bool = False,
 ) -> torch.Tensor:
     """Computes multi-head attention using PyTorch's native implementation.

comfy/ldm/wan/ar_model.py

@@ -277,6 +277,7 @@ class CausalWanModel(WanModel):
                 kv_caches=ar_state["kv_caches"],
                 crossattn_caches=ar_state["crossattn_caches"],
                 clip_fea=clip_fea,
+                transformer_options=transformer_options,
             )
 
         return super().forward(x, timestep, context, clip_fea=clip_fea,

tests-unit/comfy_test/attention_context_test.py

@@ -0,0 +1,48 @@
+import torch
+
+from comfy.ldm.cosmos import predict2
+from comfy.ldm.wan.ar_model import CausalWanModel
+
+
+def test_cosmos_attention_passes_self_attention_context(monkeypatch):
+    captured = {}
+
+    def capture_attention(q, k, v, heads, **kwargs):
+        captured.update(kwargs)
+        return q
+
+    monkeypatch.setattr(predict2, "optimized_attention", capture_attention)
+    q = torch.zeros((1, 2, 3, 4))
+
+    predict2.torch_attention_op(q, q, q, is_self_attention=True)
+
+    assert captured["is_self_attention"] is True
+    assert captured["attention_token_shape"] == (2,)
+
+
+def test_causal_wan_forward_passes_transformer_options_to_ar_block():
+    transformer_options = {
+        "ar_state": {
+            "start_frame": 2,
+            "kv_caches": [object()],
+            "crossattn_caches": [object()],
+        },
+        "optimized_attention_override": object(),
+    }
+    captured = {}
+
+    class FakeCausalWan:
+        def forward_block(self, **kwargs):
+            captured.update(kwargs)
+            return "result"
+
+    result = CausalWanModel.forward(
+        FakeCausalWan(),
+        x=torch.zeros((1, 4, 3, 2, 2)),
+        timestep=torch.zeros(1),
+        context=torch.zeros((1, 1, 1)),
+        transformer_options=transformer_options,
+    )
+
+    assert result == "result"
+    assert captured["transformer_options"] is transformer_options