import argparse import ctypes import enum import os import threading from typing import Optional, Sequence, Tuple import numpy as np # Define constants from the header CPU0 = (1 << 0) # 0x01 CPU1 = (1 << 1) # 0x02 CPU2 = (1 << 2) # 0x04 CPU3 = (1 << 3) # 0x08 CPU4 = (1 << 4) # 0x10 CPU5 = (1 << 5) # 0x20 CPU6 = (1 << 6) # 0x40 CPU7 = (1 << 7) # 0x80 # --- Enums --- class LLMCallState(enum.IntEnum): RKLLM_RUN_NORMAL = 0 RKLLM_RUN_WAITING = 1 RKLLM_RUN_FINISH = 2 RKLLM_RUN_ERROR = 3 class RKLLMInputType(enum.IntEnum): RKLLM_INPUT_PROMPT = 0 RKLLM_INPUT_TOKEN = 1 RKLLM_INPUT_EMBED = 2 RKLLM_INPUT_MULTIMODAL = 3 class RKLLMInferMode(enum.IntEnum): RKLLM_INFER_GENERATE = 0 RKLLM_INFER_GET_LAST_HIDDEN_LAYER = 1 RKLLM_INFER_GET_LOGITS = 2 # --- Structures --- class RKLLMExtendParam(ctypes.Structure): base_domain_id: ctypes.c_int32 embed_flash: ctypes.c_int8 enabled_cpus_num: ctypes.c_int8 enabled_cpus_mask: ctypes.c_uint32 n_batch: ctypes.c_uint8 use_cross_attn: ctypes.c_int8 reserved: ctypes.c_uint8 * 104 _fields_ = [ ("base_domain_id", ctypes.c_int32), # 基础域ID ("embed_flash", ctypes.c_int8), # 是否从闪存查询词嵌入向量（1启用，0禁用） ("enabled_cpus_num", ctypes.c_int8), # 推理启用的CPU数量 ("enabled_cpus_mask", ctypes.c_uint32), # 指示启用哪些CPU的位掩码 ("n_batch", ctypes.c_uint8), # 一次前向传播中并发处理的输入样本数，设置>1启用批量推理，默认为1 ("use_cross_attn", ctypes.c_int8), # 是否启用交叉注意力（非零启用，0禁用） ("reserved", ctypes.c_uint8 * 104) # 保留字段 ] class RKLLMParam(ctypes.Structure): model_path: ctypes.c_char_p max_context_len: ctypes.c_int32 max_new_tokens: ctypes.c_int32 top_k: ctypes.c_int32 n_keep: ctypes.c_int32 top_p: ctypes.c_float temperature: ctypes.c_float repeat_penalty: ctypes.c_float frequency_penalty: ctypes.c_float presence_penalty: ctypes.c_float mirostat: ctypes.c_int32 mirostat_tau: ctypes.c_float mirostat_eta: ctypes.c_float skip_special_token: ctypes.c_bool is_async: ctypes.c_bool img_start: ctypes.c_char_p img_end: ctypes.c_char_p img_content: ctypes.c_char_p extend_param: RKLLMExtendParam _fields_ = [ ("model_path", ctypes.c_char_p), # 模型文件路径 ("max_context_len", ctypes.c_int32), # 上下文窗口最大token数 ("max_new_tokens", ctypes.c_int32), # 最大生成新token数 ("top_k", ctypes.c_int32), # Top-K采样参数 ("n_keep", ctypes.c_int32), # 上下文窗口移动时保留的kv缓存数量 ("top_p", ctypes.c_float), # Top-P（nucleus）采样参数 ("temperature", ctypes.c_float), # 采样温度，影响token选择的随机性 ("repeat_penalty", ctypes.c_float), # 重复token惩罚 ("frequency_penalty", ctypes.c_float), # 频繁token惩罚 ("presence_penalty", ctypes.c_float), # 输入中已存在token的惩罚 ("mirostat", ctypes.c_int32), # Mirostat采样策略标志（0表示禁用） ("mirostat_tau", ctypes.c_float), # Mirostat采样Tau参数 ("mirostat_eta", ctypes.c_float), # Mirostat采样Eta参数 ("skip_special_token", ctypes.c_bool), # 是否跳过特殊token ("is_async", ctypes.c_bool), # 是否异步推理 ("img_start", ctypes.c_char_p), # 多模态输入中图像的起始位置 ("img_end", ctypes.c_char_p), # 多模态输入中图像的结束位置 ("img_content", ctypes.c_char_p), # 图像内容指针 ("extend_param", RKLLMExtendParam) # 扩展参数 ] class RKLLMLoraAdapter(ctypes.Structure): lora_adapter_path: ctypes.c_char_p lora_adapter_name: ctypes.c_char_p scale: ctypes.c_float _fields_ = [ ("lora_adapter_path", ctypes.c_char_p), ("lora_adapter_name", ctypes.c_char_p), ("scale", ctypes.c_float) ] class RKLLMEmbedInput(ctypes.Structure): embed: ctypes.POINTER(ctypes.c_float) n_tokens: ctypes.c_size_t _fields_ = [ ("embed", ctypes.POINTER(ctypes.c_float)), ("n_tokens", ctypes.c_size_t) ] class RKLLMTokenInput(ctypes.Structure): input_ids: ctypes.POINTER(ctypes.c_int32) n_tokens: ctypes.c_size_t _fields_ = [ ("input_ids", ctypes.POINTER(ctypes.c_int32)), ("n_tokens", ctypes.c_size_t) ] class RKLLMMultiModelInput(ctypes.Structure): prompt: ctypes.c_char_p image_embed: ctypes.POINTER(ctypes.c_float) n_image_tokens: ctypes.c_size_t n_image: ctypes.c_size_t image_width: ctypes.c_size_t image_height: ctypes.c_size_t _fields_ = [ ("prompt", ctypes.c_char_p), ("image_embed", ctypes.POINTER(ctypes.c_float)), ("n_image_tokens", ctypes.c_size_t), ("n_image", ctypes.c_size_t), ("image_width", ctypes.c_size_t), ("image_height", ctypes.c_size_t) ] class RKLLMCrossAttnParam(ctypes.Structure): """ 交叉注意力参数结构体 该结构体用于在解码器中执行交叉注意力时使用。 它提供编码器输出（键/值缓存）、位置索引和注意力掩码。 - encoder_k_cache必须存储在连续内存中，布局为： [num_layers][num_tokens][num_kv_heads][head_dim] - encoder_v_cache必须存储在连续内存中，布局为： [num_layers][num_kv_heads][head_dim][num_tokens] """ encoder_k_cache: ctypes.POINTER(ctypes.c_float) encoder_v_cache: ctypes.POINTER(ctypes.c_float) encoder_mask: ctypes.POINTER(ctypes.c_float) encoder_pos: ctypes.POINTER(ctypes.c_int32) num_tokens: ctypes.c_int _fields_ = [ ("encoder_k_cache", ctypes.POINTER(ctypes.c_float)), # 编码器键缓存指针（大小：num_layers * num_tokens * num_kv_heads * head_dim） ("encoder_v_cache", ctypes.POINTER(ctypes.c_float)), # 编码器值缓存指针（大小：num_layers * num_kv_heads * head_dim * num_tokens） ("encoder_mask", ctypes.POINTER(ctypes.c_float)), # 编码器注意力掩码指针（大小：num_tokens的数组） ("encoder_pos", ctypes.POINTER(ctypes.c_int32)), # 编码器token位置指针（大小：num_tokens的数组） ("num_tokens", ctypes.c_int) # 编码器序列中的token数量 ] class RKLLMPerfStat(ctypes.Structure): """ 性能统计结构体 用于保存预填充和生成阶段的性能统计信息。 """ prefill_time_ms: ctypes.c_float prefill_tokens: ctypes.c_int generate_time_ms: ctypes.c_float generate_tokens: ctypes.c_int memory_usage_mb: ctypes.c_float _fields_ = [ ("prefill_time_ms", ctypes.c_float), # 预填充阶段总耗时（毫秒） ("prefill_tokens", ctypes.c_int), # 预填充阶段处理的token数量 ("generate_time_ms", ctypes.c_float), # 生成阶段总耗时（毫秒） ("generate_tokens", ctypes.c_int), # 生成阶段处理的token数量 ("memory_usage_mb", ctypes.c_float) # 推理期间VmHWM常驻内存使用量（MB） ] class _RKLLMInputUnion(ctypes.Union): prompt_input: ctypes.c_char_p embed_input: RKLLMEmbedInput token_input: RKLLMTokenInput multimodal_input: RKLLMMultiModelInput _fields_ = [ ("prompt_input", ctypes.c_char_p), ("embed_input", RKLLMEmbedInput), ("token_input", RKLLMTokenInput), ("multimodal_input", RKLLMMultiModelInput) ] class RKLLMInput(ctypes.Structure): """ LLM输入结构体 通过联合体表示不同类型的LLM输入。 """ role: ctypes.c_char_p enable_thinking: ctypes.c_bool input_type: ctypes.c_int _union_data: _RKLLMInputUnion _fields_ = [ ("role", ctypes.c_char_p), # 消息角色："user"（用户输入）、"tool"（函数结果） ("enable_thinking", ctypes.c_bool), # 控制Qwen3模型是否启用"思考模式" ("input_type", ctypes.c_int), # 枚举类型，指定输入类型（如prompt、token、embed、multimodal） ("_union_data", _RKLLMInputUnion) # 联合体数据 ] # Properties to make accessing union members easier @property def prompt_input(self) -> bytes: # Assuming c_char_p maps to bytes if self.input_type == RKLLMInputType.RKLLM_INPUT_PROMPT: return self._union_data.prompt_input raise AttributeError("Not a prompt input") @prompt_input.setter def prompt_input(self, value: bytes): # Assuming c_char_p maps to bytes if self.input_type == RKLLMInputType.RKLLM_INPUT_PROMPT: self._union_data.prompt_input = value else: raise AttributeError("Not a prompt input") @property def embed_input(self) -> RKLLMEmbedInput: if self.input_type == RKLLMInputType.RKLLM_INPUT_EMBED: return self._union_data.embed_input raise AttributeError("Not an embed input") @embed_input.setter def embed_input(self, value: RKLLMEmbedInput): if self.input_type == RKLLMInputType.RKLLM_INPUT_EMBED: self._union_data.embed_input = value else: raise AttributeError("Not an embed input") @property def token_input(self) -> RKLLMTokenInput: if self.input_type == RKLLMInputType.RKLLM_INPUT_TOKEN: return self._union_data.token_input raise AttributeError("Not a token input") @token_input.setter def token_input(self, value: RKLLMTokenInput): if self.input_type == RKLLMInputType.RKLLM_INPUT_TOKEN: self._union_data.token_input = value else: raise AttributeError("Not a token input") @property def multimodal_input(self) -> RKLLMMultiModelInput: if self.input_type == RKLLMInputType.RKLLM_INPUT_MULTIMODAL: return self._union_data.multimodal_input raise AttributeError("Not a multimodal input") @multimodal_input.setter def multimodal_input(self, value: RKLLMMultiModelInput): if self.input_type == RKLLMInputType.RKLLM_INPUT_MULTIMODAL: self._union_data.multimodal_input = value else: raise AttributeError("Not a multimodal input") class RKLLMLoraParam(ctypes.Structure): # For inference lora_adapter_name: ctypes.c_char_p _fields_ = [ ("lora_adapter_name", ctypes.c_char_p) ] class RKLLMPromptCacheParam(ctypes.Structure): # For inference save_prompt_cache: ctypes.c_int # bool-like prompt_cache_path: ctypes.c_char_p _fields_ = [ ("save_prompt_cache", ctypes.c_int), # bool-like ("prompt_cache_path", ctypes.c_char_p) ] class RKLLMInferParam(ctypes.Structure): mode: ctypes.c_int lora_params: ctypes.POINTER(RKLLMLoraParam) prompt_cache_params: ctypes.POINTER(RKLLMPromptCacheParam) keep_history: ctypes.c_int # bool-like _fields_ = [ ("mode", ctypes.c_int), # Enum will be passed as int, changed RKLLMInferMode to ctypes.c_int ("lora_params", ctypes.POINTER(RKLLMLoraParam)), ("prompt_cache_params", ctypes.POINTER(RKLLMPromptCacheParam)), ("keep_history", ctypes.c_int) # bool-like ] class RKLLMResultLastHiddenLayer(ctypes.Structure): hidden_states: ctypes.POINTER(ctypes.c_float) embd_size: ctypes.c_int num_tokens: ctypes.c_int _fields_ = [ ("hidden_states", ctypes.POINTER(ctypes.c_float)), ("embd_size", ctypes.c_int), ("num_tokens", ctypes.c_int) ] class RKLLMResultLogits(ctypes.Structure): logits: ctypes.POINTER(ctypes.c_float) vocab_size: ctypes.c_int num_tokens: ctypes.c_int _fields_ = [ ("logits", ctypes.POINTER(ctypes.c_float)), ("vocab_size", ctypes.c_int), ("num_tokens", ctypes.c_int) ] class RKLLMResult(ctypes.Structure): """ LLM推理结果结构体 表示LLM推理的结果，包含生成的文本、token ID、隐藏层状态、logits和性能统计。 """ text: ctypes.c_char_p token_id: ctypes.c_int32 last_hidden_layer: RKLLMResultLastHiddenLayer logits: RKLLMResultLogits perf: RKLLMPerfStat _fields_ = [ ("text", ctypes.c_char_p), # 生成的文本结果 ("token_id", ctypes.c_int32), # 生成的token ID ("last_hidden_layer", RKLLMResultLastHiddenLayer), # 最后一层的隐藏状态（如果请求的话） ("logits", RKLLMResultLogits), # 模型输出的logits ("perf", RKLLMPerfStat) # 性能统计（预填充和生成） ] # --- Typedefs --- LLMHandle = ctypes.c_void_p # --- Callback Function Type --- LLMResultCallback = ctypes.CFUNCTYPE( ctypes.c_int, # 返回类型：int，表示处理状态 ctypes.POINTER(RKLLMResult), # LLM结果指针 ctypes.c_void_p, # 用户数据指针 ctypes.c_int # LLM调用状态（LLMCallState枚举值） ) """ 回调函数类型定义 用于处理LLM结果的回调函数。 参数： - result: 指向LLM结果的指针 - userdata: 回调的用户数据指针 - state: LLM调用状态（例如：完成、错误） 返回值： - 0: 正常继续推理 - 1: 暂停推理。如果用户想要修改或干预结果（例如编辑输出、注入新提示）， 返回1以暂停当前推理。稍后，使用更新的内容调用rkllm_run来恢复推理。 """ class RKLLMRuntime: def __init__(self, library_path="./librkllmrt.so"): try: self.lib = ctypes.CDLL(library_path) except OSError as e: raise OSError(f"Failed to load RKLLM library from {library_path}. " f"Ensure it's in your LD_LIBRARY_PATH or provide the full path. Error: {e}") self._setup_functions() self.llm_handle = LLMHandle() self._c_callback = None # To keep the callback object alive self._user_callback = None def _setup_functions(self): # RKLLMParam rkllm_createDefaultParam(); self.lib.rkllm_createDefaultParam.restype = RKLLMParam self.lib.rkllm_createDefaultParam.argtypes = [] # int rkllm_init(LLMHandle* handle, RKLLMParam* param, LLMResultCallback callback); self.lib.rkllm_init.restype = ctypes.c_int self.lib.rkllm_init.argtypes = [ ctypes.POINTER(LLMHandle), ctypes.POINTER(RKLLMParam), LLMResultCallback ] # int rkllm_load_lora(LLMHandle handle, RKLLMLoraAdapter* lora_adapter); self.lib.rkllm_load_lora.restype = ctypes.c_int self.lib.rkllm_load_lora.argtypes = [LLMHandle, ctypes.POINTER(RKLLMLoraAdapter)] # int rkllm_load_prompt_cache(LLMHandle handle, const char* prompt_cache_path); self.lib.rkllm_load_prompt_cache.restype = ctypes.c_int self.lib.rkllm_load_prompt_cache.argtypes = [LLMHandle, ctypes.c_char_p] # int rkllm_release_prompt_cache(LLMHandle handle); self.lib.rkllm_release_prompt_cache.restype = ctypes.c_int self.lib.rkllm_release_prompt_cache.argtypes = [LLMHandle] # int rkllm_destroy(LLMHandle handle); self.lib.rkllm_destroy.restype = ctypes.c_int self.lib.rkllm_destroy.argtypes = [LLMHandle] # int rkllm_run(LLMHandle handle, RKLLMInput* rkllm_input, RKLLMInferParam* rkllm_infer_params, void* userdata); self.lib.rkllm_run.restype = ctypes.c_int self.lib.rkllm_run.argtypes = [ LLMHandle, ctypes.POINTER(RKLLMInput), ctypes.POINTER(RKLLMInferParam), ctypes.c_void_p # userdata ] # int rkllm_run_async(LLMHandle handle, RKLLMInput* rkllm_input, RKLLMInferParam* rkllm_infer_params, void* userdata); # Assuming async also takes userdata for the callback context self.lib.rkllm_run_async.restype = ctypes.c_int self.lib.rkllm_run_async.argtypes = [ LLMHandle, ctypes.POINTER(RKLLMInput), ctypes.POINTER(RKLLMInferParam), ctypes.c_void_p # userdata ] # int rkllm_abort(LLMHandle handle); self.lib.rkllm_abort.restype = ctypes.c_int self.lib.rkllm_abort.argtypes = [LLMHandle] # int rkllm_is_running(LLMHandle handle); self.lib.rkllm_is_running.restype = ctypes.c_int # 0 if running, non-zero otherwise self.lib.rkllm_is_running.argtypes = [LLMHandle] # int rkllm_clear_kv_cache(LLMHandle handle, int keep_system_prompt, int* start_pos, int* end_pos); self.lib.rkllm_clear_kv_cache.restype = ctypes.c_int self.lib.rkllm_clear_kv_cache.argtypes = [ LLMHandle, ctypes.c_int, ctypes.POINTER(ctypes.c_int), # start_pos ctypes.POINTER(ctypes.c_int) # end_pos ] # int rkllm_get_kv_cache_size(LLMHandle handle, int* cache_sizes); self.lib.rkllm_get_kv_cache_size.restype = ctypes.c_int self.lib.rkllm_get_kv_cache_size.argtypes = [LLMHandle, ctypes.POINTER(ctypes.c_int)] # int rkllm_set_chat_template(LLMHandle handle, const char* system_prompt, const char* prompt_prefix, const char* prompt_postfix); self.lib.rkllm_set_chat_template.restype = ctypes.c_int self.lib.rkllm_set_chat_template.argtypes = [ LLMHandle, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p ] # int rkllm_set_function_tools(LLMHandle handle, const char* system_prompt, const char* tools, const char* tool_response_str); self.lib.rkllm_set_function_tools.restype = ctypes.c_int self.lib.rkllm_set_function_tools.argtypes = [ LLMHandle, ctypes.c_char_p, # system_prompt ctypes.c_char_p, # tools ctypes.c_char_p # tool_response_str ] # int rkllm_set_cross_attn_params(LLMHandle handle, RKLLMCrossAttnParam* cross_attn_params); self.lib.rkllm_set_cross_attn_params.restype = ctypes.c_int self.lib.rkllm_set_cross_attn_params.argtypes = [LLMHandle, ctypes.POINTER(RKLLMCrossAttnParam)] def create_default_param(self) -> RKLLMParam: """Creates a default RKLLMParam structure.""" return self.lib.rkllm_createDefaultParam() def init(self, param: RKLLMParam, callback_func) -> int: """ Initializes the LLM. :param param: RKLLMParam structure. :param callback_func: A Python function that matches the signature: def my_callback(result_ptr, userdata_ptr, state_enum): result = result_ptr.contents # RKLLMResult # Process result # userdata can be retrieved if passed during run, or ignored # state = LLMCallState(state_enum) :return: 0 for success, non-zero for failure. """ if not callable(callback_func): raise ValueError("callback_func must be a callable Python function.") self._user_callback = callback_func # Keep a reference to the ctypes callback object to prevent it from being garbage collected. # Always register a trampoline so we can swap the Python-level handler when needed. self._c_callback = LLMResultCallback(self._callback_trampoline) ret = self.lib.rkllm_init(ctypes.byref(self.llm_handle), ctypes.byref(param), self._c_callback) if ret != 0: raise RuntimeError(f"rkllm_init failed with error code {ret}") return ret def load_lora(self, lora_adapter: RKLLMLoraAdapter) -> int: """Loads a Lora adapter.""" ret = self.lib.rkllm_load_lora(self.llm_handle, ctypes.byref(lora_adapter)) if ret != 0: raise RuntimeError(f"rkllm_load_lora failed with error code {ret}") return ret def load_prompt_cache(self, prompt_cache_path: str) -> int: """Loads a prompt cache from a file.""" c_path = prompt_cache_path.encode('utf-8') ret = self.lib.rkllm_load_prompt_cache(self.llm_handle, c_path) if ret != 0: raise RuntimeError(f"rkllm_load_prompt_cache failed for {prompt_cache_path} with error code {ret}") return ret def release_prompt_cache(self) -> int: """Releases the prompt cache from memory.""" ret = self.lib.rkllm_release_prompt_cache(self.llm_handle) if ret != 0: raise RuntimeError(f"rkllm_release_prompt_cache failed with error code {ret}") return ret def destroy(self) -> int: """Destroys the LLM instance and releases resources.""" if self.llm_handle and self.llm_handle.value: # Check if handle is not NULL ret = self.lib.rkllm_destroy(self.llm_handle) self.llm_handle = LLMHandle() # Reset handle if ret != 0: # Don't raise here as it might be called in __del__ print(f"Warning: rkllm_destroy failed with error code {ret}") return ret return 0 # Already destroyed or not initialized def run(self, rkllm_input: RKLLMInput, rkllm_infer_params: RKLLMInferParam, userdata=None) -> int: """Runs an LLM inference task synchronously.""" # userdata can be a ctypes.py_object if you want to pass Python objects, # then cast to c_void_p. Or simply None. if userdata is not None: # Store the userdata object to keep it alive during the call self._userdata_ref = userdata c_userdata = ctypes.cast(ctypes.pointer(ctypes.py_object(userdata)), ctypes.c_void_p) else: c_userdata = None ret = self.lib.rkllm_run(self.llm_handle, ctypes.byref(rkllm_input), ctypes.byref(rkllm_infer_params), c_userdata) if ret != 0: raise RuntimeError(f"rkllm_run failed with error code {ret}") return ret def run_async(self, rkllm_input: RKLLMInput, rkllm_infer_params: RKLLMInferParam, userdata=None) -> int: """Runs an LLM inference task asynchronously.""" if userdata is not None: # Store the userdata object to keep it alive during the call self._userdata_ref = userdata c_userdata = ctypes.cast(ctypes.pointer(ctypes.py_object(userdata)), ctypes.c_void_p) else: c_userdata = None ret = self.lib.rkllm_run_async(self.llm_handle, ctypes.byref(rkllm_input), ctypes.byref(rkllm_infer_params), c_userdata) if ret != 0: raise RuntimeError(f"rkllm_run_async failed with error code {ret}") return ret def abort(self) -> int: """Aborts an ongoing LLM task.""" ret = self.lib.rkllm_abort(self.llm_handle) if ret != 0: raise RuntimeError(f"rkllm_abort failed with error code {ret}") return ret def is_running(self) -> bool: """Checks if an LLM task is currently running. Returns True if running.""" # The C API returns 0 if running, non-zero otherwise. # This is a bit counter-intuitive for a boolean "is_running". return self.lib.rkllm_is_running(self.llm_handle) == 0 def clear_kv_cache(self, keep_system_prompt: bool, start_pos: list = None, end_pos: list = None) -> int: """ 清除键值缓存 此函数用于清除部分或全部KV缓存。 参数： - keep_system_prompt: 是否在缓存中保留系统提示（True保留，False清除） 如果提供了特定范围[start_pos, end_pos)，此标志将被忽略 - start_pos: 要清除的KV缓存范围的起始位置数组（包含），每个批次一个 - end_pos: 要清除的KV缓存范围的结束位置数组（不包含），每个批次一个 如果start_pos和end_pos都设置为None，将清除整个缓存，keep_system_prompt将生效 如果start_pos[i] < end_pos[i]，只有指定的范围会被清除，keep_system_prompt将被忽略 注意：start_pos或end_pos只有在keep_history == 0且生成已通过在回调中返回1暂停时才有效 返回：0表示缓存清除成功，非零表示失败 """ # 准备C数组参数 c_start_pos = None c_end_pos = None if start_pos is not None and end_pos is not None: if len(start_pos) != len(end_pos): raise ValueError("start_pos和end_pos数组长度必须相同") # 创建C数组 c_start_pos = (ctypes.c_int * len(start_pos))(*start_pos) c_end_pos = (ctypes.c_int * len(end_pos))(*end_pos) ret = self.lib.rkllm_clear_kv_cache( self.llm_handle, ctypes.c_int(1 if keep_system_prompt else 0), c_start_pos, c_end_pos ) if ret != 0: raise RuntimeError(f"rkllm_clear_kv_cache失败，错误代码：{ret}") return ret def set_chat_template(self, system_prompt: str, prompt_prefix: str, prompt_postfix: str) -> int: """Sets the chat template for the LLM.""" c_system = system_prompt.encode('utf-8') if system_prompt else b"" c_prefix = prompt_prefix.encode('utf-8') if prompt_prefix else b"" c_postfix = prompt_postfix.encode('utf-8') if prompt_postfix else b"" ret = self.lib.rkllm_set_chat_template(self.llm_handle, c_system, c_prefix, c_postfix) if ret != 0: raise RuntimeError(f"rkllm_set_chat_template failed with error code {ret}") return ret def get_kv_cache_size(self, n_batch: int) -> list: """ 获取给定LLM句柄的键值缓存当前大小 此函数返回当前存储在模型KV缓存中的位置总数。 参数： - n_batch: 批次数量，用于确定返回数组的大小 返回： - list: 每个批次的缓存大小列表 """ # 预分配数组以存储每个批次的缓存大小 cache_sizes = (ctypes.c_int * n_batch)() ret = self.lib.rkllm_get_kv_cache_size(self.llm_handle, cache_sizes) if ret != 0: raise RuntimeError(f"rkllm_get_kv_cache_size失败，错误代码：{ret}") # 转换为Python列表 return [cache_sizes[i] for i in range(n_batch)] def set_function_tools(self, system_prompt: str, tools: str, tool_response_str: str) -> int: """ 为LLM设置函数调用配置，包括系统提示、工具定义和工具响应token 参数： - system_prompt: 定义语言模型上下文或行为的系统提示 - tools: JSON格式的字符串，定义可用的函数，包括它们的名称、描述和参数 - tool_response_str: 用于识别对话中函数调用结果的唯一标签。它作为标记标签， 允许分词器将工具输出与正常对话轮次分开识别 返回：0表示配置设置成功，非零表示错误 """ c_system = system_prompt.encode('utf-8') if system_prompt else b"" c_tools = tools.encode('utf-8') if tools else b"" c_tool_response = tool_response_str.encode('utf-8') if tool_response_str else b"" ret = self.lib.rkllm_set_function_tools(self.llm_handle, c_system, c_tools, c_tool_response) if ret != 0: raise RuntimeError(f"rkllm_set_function_tools失败，错误代码：{ret}") return ret def set_cross_attn_params(self, cross_attn_params: RKLLMCrossAttnParam) -> int: """ 为LLM解码器设置交叉注意力参数 参数： - cross_attn_params: 包含用于交叉注意力的编码器相关输入数据的结构体 （详见RKLLMCrossAttnParam说明） 返回：0表示参数设置成功，非零表示错误 """ ret = self.lib.rkllm_set_cross_attn_params(self.llm_handle, ctypes.byref(cross_attn_params)) if ret != 0: raise RuntimeError(f"rkllm_set_cross_attn_params失败，错误代码：{ret}") return ret def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.destroy() def __del__(self): self.destroy() # Ensure resources are freed if object is garbage collected def _callback_trampoline(self, result_ptr, userdata_ptr, state_enum): """ Bridge callback that forwards to the currently active Python handler. This keeps the C callback pointer stable while allowing per-call overrides. """ handler = self._user_callback if handler is None: return 0 try: return handler(result_ptr, userdata_ptr, state_enum) except Exception as exc: # Avoid propagating exceptions through the C callback boundary. print(f"[rkllm_binding] Callback raised an exception: {exc}") return 0 def forward_embed( self, embeds: np.ndarray, *, keep_history: bool = False, timeout: Optional[float] = None, return_last_only: bool = False, ) -> np.ndarray: """ Run a single forward pass with embedding input and return the last hidden layer. Args: embeds: Float32 embeddings shaped (T, H) or (1, T, H). Batch>1 is not supported. keep_history: When False, KV cache will be cleared after the call. When True, cache is kept; call clear_kv_cache() manually if needed. timeout: Optional timeout (seconds) for waiting on the callback. return_last_only: If True, return the last token vector shape (H,). Returns: np.ndarray containing hidden states (T, H) or the last token (H,). """ if embeds is None: raise ValueError("embeds must not be None.") np_embeds = np.asarray(embeds, dtype=np.float32) if np_embeds.ndim == 3: if np_embeds.shape[0] != 1: raise ValueError("Only batch size 1 is supported for forward_embed.") num_tokens = np_embeds.shape[1] flat = np_embeds.reshape(-1) elif np_embeds.ndim == 2: num_tokens = np_embeds.shape[0] flat = np_embeds.reshape(-1) else: raise ValueError("embeds must have shape (T, H) or (1, T, H).") flat = np.ascontiguousarray(flat, dtype=np.float32) embed_buffer = (ctypes.c_float * flat.size)(*flat) rk_input = RKLLMInput() rk_input.input_type = RKLLMInputType.RKLLM_INPUT_EMBED embed_input = RKLLMEmbedInput() embed_input.embed = embed_buffer embed_input.n_tokens = num_tokens rk_input._union_data.embed_input = embed_input infer_params = RKLLMInferParam() infer_params.mode = RKLLMInferMode.RKLLM_INFER_GET_LAST_HIDDEN_LAYER infer_params.keep_history = 1 if keep_history else 0 infer_params.lora_params = None infer_params.prompt_cache_params = None done = threading.Event() result_holder = {"hidden": None, "error": None} def _capture_hidden(result_ptr, userdata_ptr, state_enum): state = LLMCallState(state_enum) if state == LLMCallState.RKLLM_RUN_ERROR: result_holder["error"] = "RKLLM reported an error state." done.set() return 0 if not result_ptr: result_holder["error"] = "Empty result pointer received." done.set() return 0 result = result_ptr.contents if result.last_hidden_layer.hidden_states and result.last_hidden_layer.embd_size > 0: hidden = np.ctypeslib.as_array( result.last_hidden_layer.hidden_states, shape=(1, result.last_hidden_layer.num_tokens, result.last_hidden_layer.embd_size), ).copy() result_holder["hidden"] = hidden[-1].copy() if return_last_only else hidden done.set() return 1 # Pause further work; we already have the hidden states. if state == LLMCallState.RKLLM_RUN_FINISH: done.set() return 0 previous_callback = self._user_callback self._user_callback = _capture_hidden try: self.run(rk_input, infer_params) if not done.wait(timeout): raise TimeoutError("forward_embed timed out waiting for hidden states.") finally: self._user_callback = previous_callback if result_holder["error"]: raise RuntimeError(result_holder["error"]) if result_holder["hidden"] is None: raise RuntimeError("forward_embed did not receive hidden states.") try: if not keep_history: self.clear_kv_cache(True) except Exception: # Cache clearing best-effort; keep the forward result usable even if clearing fails. pass return result_holder["hidden"] # --- Demo CLI --- def _cli_parse_arguments() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Demo application showcasing rkllm_binding usage." ) parser.add_argument( "model", help="Path to the .rkllm model file used for inference." ) parser.add_argument( "--lib", default="./librkllmrt.so", help="Path to librkllmrt.so. Defaults to ./librkllmrt.so." ) # Core generation parameters parser.add_argument("--max-context-len", type=int, default=512, help="Maximum context length.") parser.add_argument("--max-new-tokens", type=int, default=128, help="Maximum number of new tokens to generate.") parser.add_argument("--top-k", type=int, default=1, help="Top-K sampling parameter.") parser.add_argument("--top-p", type=float, default=0.0, help="Top-P (nucleus) sampling parameter.") parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature.") parser.add_argument("--repeat-penalty", type=float, default=1.1, help="Penalty applied to repeated tokens.") parser.add_argument("--n-keep", type=int, default=0, help="Number of tokens to keep when context slides.") parser.add_argument("--mirostat", type=int, default=0, help="Enable Mirostat sampling (0 disables).") parser.add_argument("--mirostat-tau", type=float, default=5.0, help="Mirostat tau parameter.") parser.add_argument("--mirostat-eta", type=float, default=0.1, help="Mirostat eta parameter.") parser.add_argument( "--skip-special-token", action="store_true", help="Skip special tokens when generating output." ) # Input management parser.add_argument( "--input-type", choices=("prompt", "token", "multimodal"), default="prompt", help="Select prompt, raw token, or multimodal (image + prompt) input." ) parser.add_argument("--prompt", help="Prompt text to send to the model.") parser.add_argument("--prompt-file", help="Path to a UTF-8 text file containing the prompt.") parser.add_argument( "--token-ids", type=int, nargs="+", help="Raw token IDs (space separated). Only valid when --input-type token." ) parser.add_argument("--role", default="user", help="Role metadata for the input message (e.g., user/system).") parser.add_argument( "--enable-thinking", action="store_true", help="Enable thinking mode for supported models." ) parser.add_argument("--image", help="Path to an image file used when --input-type multimodal.") parser.add_argument("--vision-encoder", help="Path to the ONNX vision encoder model.") parser.add_argument( "--encoder-provider", help="Comma separated ONNX Runtime providers (e.g., 'CPUExecutionProvider')." ) parser.add_argument( "--encoder-threads", type=int, help="Thread count hint for ONNX Runtime session." ) parser.add_argument( "--encoder-input-shape", help="Override encoder input spatial size as HxW or H,W (e.g., 392x392)." ) parser.add_argument( "--norm", choices=("imagenet", "divide_255", "divide_128_sub_1"), default="imagenet", help="Image normalization preset." ) parser.add_argument( "--norm-mean", type=float, nargs=3, metavar=("R", "G", "B"), help="Override normalization mean (RGB order)." ) parser.add_argument( "--norm-std", type=float, nargs=3, metavar=("R", "G", "B"), help="Override normalization std (RGB order)." ) parser.add_argument( "--image-background", type=int, nargs=3, metavar=("R", "G", "B"), default=(128, 128, 128), help="Background color used when padding image to target size." ) parser.add_argument("--img-start-token", help="Override image start token string passed to the model.") parser.add_argument("--img-end-token", help="Override image end token string passed to the model.") parser.add_argument("--img-content-token", help="Override image content token string passed to the model.") # Inference options parser.add_argument( "--mode", choices=("generate", "hidden", "logits"), default="generate", help="Inference mode: generate tokens, return last hidden layer, or logits." ) parser.add_argument( "--no-keep-history", action="store_true", help="Do not keep dialogue history on the device." ) # Output options parser.add_argument( "--stream", action="store_true", default=True, help="Stream tokens to stdout as they arrive from the callback." ) parser.add_argument( "--hide-stats", action="store_true", help="Suppress performance statistics after inference." ) args = parser.parse_args() if args.prompt and args.prompt_file: parser.error("Arguments --prompt and --prompt-file cannot be used together.") if args.input_type == "prompt": if not args.prompt and not args.prompt_file: parser.error("Provide --prompt or --prompt-file when --input-type is prompt.") if args.token_ids: parser.error("--token-ids is only valid when --input-type token.") elif args.input_type == "token": if not args.token_ids: parser.error("--token-ids is required when --input-type token.") if args.prompt or args.prompt_file: parser.error("--prompt/--prompt-file cannot be combined with --input-type token.") else: # multimodal if args.token_ids: parser.error("--token-ids cannot be used with --input-type multimodal.") if not args.prompt and not args.prompt_file: parser.error("Provide --prompt or --prompt-file when --input-type is multimodal.") if not args.image: parser.error("--image is required when --input-type multimodal.") if not args.vision_encoder: parser.error("--vision-encoder is required when --input-type multimodal.") if args.image_background: for component in args.image_background: if component < 0 or component > 255: parser.error("--image-background values must be in the range [0, 255].") return args def _load_prompt_from_args(args: argparse.Namespace) -> str: if args.prompt: return args.prompt if args.prompt_file: try: with open(args.prompt_file, "r", encoding="utf-8") as fp: return fp.read() except OSError as exc: raise RuntimeError(f"Failed to read prompt file '{args.prompt_file}': {exc}") from exc raise RuntimeError("Prompt text is required but not provided.") def _mode_to_enum(mode: str) -> int: mapping = { "generate": RKLLMInferMode.RKLLM_INFER_GENERATE, "hidden": RKLLMInferMode.RKLLM_INFER_GET_LAST_HIDDEN_LAYER, "logits": RKLLMInferMode.RKLLM_INFER_GET_LOGITS, } return mapping[mode] def _parse_hw_string(value: str) -> Tuple[int, int]: separators = ("x", "X", ",", " ") token = value.strip() for sep in separators: if sep in token: parts = [p for p in token.split(sep) if p] break else: parts = [token] if len(parts) != 2: raise ValueError(f"Unable to parse height/width from '{value}'. Expected format like 392x392.") try: height = int(parts[0]) width = int(parts[1]) except ValueError as exc: raise ValueError(f"Height/width must be integers, got '{value}'.") from exc if height <= 0 or width <= 0: raise ValueError("Height and width must be positive integers.") return height, width def _infer_hw_from_onnx_shape(shape: Sequence) -> Tuple[Optional[int], Optional[int]]: if shape is None or len(shape) < 4: return None, None height = shape[-2] width = shape[-1] if isinstance(height, str) or height is None: height = None if isinstance(width, str) or width is None: width = None return height, width def _parse_providers(provider_str: Optional[str]) -> Optional[list]: if not provider_str: return None providers = [item.strip() for item in provider_str.split(",") if item.strip()] return providers or None def _load_vision_encoder_session(encoder_path: str, providers: Optional[list], threads: Optional[int]): try: import onnxruntime as ort except ImportError as exc: raise RuntimeError("onnxruntime is required for multimodal inference. Please install onnxruntime.") from exc sess_options = ort.SessionOptions() if threads and threads > 0: sess_options.intra_op_num_threads = threads try: if providers: session = ort.InferenceSession(encoder_path, sess_options=sess_options, providers=providers) else: session = ort.InferenceSession(encoder_path, sess_options=sess_options) except Exception as exc: raise RuntimeError(f"Failed to load vision encoder '{encoder_path}': {exc}") from exc return session def _letterbox_resize(image, target_hw: Tuple[int, int], background_color: Sequence[int]): try: import cv2 import numpy as np except ImportError as exc: raise RuntimeError("OpenCV (cv2) and numpy are required for multimodal preprocessing.") from exc target_h, target_w = target_hw if image.ndim != 3 or image.shape[2] != 3: raise RuntimeError("Expected RGB image with 3 channels.") src_h, src_w = image.shape[:2] if src_h == 0 or src_w == 0: raise RuntimeError("Loaded image has invalid dimensions.") scale = min(target_w / src_w, target_h / src_h) resized_w = max(1, int(round(src_w * scale))) resized_h = max(1, int(round(src_h * scale))) resized = cv2.resize(image, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR) canvas = np.full((target_h, target_w, 3), background_color, dtype=resized.dtype) top = (target_h - resized_h) // 2 left = (target_w - resized_w) // 2 canvas[top:top + resized_h, left:left + resized_w] = resized return canvas, resized_h, resized_w def _normalize_image(image, method: str, mean: Optional[Sequence[float]], std: Optional[Sequence[float]]): import numpy as np img = image.astype(np.float32) mean_arr = np.array(mean, dtype=np.float32) if mean else None std_arr = np.array(std, dtype=np.float32) if std else None if method == "imagenet": img = img / 255.0 if mean_arr is None: mean_arr = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32) if std_arr is None: std_arr = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32) img = (img - mean_arr) / std_arr elif method == "divide_255": img = img / 255.0 if mean_arr is not None: img = img - mean_arr if std_arr is not None: img = img / std_arr elif method == "divide_128_sub_1": img = img / 128.0 - 1.0 if mean_arr is not None: img = img - mean_arr if std_arr is not None: img = img / std_arr else: raise RuntimeError(f"Unsupported normalization method '{method}'.") return img def _encode_image_to_embedding( session, image_path: str, input_name: str, output_name: str, target_hw: Tuple[int, int], background_color: Sequence[int], norm_method: str, norm_mean: Optional[Sequence[float]], norm_std: Optional[Sequence[float]] ): try: import cv2 import numpy as np except ImportError as exc: raise RuntimeError("OpenCV (cv2) and numpy are required for multimodal preprocessing.") from exc image = cv2.imread(image_path, cv2.IMREAD_COLOR) if image is None: raise RuntimeError(f"Failed to read image from '{image_path}'.") image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) padded, resized_h, resized_w = _letterbox_resize(image, target_hw, background_color) normalized = _normalize_image(padded, norm_method, norm_mean, norm_std) tensor = np.transpose(normalized, (2, 0, 1)) # HWC -> CHW tensor = np.expand_dims(tensor, axis=0) # Add batch dimension tensor = np.ascontiguousarray(tensor, dtype=np.float32) try: output_list = session.run([output_name], {input_name: tensor}) except Exception as exc: raise RuntimeError(f"Vision encoder inference failed: {exc}") from exc if not output_list: raise RuntimeError("Vision encoder returned no outputs.") embedding = output_list[0] if embedding.ndim == 3: if embedding.shape[0] != 1: raise RuntimeError("Vision encoder output batch dimension must be 1 for a single image.") n_tokens = embedding.shape[1] elif embedding.ndim == 2: n_tokens = embedding.shape[0] else: raise RuntimeError(f"Unsupported vision encoder output shape {embedding.shape}.") flat_embedding = embedding.reshape(-1).astype(np.float32, copy=False) flat_embedding = np.ascontiguousarray(flat_embedding) return flat_embedding, n_tokens, target_hw if __name__ == "__main__": import os os.environ["RKLLM_LOG_LEVEL"] = "1" args = _cli_parse_arguments() prompt_text = None if args.input_type == "prompt": prompt_text = _load_prompt_from_args(args) token_id_array = None token_input_struct = None generated_chunks = [] perf_snapshot = { "prefill_tokens": 0, "prefill_time_ms": 0.0, "generate_tokens": 0, "generate_time_ms": 0.0, "memory_usage_mb": 0.0, } def demo_callback(result_ptr, userdata_ptr, state_enum): state = LLMCallState(state_enum) result = result_ptr.contents current_text = "" if result.text: current_text = result.text.decode("utf-8", errors="ignore") generated_chunks.append(current_text) if args.stream and current_text: print(current_text, end="", flush=True) perf_snapshot.update( prefill_tokens=result.perf.prefill_tokens, prefill_time_ms=result.perf.prefill_time_ms, generate_tokens=result.perf.generate_tokens, generate_time_ms=result.perf.generate_time_ms, memory_usage_mb=result.perf.memory_usage_mb, ) if state == LLMCallState.RKLLM_RUN_ERROR: print("\n[Callback] 推理过程中出现错误。") return 0 try: with RKLLMRuntime(library_path=args.lib) as rk_llm: params = rk_llm.create_default_param() params.model_path = os.path.abspath(args.model).encode("utf-8") params.max_context_len = args.max_context_len params.max_new_tokens = args.max_new_tokens params.top_k = args.top_k params.top_p = float(args.top_p) params.temperature = float(args.temperature) params.repeat_penalty = float(args.repeat_penalty) params.n_keep = args.n_keep params.mirostat = args.mirostat params.mirostat_tau = float(args.mirostat_tau) params.mirostat_eta = float(args.mirostat_eta) params.skip_special_token = bool(args.skip_special_token) params.is_async = False rk_llm.init(params, demo_callback) rk_input = RKLLMInput() rk_input.role = args.role.encode("utf-8") rk_input.enable_thinking = bool(args.enable_thinking) if args.input_type == "prompt": rk_input.input_type = RKLLMInputType.RKLLM_INPUT_PROMPT rk_input._union_data.prompt_input = prompt_text.encode("utf-8") else: rk_input.input_type = RKLLMInputType.RKLLM_INPUT_TOKEN token_id_array = (ctypes.c_int32 * len(args.token_ids))(*args.token_ids) token_input_struct = RKLLMTokenInput() token_input_struct.input_ids = token_id_array token_input_struct.n_tokens = len(args.token_ids) rk_input._union_data.token_input = token_input_struct infer_params = RKLLMInferParam() infer_params.mode = _mode_to_enum(args.mode) infer_params.keep_history = 0 if args.no_keep_history else 1 infer_params.lora_params = None infer_params.prompt_cache_params = None if args.stream: print("=== Streaming Output ===") rk_llm.run(rk_input, infer_params) except OSError as exc: print(f"无法加载 RKLLM 运行时库：{exc}") except RuntimeError as exc: print(f"推理失败：{exc}") except Exception as exc: print(f"发生未预期的错误：{exc}") else: if args.stream: print() # Ensure newline after streaming output final_text = "".join(generated_chunks) if final_text: print("=== 生成结果 ===") print(final_text) else: print("未收到生成文本。") if not args.hide_stats: print("=== 性能统计 ===") print( f"预填充: {perf_snapshot['prefill_tokens']} tokens / {perf_snapshot['prefill_time_ms']:.2f} ms" ) print( f"生成: {perf_snapshot['generate_tokens']} tokens / {perf_snapshot['generate_time_ms']:.2f} ms" ) print(f"最大常驻内存: {perf_snapshot['memory_usage_mb']:.2f} MB")