Fix max allowed lengths for prompts larger than 4096

aiu_fms_testing_utils/testing/validation.py

@@ -188,15 +188,14 @@ def load_validation_information(validation_path, validation_files_type, batch_si
     return ValidationInfo(validation_info)
 
 def extract_validation_information(model, input_ids, max_new_tokens, post_iteration_hook, attn_algorithm=None, eos_token_id = None, only_last_token=False, timing="", **extra_kwargs):
-    max_seq_len = model.config.max_expected_seq_len
     attention_specific_kwargs = {}
     if "paged" in extra_kwargs["attn_name"]:
         from aiu_fms_testing_utils.utils.paged import generate
     else:
         # TODO: Add a unified generation dependent on attn_type
         from fms.utils.generation import generate
         attention_specific_kwargs["contiguous_cache"] = True
-        attention_specific_kwargs["max_seq_len"] = max_seq_len
+        attention_specific_kwargs["max_seq_len"] = input_ids.shape[1] + max_new_tokens
 
     # Add only_last_token optimization
     extra_generation_kwargs = {**extra_kwargs}

aiu_fms_testing_utils/utils/__init__.py

@@ -30,6 +30,7 @@ def warmup_model(
         # TODO: Add a unified generation dependent on attn_type
         from fms.utils.generation import generate
         attention_specific_kwargs["contiguous_cache"] = True
+        attention_specific_kwargs["max_seq_len"] = input_ids.shape[1] + max_new_tokens
 
     dprint("AIU warmup")
     pt_compile_model_time = time.time()

aiu_fms_testing_utils/utils/paged.py

@@ -54,7 +54,6 @@ def generate(
         model: A function or nn.Module that takes a batch of input_ids and
             returns logits
         input_ids: a rectangular tensor of input_ids (batch x seq)
-        max_seq_len: the sequence length of the model
         max_new_tokens: max tokens to generate
         temperature: temperature of softmax when sampling
         top_k: only search among top k tokens

scripts/inference.py

@@ -713,6 +713,7 @@ def infer(use_cache, do_sample, warmup):
     attention_specific_kwargs = {}
     if attn_name == "sdpa_causal":
         attention_specific_kwargs["contiguous_cache"] = True
+        attention_specific_kwargs["max_seq_len"] = ids.shape[1] + args.max_new_tokens
 
     result = generate(
         model,

tests/utils/test_paged.py

@@ -29,6 +29,7 @@ def test_paged_equivalence():
         result = generate(
             _model_mock,
             ids,
+            max_seq_len=ids.shape[1] + 5,
             max_new_tokens=5,
             do_sample=False,
             use_cache=True,