[Fix] generate() prefills the full prompt instead of only the last token

fla/models/utils.py

@@ -511,7 +511,7 @@ def prepare_inputs_for_generation(
                     # Fallback: manually slice using cache_position
                     if input_ids is not None and input_ids.shape[1] != cache_position.shape[0]:
                         input_ids = input_ids[:, cache_position]
-                elif hasattr(past_key_values, '__len__') and len(past_key_values) > 0:
+                elif past_key_values.get_seq_length() > 0:
                     # Ultimate fallback to old behavior
                     input_ids = input_ids[:, -1:]
 
@@ -529,7 +529,7 @@ def prepare_inputs_for_generation(
             # For older transformers versions, use the original logic
             model_inputs = {}
             # only last token for `inputs_ids` if the `past_key_values` is not empty.
-            if past_key_values is not None and hasattr(past_key_values, '__len__') and len(past_key_values) > 0:
+            if past_key_values is not None and past_key_values.get_seq_length() > 0:
                 input_ids = input_ids[:, -1:]
             # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
             if inputs_embeds is not None and hasattr(past_key_values, '__len__') and len(past_key_values) == 0:

tests/models/test_modeling_base.py

@@ -131,3 +131,59 @@ def run_test_generation(
     gen = torch.cat(logits, 1)
     gen = torch.cat([gen[i:i+1, start:] for i, start in enumerate(seq_start)], 1)
     assert_close('logits', ref, gen, tol)
+
+
+# ===================================================================================
+# REGRESSION TEST FOR FULL-PROMPT PREFILL IN GENERATION
+# ===================================================================================
+def run_test_generate_matches_forward(
+    L: int,
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    config_class: type,
+    dtype: torch.dtype,
+    num_new_tokens: int = 8,
+):
+    """
+    A regression test that `generate()` conditions on the whole prompt, not just the last token.
+    """
+    torch.manual_seed(42)
+    os.environ['FLA_CONV_BACKEND'] = 'triton'
+    if config_class.__name__ in GENERATION_UNSUPPORTED:
+        pytest.skip(f"Generation test not supported for {config_class.__name__}.")
+    if config_class.__name__ in NOT_READY_FOR_TESTING:
+        pytest.skip(f"{config_class.__name__} is not yet ready for testing.")
+
+    model, config = create_model_and_config(config_class, L, H, D, dtype=dtype)
+    model.eval()
+    model = model.to(dtype).to(device)
+    # avoid the pad/eos id so attention_mask is all ones and tokens are unambiguous
+    input_ids = torch.randint(low=1, high=config.vocab_size, size=(B, T)).to(device)
+    attention_mask = torch.ones_like(input_ids)
+
+    with torch.no_grad():
+        out = model(input_ids=input_ids, use_cache=True)
+        past_key_values = out.past_key_values
+        next_token = out.logits[:, -1:].argmax(dim=-1)
+        manual = [next_token]
+        for _ in range(num_new_tokens - 1):
+            out = model(input_ids=next_token, use_cache=True, past_key_values=past_key_values)
+            past_key_values = out.past_key_values
+            next_token = out.logits[:, -1:].argmax(dim=-1)
+            manual.append(next_token)
+        manual = torch.cat(manual, dim=1)
+
+        generated = model.generate(
+            input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=num_new_tokens,
+            do_sample=False,
+        )[:, T:]
+
+    assert torch.equal(generated, manual), (
+        "generate() output diverges from a manual greedy decode that prefills the full prompt; "
+        "prepare_inputs_for_generation likely dropped the prompt on the first step.\n"
+        f"generate={generated.tolist()}\nmanual={manual.tolist()}"
+    )

tests/models/test_modeling_gated_deltanet.py

@@ -10,7 +10,11 @@
 
 from fla.models import GatedDeltaNetConfig
 
-from .test_modeling_base import run_test_generation, run_test_model_forward_backward
+from .test_modeling_base import (
+    run_test_generate_matches_forward,
+    run_test_generation,
+    run_test_model_forward_backward,
+)
 
 
 # ===================================================================================
@@ -72,3 +76,23 @@ def test_generation(
     dtype: torch.dtype,
 ):
     run_test_generation(L, B, T, H, D, GatedDeltaNetConfig, dtype)
+
+
+@pytest.mark.parametrize(
+    ['L', 'B', 'T', 'H', 'D', 'dtype'],
+    [
+        pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test))
+        for test in [
+            (2, 2, 64, 8, 64, torch.float32),
+        ]
+    ],
+)
+def test_generate_prefill(
+    L: int,
+    B: int,
+    T: int,
+    H: int,
+    D: int,
+    dtype: torch.dtype,
+):
+    run_test_generate_matches_forward(L, B, T, H, D, GatedDeltaNetConfig, dtype)