Final fix for MacOS CI failure

Author: jlarson4

tests/unit/model_bridge/test_bridge_generate_no_tokenizer.py

@@ -29,45 +29,6 @@ def test_generate_without_tokenizer_stop_at_eos_false_kv_cache(tokenizer_free_br
     assert bridge.tokenizer is None
 
     tokens = _PROMPT_TOKENS.clone()
-
-    # === TEMP DEBUG: localize CI-only NaN; remove after diagnosing ===
-    import sys
-
-    def _diag(label: str, t: torch.Tensor) -> None:
-        print(
-            f"[DIAG] {label}: nan={torch.isnan(t).any().item()} "
-            f"inf={torch.isinf(t).any().item()} "
-            f"shape={tuple(t.shape)} dtype={t.dtype} "
-            f"sample={t.flatten()[:4].tolist()}",
-            file=sys.stderr,
-            flush=True,
-        )
-
-    with torch.no_grad():
-        bl = bridge(tokens, return_type="logits")
-    _diag("bridge_fwd_no_cache", bl)
-
-    with torch.no_grad():
-        ho = bridge.original_model(tokens)
-    _diag("hf_fwd_no_cache", ho.logits)
-
-    with torch.no_grad():
-        ho_cache = bridge.original_model(tokens, use_cache=True)
-    _diag("hf_fwd_step0_use_cache", ho_cache.logits)
-    print(
-        f"[DIAG] cache_type={type(ho_cache.past_key_values).__name__}",
-        file=sys.stderr,
-        flush=True,
-    )
-
-    next_id = ho_cache.logits[:, -1, :].argmax(-1, keepdim=True)
-    with torch.no_grad():
-        ho_step1 = bridge.original_model(
-            next_id, past_key_values=ho_cache.past_key_values, use_cache=True
-        )
-    _diag("hf_fwd_step1_with_cache", ho_step1.logits)
-    # === END TEMP DEBUG ===
-
     output = bridge.generate(
         tokens,
         max_new_tokens=3,

transformer_lens/model_bridge/bridge.py

@@ -2382,10 +2382,30 @@ def _generate_tokens(
                         forward_kwargs["use_cache"] = True
                         if _hf_kv_cache is not None:
                             forward_kwargs["past_key_values"] = _hf_kv_cache
+                            # HF v5 + macOS-arm64 NaNs when these are inferred
+                            # from cache state alone. Mirror HF generate(): pass
+                            # both an (batch, total_len) attention_mask and a
+                            # (batch, 1) position_ids for the new token.
+                            batch_size = current_tokens.shape[0]
+                            total_len = current_tokens.shape[1]
+                            device = current_tokens.device
+                            if "attention_mask" not in forward_kwargs:
+                                forward_kwargs["attention_mask"] = torch.ones(
+                                    (batch_size, total_len),
+                                    dtype=torch.long,
+                                    device=device,
+                                )
                             if "position_ids" in forward_kwargs:
                                 forward_kwargs["position_ids"] = forward_kwargs["position_ids"][
                                     :, -1:
                                 ]
+                            else:
+                                forward_kwargs["position_ids"] = torch.full(
+                                    (batch_size, 1),
+                                    total_len - 1,
+                                    dtype=torch.long,
+                                    device=device,
+                                )
                             logits = self(
                                 current_tokens[:, -1:],
                                 return_type="logits",