fix: preserve sampling penalties in stateful generation fallback

tests/unit/model_bridge/test_generate_fallback_flags.py

@@ -1,7 +1,8 @@
-"""generate()'s stateful hf_generate fallback must forward its stop flags."""
+"""Stateful hf_generate fallback preserves stop flags and sampling penalties."""
 
 from types import SimpleNamespace
 
+import pytest
 import torch
 from transformers import GPT2Config, GPT2LMHeadModel
 
@@ -108,3 +109,59 @@ def spy(input_ids, **kwargs):
     )
     # HF takes int or list; any other sequence must arrive normalized.
     assert captured["eos_token_id"] == [5, 7]
+
+
+@pytest.mark.parametrize("penalty", [0.8, 1.0, 1.5])
+def test_stateful_fallback_forwards_repetition_penalty(monkeypatch, penalty):
+    bridge = _tiny_bridge()
+    bridge.cfg.is_stateful = True
+    bridge.tokenizer = SimpleNamespace(eos_token_id=0)
+    captured: dict = {}
+
+    def spy(input_ids, **kwargs):
+        captured.update(kwargs)
+        return input_ids
+
+    monkeypatch.setattr(bridge.original_model, "generate", spy)
+    bridge.generate(
+        torch.tensor([[1, 2, 3]]),
+        use_past_kv_cache=False,
+        repetition_penalty=penalty,
+        do_sample=False,
+        verbose=False,
+    )
+    assert captured["repetition_penalty"] == penalty
+
+
+@pytest.mark.parametrize("penalty", [-0.5, 0.5])
+def test_stateful_fallback_rejects_frequency_penalty(monkeypatch, penalty):
+    bridge = _tiny_bridge()
+    bridge.cfg.is_stateful = True
+    captured = _capture_hf_generate(bridge, monkeypatch)
+    with pytest.raises(NotImplementedError, match="freq_penalty.*use_past_kv_cache=True"):
+        bridge.generate(
+            torch.tensor([[1, 2, 3]]),
+            use_past_kv_cache=False,
+            freq_penalty=penalty,
+            verbose=False,
+        )
+    assert not captured
+
+
+def test_stateful_fallback_repetition_penalty_matches_hf_generation():
+    bridge = _tiny_bridge()
+    bridge.cfg.is_stateful = True
+    bridge.tokenizer = SimpleNamespace(eos_token_id=0)
+    tokens = torch.tensor([[1, 2, 3]])
+    kwargs = dict(max_new_tokens=4, do_sample=False, stop_at_eos=False)
+    expected = bridge.hf_generate(tokens, repetition_penalty=2.0, **kwargs)
+    unpenalized = bridge.hf_generate(tokens, repetition_penalty=1.0, **kwargs)
+    actual = bridge.generate(
+        tokens,
+        repetition_penalty=2.0,
+        use_past_kv_cache=False,
+        verbose=False,
+        **kwargs,
+    )
+    torch.testing.assert_close(actual, expected)
+    assert not torch.equal(actual, unpenalized)

transformer_lens/model_bridge/transformer_bridge.py

@@ -2998,7 +2998,9 @@ def generate(
             top_k: Number of tokens to sample from. If None, sample from all tokens
             top_p: Probability mass to sample from. If 1.0, sample from all tokens
             temperature: Temperature for sampling. Higher values will make the model more random
-            freq_penalty: Frequency penalty for sampling - how much to penalise previous tokens
+            freq_penalty: Frequency penalty for sampling - how much to penalise previous tokens.
+                Nonzero values require the hooked generation loop; the stateful/SSM
+                HF fallback with use_past_kv_cache=False raises NotImplementedError.
             repetition_penalty: HuggingFace-style repetition penalty. Values > 1.0 discourage
                 repetition by dividing positive logits and multiplying negative logits for
                 previously seen tokens. Default 1.0 (no penalty).
@@ -3367,10 +3369,16 @@ def generate(
                     )
 
         if is_stateful_model and not use_stateful_cache:
+            if freq_penalty != 0.0:
+                raise NotImplementedError(
+                    "freq_penalty requires use_past_kv_cache=True for stateful/SSM models. "
+                    "The HF generation fallback has no equivalent frequency penalty."
+                )
             hf_kwargs: dict[str, Any] = {
                 "max_new_tokens": max_new_tokens,
                 "do_sample": do_sample,
                 "temperature": temperature,
+                "repetition_penalty": repetition_penalty,
                 # hf_generate defaults this to True, so omitting it would make
                 # generate(stop_at_eos=False) silently stop at EOS on this path.
                 "stop_at_eos": stop_at_eos,