Merge pull request #10 from foundation-model-stack/contiguous-cache

Force kv cache to be contiguous to reduce number of graph traces

Author: ani300

.gitignore

@@ -7,4 +7,3 @@ aiu-fms-testing-utils.egg-info
 */**/*.pyc
 .vscode
 aiu-fms-testing-utils.egg-info
-

aiu_fms_testing_utils/testing/validation.py

@@ -2,7 +2,6 @@
 from typing import List, Tuple, Callable, MutableMapping, Any, Optional
 
 import torch
-import torch.nn as nn
 from fms.utils.generation import generate
 from aiu_fms_testing_utils.utils import ids_for_prompt
 from aiu_fms_testing_utils.utils.aiu_setup import dprint
@@ -205,6 +204,7 @@ def extract_validation_information(model, input_ids, max_new_tokens, post_iterat
         post_iteration_hook=post_iteration_hook,
         eos_token_id=eos_token_id,
         timing=timing,
+        contiguous_cache=True,
         extra_kwargs=extra_generation_kwargs,
     )
 

aiu_fms_testing_utils/utils/__init__.py

@@ -15,7 +15,7 @@ def warmup_model(model: nn.Module, input_ids: torch.Tensor, max_new_tokens: int,
     dprint("AIU warmup")
     pt_compile_model_time = time.time()
     extra_kwargs = {**padding_kwargs, "only_last_token": True}
-    generate(model, input_ids, max_new_tokens=max_new_tokens, max_seq_len=model.config.max_expected_seq_len, use_cache=True, do_sample=False, extra_kwargs=extra_kwargs)
+    generate(model, input_ids, max_new_tokens=max_new_tokens, max_seq_len=model.config.max_expected_seq_len, use_cache=True, do_sample=False, contiguous_cache=True, extra_kwargs=extra_kwargs)
     pt_compile_model_time = time.time() - pt_compile_model_time
     dprint(f"PT compile complete, took {pt_compile_model_time:.3f}s")
 

scripts/inference.py

@@ -589,6 +589,7 @@ def infer(use_cache, do_sample, warmup):
         max_seq_len=max_seq_len,
         timing=args.timing,
         eos_token_id=eos_token_id,
+        contiguous_cache=True,
         extra_kwargs=extra_generation_kwargs,
     )
     if args.timing != "":

scripts/validation.py

@@ -1,22 +1,18 @@
 import argparse
-import itertools
 import json
 import os
 import random
-import sys
 import time
 from pathlib import Path
-from typing import Any, Callable, MutableMapping, Optional, Tuple
-import sys
 import ast
 
 import numpy as np
 import torch
 import torch._inductor.config
 from fms.models import get_model, register_model
 from fms.models.llama import LLaMAConfig, _llama_factory_factory
-from fms.utils import fusion, generation, tokenizers
-from fms.utils.generation import generate, pad_input_ids
+from fms.utils import generation, tokenizers
+from fms.utils.generation import pad_input_ids
 from torch import distributed as dist
 from aiu_fms_testing_utils.utils import warmup_model
 from aiu_fms_testing_utils.testing.validation import LogitsExtractorHook, capture_level_1_metrics, extract_validation_information, StaticTokenInjectorHook, GoldenTokenHook, filter_failed_level_1_cases, validate_level_0, load_validation_information, print_failed_cases