API 参考

AutoModel

加载和使用模型的主要入口。

from_pretrained

AutoModel.from_pretrained(
    model_name: str,
    device_map: str = "auto",
    ktransformers_config: str = None,
    **kwargs
) -> Model

参数：

• model_name: HuggingFace 模型名称或本地路径
• device_map: 设备放置策略（"auto"、"cuda:0" 等）
• ktransformers_config: YAML 配置文件路径

返回： Model 实例

generate

model.generate(
    prompt: str,
    max_new_tokens: int = 512,
    temperature: float = 0.7,
    top_p: float = 0.9,
    stream: bool = False
) -> str | Generator

参数：

• prompt: 输入文本
• max_new_tokens: 最大生成 token 数
• temperature: 采样温度
• top_p: 核采样参数
• stream: 启用流式输出

配置 API

load_config

from ktransformers import load_config

config = load_config("config.yaml")

merge_configs

from ktransformers import merge_configs

config = merge_configs(base_config, override_config)

工具函数

get_device_info

from ktransformers.utils import get_device_info

info = get_device_info()
# {'gpu_count': 1, 'gpu_memory': 24576, 'cpu_memory': 131072}