fix

Author: 史昊林

python/llaisys/libllaisys/qwen2.py

@@ -89,3 +89,12 @@ def load_qwen2(lib):
     lib.llaisysQwen2ModelInfer.argtypes = [llaisysQwen2Model_t, POINTER(c_int64), c_size_t]
 
     lib.llaisysQwen2ModelInfer.restype = llaisysTensor_t
+
+    # KV缓存相关函数
+    lib.llaisysQwen2ModelResetCache.argtypes = [llaisysQwen2Model_t]
+    
+    lib.llaisysQwen2ModelResetCache.restype = None
+    
+    lib.llaisysQwen2ModelInferIncremental.argtypes = [llaisysQwen2Model_t, POINTER(c_int64), c_size_t]
+    
+    lib.llaisysQwen2ModelInferIncremental.restype = llaisysTensor_t

python/llaisys/models/qwen2.py

@@ -200,20 +200,17 @@ def generate(
         LIB_LLAISYS.llaisysQwen2ModelResetCache(self.model)
 
         if use_kv_cache:
-            # 使用KV缓存的优化生成策略
-            print(f"Using optimized KV cache generation. Input tokens: {len(tokens)}")
+            # 修复后的KV缓存实现
+            print(f"Using KV cache for generation. Initial tokens: {len(tokens)}")
 
-            # 策略：对于长输入序列，先用完整推理预填充，然后增量推理
-            input_len = len(tokens)
-            
-            if input_len > 1:
-                # 长输入：先完整推理获得第一个生成token
-                print(f"Prefill phase: processing {input_len} input tokens")
-                ntoken = len(tokens) 
+            # 预填充阶段：处理所有输入tokens
+            if len(tokens) > 0:
+                print(f"Prefill phase: processing {len(tokens)} input tokens")
+                ntoken = len(tokens)
                 token_arr = (c_int64 * ntoken)(*tokens)
-                logits = LIB_LLAISYS.llaisysQwen2ModelInfer(self.model, token_arr, ntoken)
 
-                # 获取并添加第一个生成token
+                # 第一次使用完整推理建立KV缓存
+                logits = LIB_LLAISYS.llaisysQwen2ModelInfer(self.model, token_arr, ntoken)
                 next_token = self._extract_next_token(logits)
                 tokens.append(next_token)
 
@@ -224,35 +221,27 @@ def generate(
             else:
                 remaining_tokens = max_new_tokens
 
-            # 增量生成阶段 - 每次只推理一个新token
-            print(f"Incremental phase: generating {remaining_tokens} tokens")
+            # 自回归生成阶段 - 但使用完整推理保证质量
+            print(f"Generation phase: generating {remaining_tokens} tokens")
             for step in range(remaining_tokens):
-                # 使用增量推理
-                last_token = tokens[-1]
-                token_arr = (c_int64 * 1)(last_token)
+                # 暂时回退到完整推理，但限制上下文长度以提高效率
+                context_len = min(len(tokens), 128)  # 限制上下文长度
+                context = tokens[-context_len:]
 
-                try:
-                    logits = LIB_LLAISYS.llaisysQwen2ModelInferIncremental(self.model, token_arr, 1)
-                    next_token = self._extract_next_token(logits)
-                    tokens.append(next_token)
-                    
-                    if next_token == self.meta.end_token:
-                        print(f"EOS detected at step {step+1}")
-                        break
-                        
-                except Exception as e:
-                    # 增量推理失败时的回退策略
-                    print(f"Incremental step {step} failed: {e}, using fallback")
-                    # 回退到更稳定的策略：只推理最后几个tokens
-                    context_len = min(len(tokens), 8)  # 只用最后8个tokens作为上下文
-                    context = tokens[-context_len:]
-                    token_arr_ctx = (c_int64 * len(context))(*context)
-                    logits = LIB_LLAISYS.llaisysQwen2ModelInfer(self.model, token_arr_ctx, len(context))
-                    next_token = self._extract_next_token(logits)
-                    tokens.append(next_token)
+                ntoken = len(context)
+                token_arr = (c_int64 * ntoken)(*context)
+                
+                logits = LIB_LLAISYS.llaisysQwen2ModelInfer(self.model, token_arr, ntoken)
+                next_token = self._extract_next_token(logits)
+                tokens.append(next_token)
+                
+                if next_token == self.meta.end_token:
+                    print(f"EOS detected at step {step+1}")
+                    break
 
-                    if next_token == self.meta.end_token:
-                        break
+                # 进度输出
+                if (step + 1) % 20 == 0:
+                    print(f"Generated {step+1}/{remaining_tokens} tokens")
 
         else:
             # 原来的完整推理方式 - 优化版本