Skip to content

[Feat] toolkit: add ttft-analyze tool for UCM prefix-cache TTFT estimation - #1308

Draft
student-jhz wants to merge 5 commits into
ModelEngine-Group:developfrom
student-jhz:toolkit
Draft

[Feat] toolkit: add ttft-analyze tool for UCM prefix-cache TTFT estimation#1308
student-jhz wants to merge 5 commits into
ModelEngine-Group:developfrom
student-jhz:toolkit

Conversation

@student-jhz

Copy link
Copy Markdown

Purpose
新增 ttft-analyze 工具,用于在部署前预估 KV Cache 完全命中 UCM SSD/DRAM Prefix Cache 时的 TTFT(首 Token 时间)。基于用户给定(或实测)的 Full Prefill / Full HBM Prefix Cache TTFT,结合存储读带宽、H2D 带宽与输入长度,输出预估 TTFT、相对增益/损失与瓶颈拆分,辅助部署前的 TTFT 预判。纯计算,无外部依赖,不运行真实推理。
Modifications

  • 新增 ttft-analyze 工具(toolkit/ucm_toolkit/tools/ttft_analyze/):
  • adapter.py:CLI 入口,解析 --model-dir、--posix-bw、--h2d-bw、--input-len、--ttft-prefill、--ttft-hbm、--tp 参数,渲染结果表格
  • model.py:TTFT 计算模型,支持 layered(边加载边计算,TTFT = max(compute, load))和 full(先整段加载再计算,TTFT = compute + load)两种模式
  • kv_size.py:根据 config.json 自动识别 DSA / MLA / GQA(含 MHA、MQA)架构并推导 KV cache 字节数,口径与 docs/source/getting-started/kv_cache_calculator.md 一致
  • README.md:建模原理、用法、参数说明
  • 注册工具(toolkit/ucm_toolkit/registry.py):在 init_builtin_tools() 中导入并注册 TtftAnalyzeTool
  • 更新工具文档(toolkit/README.md):工具列表中新增 ttft-analyze 行
    建模要点:
  • GQA/MHA:每卡加载 cache_total / tp,存储读带宽多卡均分
  • MLA/DSA:latent 不切分,每卡加载 cache_total,单卡独享 H2D 带宽
    Test
    新增 toolkit/tests/test_ttft_analyze.py(329 行),覆盖:
  • KV cache 大小推导(GQA / MLA / DSA 各架构)
  • TTFT 计算模型(layered / full 两种模式)
  • 瓶颈识别与增益/损失计算
  • 边界输入校验(带宽 ≤ 0、输入长度 ≤ 0、tp < 1、缺模型配置字段)
    运行测试:
    cd toolkit && python -m pytest tests/test_ttft_analyze.py -v

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant