Prefix Cache 数据集工具

本项目提供两个可串联使用的 JSONL 工具：

工具	功能	输入	输出
distribution_dataset.py	生成用于前缀缓存（prefix cache）压测的完整数据集，并保留审计字段。	语料、tokenizer 与长度/缓存参数。	包含 prompt、缓存信息和 max_tokens 的 JSONL，以及审计文件。
extract_qa.py	精简生成器产物，只保留推理请求所需的三个字段。	distribution_dataset.py 生成的 JSONL。	仅含 question、answer、max_tokens 的 JSONL。

典型流程：先运行 distribution_dataset.py 生成并校验缓存数据，再运行 extract_qa.py 导出精简的请求文件。

依赖与输入

必需运行环境

• Python 3.10 或更高版本；
• transformers：运行主程序时加载 Hugging Face tokenizer 所必需；
• 可用的 tokenizer：本地 tokenizer 目录，或可访问的 Hugging Face 模型名。

安装主程序必需依赖：

python -m pip install transformers

按需依赖

• numpy：仅在 --output-length-mode truncated-normal 时必需；fixed、uniform 和 custom 模式不需要它。
• extract_qa.py：不需要第三方包，仅使用 Python 标准库。

python -m pip install numpy

语料为 UTF-8 JSONL，每行至少应有非空的 question 或 text 字段。默认语料路径为 GSM8K.jsonl。

--tokenizer 可以是本地 tokenizer 目录，也可以是 Hugging Face 模型名。

快速开始

下面的示例按输入长度分布生成数据，并将输入按递增顺序排列：

python .\distribution_dataset.py `
  --tokenizer <本地-tokenizer-目录或模型名> `
  --corpus .\GSM8K.jsonl `
  --output .\distribution.jsonl `
  --distribution "3600-3800:56,3800-4000:267,4000-4200:411,4200-4400:219,4400-4600:47" `
  --repeat-rate 0.5 `
  --blocksize 256 `
  --prefix-count 20 `
  --prefix-algorithm zipf `
  --zipf-exponent 1.0 `
  --order increasing `
  --output-length-mode fixed `
  --output-len 512 `
  --seed 42

运行结束会打印主数据文件和对应的 .prefixes.jsonl 文件路径。

输入长度来源

三种来源互斥，必须选择其一。

方式	参数	说明
分布采样	--distribution LOW-HIGH:COUNT,...	每个区间内均匀采样 COUNT 次；可直接传入存在的 .txt 文件，例如 fenbu.txt。
CSV	--input-token-csv FILE.csv	读取必需列 input_prompt_tokens；每行对应一个 prompt。
固定长度	--input-len N --num-prompts M	生成 M 条、每条输入长度均为 N 的 prompt。

--order random（默认）会打乱分布采样的结果；--order increasing 会将其排序。CSV 与固定长度模式不受 --order 影响。

前缀缓存生成规则

• 首条 prompt 是冷启动：prefix_token_length 固定为 0，其完整文本由语料拼接并补齐到目标长度。
• 从第二条起，脚本选择一条更早且长度足够的 prompt 作为 parent_prompt_index，直接复制其前 prefix_token_length 个 token。
• --repeat-rate 指定全局前缀 token 的目标比例。实际总量会受首条冷启动、可用历史长度和 block 对齐限制。
• 对长度不小于 --blocksize 的 prompt，缓存前缀按 block 对齐；短于一个 block 的 prompt 可按原始 token 长度处理。
• 前缀之后先放置最多一个 block 的 seed，剩余部分为随机 tail。非空 seed 的第一个 token 使用唯一、可安全重编码的空格开头 token，便于区分请求。
• --prefix-algorithm uniform（默认）以均匀策略分配前缀预算；zipf 以 Zipf 权重分配。--zipf-exponent 仅对 Zipf 生效。
• --prefix-count 不是独立的“共享前缀池”大小；它限制选择父 prompt 时可考虑的最近历史 prompt 数量。

生成过程中会重新编码每条文本，并验证子 prompt 的缓存前缀确实与其父 prompt 一致。

输出长度（max_tokens）

每条输出记录都包含 max_tokens。使用 --output-length-mode 选择策略：

模式	需要的参数	行为
fixed（默认）	--output-len N	每条记录的 max_tokens 都为 N。
truncated-normal	--min-output-len A --max-output-len B	在闭区间 [A, B] 内拒绝采样截断正态分布；随机种子固定为 0。
uniform	--min-output-len A --max-output-len B	在闭区间 [A, B] 内均匀整数采样，使用 --seed。
custom	--input-token-csv FILE.csv	逐行读取 CSV 的 output_tokens 列，并写入对应行的 max_tokens。

指定每条请求的输出长度

使用 custom 模式时，输入长度也必须来自同一个 CSV。CSV 必须同时包含以下两列，并且每行代表一条请求：

input_prompt_tokens,output_tokens
4096,512
2048,256
8192,1024

input_prompt_tokens 用于构造 prompt，output_tokens 原样写入该行 JSONL 的 max_tokens。因此不能将 custom 与 --input-len 或 --distribution 组合使用；缺少 output_tokens 列时会报错。

python .\distribution_dataset.py `
  --tokenizer <TOKENIZER> `
  --corpus .\GSM8K.jsonl `
  --input-token-csv .\input_prompt_tokens_output_tokens.csv `
  --output .\custom-output.jsonl `
  --repeat-rate 0.3 `
  --blocksize 64 `
  --output-length-mode custom `
  --seed 42

固定输入、均匀输出长度的示例：

python .\distribution_dataset.py `
  --tokenizer <TOKENIZER> `
  --output .\uniform.jsonl `
  --input-len 4096 `
  --num-prompts 100 `
  --repeat-rate 0.3 `
  --blocksize 64 `
  --output-length-mode uniform `
  --min-output-len 128 `
  --max-output-len 1024 `
  --seed 42

输出文件与字段

若指定的 --output 已存在，脚本会自动追加编号，例如 distribution.jsonl 会生成 distribution_1.jsonl，避免覆盖原有数据。

主输出为 JSONL，每行包含：

字段	含义
question / answer	生成的 prompt 文本；answer 固定为 "none"。
target_token_length / actual_token_length	期望输入 token 数与文本重新编码后的实际 token 数。
max_tokens	本条请求的目标输出长度。
prefix_token_length	可命中缓存的真实复制前缀长度；首条为 0。
constructed_prefix_token_length	构造 prompt 时使用的前缀长度；首条等于其完整输入长度。
seed_token_length / tail_token_length	前缀后的 seed 与尾部随机 token 数。
prefix_id / parent_prompt_index	父 prompt 的行号；首条为 -1。
seed_marker_token_id	非空 seed 的唯一首 token ID；没有 seed 时为 null。
prefix_candidate_count	本次运行的 --prefix-count 值。
prefix_pool_path	配套审计文件路径。

配套文件命名为 <输出文件名>.prefixes.jsonl（例如 distribution.prefixes.jsonl）。它逐行记录每条生成 prompt 的 prompt_index、cached_input_token_length，以及其来源：首条为 GSM8K，后续为 prompt:<父行号>。

示例：从 CSV 读取输入长度

python .\distribution_dataset.py `
  --tokenizer Qwen/Qwen3.6-27B `
  --corpus .\GSM8K.jsonl `
  --input-token-csv .\input_prompt_tokens_output_tokens.csv `
  --output .\output.jsonl `
  --repeat-rate 0.5 `
  --blocksize 16 `
  --prefix-algorithm zipf `
  --output-length-mode custom `
  --seed 42

功能二：提取精简请求文件（extract_qa.py）

extract_qa.py 用于处理 distribution_dataset.py 生成的主 JSONL。它逐行读取输入，并仅保留以下字段：

{"question": "…", "answer": "none", "max_tokens": 512}

它不会修改 prompt 文本或重新计算 token 长度；前缀、seed、tail、父 prompt 等审计字段会被移除。输入 JSONL 的每行必须是有效 JSON，且必须包含 question、answer 与 max_tokens，否则程序会报错并停止。

用法

指定输入文件后，默认输出到同一目录，并在原文件名后追加 _extracted：

python .\extract_qa.py .\output.jsonl
# 输出：.\output_extracted.jsonl

使用 -o 或 --output 指定输出文件：

python .\extract_qa.py .\output.jsonl -o .\requests.jsonl

完整串联示例：

python .\distribution_dataset.py `
  --tokenizer <TOKENIZER> `
  --input-len 4096 `
  --num-prompts 100 `
  --output .\full-dataset.jsonl `
  --repeat-rate 0.3 `
  --blocksize 64 `
  --output-len 512

python .\extract_qa.py .\full-dataset.jsonl -o .\inference-requests.jsonl