From b18934217010f69ce9a03d68fe3b043cb2217c5b Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Mon, 19 Feb 2024 10:06:56 +0000 Subject: [PATCH] Let it work. --- server/text_generation_server/models/flash_mistral.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/server/text_generation_server/models/flash_mistral.py b/server/text_generation_server/models/flash_mistral.py index 75c3f838..5df4e214 100644 --- a/server/text_generation_server/models/flash_mistral.py +++ b/server/text_generation_server/models/flash_mistral.py @@ -460,8 +460,11 @@ class BaseFlashMistral(FlashCausalLM): max_s = batch.max_seqlen lm_head_indices = batch.prefill_head_indices - # if self.model.max_past is not None: - # max_s = min(self.model.max_past, max_s) + if cu_seqlen_prefill is None and self.model.max_past is not None: + # In decode, not prefill, we're actually overwriting the KV-cache + # in a circular buffer mode. + # This makes sure the max_s for the decode pass is correct. + max_s = min(self.model.max_past, max_s) bs = input_ids.shape[0] padded_bs = bs