fix: fix benchmark_multiturn.py

Author: niushengxiao

lightllm/common/linear_att_cache_manager/linear_att_buffer_manager.py

@@ -26,15 +26,15 @@ def __init__(
             dtype=self.linear_config.conv_state_dtype,
             shape=self.linear_config.get_conv_state_shape(),
             layer_num=self.linear_config.linear_layer_num,
-            device="cuda",
+            device="cpu",
             size_first=True,
         )
         self.ssm_state_cache = LayerCache(
             size=self.size,
             dtype=self.linear_config.ssm_state_dtype,
             shape=self.linear_config.get_ssm_state_shape(),
             layer_num=self.linear_config.linear_layer_num,
-            device="cuda",
+            device="cpu",
             size_first=True,
         )
         self.clear_to_init_state()

requirements.txt

@@ -98,4 +98,4 @@ nixl==1.1.0
 xformers==0.0.35
 redis==7.3.0
 litellm>=1.52.0,<1.85
-flash-attn-4[13]==4.0.0b13
+flash-attn-4[13]==4.0.0b14

test/benchmark/service/benchmark_multiturn.py

@@ -50,6 +50,7 @@
 _STREAM_READ_BUFSIZE = 1 << 20
 _STREAM_MAX_LINE_SIZE = 1 << 20
 _DEFAULT_TRANSIENT_RETRIES = 2
+_PROMPT_LEN_OVERLAP_CHARS = 512
 _TRANSIENT_STREAM_ERRORS = (
     aiohttp.ServerDisconnectedError,
     aiohttp.ClientPayloadError,
@@ -177,6 +178,7 @@ def gen_session_initial_prompt(
 def append_turn_input(
     tokenizer,
     prompt: str,
+    prompt_token_len: int,
     generated_text: str,
     turn_input_increment: int,
     rng: random.Random,
@@ -188,8 +190,22 @@ def append_turn_input(
         new_text = decode_ids(tokenizer, new_ids)
     else:
         new_text = ""
-    new_prompt = prompt + generated_text + new_text
-    new_len = len(tokenizer.encode(new_prompt, add_special_tokens=False))
+
+    appended_text = generated_text + new_text
+    new_prompt = prompt + appended_text
+    if not appended_text:
+        return new_prompt, prompt_token_len
+
+    # Token merges only depend on a small boundary window, so avoid
+    # re-encoding the entire prompt on every turn.
+    overlap_text = prompt[-_PROMPT_LEN_OVERLAP_CHARS:]
+    if overlap_text:
+        overlap_token_len = len(tokenizer.encode(overlap_text, add_special_tokens=False))
+        merged_token_len = len(tokenizer.encode(overlap_text + appended_text, add_special_tokens=False))
+        appended_token_len = max(merged_token_len - overlap_token_len, 0)
+    else:
+        appended_token_len = len(tokenizer.encode(appended_text, add_special_tokens=False))
+    new_len = prompt_token_len + appended_token_len
     return new_prompt, new_len
 
 
@@ -352,7 +368,12 @@ async def run_session(
     """Run a single multi-turn dialogue session. Returns a list of per-turn
     stat dicts (same schema as stream_one_turn output)."""
     rng = random.Random(base_seed + session_id)
-    prompt, prompt_len = gen_session_initial_prompt(tokenizer, start_input_len, base_seed + session_id)
+    prompt, prompt_len = await asyncio.to_thread(
+        gen_session_initial_prompt,
+        tokenizer,
+        start_input_len,
+        base_seed + session_id,
+    )
 
     per_turn: List[Dict] = []
     turn_idx = 0
@@ -370,9 +391,11 @@ async def run_session(
             end="",
         )
         turn_input_len = rng.randint(min_turn_input_increment, turn_input_increment)
-        prompt, prompt_len = append_turn_input(
+        prompt, prompt_len = await asyncio.to_thread(
+            append_turn_input,
             tokenizer,
             prompt,
+            result["prompt_tokens"] or prompt_len,
             result["generated_text"],
             turn_input_len,
             rng,