From df8033d91bb4e0e55996c0c3d7489237c7b0170a Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Tue, 25 Aug 2026 20:02:10 +0800 Subject: [PATCH] feat: add AFD prefill performance simulator Signed-off-by: jiangkuaixue123 --- simulator/.gitignore | 4 + simulator/README.md | 289 +++++ simulator/__init__.py | 12 + simulator/__main__.py | 129 ++ simulator/config.py | 386 ++++++ simulator/engine.py | 1075 +++++++++++++++++ .../examples/continuous-prefix-cache.json | 60 + simulator/examples/fixed.json | 40 + simulator/examples/requests.csv | 5 + simulator/profile_builder.py | 438 +++++++ simulator/profiles.py | 354 ++++++ simulator/profiles/README.md | 6 + simulator/server.py | 127 ++ simulator/tests/__init__.py | 1 + simulator/tests/helpers.py | 90 ++ simulator/tests/test_engine.py | 372 ++++++ simulator/tests/test_profiles.py | 123 ++ simulator/tests/test_server.py | 56 + simulator/tests/test_workload.py | 91 ++ simulator/web/index.html | 652 ++++++++++ simulator/workload.py | 261 ++++ 21 files changed, 4571 insertions(+) create mode 100644 simulator/.gitignore create mode 100644 simulator/README.md create mode 100644 simulator/__init__.py create mode 100644 simulator/__main__.py create mode 100644 simulator/config.py create mode 100644 simulator/engine.py create mode 100644 simulator/examples/continuous-prefix-cache.json create mode 100644 simulator/examples/fixed.json create mode 100644 simulator/examples/requests.csv create mode 100644 simulator/profile_builder.py create mode 100644 simulator/profiles.py create mode 100644 simulator/profiles/README.md create mode 100644 simulator/server.py create mode 100644 simulator/tests/__init__.py create mode 100644 simulator/tests/helpers.py create mode 100644 simulator/tests/test_engine.py create mode 100644 simulator/tests/test_profiles.py create mode 100644 simulator/tests/test_server.py create mode 100644 simulator/tests/test_workload.py create mode 100644 simulator/web/index.html create mode 100644 simulator/workload.py diff --git a/simulator/.gitignore b/simulator/.gitignore new file mode 100644 index 00000000..7806ed66 --- /dev/null +++ b/simulator/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.py[cod] +profiles/*.json +profiles/traces/ diff --git a/simulator/README.md b/simulator/README.md new file mode 100644 index 00000000..a211a2d4 --- /dev/null +++ b/simulator/README.md @@ -0,0 +1,289 @@ +# DSV4-Flash Prefill 性能仿真器 + +该工具在相同 16-die 预算下比较两种 Prefill 执行语义： + +- CAMAsync AFD：Attention `DP2×TP4`（8 die）与 FFN `EP8`（8 die），两个 Attention DP 独立调度，FFN 使用共享 FCFS 队列； +- 合并部署：`DP4×TP4/SP4/global EP16`，各 DP 每层等待全局 wave；全局 dispatch/routed expert/combine collective 按四个 DP 的 query token 总量建模，combine 后的本地尾段按最重 DP 建模。 + +Python 后端是唯一仿真实现。CLI、HTTP API 和页面都调用同一套调度、离散事件和指标计算逻辑，前端不重复实现模型。 + +## 1. 快速开始 + +要求 Python 3.10 或更高版本。仿真器本身只使用标准库。 + +### 1.1 生成 msModeling profile + +运行时只读取归一化 profile JSON，不会在每次页面操作时启动 msModeling： + +```bash +python -m simulator profiles build \ + --msmodeling-root /path/to/msmodeling \ + --python /path/to/msmodeling/python \ + --output simulator/profiles/dsv4-flash-910c.json +``` + +默认网格： + +- query anchors：`1,128,512,2K,4K,8K,16K,32K,64K,128K`； +- prefix anchors：`0,8K,32K,64K,96K,120K`； +- AFD Attention profile：8-device `DP2×TP4`，保留 `attention_router/afd_post`； +- AFD 单 FFN job profile：8-device `DP8×TP1×EP8`，关闭 SP，每 rank 输入 `ceil(stage_query_tokens/8)`，保留 `routed_experts/shared_expert`； +- 合并 profile：16-device `DP4×TP4×EP16`； +- `DeepSeek-V4-Flash`、msModeling `analytic`、sequence parallel、compile 路径。 + +最大的 query anchor 同时定义最大 context（默认 128K）。生成器会自动补 `prefix=0`、`query=1`、每个 prefix 的 `query=max_context-prefix` 边界点，以及 `prefix=max_context-1, query=1`，因此整个 `prefix+query<=max_context` 三角域都可插值。可通过 `--query-anchors`、`--prefix-anchors`、`--model-id`、`--device` 修改。非默认模型还应通过 `--hidden-size` 和 `--moe-top-k` 记录正确的模型 provenance；tooltip 的 Shape 直接来自 trace，TopK 来自 `--moe-top-k`。超出生成域的输入会报错，不做静默外推。`--keep-traces DIR` 可保留中间 Chrome trace。 + +### 1.2 启动页面 + +```bash +python -m simulator serve \ + --profiles simulator/profiles/dsv4-flash-910c.json \ + --host 127.0.0.1 \ + --port 8765 +``` + +浏览器访问 `http://127.0.0.1:8765`。 + +逐层关键路径时间线支持交互浏览：鼠标悬停或聚焦时间线后，按 +`W` / `S` 缩放，按 `A` / `D` 左右平移，按 `R` 复位；也可以使用 +鼠标滚轮缩放、按住拖拽平移、双击或点击“重置视图”复位。当前可见 +时间范围和缩放倍数显示在时间线右上角。鼠标悬停在事件色块上会显示 +阶段、资源、层号、起止时间、持续时间、token 数、批次和 uBatch 信息。 +`W` / `S` / `A` / `D` 支持长按连续移动，动画速度与浏览器刷新率同步。 +放大到事件色块有足够空间时，色块内会直接显示 Attention、Router、 +Dispatch、FFN、Combine、Barrier 等阶段标签；空间不足时自动隐藏文字。 +时间线按执行路径合并展示泳道：AFD 的 CAM dispatch/combine 事件并入对应 +DP Attention 泳道，merged 的全局 EP16 阶段复制显示在四条 DP +Attention/FFN 泳道中。这里只改变前端布局，不改变后端事件和仿真结果。 +merged 时间线还会把相邻的 combine collective 与本地 +unpermute/TopK-weight 阶段合并成一个 `Combine` 色块，并将 SP 收尾阶段 +显示为 `TP AllGather + HC Post`；profile 中的原始分项时延仍保持独立。 +悬停 AFD 的 dispatch/combine 通信色块时，时间线会按相同的 +layer、batch 和 uBatch 配对 Attention/CAM 与 FFN EP8 两端，并用带方向的 +箭头显示数据流；横向跨度同时反映发送完成到接收开始之间的排队时间。 +悬停 `FFN Compute`、`Routed Experts` 或 `Shared Expert` 时还会显示 +msModeling trace 中 EP rank 的总输入 Shape、每个本地 Expert 的 GMM 实采 +Shape、MoE TopK 和对应架构的 EP 数。Shape 直接从生成 profile 所用 trace +提取，不再根据请求 token 数推导。相同 Shape 的本地 Expert 会压缩显示为 +`count × [tokens, hidden]`。运行点落在两个 query anchor 之间时，页面会同时 +列出两个实采 Shape、各自的采样 Query 和线性插值权重。Shared Expert 不经过 +TopK 路由，因此 TopK 显示“不适用”。Expert Shape 只随 query 变化，profile +仅在 `prefix=0` 的 query anchor 保存一份，避免随 prefix 重复数据。 + +### 1.3 CLI 仿真 + +```bash +python -m simulator simulate \ + --profiles simulator/profiles/dsv4-flash-910c.json \ + --config simulator/examples/continuous-prefix-cache.json \ + --output /tmp/dsv4-result.json + +python -m simulator sweep \ + --profiles simulator/profiles/dsv4-flash-910c.json \ + --config simulator/examples/continuous-prefix-cache.json \ + --output /tmp/dsv4-sweep.json +``` + +## 2. 全部配置字段 + +配置文件顶层是 JSON object。没有提供的字段使用表中默认值。 + +### 2.1 Workload + +| 字段 | 类型 / 默认值 | 说明 | +| --- | --- | --- | +| `mode` | `fixed` / `"fixed"` | `fixed`：所有请求在 `t=0` 可调度；`continuous`：按 arrival 配置持续生成请求。 | +| `fixed_lengths` | `int[]` / `[512,8192,2048,6144]` | 固定模式的完整 Prompt 长度。未使用 CSV 时生效。 | +| `length_mix` | `{tokens,weight}[]` | 持续模式的离散长度分布；`weight` 只需为正数，不要求预归一化。 | +| `csv_path` | `string|null` | CLI 读取的 CSV 路径；与 `csv_text` 互斥。页面上传会使用 `csv_text`。 | +| `csv_text` | `string|null` | CSV 原始内容；与 `csv_path` 互斥。 | +| `csv_sampling` | `cycle|sample` / `"cycle"` | 无时间戳 CSV 在持续模式下的选取方式：按行循环或有放回随机采样。 | + +优先级：CSV > `fixed_lengths`/`length_mix`。固定模式下 CSV 每行回放一次；持续模式下无时间戳 CSV 提供经验长度分布。 + +### 2.2 `arrival` + +| 字段 | 类型 / 默认值 | 说明 | +| --- | --- | --- | +| `arrival.kind` | `constant|poisson|trace` / `constant` | 固定间隔、泊松到达或回放 CSV 时间戳。 | +| `arrival.qps` | `float` / `1.0` | `constant`/`poisson` 的 offered QPS；trace 模式忽略。 | +| `arrival.duration_s` | `float` / `60` | warmup 后的统计窗口和请求生成时长。 | +| `arrival.warmup_s` | `float` / `10` | 持续负载预热时长；该区间到达的请求不进入最终指标。 | +| `arrival.seed` | `int` / `1024` | 泊松间隔随机种子；长度采样使用 `seed+1`。 | + +`trace` 只支持持续模式，并要求 CSV 每一行都有 `arrival_time_ms`。第一条时间戳归零后按原间隔回放，只保留 `[0, warmup+duration)` 的请求；指标只统计 `[warmup, warmup+duration)` 内到达的请求，窗口内的空闲时间也进入吞吐分母。精确 trace 模式不支持自动 QPS 扫描。 + +### 2.3 `scheduler` + +| 字段 | 类型 / 默认值 | 说明 | +| --- | --- | --- | +| `scheduler.policy` | `round_robin|vllm_queue_aware` / `round_robin` | `round_robin` 按到达顺序严格轮询；`vllm_queue_aware` 在请求到达时选择尚未完成请求数最少的 DP，同分时轮转。 | +| `scheduler.max_num_seqs` | `int` / `64` | 一个 DP scheduler batch 最多包含的请求/请求 chunk 数。 | +| `scheduler.max_num_batched_tokens` | `int` / `8192` | 一个 DP batch 的未缓存 query token 预算。 | +| `scheduler.chunked_prefill` | `bool` / `false` | 是否允许长 Prompt 分多个 Prefill batch。AFD chunked 路径是敏感性假设，不代表当前运行时已支持。 | +| `scheduler.chunk_size` | `int` / `max_num_batched_tokens` | 单请求每次最多调度的 query token；不得超过 batch token 预算。 | + +每个 DP 使用 FIFO 装箱。non-chunked 请求的未缓存长度若超过 token 预算会直接报错。chunked 请求保留已计算 prefix，最终 chunk 完成才算请求完成。 + +`vllm_queue_aware` 对齐当前 vLLM 内置 DP 负载均衡的请求数口径和轮转平局规则。vLLM 还会在存在 waiting 请求且 KV 使用率超过 50% 时增加压力惩罚；本仿真器没有 Decode 与 KV 容量/驻留模型，因此不模拟该项，也不将 Prefix Cache 命中率误作 KV 容量压力。请求长度本身不参与该策略评分。 + +### 2.4 `prefix_cache` + +| 字段 | 类型 / 默认值 | 说明 | +| --- | --- | --- | +| `prefix_cache.enabled` | `bool` / `false` | Prefix Cache 总开关。关闭时忽略 CSV cache 字段和全局命中参数。 | +| `prefix_cache.request_hit_rate` | `[0,1]` / `0` | 未提供请求级 cache 数据时，一个请求发生命中的概率。 | +| `prefix_cache.matched_prefix_ratio` | `[0,1]` / `0` | 命中请求中被缓存的前缀 token 比例。 | +| `prefix_cache.block_size` | `int` / `32` | 采样的缓存长度向下对齐到该 block；CSV 的真实 `cached_prefix_tokens` 不再对齐。 | +| `prefix_cache.lookup_fixed_ms` | `float` / `0` | 每个请求的固定缓存查找开销。 | +| `prefix_cache.lookup_per_block_ms` | `float` / `0` | 每个已缓存 block 的附加查找开销。 | +| `prefix_cache.seed` | `int` / `1024` | 命中与否的随机种子。两种架构复用同一采样结果。 | + +计算语义： + +```text +prefix_tokens = cached_prefix_tokens + 已完成的 chunk tokens +query_tokens = 当前实际 Prefill tokens +``` + +Scheduler、FFN、Router 和 CAM 只处理 query tokens；Attention 用 `(prefix_tokens, query_tokens)` 查询 msModeling profile。输出同时报告逻辑输入 token 与实际计算 token。 + +merged 的 Attention 仍按每个 DP 的真实 query tokens 分别查表，barrier 等待最慢 +DP。`merged_dispatch`、`routed_experts`、`merged_combine` 先求当前 wave 的 +四 DP query tokens 总和，再用 `global_query_tokens/4` 查询对称 DP4 profile； +不足 1 token 时按最小 anchor 1 处理。combine 后的 `merged_combine_local`、 +`shared_expert`、`merged_sp_post` 属于 per-DP 本地尾段，使用当前 wave 的 +`max(dp_query_tokens)` 查询 profile。时间线的 `Token 数` 是该 phase 的 workload +口径：前三段显示全局总量，后三段显示最重 DP token；`Profile Query` 是实际 +用于查表的 token 数。这里按需求将 dispatch 也视作全局 payload 的对称等效 +近似；它不刻画发送端 DP token skew。若要研究 sender-side dispatch 尾延迟， +需要另行使用 `max(dp_query_tokens)` 或建立非对称通信 profile。 + +### 2.5 `afd` + +| 字段 | 类型 / 默认值 | 说明 | +| --- | --- | --- | +| `afd.ubatch_split` | `request|token` / `request` | 两个 CAM MoE stage 的切分方式。请求模式按 token 总量选择最接近均衡的请求边界；token 模式可切开单个请求 chunk。无法形成两个非空 stage 时回退单 stage。 | + +### 2.6 `cam` + +每个 CAM leg 使用： + +```text +latency_ms = fixed_ms + per_token_ms × stage_query_tokens +``` + +| 字段 | 默认值 `(fixed_ms, per_token_ms)` | 说明 | +| --- | --- | --- | +| `cam.calibrated` | `false` | 仅作结果可信度标记，不改变计算。 | +| `cam.dispatch_send` | `(0.11, 1/52000)` | Attention 侧 dispatch send。 | +| `cam.dispatch_recv` | `(0.10, 1/68000)` | FFN 侧 dispatch recv。 | +| `cam.combine_send` | `(0.10, 1/70000)` | FFN 侧 combine send。 | +| `cam.combine_recv` | `(0.12, 1/58000)` | Attention 侧 combine recv。 | +| `cam..fixed_ms` | 见上 | 该 leg 固定启动时延。 | +| `cam..per_token_ms` | 见上 | 每 query token 的线性时延。 | + +默认值来自旧流水页面，只是未校准占位值。正式结论应使用 CAM microbenchmark 拟合值并设置 `calibrated=true`。 + +### 2.7 `slo`、`sweep` 和 TTFT + +| 字段 | 类型 / 默认值 | 说明 | +| --- | --- | --- | +| `slo.ttft_limit_ms` | `float` / `1000` | Prefill TTFT 代理的 SLO 上限。 | +| `slo.target_ratio` | `(0,1]` / `0.99` | 自动容量判断要求的请求达标比例。 | +| `fixed_ttft_overhead_ms` | `float` / `0` | 加到每个请求 TTFT 上的 tokenizer/HTTP 等固定外部开销；不占用模拟计算资源。 | +| `sweep.min_qps` | `float` / `0.5` | 粗扫下界。 | +| `sweep.max_qps` | `float` / `64` | 粗扫上界。 | +| `sweep.coarse_points` | `int` / `10` | 对数间隔粗扫点数。 | +| `sweep.refinement_steps` | `int` / `7` | 最后一个 PASS 与第一个 FAIL 之间的二分轮数。 | +| `sweep.throughput_tolerance_ratio` | `(0,1]` / `0.99` | 除 SLO 外，要求 achieved throughput 至少达到 offered QPS 的比例。 | + +QPS 扫描只支持 `mode="continuous"`，因为固定请求集没有 offered QPS；精确 CSV timestamp trace 也不支持改变 QPS。 + +TTFT 定义： + +```text +TTFT_proxy = Prefill完成时间 - arrival_time + fixed_ttft_overhead_ms +``` + +它不包含首个 Decode step。持续模式停止到达后会 drain 完成，避免遗漏积压请求。 + +### 2.8 `output` + +| 字段 | 类型 / 默认值 | 说明 | +| --- | --- | --- | +| `output.include_timeline` | `bool` / `true` | 是否返回逐层阶段事件。QPS 扫描内部会关闭。 | +| `output.timeline_max_events` | `int` / `20000` | 最大时间线事件数，超过后截断并设置 `timeline_truncated=true`。 | +| `output.include_requests` | `bool` / `true` | 是否返回每个请求的到达、DP、cache、完成时间和 TTFT。 | + +## 3. CSV 格式 + +### 3.1 只有线上长度列表 + +```csv +input_length +512 +8192 +8192 +32768 +``` + +`input_length` 必填。重复行会保留，因此自然构成经验分布。 + +### 3.2 完整 trace + +```csv +request_id,arrival_time_ms,input_length,cached_prefix_tokens +r001,0,8192,4096 +r002,17,512,0 +r003,21,32768,24576 +``` + +| 列 | 必填 | 说明 | +| --- | --- | --- | +| `input_length` | 是 | 完整 Prompt token 数，必须为正整数。 | +| `request_id` | 否 | 原始请求标识；缺失时生成 `r1`、`r2`。 | +| `arrival_time_ms` | 否 | 线上到达时间。必须全部行都有或全部没有。 | +| `cached_prefix_tokens` | 否 | 真实缓存前缀长度，必须满足 `0 <= cached < input_length`。Prefix Cache 开启时覆盖全局采样。 | + +未知列会被忽略。解析错误会报告 CSV 行号。 + +## 4. HTTP API + +| 接口 | 说明 | +| --- | --- | +| `GET /api/defaults` | 返回全部默认配置与当前 profile 元数据。 | +| `POST /api/simulate` | Body 为完整配置 JSON；返回 AFD、合并结果和对比倍率。 | +| `POST /api/sweep` | Body 为完整配置 JSON；返回两种架构的 QPS 曲线和最大 SLO QPS。 | + +页面与 API 同源，默认只监听 `127.0.0.1`。请求 body 上限为 10 MiB。 + +## 5. 输出指标 + +- `throughput_rps`：统计请求数除以包含 drain 的有效时长； +- `input_tokens_per_s`：逻辑完整 Prompt 吞吐； +- `compute_tokens_per_s`：扣除缓存前缀后的实际 Prefill token 吞吐； +- `ttft_mean/p50/p90/p99_ms`：Prefill TTFT 代理； +- `slo_attainment`：TTFT 不超过上限的请求比例； +- `slo_goodput_rps`：SLO 达标请求数除以有效时长； +- `utilization`：各 Attention DP 和 FFN/EP 关键资源的忙时比例； +- `barrier_wait_ms`：合并路径所有 DP 的同步等待总和； +- `attention_wait_ms`：AFD Attention 等待 FFN 返回的总和。 + +## 6. 当前模型边界 + +- 算子时延来自 msModeling analytic trace；CAM 时延来自独立参数模型。AFD profile 在归一化阶段按 phase 合成 Attention 与单 FFN job 两种 trace，并在 JSON metadata 中保存来源和命令。 +- Attention batch 时延按各请求/chunk profile 求和；AFD FFN/MoE 使用 stage 总 query tokens 查表。FFN 侧假设 EP8 各 rank 均分单 job token，shared expert 也按每 rank `ceil(tokens/8)` 的 DP8×TP1 口径建模。 +- merged 的 routed/combine 以及按需求近似的 dispatch 使用 `global_query_tokens/4`；combine-local/shared/SP 使用 `max(dp_query_tokens)`。 +- 不模拟 Decode、KV transfer、prefix cache 容量/淘汰算法、MTP、graph、prefix cache lookup 并发、HBM OOM、EPLB 或真实专家负载偏斜。 +- AFD DSV4 NPU 与 AFD chunked prefill 都是架构性能假设，不代表当前 afd-plugin 已完成对应 E2E 支持。 + +## 7. 测试 + +无额外测试依赖： + +```bash +python -m unittest discover -s simulator/tests -v +``` + +测试覆盖 profile 插值/越界、CSV、Prefix Cache、固定/持续负载、chunked prefill、AFD 双 uBatch、合并屏障和 QPS 扫描复现性。 diff --git a/simulator/__init__.py b/simulator/__init__.py new file mode 100644 index 00000000..5bc94dda --- /dev/null +++ b/simulator/__init__.py @@ -0,0 +1,12 @@ +"""DeepSeek-V4-Flash Prefill performance simulator.""" + +from simulator.config import SimulationConfig +from simulator.engine import compare_architectures, sweep_qps +from simulator.profiles import ProfileBundle + +__all__ = [ + "ProfileBundle", + "SimulationConfig", + "compare_architectures", + "sweep_qps", +] diff --git a/simulator/__main__.py b/simulator/__main__.py new file mode 100644 index 00000000..9a580b7f --- /dev/null +++ b/simulator/__main__.py @@ -0,0 +1,129 @@ +"""Command-line entry point for the DSV4 Prefill simulator.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +from simulator.config import SimulationConfig +from simulator.engine import compare_architectures, sweep_qps +from simulator.profile_builder import ( + DEFAULT_HIDDEN_SIZE, + DEFAULT_MOE_TOP_K, + build_profile_bundle, +) +from simulator.profiles import ProfileBundle +from simulator.server import serve + +DEFAULT_QUERY_ANCHORS = "1,128,512,2048,4096,8192,16384,32768,65536,131072" +DEFAULT_PREFIX_ANCHORS = "0,8192,32768,65536,98304,122880" + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + + serve_parser = commands.add_parser("serve", help="start the local web UI") + _add_profile_argument(serve_parser) + serve_parser.add_argument("--host", default="127.0.0.1") + serve_parser.add_argument("--port", type=int, default=8765) + + simulate_parser = commands.add_parser("simulate", help="run one simulation") + _add_profile_argument(simulate_parser) + _add_config_arguments(simulate_parser) + + sweep_parser = commands.add_parser("sweep", help="run an SLO QPS sweep") + _add_profile_argument(sweep_parser) + _add_config_arguments(sweep_parser) + + profiles_parser = commands.add_parser("profiles", help="profile operations") + profile_commands = profiles_parser.add_subparsers( + dest="profile_command", required=True + ) + build_parser = profile_commands.add_parser( + "build", help="generate an analytic profile grid with msModeling" + ) + build_parser.add_argument("--msmodeling-root", required=True) + build_parser.add_argument("--python", dest="python_executable", required=True) + build_parser.add_argument("--output", required=True) + build_parser.add_argument("--model-id", default=None) + build_parser.add_argument("--device", default=None) + build_parser.add_argument("--hidden-size", type=int, default=DEFAULT_HIDDEN_SIZE) + build_parser.add_argument("--moe-top-k", type=int, default=DEFAULT_MOE_TOP_K) + build_parser.add_argument("--query-anchors", default=DEFAULT_QUERY_ANCHORS) + build_parser.add_argument("--prefix-anchors", default=DEFAULT_PREFIX_ANCHORS) + build_parser.add_argument("--keep-traces") + + args = parser.parse_args(argv) + if args.command == "serve": + serve(ProfileBundle.load(args.profiles), host=args.host, port=args.port) + return 0 + if args.command in {"simulate", "sweep"}: + config = _load_config(args.config) + profiles = ProfileBundle.load(args.profiles) + result = ( + compare_architectures(config, profiles) + if args.command == "simulate" + else sweep_qps(config, profiles) + ) + _write_result(result, args.output) + return 0 + if args.command == "profiles" and args.profile_command == "build": + kwargs = { + "msmodeling_root": args.msmodeling_root, + "python_executable": args.python_executable, + "output_path": args.output, + "query_anchors": _parse_anchors(args.query_anchors), + "prefix_anchors": _parse_anchors(args.prefix_anchors), + "hidden_size": args.hidden_size, + "moe_top_k": args.moe_top_k, + "keep_traces": args.keep_traces, + } + if args.model_id: + kwargs["model_id"] = args.model_id + if args.device: + kwargs["device"] = args.device + build_profile_bundle(**kwargs) + print(args.output) + return 0 + parser.error("unsupported command") + return 2 + + +def _add_profile_argument(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--profiles", required=True, help="normalized profile JSON") + + +def _add_config_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--config", required=True, help="simulation config JSON") + parser.add_argument("--output", help="result JSON; stdout when omitted") + + +def _load_config(path: str) -> SimulationConfig: + with Path(path).open(encoding="utf-8") as handle: + raw = json.load(handle) + if not isinstance(raw, dict): + raise ValueError("config JSON must be an object") + return SimulationConfig.from_mapping(raw) + + +def _write_result(result: dict, output: str | None) -> None: + text = json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False) + if output: + Path(output).write_text(text + "\n", encoding="utf-8") + print(output) + else: + sys.stdout.write(text + "\n") + + +def _parse_anchors(raw: str) -> tuple[int, ...]: + values = tuple(sorted({int(value.strip()) for value in raw.split(",")})) + if not values or any(value < 0 for value in values): + raise ValueError("anchors must be comma-separated non-negative integers") + return values + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/simulator/config.py b/simulator/config.py new file mode 100644 index 00000000..50a72128 --- /dev/null +++ b/simulator/config.py @@ -0,0 +1,386 @@ +"""Typed configuration and validation for the Prefill simulator.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any + +DEFAULT_PROFILE_MIN_TOKENS = 128 +DEFAULT_PROFILE_MAX_TOKENS = 131_072 +DEFAULT_MAX_NUM_SEQS = 64 +DEFAULT_MAX_BATCHED_TOKENS = 8_192 +DEFAULT_SIMULATION_DURATION_S = 60.0 +DEFAULT_WARMUP_S = 10.0 +DEFAULT_RANDOM_SEED = 1_024 +DEFAULT_SLO_LIMIT_MS = 1_000.0 +DEFAULT_SLO_TARGET_RATIO = 0.99 +DEFAULT_TIMELINE_MAX_EVENTS = 20_000 + + +def _mapping(value: Any, name: str) -> dict[str, Any]: + if value is None: + return {} + if not isinstance(value, dict): + raise ValueError(f"{name} must be an object") + return value + + +def _positive(value: float | int, name: str) -> None: + if value <= 0: + raise ValueError(f"{name} must be positive") + + +@dataclass(frozen=True) +class LengthBucket: + tokens: int + weight: float = 1.0 + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> LengthBucket: + item = cls(tokens=int(raw["tokens"]), weight=float(raw.get("weight", 1.0))) + _positive(item.tokens, "length_mix.tokens") + _positive(item.weight, "length_mix.weight") + return item + + +@dataclass(frozen=True) +class ArrivalConfig: + kind: str = "constant" + qps: float = 1.0 + duration_s: float = DEFAULT_SIMULATION_DURATION_S + warmup_s: float = DEFAULT_WARMUP_S + seed: int = DEFAULT_RANDOM_SEED + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> ArrivalConfig: + item = cls( + kind=str(raw.get("kind", "constant")), + qps=float(raw.get("qps", 1.0)), + duration_s=float(raw.get("duration_s", DEFAULT_SIMULATION_DURATION_S)), + warmup_s=float(raw.get("warmup_s", DEFAULT_WARMUP_S)), + seed=int(raw.get("seed", DEFAULT_RANDOM_SEED)), + ) + if item.kind not in {"constant", "poisson", "trace"}: + raise ValueError("arrival.kind must be constant, poisson, or trace") + if item.kind != "trace": + _positive(item.qps, "arrival.qps") + _positive(item.duration_s, "arrival.duration_s") + if item.warmup_s < 0: + raise ValueError("arrival.warmup_s must be non-negative") + return item + + +@dataclass(frozen=True) +class SchedulerConfig: + policy: str = "round_robin" + max_num_seqs: int = DEFAULT_MAX_NUM_SEQS + max_num_batched_tokens: int = DEFAULT_MAX_BATCHED_TOKENS + chunked_prefill: bool = False + chunk_size: int = DEFAULT_MAX_BATCHED_TOKENS + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> SchedulerConfig: + max_batched_tokens = int( + raw.get("max_num_batched_tokens", DEFAULT_MAX_BATCHED_TOKENS) + ) + item = cls( + policy=str(raw.get("policy", "round_robin")), + max_num_seqs=int(raw.get("max_num_seqs", DEFAULT_MAX_NUM_SEQS)), + max_num_batched_tokens=max_batched_tokens, + chunked_prefill=bool(raw.get("chunked_prefill", False)), + chunk_size=int(raw.get("chunk_size", max_batched_tokens)), + ) + if item.policy not in {"round_robin", "vllm_queue_aware"}: + raise ValueError("scheduler.policy must be round_robin or vllm_queue_aware") + _positive(item.max_num_seqs, "scheduler.max_num_seqs") + _positive(item.max_num_batched_tokens, "scheduler.max_num_batched_tokens") + _positive(item.chunk_size, "scheduler.chunk_size") + if item.chunk_size > item.max_num_batched_tokens: + raise ValueError( + "scheduler.chunk_size cannot exceed max_num_batched_tokens" + ) + return item + + +@dataclass(frozen=True) +class PrefixCacheConfig: + enabled: bool = False + request_hit_rate: float = 0.0 + matched_prefix_ratio: float = 0.0 + block_size: int = 32 + lookup_fixed_ms: float = 0.0 + lookup_per_block_ms: float = 0.0 + seed: int = DEFAULT_RANDOM_SEED + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> PrefixCacheConfig: + item = cls( + enabled=bool(raw.get("enabled", False)), + request_hit_rate=float(raw.get("request_hit_rate", 0.0)), + matched_prefix_ratio=float(raw.get("matched_prefix_ratio", 0.0)), + block_size=int(raw.get("block_size", 32)), + lookup_fixed_ms=float(raw.get("lookup_fixed_ms", 0.0)), + lookup_per_block_ms=float(raw.get("lookup_per_block_ms", 0.0)), + seed=int(raw.get("seed", DEFAULT_RANDOM_SEED)), + ) + for name, value in ( + ("request_hit_rate", item.request_hit_rate), + ("matched_prefix_ratio", item.matched_prefix_ratio), + ): + if not 0 <= value <= 1: + raise ValueError(f"prefix_cache.{name} must be in [0, 1]") + _positive(item.block_size, "prefix_cache.block_size") + if item.lookup_fixed_ms < 0 or item.lookup_per_block_ms < 0: + raise ValueError("prefix cache lookup latency cannot be negative") + return item + + def lookup_latency_ms(self, cached_tokens: int) -> float: + if not self.enabled: + return 0.0 + blocks = cached_tokens // self.block_size + return self.lookup_fixed_ms + blocks * self.lookup_per_block_ms + + +@dataclass(frozen=True) +class CamLegConfig: + fixed_ms: float + per_token_ms: float + + @classmethod + def from_mapping( + cls, + raw: dict[str, Any], + *, + default_fixed_ms: float, + default_per_token_ms: float, + ) -> CamLegConfig: + item = cls( + fixed_ms=float(raw.get("fixed_ms", default_fixed_ms)), + per_token_ms=float(raw.get("per_token_ms", default_per_token_ms)), + ) + if item.fixed_ms < 0 or item.per_token_ms < 0: + raise ValueError("CAM latency parameters cannot be negative") + return item + + def latency_ms(self, tokens: int) -> float: + return self.fixed_ms + self.per_token_ms * tokens + + +@dataclass(frozen=True) +class CamConfig: + calibrated: bool = False + dispatch_send: CamLegConfig = field( + default_factory=lambda: CamLegConfig(0.11, 1.0 / 52_000) + ) + dispatch_recv: CamLegConfig = field( + default_factory=lambda: CamLegConfig(0.10, 1.0 / 68_000) + ) + combine_send: CamLegConfig = field( + default_factory=lambda: CamLegConfig(0.10, 1.0 / 70_000) + ) + combine_recv: CamLegConfig = field( + default_factory=lambda: CamLegConfig(0.12, 1.0 / 58_000) + ) + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> CamConfig: + return cls( + calibrated=bool(raw.get("calibrated", False)), + dispatch_send=CamLegConfig.from_mapping( + _mapping(raw.get("dispatch_send"), "cam.dispatch_send"), + default_fixed_ms=0.11, + default_per_token_ms=1.0 / 52_000, + ), + dispatch_recv=CamLegConfig.from_mapping( + _mapping(raw.get("dispatch_recv"), "cam.dispatch_recv"), + default_fixed_ms=0.10, + default_per_token_ms=1.0 / 68_000, + ), + combine_send=CamLegConfig.from_mapping( + _mapping(raw.get("combine_send"), "cam.combine_send"), + default_fixed_ms=0.10, + default_per_token_ms=1.0 / 70_000, + ), + combine_recv=CamLegConfig.from_mapping( + _mapping(raw.get("combine_recv"), "cam.combine_recv"), + default_fixed_ms=0.12, + default_per_token_ms=1.0 / 58_000, + ), + ) + + +@dataclass(frozen=True) +class AfdConfig: + ubatch_split: str = "request" + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> AfdConfig: + item = cls(ubatch_split=str(raw.get("ubatch_split", "request"))) + if item.ubatch_split not in {"request", "token"}: + raise ValueError("afd.ubatch_split must be request or token") + return item + + +@dataclass(frozen=True) +class SloConfig: + ttft_limit_ms: float = DEFAULT_SLO_LIMIT_MS + target_ratio: float = DEFAULT_SLO_TARGET_RATIO + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> SloConfig: + item = cls( + ttft_limit_ms=float(raw.get("ttft_limit_ms", DEFAULT_SLO_LIMIT_MS)), + target_ratio=float(raw.get("target_ratio", DEFAULT_SLO_TARGET_RATIO)), + ) + _positive(item.ttft_limit_ms, "slo.ttft_limit_ms") + if not 0 < item.target_ratio <= 1: + raise ValueError("slo.target_ratio must be in (0, 1]") + return item + + +@dataclass(frozen=True) +class SweepConfig: + min_qps: float = 0.5 + max_qps: float = 64.0 + coarse_points: int = 10 + refinement_steps: int = 7 + throughput_tolerance_ratio: float = 0.99 + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> SweepConfig: + item = cls( + min_qps=float(raw.get("min_qps", 0.5)), + max_qps=float(raw.get("max_qps", 64.0)), + coarse_points=int(raw.get("coarse_points", 10)), + refinement_steps=int(raw.get("refinement_steps", 7)), + throughput_tolerance_ratio=float( + raw.get("throughput_tolerance_ratio", 0.99) + ), + ) + _positive(item.min_qps, "sweep.min_qps") + if item.max_qps <= item.min_qps: + raise ValueError("sweep.max_qps must exceed min_qps") + if item.coarse_points < 2 or item.refinement_steps < 0: + raise ValueError("invalid sweep point counts") + if not 0 < item.throughput_tolerance_ratio <= 1: + raise ValueError("sweep.throughput_tolerance_ratio must be in (0, 1]") + return item + + +@dataclass(frozen=True) +class OutputConfig: + include_timeline: bool = True + timeline_max_events: int = DEFAULT_TIMELINE_MAX_EVENTS + include_requests: bool = True + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> OutputConfig: + item = cls( + include_timeline=bool(raw.get("include_timeline", True)), + timeline_max_events=int( + raw.get("timeline_max_events", DEFAULT_TIMELINE_MAX_EVENTS) + ), + include_requests=bool(raw.get("include_requests", True)), + ) + _positive(item.timeline_max_events, "output.timeline_max_events") + return item + + +@dataclass(frozen=True) +class SimulationConfig: + mode: str = "fixed" + fixed_lengths: tuple[int, ...] = (512, 8_192, 2_048, 6_144) + length_mix: tuple[LengthBucket, ...] = ( + LengthBucket(512, 0.5), + LengthBucket(8_192, 0.5), + ) + csv_path: str | None = None + csv_text: str | None = None + csv_sampling: str = "cycle" + arrival: ArrivalConfig = field(default_factory=ArrivalConfig) + scheduler: SchedulerConfig = field(default_factory=SchedulerConfig) + prefix_cache: PrefixCacheConfig = field(default_factory=PrefixCacheConfig) + afd: AfdConfig = field(default_factory=AfdConfig) + cam: CamConfig = field(default_factory=CamConfig) + slo: SloConfig = field(default_factory=SloConfig) + sweep: SweepConfig = field(default_factory=SweepConfig) + output: OutputConfig = field(default_factory=OutputConfig) + fixed_ttft_overhead_ms: float = 0.0 + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> SimulationConfig: + fixed_lengths = tuple(int(value) for value in raw.get("fixed_lengths", [])) + length_mix = tuple( + LengthBucket.from_mapping(_mapping(value, "length_mix item")) + for value in raw.get("length_mix", []) + ) + item = cls( + mode=str(raw.get("mode", "fixed")), + fixed_lengths=( + fixed_lengths or cls.__dataclass_fields__["fixed_lengths"].default + ), + length_mix=length_mix or cls.__dataclass_fields__["length_mix"].default, + csv_path=raw.get("csv_path"), + csv_text=raw.get("csv_text"), + csv_sampling=str(raw.get("csv_sampling", "cycle")), + arrival=ArrivalConfig.from_mapping(_mapping(raw.get("arrival"), "arrival")), + scheduler=SchedulerConfig.from_mapping( + _mapping(raw.get("scheduler"), "scheduler") + ), + prefix_cache=PrefixCacheConfig.from_mapping( + _mapping(raw.get("prefix_cache"), "prefix_cache") + ), + afd=AfdConfig.from_mapping(_mapping(raw.get("afd"), "afd")), + cam=CamConfig.from_mapping(_mapping(raw.get("cam"), "cam")), + slo=SloConfig.from_mapping(_mapping(raw.get("slo"), "slo")), + sweep=SweepConfig.from_mapping(_mapping(raw.get("sweep"), "sweep")), + output=OutputConfig.from_mapping(_mapping(raw.get("output"), "output")), + fixed_ttft_overhead_ms=float(raw.get("fixed_ttft_overhead_ms", 0.0)), + ) + item.validate() + return item + + def validate(self) -> None: + if self.mode not in {"fixed", "continuous"}: + raise ValueError("mode must be fixed or continuous") + if self.mode == "fixed" and self.arrival.kind == "trace": + raise ValueError("arrival.kind=trace requires mode='continuous'") + if self.csv_sampling not in {"cycle", "sample"}: + raise ValueError("csv_sampling must be cycle or sample") + if self.fixed_ttft_overhead_ms < 0: + raise ValueError("fixed_ttft_overhead_ms cannot be negative") + for tokens in self.fixed_lengths: + _positive(tokens, "fixed_lengths item") + if ( + self.mode == "fixed" + and not self.scheduler.chunked_prefill + and any( + tokens > self.scheduler.max_num_batched_tokens + for tokens in self.fixed_lengths + ) + and not (self.csv_path or self.csv_text) + ): + raise ValueError( + "non-chunked fixed request exceeds scheduler.max_num_batched_tokens" + ) + + def to_mapping(self) -> dict[str, Any]: + return asdict(self) + + +def default_config_mapping() -> dict[str, Any]: + return SimulationConfig().to_mapping() + + +__all__ = [ + "AfdConfig", + "ArrivalConfig", + "CamConfig", + "LengthBucket", + "PrefixCacheConfig", + "SchedulerConfig", + "SimulationConfig", + "SloConfig", + "SweepConfig", + "default_config_mapping", +] diff --git a/simulator/engine.py b/simulator/engine.py new file mode 100644 index 00000000..d2614912 --- /dev/null +++ b/simulator/engine.py @@ -0,0 +1,1075 @@ +"""Discrete-event simulation for merged and CAMAsync AFD Prefill.""" + +from __future__ import annotations + +import math +from collections import deque +from collections.abc import Iterable +from dataclasses import dataclass, replace +from typing import Any + +from simulator.config import OutputConfig, SimulationConfig +from simulator.profiles import ProfileBundle +from simulator.workload import RequestSpec, RuntimeRequest, generate_workload + +AFD_DP_COUNT = 2 +MERGED_DP_COUNT = 4 + + +@dataclass(frozen=True) +class BatchSegment: + request: RuntimeRequest + prefix_tokens: int + query_tokens: int + + +@dataclass(frozen=True) +class SchedulerBatch: + batch_id: int + dp: int + segments: tuple[BatchSegment, ...] + + @property + def query_tokens(self) -> int: + return sum(segment.query_tokens for segment in self.segments) + + +@dataclass(frozen=True) +class AfdStage: + segments: tuple[BatchSegment, ...] + + @property + def query_tokens(self) -> int: + return sum(segment.query_tokens for segment in self.segments) + + +@dataclass +class AfdDpState: + dp: int + batch: SchedulerBatch + stages: tuple[AfdStage, ...] + ops: list[tuple[str, int, int]] + time_ms: float + op_index: int = 0 + + +@dataclass(frozen=True) +class FfnJob: + end_ms: float + + +@dataclass(frozen=True) +class PendingFfnJob: + key: str + arrival_ms: float + layer: int + batch: int + stage: int + tokens: int + + +class RequestDispatcher: + """Route requests to DP queues when they arrive.""" + + def __init__( + self, + specs: tuple[RequestSpec, ...], + dp_count: int, + policy: str, + ) -> None: + self.dp_count = dp_count + self.policy = policy + self.requests = [RuntimeRequest(spec=spec, assigned_dp=-1) for spec in specs] + self.pending = sorted( + enumerate(self.requests), + key=lambda item: (item[1].spec.arrival_ms, item[0]), + ) + self.pending_index = 0 + self.round_robin_index = 0 + self.tie_start_index = 0 + self.queues: list[list[RuntimeRequest]] = [[] for _ in range(dp_count)] + self.active: list[deque[RuntimeRequest]] = [deque() for _ in range(dp_count)] + + @property + def has_pending(self) -> bool: + return self.pending_index < len(self.pending) + + @property + def next_arrival_ms(self) -> float: + if not self.has_pending: + return math.inf + return self.pending[self.pending_index][1].spec.arrival_ms + + def dispatch_until(self, time_ms: float) -> None: + while self.next_arrival_ms <= time_ms: + request = self.pending[self.pending_index][1] + self.pending_index += 1 + dp = self._select_dp(request.spec.arrival_ms) + request.assigned_dp = dp + self.queues[dp].append(request) + if self.policy == "vllm_queue_aware": + self.active[dp].append(request) + + def _select_dp(self, arrival_ms: float) -> int: + if self.policy == "round_robin": + dp = self.round_robin_index + self.round_robin_index = (self.round_robin_index + 1) % self.dp_count + return dp + + loads = [] + for dp in range(self.dp_count): + active = self.active[dp] + while ( + active + and active[0].completion_ms is not None + and active[0].completion_ms <= arrival_ms + ): + active.popleft() + loads.append(len(active)) + min_load = min(loads) + for offset in range(self.dp_count): + dp = (self.tie_start_index + offset) % self.dp_count + if loads[dp] == min_load: + self.tie_start_index = (self.tie_start_index + 1) % self.dp_count + return dp + raise RuntimeError("DP dispatcher failed to select an engine") + + +def compare_architectures( + config: SimulationConfig, + profiles: ProfileBundle, + requests: tuple[RequestSpec, ...] | None = None, +) -> dict[str, Any]: + """Run one identical workload through AFD and merged architectures.""" + + workload = requests if requests is not None else generate_workload(config) + if not workload: + raise ValueError("workload contains no requests") + _validate_workload(config, workload) + afd = simulate_afd(config, profiles, workload) + merged = simulate_merged(config, profiles, workload) + return { + "config": config.to_mapping(), + "profile": profiles.summary(), + "workload": { + "request_count": len(workload), + "input_tokens": sum(item.input_tokens for item in workload), + "cached_tokens": sum(item.cached_prefix_tokens for item in workload), + }, + "afd": afd, + "merged": merged, + "comparison": _comparison(afd["summary"], merged["summary"]), + } + + +def simulate_merged( + config: SimulationConfig, + profiles: ProfileBundle, + specs: tuple[RequestSpec, ...], +) -> dict[str, Any]: + dispatcher = RequestDispatcher(specs, MERGED_DP_COUNT, config.scheduler.policy) + requests = dispatcher.requests + queues = dispatcher.queues + timeline: list[dict[str, Any]] = [] + wave_start = 0.0 + batch_id = 0 + total_barrier_ms = 0.0 + busy_ms = [0.0 for _ in range(MERGED_DP_COUNT)] + ep_busy_ms = 0.0 + + while dispatcher.has_pending or any(queues): + dispatcher.dispatch_until(wave_start) + next_ready = min( + (_request_ready_ms(queue[0]) for queue in queues if queue), + default=math.inf, + ) + if dispatcher.next_arrival_ms < next_ready: + wave_start = max(wave_start, dispatcher.next_arrival_ms) + dispatcher.dispatch_until(wave_start) + continue + wave_start = max(wave_start, next_ready) + dispatcher.dispatch_until(wave_start) + batches = [] + for dp, queue in enumerate(queues): + batch = _pack_batch( + config, + queue, + dp=dp, + batch_id=batch_id, + start_ms=wave_start, + ) + batches.append(batch) + if batch is not None: + batch_id += 1 + if not any(batches): + wave_start = next_ready + continue + + layer_start = wave_start + for layer_idx in range(profiles.layer_count): + attention_ends: list[float | None] = [] + for dp, batch in enumerate(batches): + if batch is None: + attention_ends.append(None) + continue + duration = _segments_duration( + profiles, + "merged", + layer_idx, + "attention_router", + batch.segments, + ) + end = layer_start + duration + attention_ends.append(end) + busy_ms[dp] += duration + _append_timeline( + timeline, + config, + architecture="merged", + resource=f"DP{dp} Attention", + phase="attention_router", + start_ms=layer_start, + end_ms=end, + layer=layer_idx, + batch=batch.batch_id, + tokens=batch.query_tokens, + ) + active_ends = [value for value in attention_ends if value is not None] + barrier_end = max(active_ends, default=layer_start) + for dp, end in enumerate(attention_ends): + wait_start = layer_start if end is None else end + if wait_start < barrier_end: + total_barrier_ms += barrier_end - wait_start + _append_timeline( + timeline, + config, + architecture="merged", + resource=f"DP{dp} Attention", + phase="barrier", + start_ms=wait_start, + end_ms=barrier_end, + layer=layer_idx, + batch=None, + tokens=0, + ) + + global_ep_tokens = sum( + batch.query_tokens for batch in batches if batch is not None + ) + global_profile_query = max(1.0, global_ep_tokens / MERGED_DP_COUNT) + max_dp_tokens = max( + batch.query_tokens for batch in batches if batch is not None + ) + cursor = barrier_end + for phase in ( + "merged_dispatch", + "routed_experts", + "merged_combine", + "merged_combine_local", + "shared_expert", + "merged_sp_post", + ): + is_global_phase = phase in { + "merged_dispatch", + "routed_experts", + "merged_combine", + } + phase_tokens = global_ep_tokens if is_global_phase else max_dp_tokens + profile_query_tokens = ( + global_profile_query if is_global_phase else max_dp_tokens + ) + duration = profiles.duration_ms( + "merged", layer_idx, phase, 0, profile_query_tokens + ) + end = cursor + duration + ep_busy_ms += duration + _append_timeline( + timeline, + config, + architecture="merged", + resource="Global EP16", + phase=phase, + start_ms=cursor, + end_ms=end, + layer=layer_idx, + batch=None, + tokens=phase_tokens, + profile_query_tokens=profile_query_tokens, + expert_details=_expert_event_details( + profiles, + "merged", + layer_idx, + phase, + profile_query_tokens, + ), + ) + cursor = end + layer_start = cursor + + wave_start = layer_start + for batch in batches: + if batch is not None: + _complete_batch(batch, wave_start) + + makespan = max(request.completion_ms or 0.0 for request in requests) + utilization = { + f"attention_dp{dp}": value / makespan if makespan else 0.0 + for dp, value in enumerate(busy_ms) + } + utilization["global_ep16"] = ep_busy_ms / makespan if makespan else 0.0 + result = _build_result( + "merged", + config, + requests, + timeline, + makespan, + utilization, + ) + result["summary"]["barrier_wait_ms"] = total_barrier_ms + return result + + +def simulate_afd( + config: SimulationConfig, + profiles: ProfileBundle, + specs: tuple[RequestSpec, ...], +) -> dict[str, Any]: + dispatcher = RequestDispatcher(specs, AFD_DP_COUNT, config.scheduler.policy) + dispatcher.dispatch_until(0.0) + requests = dispatcher.requests + queues = dispatcher.queues + timeline: list[dict[str, Any]] = [] + states: dict[int, AfdDpState] = {} + next_batch_id = 0 + jobs: dict[str, FfnJob] = {} + pending_jobs: list[PendingFfnJob] = [] + ffn_available_ms = 0.0 + ffn_busy_ms = 0.0 + attention_busy_ms = [0.0 for _ in range(AFD_DP_COUNT)] + attention_wait_ms = 0.0 + dp_available_ms = [0.0 for _ in range(AFD_DP_COUNT)] + + while states or dispatcher.has_pending or any(queues): + candidates = [] + for state in states.values(): + op_type, stage_idx, layer_idx = state.ops[state.op_index] + ready_ms = state.time_ms + if op_type == "combine": + job = jobs.get(_job_key(state.batch.batch_id, stage_idx, layer_idx)) + if job is None: + continue + ready_ms = max(ready_ms, job.end_ms) + candidates.append((ready_ms, state.dp, op_type, stage_idx, layer_idx)) + + next_state_op = min(candidates) if candidates else None + next_ffn_job = min( + pending_jobs, + key=lambda item: (item.arrival_ms, item.batch, item.stage, item.layer), + default=None, + ) + ffn_ready_ms = ( + max(ffn_available_ms, next_ffn_job.arrival_ms) + if next_ffn_job is not None + else math.inf + ) + idle_ready = [ + ( + max(dp_available_ms[dp], _request_ready_ms(queues[dp][0])), + dp, + ) + for dp in range(AFD_DP_COUNT) + if dp not in states and queues[dp] + ] + next_idle = min(idle_ready, default=(math.inf, -1)) + next_state_ms = next_state_op[0] if next_state_op is not None else math.inf + next_operation_ms = min(next_state_ms, ffn_ready_ms, next_idle[0]) + if dispatcher.next_arrival_ms <= next_operation_ms: + dispatcher.dispatch_until(dispatcher.next_arrival_ms) + continue + if next_idle[0] <= min(next_state_ms, ffn_ready_ms): + start_ms, idle_dp = next_idle + state, next_batch_id = _load_afd_state( + config, + queues[idle_dp], + idle_dp, + current_time_ms=start_ms, + next_batch_id=next_batch_id, + layer_count=profiles.layer_count, + ) + if state is None: + raise RuntimeError("AFD scheduler failed to load a ready batch") + states[idle_dp] = state + continue + if next_ffn_job is not None and ( + next_state_op is None or next_state_op[0] > ffn_ready_ms + ): + pending_jobs.remove(next_ffn_job) + recv_start = ffn_ready_ms + recv_end = recv_start + config.cam.dispatch_recv.latency_ms( + next_ffn_job.tokens + ) + compute_duration = sum( + profiles.duration_ms( + "afd", + next_ffn_job.layer, + phase, + 0, + next_ffn_job.tokens, + ) + for phase in ("routed_experts", "shared_expert") + ) + compute_end = recv_end + compute_duration + combine_send_end = compute_end + config.cam.combine_send.latency_ms( + next_ffn_job.tokens + ) + for phase, start, end in ( + ("dispatch_recv", recv_start, recv_end), + ("ffn_compute", recv_end, compute_end), + ("combine_send", compute_end, combine_send_end), + ): + _append_timeline( + timeline, + config, + architecture="afd", + resource="FFN EP8", + phase=phase, + start_ms=start, + end_ms=end, + layer=next_ffn_job.layer, + batch=next_ffn_job.batch, + stage=next_ffn_job.stage, + tokens=next_ffn_job.tokens, + profile_query_tokens=( + next_ffn_job.tokens if phase == "ffn_compute" else None + ), + expert_details=_expert_event_details( + profiles, + "afd", + next_ffn_job.layer, + phase, + next_ffn_job.tokens, + ), + ) + ffn_busy_ms += combine_send_end - recv_start + ffn_available_ms = combine_send_end + jobs[next_ffn_job.key] = FfnJob(end_ms=combine_send_end) + continue + + if next_state_op is None: + raise RuntimeError("AFD event loop has no runnable operation") + _, dp, op_type, stage_idx, layer_idx = next_state_op + state = states[dp] + stage = state.stages[stage_idx] + + if op_type == "attention": + start = state.time_ms + duration = _segments_duration( + profiles, + "afd", + layer_idx, + "attention_router", + stage.segments, + ) + end = start + duration + attention_busy_ms[dp] += duration + _append_timeline( + timeline, + config, + architecture="afd", + resource=f"DP{dp} Attention", + phase="attention_router", + start_ms=start, + end_ms=end, + layer=layer_idx, + batch=state.batch.batch_id, + stage=stage_idx, + tokens=stage.query_tokens, + ) + state.time_ms = end + elif op_type == "dispatch": + send_start = state.time_ms + send_end = send_start + config.cam.dispatch_send.latency_ms( + stage.query_tokens + ) + _append_timeline( + timeline, + config, + architecture="afd", + resource=f"DP{dp} CAM", + phase="dispatch_send", + start_ms=send_start, + end_ms=send_end, + layer=layer_idx, + batch=state.batch.batch_id, + stage=stage_idx, + tokens=stage.query_tokens, + ) + key = _job_key(state.batch.batch_id, stage_idx, layer_idx) + pending_jobs.append( + PendingFfnJob( + key=key, + arrival_ms=send_end, + layer=layer_idx, + batch=state.batch.batch_id, + stage=stage_idx, + tokens=stage.query_tokens, + ) + ) + state.time_ms = send_end + else: + job = jobs.pop(_job_key(state.batch.batch_id, stage_idx, layer_idx)) + if job.end_ms > state.time_ms: + _append_timeline( + timeline, + config, + architecture="afd", + resource=f"DP{dp} Attention", + phase="wait_ffn", + start_ms=state.time_ms, + end_ms=job.end_ms, + layer=layer_idx, + batch=state.batch.batch_id, + stage=stage_idx, + tokens=stage.query_tokens, + ) + attention_wait_ms += job.end_ms - state.time_ms + state.time_ms = job.end_ms + recv_start = state.time_ms + post_duration = profiles.duration_ms( + "afd", layer_idx, "afd_post", 0, stage.query_tokens + ) + recv_end = ( + recv_start + + config.cam.combine_recv.latency_ms(stage.query_tokens) + + post_duration + ) + _append_timeline( + timeline, + config, + architecture="afd", + resource=f"DP{dp} CAM", + phase="combine_recv_post", + start_ms=recv_start, + end_ms=recv_end, + layer=layer_idx, + batch=state.batch.batch_id, + stage=stage_idx, + tokens=stage.query_tokens, + ) + state.time_ms = recv_end + + state.op_index += 1 + if state.op_index == len(state.ops): + _complete_batch(state.batch, state.time_ms) + dp_available_ms[dp] = state.time_ms + del states[dp] + + makespan = max(request.completion_ms or 0.0 for request in requests) + utilization = { + f"attention_dp{dp}": value / makespan if makespan else 0.0 + for dp, value in enumerate(attention_busy_ms) + } + utilization["ffn_ep8"] = ffn_busy_ms / makespan if makespan else 0.0 + result = _build_result("afd", config, requests, timeline, makespan, utilization) + result["summary"]["attention_wait_ms"] = attention_wait_ms + result["summary"]["cam_calibrated"] = config.cam.calibrated + return result + + +def sweep_qps(config: SimulationConfig, profiles: ProfileBundle) -> dict[str, Any]: + """Find the highest SLO-feasible QPS with coarse search and refinement.""" + + if config.mode != "continuous": + raise ValueError("QPS sweep requires mode='continuous'") + if config.arrival.kind == "trace": + raise ValueError("QPS sweep is unavailable for exact timestamp trace replay") + sweep = config.sweep + ratio = sweep.max_qps / sweep.min_qps + coarse = [ + sweep.min_qps * ratio ** (index / (sweep.coarse_points - 1)) + for index in range(sweep.coarse_points) + ] + cache: dict[float, dict[str, Any]] = {} + compact_output = OutputConfig( + include_timeline=False, + timeline_max_events=config.output.timeline_max_events, + include_requests=False, + ) + + def evaluate(qps: float) -> dict[str, Any]: + key = round(qps, 9) + if key not in cache: + scenario = replace( + config, + arrival=replace(config.arrival, qps=qps), + output=compact_output, + ) + cache[key] = compare_architectures(scenario, profiles) + return cache[key] + + for qps in coarse: + evaluate(qps) + + capacities = {} + for architecture in ("afd", "merged"): + passed = [ + qps for qps in coarse if _sweep_pass(evaluate(qps), architecture, config) + ] + failed = [ + qps + for qps in coarse + if not _sweep_pass(evaluate(qps), architecture, config) + ] + lower = max(passed, default=0.0) + upper_candidates = [qps for qps in failed if qps > lower] + upper = min(upper_candidates, default=sweep.max_qps) + if lower > 0 and upper > lower: + for _ in range(sweep.refinement_steps): + midpoint = (lower + upper) / 2.0 + if _sweep_pass(evaluate(midpoint), architecture, config): + lower = midpoint + else: + upper = midpoint + capacities[architecture] = lower + + series = {"afd": [], "merged": []} + for qps in sorted(cache): + result = cache[qps] + for architecture in series: + summary = result[architecture]["summary"] + series[architecture].append( + { + "qps": qps, + "throughput_rps": summary["throughput_rps"], + "ttft_p50_ms": summary["ttft_p50_ms"], + "ttft_p90_ms": summary["ttft_p90_ms"], + "ttft_p99_ms": summary["ttft_p99_ms"], + "slo_attainment": summary["slo_attainment"], + "passed": _sweep_pass(result, architecture, config), + } + ) + return { + "capacity_qps": capacities, + "series": series, + "slo": { + "ttft_limit_ms": config.slo.ttft_limit_ms, + "target_ratio": config.slo.target_ratio, + "throughput_tolerance_ratio": sweep.throughput_tolerance_ratio, + }, + } + + +def _sweep_pass( + result: dict[str, Any], architecture: str, config: SimulationConfig +) -> bool: + summary = result[architecture]["summary"] + offered = summary["offered_qps"] + return ( + summary["slo_attainment"] >= config.slo.target_ratio + and summary["throughput_rps"] + >= offered * config.sweep.throughput_tolerance_ratio + ) + + +def _request_ready_ms(request: RuntimeRequest) -> float: + return request.spec.arrival_ms + request.spec.cache_lookup_ms + + +def _pack_batch( + config: SimulationConfig, + queue: list[RuntimeRequest], + *, + dp: int, + batch_id: int, + start_ms: float, +) -> SchedulerBatch | None: + budget = config.scheduler.max_num_batched_tokens + segments = [] + while queue and len(segments) < config.scheduler.max_num_seqs and budget > 0: + request = queue[0] + if _request_ready_ms(request) > start_ms: + break + remaining = request.remaining_query_tokens + if not config.scheduler.chunked_prefill: + if remaining > budget: + if not segments: + raise ValueError( + f"request {request.spec.request_id} has {remaining} " + "uncached tokens, " + "exceeding non-chunked max_num_batched_tokens" + ) + break + query_tokens = remaining + else: + query_tokens = min(remaining, config.scheduler.chunk_size, budget) + prefix_tokens = request.current_prefix_tokens + request.computed_query_tokens += query_tokens + if request.first_scheduled_ms is None: + request.first_scheduled_ms = start_ms + segments.append( + BatchSegment( + request=request, + prefix_tokens=prefix_tokens, + query_tokens=query_tokens, + ) + ) + budget -= query_tokens + if request.remaining_query_tokens == 0: + queue.pop(0) + else: + break + if not segments: + return None + return SchedulerBatch( + batch_id=batch_id, + dp=dp, + segments=tuple(segments), + ) + + +def _split_afd_stages(batch: SchedulerBatch, split: str) -> tuple[AfdStage, ...]: + segments = batch.segments + if split == "request": + if len(segments) < 2: + return (AfdStage(segments),) + cumulative = 0 + total = batch.query_tokens + best_index = 1 + best_distance = math.inf + for index in range(1, len(segments)): + cumulative += segments[index - 1].query_tokens + distance = abs(cumulative * 2 - total) + if distance < best_distance: + best_distance = distance + best_index = index + return ( + AfdStage(segments[:best_index]), + AfdStage(segments[best_index:]), + ) + + if batch.query_tokens < 2: + return (AfdStage(segments),) + split_token = (batch.query_tokens + 1) // 2 + stage_segments: list[list[BatchSegment]] = [[], []] + cursor = 0 + for segment in segments: + segment_start = cursor + segment_end = cursor + segment.query_tokens + for stage_idx, (start, end) in enumerate( + ((0, split_token), (split_token, batch.query_tokens)) + ): + overlap_start = max(segment_start, start) + overlap_end = min(segment_end, end) + if overlap_end > overlap_start: + offset = overlap_start - segment_start + stage_segments[stage_idx].append( + BatchSegment( + request=segment.request, + prefix_tokens=segment.prefix_tokens + offset, + query_tokens=overlap_end - overlap_start, + ) + ) + cursor = segment_end + return tuple(AfdStage(tuple(items)) for items in stage_segments if items) + + +def _build_afd_ops(stage_count: int, layer_count: int) -> list[tuple[str, int, int]]: + if stage_count == 1: + return [ + operation + for layer in range(layer_count) + for operation in ( + ("attention", 0, layer), + ("dispatch", 0, layer), + ("combine", 0, layer), + ) + ] + ops = [("attention", 0, 0), ("dispatch", 0, 0)] + for layer in range(layer_count - 1): + ops.extend( + ( + ("attention", 1, layer), + ("combine", 0, layer), + ("dispatch", 1, layer), + ("attention", 0, layer + 1), + ("combine", 1, layer), + ("dispatch", 0, layer + 1), + ) + ) + last = layer_count - 1 + ops.extend( + ( + ("attention", 1, last), + ("combine", 0, last), + ("dispatch", 1, last), + ("combine", 1, last), + ) + ) + return ops + + +def _load_afd_state( + config: SimulationConfig, + queue: list[RuntimeRequest], + dp: int, + *, + current_time_ms: float, + next_batch_id: int, + layer_count: int, +) -> tuple[AfdDpState | None, int]: + if not queue: + return None, next_batch_id + start_ms = max(current_time_ms, _request_ready_ms(queue[0])) + batch = _pack_batch( + config, + queue, + dp=dp, + batch_id=next_batch_id, + start_ms=start_ms, + ) + if batch is None: + return None, next_batch_id + stages = _split_afd_stages(batch, config.afd.ubatch_split) + return ( + AfdDpState( + dp=dp, + batch=batch, + stages=stages, + ops=_build_afd_ops(len(stages), layer_count), + time_ms=start_ms, + ), + next_batch_id + 1, + ) + + +def _segments_duration( + profiles: ProfileBundle, + topology: str, + layer_idx: int, + phase: str, + segments: Iterable[BatchSegment], +) -> float: + return sum( + profiles.duration_ms( + topology, + layer_idx, + phase, + segment.prefix_tokens, + segment.query_tokens, + ) + for segment in segments + ) + + +def _complete_batch(batch: SchedulerBatch, completion_ms: float) -> None: + for segment in batch.segments: + request = segment.request + if request.remaining_query_tokens == 0: + request.completion_ms = completion_ms + + +def _job_key(batch_id: int, stage_idx: int, layer_idx: int) -> str: + return f"b{batch_id}-u{stage_idx}-l{layer_idx}" + + +def _expert_event_details( + profiles: ProfileBundle, + topology: str, + layer_idx: int, + phase: str, + tokens: int | float, +) -> dict[str, Any] | None: + if phase not in {"ffn_compute", "routed_experts", "shared_expert"}: + return None + model_config = profiles.metadata.get("model_config") + topology_spec = profiles.topology_specs[topology] + if not model_config or not topology_spec: + return None + expert_spec = topology_spec["ffn"] if topology == "afd" else topology_spec + return { + "top_k": (None if phase == "shared_expert" else int(model_config["moe_top_k"])), + "ep_size": int(expert_spec["ep_size"]), + "sampled_input_shapes": profiles.expert_shape_samples( + topology, layer_idx, phase, tokens + ), + } + + +def _append_timeline( + timeline: list[dict[str, Any]], + config: SimulationConfig, + *, + architecture: str, + resource: str, + phase: str, + start_ms: float, + end_ms: float, + layer: int, + batch: int | None, + tokens: int, + stage: int | None = None, + profile_query_tokens: int | float | None = None, + expert_details: dict[str, Any] | None = None, +) -> None: + if not config.output.include_timeline: + return + if len(timeline) > config.output.timeline_max_events: + return + event = { + "architecture": architecture, + "resource": resource, + "phase": phase, + "start_ms": start_ms, + "end_ms": end_ms, + "layer": layer, + "batch": batch, + "stage": stage, + "tokens": tokens, + } + if profile_query_tokens is not None: + event["profile_query_tokens"] = profile_query_tokens + if expert_details: + event.update(expert_details) + timeline.append(event) + + +def _build_result( + architecture: str, + config: SimulationConfig, + requests: list[RuntimeRequest], + timeline: list[dict[str, Any]], + makespan_ms: float, + utilization: dict[str, float], +) -> dict[str, Any]: + measurement_start = ( + config.arrival.warmup_s * 1_000.0 if config.mode == "continuous" else 0.0 + ) + measured = [ + request for request in requests if request.spec.arrival_ms >= measurement_start + ] + ttfts = [ + (request.completion_ms or 0.0) + - request.spec.arrival_ms + + config.fixed_ttft_overhead_ms + for request in measured + ] + if config.mode == "continuous": + nominal_end = measurement_start + config.arrival.duration_s * 1_000.0 + actual_end = max( + nominal_end, + max((request.completion_ms or 0.0) for request in measured), + ) + elapsed_s = max((actual_end - measurement_start) / 1_000.0, 1e-9) + offered_qps = len(measured) / config.arrival.duration_s + else: + first_arrival = min(request.spec.arrival_ms for request in measured) + elapsed_s = max((makespan_ms - first_arrival) / 1_000.0, 1e-9) + offered_qps = len(measured) / elapsed_s + throughput = len(measured) / elapsed_s + input_tokens = sum(request.spec.input_tokens for request in measured) + cached_tokens = sum(request.spec.cached_prefix_tokens for request in measured) + compute_tokens = input_tokens - cached_tokens + slo_hits = sum(value <= config.slo.ttft_limit_ms for value in ttfts) + summary = { + "architecture": architecture, + "request_count": len(measured), + "makespan_ms": makespan_ms, + "offered_qps": offered_qps, + "throughput_rps": throughput, + "input_tokens_per_s": input_tokens / elapsed_s, + "compute_tokens_per_s": compute_tokens / elapsed_s, + "logical_input_tokens": input_tokens, + "computed_query_tokens": compute_tokens, + "cached_prefix_tokens": cached_tokens, + "cache_token_ratio": cached_tokens / input_tokens if input_tokens else 0.0, + "ttft_mean_ms": sum(ttfts) / len(ttfts), + "ttft_p50_ms": _percentile(ttfts, 0.50), + "ttft_p90_ms": _percentile(ttfts, 0.90), + "ttft_p99_ms": _percentile(ttfts, 0.99), + "slo_attainment": slo_hits / len(ttfts), + "slo_goodput_rps": slo_hits / elapsed_s, + "utilization": utilization, + "timeline_truncated": len(timeline) > config.output.timeline_max_events, + } + request_payload = [] + if config.output.include_requests: + request_payload = [ + { + "request_id": request.spec.request_id, + "arrival_ms": request.spec.arrival_ms, + "input_tokens": request.spec.input_tokens, + "cached_prefix_tokens": request.spec.cached_prefix_tokens, + "query_tokens": request.spec.query_tokens, + "assigned_dp": request.assigned_dp, + "first_scheduled_ms": request.first_scheduled_ms, + "completion_ms": request.completion_ms, + "ttft_ms": ( + (request.completion_ms or 0.0) + - request.spec.arrival_ms + + config.fixed_ttft_overhead_ms + ), + } + for request in requests + ] + return { + "summary": summary, + "requests": request_payload, + "timeline": timeline[: config.output.timeline_max_events], + } + + +def _percentile(values: list[float], quantile: float) -> float: + ordered = sorted(values) + if len(ordered) == 1: + return ordered[0] + position = (len(ordered) - 1) * quantile + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + ratio = position - lower + return ordered[lower] + (ordered[upper] - ordered[lower]) * ratio + + +def _comparison(afd: dict[str, Any], merged: dict[str, Any]) -> dict[str, float | None]: + return { + "throughput_speedup": _safe_ratio( + afd["throughput_rps"], merged["throughput_rps"] + ), + "p99_ttft_ratio": _safe_ratio(afd["ttft_p99_ms"], merged["ttft_p99_ms"]), + "slo_goodput_speedup": _safe_ratio( + afd["slo_goodput_rps"], merged["slo_goodput_rps"] + ), + } + + +def _safe_ratio(numerator: float, denominator: float) -> float | None: + return numerator / denominator if denominator else None + + +def _validate_workload( + config: SimulationConfig, workload: tuple[RequestSpec, ...] +) -> None: + if config.mode == "continuous" and not any( + request.arrival_ms >= config.arrival.warmup_s * 1_000.0 for request in workload + ): + raise ValueError("measurement window contains no requests") + if not config.scheduler.chunked_prefill: + too_large = [ + request.request_id + for request in workload + if request.query_tokens > config.scheduler.max_num_batched_tokens + ] + if too_large: + raise ValueError( + "non-chunked requests exceed max_num_batched_tokens: " + + ", ".join(too_large[:5]) + ) + + +__all__ = [ + "compare_architectures", + "simulate_afd", + "simulate_merged", + "sweep_qps", +] diff --git a/simulator/examples/continuous-prefix-cache.json b/simulator/examples/continuous-prefix-cache.json new file mode 100644 index 00000000..9b09d1a5 --- /dev/null +++ b/simulator/examples/continuous-prefix-cache.json @@ -0,0 +1,60 @@ +{ + "mode": "continuous", + "length_mix": [ + {"tokens": 512, "weight": 0.4}, + {"tokens": 2048, "weight": 0.2}, + {"tokens": 8192, "weight": 0.3}, + {"tokens": 32768, "weight": 0.1} + ], + "csv_sampling": "cycle", + "arrival": { + "kind": "poisson", + "qps": 8, + "duration_s": 60, + "warmup_s": 10, + "seed": 1024 + }, + "scheduler": { + "policy": "round_robin", + "max_num_seqs": 64, + "max_num_batched_tokens": 8192, + "chunked_prefill": true, + "chunk_size": 8192 + }, + "prefix_cache": { + "enabled": true, + "request_hit_rate": 0.6, + "matched_prefix_ratio": 0.75, + "block_size": 32, + "lookup_fixed_ms": 0, + "lookup_per_block_ms": 0, + "seed": 1024 + }, + "afd": { + "ubatch_split": "request" + }, + "cam": { + "calibrated": false, + "dispatch_send": {"fixed_ms": 0.11, "per_token_ms": 0.000019230769}, + "dispatch_recv": {"fixed_ms": 0.10, "per_token_ms": 0.000014705882}, + "combine_send": {"fixed_ms": 0.10, "per_token_ms": 0.000014285714}, + "combine_recv": {"fixed_ms": 0.12, "per_token_ms": 0.000017241379} + }, + "slo": { + "ttft_limit_ms": 1000, + "target_ratio": 0.99 + }, + "sweep": { + "min_qps": 0.5, + "max_qps": 64, + "coarse_points": 10, + "refinement_steps": 7, + "throughput_tolerance_ratio": 0.99 + }, + "fixed_ttft_overhead_ms": 0, + "output": { + "include_timeline": true, + "timeline_max_events": 20000, + "include_requests": true + } +} diff --git a/simulator/examples/fixed.json b/simulator/examples/fixed.json new file mode 100644 index 00000000..d3954ba9 --- /dev/null +++ b/simulator/examples/fixed.json @@ -0,0 +1,40 @@ +{ + "mode": "fixed", + "fixed_lengths": [512, 8192, 2048, 16384], + "arrival": { + "kind": "constant", + "qps": 4, + "duration_s": 60, + "warmup_s": 0, + "seed": 1024 + }, + "scheduler": { + "policy": "round_robin", + "max_num_seqs": 64, + "max_num_batched_tokens": 32768, + "chunked_prefill": false, + "chunk_size": 8192 + }, + "prefix_cache": { + "enabled": false, + "request_hit_rate": 0, + "matched_prefix_ratio": 0, + "block_size": 32, + "lookup_fixed_ms": 0, + "lookup_per_block_ms": 0, + "seed": 1024 + }, + "afd": { + "ubatch_split": "request" + }, + "slo": { + "ttft_limit_ms": 1000, + "target_ratio": 0.99 + }, + "fixed_ttft_overhead_ms": 0, + "output": { + "include_timeline": true, + "timeline_max_events": 20000, + "include_requests": true + } +} diff --git a/simulator/examples/requests.csv b/simulator/examples/requests.csv new file mode 100644 index 00000000..b8a15b79 --- /dev/null +++ b/simulator/examples/requests.csv @@ -0,0 +1,5 @@ +request_id,arrival_time_ms,input_length,cached_prefix_tokens +r001,0,8192,4096 +r002,17,512,0 +r003,21,32768,24576 +r004,44,2048,0 diff --git a/simulator/profile_builder.py b/simulator/profile_builder.py new file mode 100644 index 00000000..4efeb1d0 --- /dev/null +++ b/simulator/profile_builder.py @@ -0,0 +1,438 @@ +"""Generate normalized DSV4 profiles from msModeling analytic traces.""" + +from __future__ import annotations + +import json +import math +import re +import subprocess +import tempfile +from collections import Counter +from collections.abc import Iterable +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from simulator.profiles import PROFILE_PHASES, PROFILE_SCHEMA_VERSION + +DEFAULT_MODEL_ID = "deepseek-ai/DeepSeek-V4-Flash" +DEFAULT_DEVICE = "ATLAS_800_A3_752T_128G_DIE" +DEFAULT_HIDDEN_SIZE = 4_096 +DEFAULT_MOE_TOP_K = 6 +EXPECTED_LAYER_COUNT = 43 +GATE_OPS = { + "tensor_cast.moe_gating_top_k.default", + "tensor_cast.moe_gating_top_k_hash.default", +} +ALL_TO_ALL_OP = "tensor_cast.all_to_all.default" +ALL_GATHER_OP = "tensor_cast.all_gather.default" +HC_PRE_OP = "tensor_cast.hc_pre_inv_rms.default" +HC_POST_OP = "tensor_cast.hc_post.default" +SHARED_EXPERT_START_OP = "tensor_cast.static_quant_linear.default" +SPLIT_WITH_SIZES_OP = "aten.split_with_sizes.default" +ROUTED_EXPERT_SHAPE = "routed_expert_input_shape" +ROUTED_EXPERT_SAMPLE_SHAPES = "routed_expert_sample_shapes" +SHARED_EXPERT_SHAPE = "shared_expert_input_shape" + + +@dataclass(frozen=True) +class TopologySpec: + name: str + num_devices: int + dp_size: int + tp_size: int + ep_size: int + sequence_parallel: bool = True + + +DEFAULT_TOPOLOGIES = ( + TopologySpec("afd", num_devices=8, dp_size=2, tp_size=4, ep_size=8), + TopologySpec("merged", num_devices=16, dp_size=4, tp_size=4, ep_size=16), +) +DEFAULT_AFD_FFN_TOPOLOGY = TopologySpec( + "afd_ffn", + num_devices=8, + dp_size=8, + tp_size=1, + ep_size=8, + sequence_parallel=False, +) + + +def build_profile_bundle( + *, + msmodeling_root: str | Path, + python_executable: str, + output_path: str | Path, + query_anchors: Iterable[int], + prefix_anchors: Iterable[int], + model_id: str = DEFAULT_MODEL_ID, + device: str = DEFAULT_DEVICE, + hidden_size: int = DEFAULT_HIDDEN_SIZE, + moe_top_k: int = DEFAULT_MOE_TOP_K, + topologies: tuple[TopologySpec, ...] = DEFAULT_TOPOLOGIES, + afd_ffn_topology: TopologySpec = DEFAULT_AFD_FFN_TOPOLOGY, + keep_traces: str | Path | None = None, +) -> dict[str, Any]: + """Run msModeling for a triangular `(prefix, query)` grid.""" + + if hidden_size <= 0 or moe_top_k <= 0: + raise ValueError("hidden_size and moe_top_k must be positive") + + root = Path(msmodeling_root).resolve() + if not (root / "cli" / "inference" / "text_generate.py").is_file(): + raise ValueError(f"not an msModeling checkout: {root}") + query_values = tuple(sorted({int(value) for value in query_anchors})) + prefix_values = tuple(sorted({int(value) for value in prefix_anchors})) + max_context_tokens, anchor_grid = _anchor_grid(query_values, prefix_values) + + trace_root_context = ( + tempfile.TemporaryDirectory(prefix="dsv4-profile-") + if keep_traces is None + else None + ) + trace_root = Path( + trace_root_context.name if trace_root_context else keep_traces + ).resolve() + trace_root.mkdir(parents=True, exist_ok=True) + command_log: list[list[str]] = [] + topology_payload: dict[str, Any] = {} + try: + for topology in topologies: + points = [] + ffn_cache: dict[int, list[dict[str, Any]]] = {} + for prefix_tokens, row_queries in anchor_grid.items(): + for query_tokens in row_queries: + trace_path = trace_root / ( + f"{topology.name}-p{prefix_tokens}-q{query_tokens}.json" + ) + command = _msmodeling_command( + python_executable=python_executable, + model_id=model_id, + device=device, + topology=topology, + prefix_tokens=prefix_tokens, + query_tokens=query_tokens, + trace_path=trace_path, + ) + command_log.append(command) + subprocess.run(command, cwd=root, check=True) + layers = aggregate_trace(trace_path) + if topology.name == "afd": + local_query = max( + 1, + (query_tokens + afd_ffn_topology.dp_size - 1) + // afd_ffn_topology.dp_size, + ) + if local_query not in ffn_cache: + ffn_trace_path = trace_root / ( + f"afd-ffn-q{local_query}.json" + ) + ffn_command = _msmodeling_command( + python_executable=python_executable, + model_id=model_id, + device=device, + topology=afd_ffn_topology, + prefix_tokens=0, + query_tokens=local_query, + trace_path=ffn_trace_path, + ) + command_log.append(ffn_command) + subprocess.run(ffn_command, cwd=root, check=True) + ffn_cache[local_query] = aggregate_trace(ffn_trace_path) + layers = _compose_afd_layers( + attention_layers=layers, + ffn_layers=ffn_cache[local_query], + ) + if prefix_tokens != 0: + for layer in layers: + layer.pop(ROUTED_EXPERT_SHAPE, None) + layer.pop(ROUTED_EXPERT_SAMPLE_SHAPES, None) + layer.pop(SHARED_EXPERT_SHAPE, None) + points.append( + { + "prefix_tokens": prefix_tokens, + "query_tokens": query_tokens, + "layers": layers, + } + ) + spec: dict[str, Any] = topology.__dict__ + if topology.name == "afd": + spec = { + "attention": topology.__dict__, + "ffn": afd_ffn_topology.__dict__, + "ffn_job_mapping": ("local_query=ceil(job_query/ffn.dp_size)"), + } + topology_payload[topology.name] = { + "spec": spec, + "max_context_tokens": max_context_tokens, + "points": points, + } + finally: + if trace_root_context is not None: + trace_root_context.cleanup() + + payload = { + "schema_version": PROFILE_SCHEMA_VERSION, + "layer_count": EXPECTED_LAYER_COUNT, + "metadata": { + "model": model_id, + "device": device, + "model_config": { + "hidden_size": hidden_size, + "moe_top_k": moe_top_k, + }, + "performance_model": "analytic", + "generated_at": datetime.now(timezone.utc).isoformat(), + "msmodeling_root": str(root), + "commands": command_log, + "phase_sources": { + "afd.attention_router,afd.afd_post": "afd attention trace", + "afd.routed_experts,afd.shared_expert": "afd_ffn job trace", + "merged.*": "merged trace", + }, + "notes": "CAM communication is not included in these profiles.", + }, + "topologies": topology_payload, + } + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text( + json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8" + ) + return payload + + +def _anchor_grid( + query_values: tuple[int, ...], prefix_values: tuple[int, ...] +) -> tuple[int, dict[int, tuple[int, ...]]]: + if not query_values or not prefix_values: + raise ValueError("profile anchors cannot be empty") + max_context_tokens = max(query_values) + if max_context_tokens < 2: + raise ValueError("largest query anchor must be at least 2") + if any(prefix < 0 or prefix >= max_context_tokens for prefix in prefix_values): + raise ValueError("prefix anchors must be in [0, max context)") + prefixes = sorted({0, *prefix_values, max_context_tokens - 1}) + queries_with_endpoints = {1, *query_values} + grid = {} + for prefix in prefixes: + queries = { + query + for query in queries_with_endpoints + if query > 0 and prefix + query <= max_context_tokens + } + queries.add(max_context_tokens - prefix) + grid[prefix] = tuple(sorted(queries)) + return max_context_tokens, grid + + +def _compose_afd_layers( + *, + attention_layers: list[dict[str, Any]], + ffn_layers: list[dict[str, Any]], +) -> list[dict[str, Any]]: + return [ + { + **attention, + "routed_experts": ffn["routed_experts"], + "shared_expert": ffn["shared_expert"], + ROUTED_EXPERT_SHAPE: ffn[ROUTED_EXPERT_SHAPE], + ROUTED_EXPERT_SAMPLE_SHAPES: ffn[ROUTED_EXPERT_SAMPLE_SHAPES], + SHARED_EXPERT_SHAPE: ffn[SHARED_EXPERT_SHAPE], + } + for attention, ffn in zip(attention_layers, ffn_layers, strict=True) + ] + + +def aggregate_trace(path: str | Path) -> list[dict[str, Any]]: + """Aggregate one msModeling trace into stable per-layer phases.""" + + with Path(path).open(encoding="utf-8") as handle: + raw = json.load(handle) + events = [ + event + for event in raw.get("traceEvents", []) + if event.get("ph") == "X" and event.get("cat") == "analytic" + ] + gates = [ + index for index, event in enumerate(events) if event.get("name") in GATE_OPS + ] + if len(gates) != EXPECTED_LAYER_COUNT: + raise ValueError( + f"trace {path} contains {len(gates)} MoE gates, " + f"expected {EXPECTED_LAYER_COUNT}" + ) + + layers = [] + for layer_idx, gate_index in enumerate(gates): + previous_gate = gates[layer_idx - 1] if layer_idx else 0 + next_gate = gates[layer_idx + 1] if layer_idx + 1 < len(gates) else len(events) + hcpres = [ + index + for index in range(previous_gate, gate_index) + if events[index].get("name") == HC_PRE_OP + ] + layer_start = hcpres[-2] if len(hcpres) >= 2 else previous_gate + all_to_alls = [ + index + for index in range(gate_index, next_gate) + if events[index].get("name") == ALL_TO_ALL_OP + ] + if len(all_to_alls) != 2: + raise ValueError( + f"layer {layer_idx} contains {len(all_to_alls)} " + "all_to_all ops, expected 2" + ) + dispatch_index, combine_index = all_to_alls + gather_indices = [ + index + for index in range(combine_index, next_gate) + if events[index].get("name") == ALL_GATHER_OP + ] + hcpost_indices = [ + index + for index in range(combine_index, next_gate) + if events[index].get("name") == HC_POST_OP + ] + if not hcpost_indices: + raise ValueError(f"layer {layer_idx} is missing hc_post") + hcpost_index = hcpost_indices[0] + gather_index = gather_indices[0] if gather_indices else None + shared_end = gather_index if gather_index is not None else hcpost_index + layer_end = hcpost_index + 1 + shared_starts = [ + index + for index in range(combine_index + 1, shared_end) + if events[index].get("name") == SHARED_EXPERT_START_OP + ] + shared_start = shared_starts[0] if shared_starts else shared_end + + layer = {phase: 0.0 for phase in PROFILE_PHASES} + layer["attention_router"] = _duration_ms(events[layer_start:dispatch_index]) + layer["merged_dispatch"] = _duration_ms( + events[dispatch_index : dispatch_index + 1] + ) + layer["routed_experts"] = _duration_ms( + events[dispatch_index + 1 : combine_index] + ) + layer["merged_combine"] = _duration_ms( + events[combine_index : combine_index + 1] + ) + layer["merged_combine_local"] = _duration_ms( + events[combine_index + 1 : shared_start] + ) + layer["shared_expert"] = _duration_ms(events[shared_start:shared_end]) + layer[ROUTED_EXPERT_SHAPE] = _output_shape(events[dispatch_index], path) + routed_splits = [ + event + for event in events[dispatch_index + 1 : combine_index] + if event.get("name") == SPLIT_WITH_SIZES_OP + ] + if not routed_splits: + raise ValueError( + f"trace {path} layer {layer_idx} has no routed expert split" + ) + layer[ROUTED_EXPERT_SAMPLE_SHAPES] = _output_shape_counts( + routed_splits[0], path + ) + layer[SHARED_EXPERT_SHAPE] = _input_shape(events[shared_start], path) + if gather_index is not None: + layer["merged_sp_post"] = _duration_ms(events[gather_index:layer_end]) + layer["afd_post"] = _duration_ms([events[hcpost_index]]) + layers.append(layer) + return layers + + +def _output_shape(event: dict[str, Any], path: str | Path) -> list[int]: + output = str(event.get("args", {}).get("Output", "")) + match = re.search(r"size=\((\d+(?:\s*,\s*\d+)*)\)", output) + if match is None: + raise ValueError(f"trace {path} event {event.get('name')} has no output shape") + return [int(value.strip()) for value in match.group(1).split(",")] + + +def _input_shape(event: dict[str, Any], path: str | Path) -> list[int]: + raw_shapes = event.get("args", {}).get("simulation_shapes") + shapes = json.loads(raw_shapes) if isinstance(raw_shapes, str) else raw_shapes + if not shapes or not shapes[0]: + raise ValueError(f"trace {path} event {event.get('name')} has no input shape") + shape = [int(value) for value in shapes[0]] + return [math.prod(shape[:-1]), shape[-1]] + + +def _output_shape_counts( + event: dict[str, Any], path: str | Path +) -> list[dict[str, Any]]: + output = str(event.get("args", {}).get("Output", "")) + shapes = re.findall(r"size=\((\d+(?:\s*,\s*\d+)*)\)", output) + if not shapes: + raise ValueError(f"trace {path} event {event.get('name')} has no output shapes") + counts = Counter( + tuple(int(value.strip()) for value in shape.split(",")) for shape in shapes + ) + return [{"shape": list(shape), "count": count} for shape, count in counts.items()] + + +def _duration_ms(events: Iterable[dict[str, Any]]) -> float: + return sum(float(event.get("dur", 0.0)) for event in events) / 1_000.0 + + +def _msmodeling_command( + *, + python_executable: str, + model_id: str, + device: str, + topology: TopologySpec, + prefix_tokens: int, + query_tokens: int, + trace_path: Path, +) -> list[str]: + command = [ + python_executable, + "-m", + "cli.inference.text_generate", + model_id, + "--device", + device, + "--num-devices", + str(topology.num_devices), + "--tp-size", + str(topology.tp_size), + "--dp-size", + str(topology.dp_size), + "--ep-size", + str(topology.ep_size), + ] + if topology.sequence_parallel: + command.append("--enable-sequence-parallel") + command.extend( + [ + "--compile", + "--num-queries", + str(topology.dp_size), + "--query-length", + str(query_tokens), + "--context-length", + str(prefix_tokens), + "--performance-model", + "analytic", + "--chrome-trace", + str(trace_path), + ] + ) + return command + + +__all__ = [ + "DEFAULT_DEVICE", + "DEFAULT_AFD_FFN_TOPOLOGY", + "DEFAULT_HIDDEN_SIZE", + "DEFAULT_MODEL_ID", + "DEFAULT_MOE_TOP_K", + "DEFAULT_TOPOLOGIES", + "TopologySpec", + "aggregate_trace", + "build_profile_bundle", +] diff --git a/simulator/profiles.py b/simulator/profiles.py new file mode 100644 index 00000000..601eda21 --- /dev/null +++ b/simulator/profiles.py @@ -0,0 +1,354 @@ +"""Normalized msModeling profile bundle and interpolation.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +PROFILE_SCHEMA_VERSION = 1 +PROFILE_PHASES = ( + "attention_router", + "merged_dispatch", + "routed_experts", + "merged_combine", + "merged_combine_local", + "shared_expert", + "merged_sp_post", + "afd_post", +) +EXPERT_SHAPE_KEYS = { + "routed_experts": "routed_expert_input_shape", + "ffn_compute": "routed_expert_input_shape", + "shared_expert": "shared_expert_input_shape", +} +ROUTED_EXPERT_SAMPLE_SHAPES = "routed_expert_sample_shapes" + + +@dataclass(frozen=True) +class ProfilePoint: + prefix_tokens: int + query_tokens: int + layers: tuple[dict[str, Any], ...] + + @classmethod + def from_mapping(cls, raw: dict[str, Any], layer_count: int) -> ProfilePoint: + layers = tuple(_normalize_layer(layer) for layer in raw["layers"]) + if len(layers) != layer_count: + raise ValueError( + f"profile point has {len(layers)} layers, expected {layer_count}" + ) + return cls( + prefix_tokens=int(raw["prefix_tokens"]), + query_tokens=int(raw["query_tokens"]), + layers=layers, + ) + + +def _normalize_layer(raw: dict[str, Any]) -> dict[str, Any]: + layer: dict[str, Any] = { + phase: float(raw.get(phase, 0.0)) for phase in PROFILE_PHASES + } + for shape_key in set(EXPERT_SHAPE_KEYS.values()): + if shape_key in raw: + layer[shape_key] = tuple(int(value) for value in raw[shape_key]) + if ROUTED_EXPERT_SAMPLE_SHAPES in raw: + layer[ROUTED_EXPERT_SAMPLE_SHAPES] = tuple( + { + "shape": tuple(int(value) for value in sample["shape"]), + "count": int(sample["count"]), + } + for sample in raw[ROUTED_EXPERT_SAMPLE_SHAPES] + ) + return layer + + +class TopologyProfile: + """One topology's two-dimensional `(prefix, query)` profile grid.""" + + def __init__( + self, + name: str, + points: tuple[ProfilePoint, ...], + layer_count: int, + max_context_tokens: int | None = None, + ): + if not points: + raise ValueError(f"topology profile {name!r} has no points") + self.name = name + self.points = points + self.layer_count = layer_count + self.max_context_tokens = max_context_tokens + self._grid = { + (point.prefix_tokens, point.query_tokens): point for point in points + } + if len(self._grid) != len(points): + raise ValueError(f"topology profile {name!r} contains duplicate points") + self.prefix_anchors = tuple(sorted({point.prefix_tokens for point in points})) + self.query_anchors = tuple(sorted({point.query_tokens for point in points})) + + def duration_ms( + self, + layer_idx: int, + phase: str, + prefix_tokens: int, + query_tokens: int | float, + ) -> float: + if not 0 <= layer_idx < self.layer_count: + raise ValueError(f"layer_idx {layer_idx} is outside profile") + if phase not in PROFILE_PHASES: + raise ValueError(f"unknown profile phase {phase!r}") + if prefix_tokens < 0 or query_tokens <= 0: + raise ValueError("profile lookup requires prefix>=0 and query>0") + if ( + self.max_context_tokens is not None + and prefix_tokens + query_tokens > self.max_context_tokens + ): + raise ValueError( + f"prefix+query={prefix_tokens + query_tokens} outside profile " + f"context limit {self.max_context_tokens}" + ) + + lower_prefix, upper_prefix = _bounds( + self.prefix_anchors, prefix_tokens, "prefix" + ) + lower_value = self._interpolate_query( + lower_prefix, layer_idx, phase, query_tokens + ) + if upper_prefix == lower_prefix: + return lower_value + ratio = (prefix_tokens - lower_prefix) / (upper_prefix - lower_prefix) + if ( + self.max_context_tokens is not None + and query_tokens > self.max_context_tokens - upper_prefix + ): + return self._interpolate_triangle( + lower_prefix, + upper_prefix, + ratio, + layer_idx, + phase, + query_tokens, + ) + upper_value = self._interpolate_query( + upper_prefix, layer_idx, phase, query_tokens + ) + return lower_value + (upper_value - lower_value) * ratio + + def expert_shape_samples( + self, + layer_idx: int, + phase: str, + query_tokens: int | float, + ) -> list[dict[str, Any]]: + """Return the trace anchors whose measured shapes back this lookup.""" + + shape_key = EXPERT_SHAPE_KEYS.get(phase) + if shape_key is None: + return [] + anchors = tuple(sorted(query for prefix, query in self._grid if prefix == 0)) + lower_query, upper_query = _bounds(anchors, query_tokens, "query") + weights = [(lower_query, 1.0)] + if upper_query != lower_query: + upper_weight = (query_tokens - lower_query) / (upper_query - lower_query) + weights = [ + (lower_query, 1.0 - upper_weight), + (upper_query, upper_weight), + ] + samples = [] + for sampled_query, weight in weights: + shape = self._grid[(0, sampled_query)].layers[layer_idx].get(shape_key) + if shape is not None: + sample = { + "query_tokens": sampled_query, + "weight": weight, + "shape": list(shape), + } + if phase != "shared_expert": + sample["expert_shapes"] = [ + { + "shape": list(item["shape"]), + "count": item["count"], + } + for item in self._grid[(0, sampled_query)] + .layers[layer_idx] + .get(ROUTED_EXPERT_SAMPLE_SHAPES, ()) + ] + samples.append(sample) + return samples + + def _interpolate_triangle( + self, + lower_prefix: int, + upper_prefix: int, + prefix_ratio: float, + layer_idx: int, + phase: str, + query_tokens: int | float, + ) -> float: + """Interpolate the strip next to the prefix+query context boundary.""" + + assert self.max_context_tokens is not None + lower_boundary = self.max_context_tokens - lower_prefix + upper_boundary = self.max_context_tokens - upper_prefix + lower_at_boundary = self._interpolate_query( + lower_prefix, layer_idx, phase, lower_boundary + ) + upper_at_boundary = self._interpolate_query( + upper_prefix, layer_idx, phase, upper_boundary + ) + distance = lower_boundary - upper_boundary + lower_weight = (query_tokens - upper_boundary) / distance + upper_weight = prefix_ratio + corner_weight = 1.0 - lower_weight - upper_weight + lower_corner = self._interpolate_query( + lower_prefix, layer_idx, phase, upper_boundary + ) + return ( + corner_weight * lower_corner + + lower_weight * lower_at_boundary + + upper_weight * upper_at_boundary + ) + + def _interpolate_query( + self, + prefix_tokens: int, + layer_idx: int, + phase: str, + query_tokens: int | float, + ) -> float: + anchors = tuple( + sorted(query for prefix, query in self._grid if prefix == prefix_tokens) + ) + lower_query, upper_query = _bounds( + anchors, query_tokens, f"query at prefix={prefix_tokens}" + ) + lower = self._grid[(prefix_tokens, lower_query)].layers[layer_idx][phase] + if upper_query == lower_query: + return lower + upper = self._grid[(prefix_tokens, upper_query)].layers[layer_idx][phase] + ratio = (query_tokens - lower_query) / (upper_query - lower_query) + return lower + (upper - lower) * ratio + + +class ProfileBundle: + """Versioned profiles consumed by the simulation runtime.""" + + def __init__( + self, + *, + metadata: dict[str, Any], + layer_count: int, + topologies: dict[str, TopologyProfile], + topology_specs: dict[str, dict[str, Any]], + ) -> None: + self.metadata = metadata + self.layer_count = layer_count + self.topologies = topologies + self.topology_specs = topology_specs + for required in ("afd", "merged"): + if required not in topologies: + raise ValueError(f"profile bundle is missing {required!r} topology") + + @classmethod + def from_mapping(cls, raw: dict[str, Any]) -> ProfileBundle: + if int(raw.get("schema_version", 0)) != PROFILE_SCHEMA_VERSION: + raise ValueError( + f"unsupported profile schema_version {raw.get('schema_version')!r}" + ) + layer_count = int(raw.get("layer_count", 43)) + topology_profiles = {} + topology_specs = {} + for name, topology in raw.get("topologies", {}).items(): + points = tuple( + ProfilePoint.from_mapping(point, layer_count) + for point in topology.get("points", []) + ) + max_context = topology.get("max_context_tokens") + topology_profiles[name] = TopologyProfile( + name, + points, + layer_count, + int(max_context) if max_context is not None else None, + ) + topology_specs[name] = dict(topology.get("spec", {})) + return cls( + metadata=dict(raw.get("metadata", {})), + layer_count=layer_count, + topologies=topology_profiles, + topology_specs=topology_specs, + ) + + @classmethod + def load(cls, path: str | Path) -> ProfileBundle: + with Path(path).open(encoding="utf-8") as handle: + return cls.from_mapping(json.load(handle)) + + def duration_ms( + self, + topology: str, + layer_idx: int, + phase: str, + prefix_tokens: int, + query_tokens: int | float, + ) -> float: + try: + profile = self.topologies[topology] + except KeyError as exc: + raise ValueError(f"unknown topology profile {topology!r}") from exc + return profile.duration_ms(layer_idx, phase, prefix_tokens, query_tokens) + + def expert_shape_samples( + self, + topology: str, + layer_idx: int, + phase: str, + query_tokens: int | float, + ) -> list[dict[str, Any]]: + return self.topologies[topology].expert_shape_samples( + layer_idx, phase, query_tokens + ) + + def summary(self) -> dict[str, Any]: + return { + "metadata": self.metadata, + "layer_count": self.layer_count, + "topologies": { + name: { + "point_count": len(profile.points), + "prefix_anchors": profile.prefix_anchors, + "query_anchors": profile.query_anchors, + "max_context_tokens": profile.max_context_tokens, + "spec": self.topology_specs[name], + } + for name, profile in self.topologies.items() + }, + } + + +def _bounds( + anchors: tuple[int, ...], value: int | float, label: str +) -> tuple[int, int]: + if not anchors: + raise ValueError(f"profile contains no {label} anchors") + if value < anchors[0] or value > anchors[-1]: + raise ValueError( + f"{label}={value} outside profile domain [{anchors[0]}, {anchors[-1]}]" + ) + for anchor in anchors: + if value == anchor: + return anchor, anchor + if value < anchor: + index = anchors.index(anchor) + return anchors[index - 1], anchor + return anchors[-1], anchors[-1] + + +__all__ = [ + "PROFILE_PHASES", + "PROFILE_SCHEMA_VERSION", + "ProfileBundle", + "ProfilePoint", + "TopologyProfile", +] diff --git a/simulator/profiles/README.md b/simulator/profiles/README.md new file mode 100644 index 00000000..f4c19c79 --- /dev/null +++ b/simulator/profiles/README.md @@ -0,0 +1,6 @@ +# Generated profiles + +Use `python -m simulator profiles build ...` to create the normalized +msModeling analytic profile JSON consumed by the simulator. Generated JSON and +raw traces are intentionally ignored because a full 128K grid can be large and +is specific to an msModeling revision and device profile. diff --git a/simulator/server.py b/simulator/server.py new file mode 100644 index 00000000..f87d6a9e --- /dev/null +++ b/simulator/server.py @@ -0,0 +1,127 @@ +"""Dependency-free local HTTP server for the simulator UI.""" + +from __future__ import annotations + +import json +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any + +from simulator.config import SimulationConfig, default_config_mapping +from simulator.engine import compare_architectures, sweep_qps +from simulator.profiles import ProfileBundle + +WEB_ROOT = Path(__file__).with_name("web") + + +class SimulatorServer(ThreadingHTTPServer): + allow_reuse_address = True + + def __init__( + self, + address: tuple[str, int], + profiles: ProfileBundle, + ) -> None: + self.profiles = profiles + super().__init__(address, SimulatorRequestHandler) + + +class SimulatorRequestHandler(BaseHTTPRequestHandler): + server: SimulatorServer + + def do_GET(self) -> None: # noqa: N802 + if self.path in {"/", "/index.html"}: + self._send_bytes( + (WEB_ROOT / "index.html").read_bytes(), + content_type="text/html; charset=utf-8", + ) + return + if self.path == "/api/defaults": + self._send_json( + { + "config": default_config_mapping(), + "profile": self.server.profiles.summary(), + } + ) + return + self._send_json({"error": "not found"}, status=HTTPStatus.NOT_FOUND) + + def do_POST(self) -> None: # noqa: N802 + try: + raw = self._read_json() + config = SimulationConfig.from_mapping(raw) + if self.path == "/api/simulate": + self._send_json(compare_architectures(config, self.server.profiles)) + return + if self.path == "/api/sweep": + self._send_json(sweep_qps(config, self.server.profiles)) + return + self._send_json({"error": "not found"}, status=HTTPStatus.NOT_FOUND) + except (KeyError, TypeError, ValueError) as exc: + self._send_json({"error": str(exc)}, status=HTTPStatus.UNPROCESSABLE_ENTITY) + except Exception as exc: # pragma: no cover - server safety boundary + self._send_json( + {"error": f"internal error: {exc}"}, + status=HTTPStatus.INTERNAL_SERVER_ERROR, + ) + + def log_message(self, format: str, *args: Any) -> None: + print(f"[simulator] {self.address_string()} {format % args}") + + def _read_json(self) -> dict[str, Any]: + content_length = int(self.headers.get("Content-Length", "0")) + if content_length <= 0: + raise ValueError("request body is empty") + if content_length > 10 * 1024 * 1024: + raise ValueError("request body exceeds 10 MiB") + body = self.rfile.read(content_length) + value = json.loads(body) + if not isinstance(value, dict): + raise ValueError("request body must be a JSON object") + return value + + def _send_json( + self, + payload: dict[str, Any], + *, + status: HTTPStatus = HTTPStatus.OK, + ) -> None: + self._send_bytes( + json.dumps(payload, ensure_ascii=False, allow_nan=False).encode("utf-8"), + status=status, + content_type="application/json; charset=utf-8", + ) + + def _send_bytes( + self, + payload: bytes, + *, + status: HTTPStatus = HTTPStatus.OK, + content_type: str, + ) -> None: + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(payload))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(payload) + + +def serve( + profiles: ProfileBundle, + *, + host: str = "127.0.0.1", + port: int = 8765, +) -> None: + server = SimulatorServer((host, port), profiles) + print(f"DSV4 Prefill simulator: http://{host}:{port}") + try: + server.serve_forever() + except KeyboardInterrupt: + pass + finally: + server.server_close() + + +__all__ = ["SimulatorServer", "serve"] diff --git a/simulator/tests/__init__.py b/simulator/tests/__init__.py new file mode 100644 index 00000000..001e2740 --- /dev/null +++ b/simulator/tests/__init__.py @@ -0,0 +1 @@ +"""Simulator tests.""" diff --git a/simulator/tests/helpers.py b/simulator/tests/helpers.py new file mode 100644 index 00000000..0154864a --- /dev/null +++ b/simulator/tests/helpers.py @@ -0,0 +1,90 @@ +"""Small deterministic profile fixtures.""" + +from __future__ import annotations + +from simulator.profiles import PROFILE_PHASES, ProfileBundle + +QUERY_ANCHORS = (1, 64, 128, 256, 512, 1_024, 2_048, 8_192, 32_768) +PREFIX_ANCHORS = (0, 256, 512, 4_096) + + +def make_profile(layer_count: int = 3) -> ProfileBundle: + topologies = {} + for topology, factor in (("afd", 0.9), ("merged", 1.0)): + points = [] + for prefix in PREFIX_ANCHORS: + for query in QUERY_ANCHORS: + layers = [] + for layer in range(layer_count): + scale = factor * (1.0 + layer * 0.05) + token_cost = query / 4_096.0 + prefix_cost = prefix * query / (4_096.0 * 4_096.0) + values = {phase: 0.0 for phase in PROFILE_PHASES} + values.update( + { + "attention_router": scale + * (0.6 + 0.8 * token_cost + 0.2 * prefix_cost), + "merged_dispatch": scale * (0.12 + 0.05 * token_cost), + "routed_experts": scale * (0.3 + 0.4 * token_cost), + "merged_combine": scale * (0.08 + 0.04 * token_cost), + "merged_combine_local": scale * 0.05, + "shared_expert": scale * (0.12 + 0.2 * token_cost), + "merged_sp_post": scale * 0.11, + "afd_post": scale * 0.04, + } + ) + layers.append(values) + points.append( + { + "prefix_tokens": prefix, + "query_tokens": query, + "layers": [ + { + **layer, + "routed_expert_input_shape": [ + (query * 6 + (8 if topology == "afd" else 16) - 1) + // (8 if topology == "afd" else 16), + 4_096, + ], + "routed_expert_sample_shapes": [ + { + "shape": [ + ( + query * 6 + + (8 if topology == "afd" else 16) + - 1 + ) + // (8 if topology == "afd" else 16), + 4_096, + ], + "count": 1, + } + ], + "shared_expert_input_shape": [query, 4_096], + } + for layer in layers + ], + } + ) + topologies[topology] = {"points": points} + return ProfileBundle.from_mapping( + { + "schema_version": 1, + "layer_count": layer_count, + "metadata": { + "model": "test-dsv4", + "performance_model": "analytic", + "model_config": {"hidden_size": 4_096, "moe_top_k": 6}, + }, + "topologies": { + "afd": { + **topologies["afd"], + "spec": {"ffn": {"ep_size": 8}}, + }, + "merged": { + **topologies["merged"], + "spec": {"ep_size": 16}, + }, + }, + } + ) diff --git a/simulator/tests/test_engine.py b/simulator/tests/test_engine.py new file mode 100644 index 00000000..893302a5 --- /dev/null +++ b/simulator/tests/test_engine.py @@ -0,0 +1,372 @@ +from __future__ import annotations + +import unittest +from collections import deque + +from simulator.config import SimulationConfig +from simulator.engine import RequestDispatcher, compare_architectures, sweep_qps +from simulator.tests.helpers import make_profile +from simulator.workload import RequestSpec + + +class EngineTests(unittest.TestCase): + def test_scheduler_policy_validation(self) -> None: + config = SimulationConfig.from_mapping( + {"scheduler": {"policy": "vllm_queue_aware"}} + ) + self.assertEqual(config.scheduler.policy, "vllm_queue_aware") + + with self.assertRaisesRegex(ValueError, "scheduler.policy"): + SimulationConfig.from_mapping({"scheduler": {"policy": "least_tokens"}}) + + def test_vllm_queue_aware_routes_to_the_less_loaded_dp(self) -> None: + specs = ( + RequestSpec("long", 0.0, 8_192, 0, 0.0), + RequestSpec("short", 0.0, 1, 0, 0.0), + RequestSpec("later", 5.0, 1, 0, 0.0), + ) + base = { + "mode": "fixed", + "scheduler": {"max_num_batched_tokens": 8_192}, + "cam": { + "dispatch_send": {"fixed_ms": 0, "per_token_ms": 0.001}, + "dispatch_recv": {"fixed_ms": 0, "per_token_ms": 0}, + "combine_send": {"fixed_ms": 0, "per_token_ms": 0}, + "combine_recv": {"fixed_ms": 0, "per_token_ms": 0}, + }, + } + round_robin = compare_architectures( + SimulationConfig.from_mapping(base), + make_profile(layer_count=1), + specs, + ) + base["scheduler"]["policy"] = "vllm_queue_aware" + queue_aware = compare_architectures( + SimulationConfig.from_mapping(base), + make_profile(layer_count=1), + specs, + ) + + self.assertEqual(round_robin["afd"]["requests"][2]["assigned_dp"], 0) + self.assertEqual(queue_aware["afd"]["requests"][2]["assigned_dp"], 1) + self.assertLess( + queue_aware["afd"]["requests"][2]["completion_ms"], + round_robin["afd"]["requests"][2]["completion_ms"], + ) + + def test_queue_aware_active_sets_are_pruned_in_place(self) -> None: + specs = tuple( + RequestSpec(f"request-{index}", 0.0, 1, 0, 0.0) for index in range(10_000) + ) + dispatcher = RequestDispatcher(specs, 4, "vllm_queue_aware") + active_ids = [id(active) for active in dispatcher.active] + + dispatcher.dispatch_until(0.0) + + self.assertTrue(all(isinstance(active, deque) for active in dispatcher.active)) + self.assertEqual([id(active) for active in dispatcher.active], active_ids) + self.assertEqual([len(active) for active in dispatcher.active], [2_500] * 4) + + round_robin = RequestDispatcher(specs, 4, "round_robin") + round_robin.dispatch_until(0.0) + self.assertTrue(all(not active for active in round_robin.active)) + + def test_fixed_workload_runs_both_architectures_and_exposes_barrier(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "fixed", + "fixed_lengths": [128, 512, 2048, 8192], + "scheduler": {"max_num_batched_tokens": 16384}, + } + ) + + result = compare_architectures(config, make_profile()) + + self.assertEqual(result["afd"]["summary"]["request_count"], 4) + self.assertEqual(result["merged"]["summary"]["request_count"], 4) + self.assertGreater(result["merged"]["summary"]["barrier_wait_ms"], 0) + self.assertTrue( + all(item["completion_ms"] is not None for item in result["afd"]["requests"]) + ) + + def test_expert_timeline_events_expose_shape_topk_and_ep(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "fixed", + "fixed_lengths": [128, 128, 128, 128], + "scheduler": {"max_num_batched_tokens": 128}, + } + ) + + result = compare_architectures(config, make_profile(layer_count=1)) + afd_ffn = next( + event + for event in result["afd"]["timeline"] + if event["phase"] == "ffn_compute" + ) + routed = next( + event + for event in result["merged"]["timeline"] + if event["phase"] == "routed_experts" + ) + shared = next( + event + for event in result["merged"]["timeline"] + if event["phase"] == "shared_expert" + ) + + self.assertEqual( + afd_ffn["sampled_input_shapes"], + [ + { + "query_tokens": 128, + "weight": 1.0, + "shape": [96, 4_096], + "expert_shapes": [{"shape": [96, 4_096], "count": 1}], + } + ], + ) + self.assertEqual(afd_ffn["top_k"], 6) + self.assertEqual(afd_ffn["ep_size"], 8) + self.assertEqual( + routed["sampled_input_shapes"], + [ + { + "query_tokens": 128, + "weight": 1.0, + "shape": [48, 4_096], + "expert_shapes": [{"shape": [48, 4_096], "count": 1}], + } + ], + ) + self.assertEqual(routed["top_k"], 6) + self.assertEqual(routed["ep_size"], 16) + self.assertEqual(routed["tokens"], 512) + self.assertEqual(routed["profile_query_tokens"], 128) + self.assertIsNone(shared["top_k"]) + + def test_merged_ep_uses_global_tokens_with_dp4_equivalent_query(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "fixed", + "fixed_lengths": [512, 8_192, 2_048, 6_144], + "scheduler": {"max_num_batched_tokens": 8_192}, + } + ) + profiles = make_profile(layer_count=1) + + result = compare_architectures(config, profiles) + routed = next( + event + for event in result["merged"]["timeline"] + if event["phase"] == "routed_experts" + ) + + self.assertEqual(routed["tokens"], 16_896) + self.assertEqual(routed["profile_query_tokens"], 4_224) + global_ep_events = [ + event + for event in result["merged"]["timeline"] + if event["phase"] in {"merged_dispatch", "routed_experts", "merged_combine"} + ] + self.assertTrue(all(event["tokens"] == 16_896 for event in global_ep_events)) + self.assertTrue( + all(event["profile_query_tokens"] == 4_224 for event in global_ep_events) + ) + local_tail_events = [ + event + for event in result["merged"]["timeline"] + if event["phase"] + in { + "merged_combine_local", + "shared_expert", + "merged_sp_post", + } + ] + self.assertTrue(all(event["tokens"] == 8_192 for event in local_tail_events)) + self.assertTrue( + all(event["profile_query_tokens"] == 8_192 for event in local_tail_events) + ) + self.assertAlmostEqual( + routed["end_ms"] - routed["start_ms"], + profiles.duration_ms("merged", 0, "routed_experts", 0, 4_224), + ) + + def test_merged_ep_preserves_fractional_dp4_equivalent_query(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "fixed", + "fixed_lengths": [128, 129, 130, 132], + "scheduler": {"max_num_batched_tokens": 132}, + } + ) + + result = compare_architectures(config, make_profile(layer_count=1)) + routed = next( + event + for event in result["merged"]["timeline"] + if event["phase"] == "routed_experts" + ) + + self.assertEqual(routed["tokens"], 519) + self.assertEqual(routed["profile_query_tokens"], 129.75) + + def test_chunked_prefill_and_token_ubatch_preserve_request_completion(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "fixed", + "fixed_lengths": [2048], + "scheduler": { + "max_num_batched_tokens": 1024, + "chunked_prefill": True, + "chunk_size": 1024, + }, + "afd": {"ubatch_split": "token"}, + } + ) + + result = compare_architectures(config, make_profile()) + + request = result["afd"]["requests"][0] + self.assertEqual(request["query_tokens"], 2048) + self.assertGreater(request["completion_ms"], request["first_scheduled_ms"]) + self.assertTrue(any(event["stage"] == 1 for event in result["afd"]["timeline"])) + + def test_prefix_cache_reduces_computed_tokens_but_not_logical_tokens(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "fixed", + "csv_text": "input_length,cached_prefix_tokens\n1024,512\n", + "scheduler": {"max_num_batched_tokens": 2048}, + "prefix_cache": {"enabled": True}, + } + ) + + result = compare_architectures(config, make_profile()) + summary = result["afd"]["summary"] + + self.assertEqual(summary["logical_input_tokens"], 1024) + self.assertEqual(summary["computed_query_tokens"], 512) + self.assertEqual(summary["cache_token_ratio"], 0.5) + + def test_continuous_poisson_and_sweep_are_reproducible(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "continuous", + "length_mix": [{"tokens": 128, "weight": 1}], + "arrival": { + "kind": "poisson", + "qps": 1, + "duration_s": 2, + "warmup_s": 0, + "seed": 11, + }, + "scheduler": {"max_num_batched_tokens": 1024}, + "sweep": { + "min_qps": 0.5, + "max_qps": 2, + "coarse_points": 3, + "refinement_steps": 1, + }, + } + ) + + first = sweep_qps(config, make_profile(layer_count=1)) + second = sweep_qps(config, make_profile(layer_count=1)) + + self.assertEqual(first, second) + self.assertTrue(first["series"]["afd"]) + + def test_afd_ffn_is_fcfs_by_dispatch_arrival(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "fixed", + "fixed_lengths": [100, 1], + "scheduler": {"max_num_batched_tokens": 100}, + "cam": { + "dispatch_send": {"fixed_ms": 0, "per_token_ms": 1}, + "dispatch_recv": {"fixed_ms": 0, "per_token_ms": 0}, + "combine_send": {"fixed_ms": 0, "per_token_ms": 0}, + "combine_recv": {"fixed_ms": 0, "per_token_ms": 0}, + }, + } + ) + + result = compare_architectures(config, make_profile(layer_count=1)) + receives = sorted( + ( + event + for event in result["afd"]["timeline"] + if event["phase"] == "dispatch_recv" + ), + key=lambda event: event["start_ms"], + ) + + self.assertEqual([event["tokens"] for event in receives], [1, 100]) + + def test_sweep_rejects_fixed_workload(self) -> None: + config = SimulationConfig.from_mapping( + {"mode": "fixed", "fixed_lengths": [128]} + ) + + with self.assertRaisesRegex(ValueError, "continuous"): + sweep_qps(config, make_profile(layer_count=1)) + + def test_timeline_truncated_only_when_an_event_is_dropped(self) -> None: + exact = SimulationConfig.from_mapping( + { + "mode": "fixed", + "fixed_lengths": [128], + "output": {"timeline_max_events": 7}, + } + ) + truncated = SimulationConfig.from_mapping( + { + "mode": "fixed", + "fixed_lengths": [128], + "output": {"timeline_max_events": 6}, + } + ) + + exact_result = compare_architectures(exact, make_profile(layer_count=1)) + truncated_result = compare_architectures(truncated, make_profile(layer_count=1)) + + self.assertFalse(exact_result["afd"]["summary"]["timeline_truncated"]) + self.assertTrue(truncated_result["afd"]["summary"]["timeline_truncated"]) + self.assertEqual(len(truncated_result["afd"]["timeline"]), 6) + + def test_continuous_window_requires_a_measured_request(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "continuous", + "csv_text": "arrival_time_ms,input_length\n0,128\n", + "arrival": { + "kind": "trace", + "warmup_s": 1, + "duration_s": 1, + }, + } + ) + + with self.assertRaisesRegex(ValueError, "measurement window"): + compare_architectures(config, make_profile(layer_count=1)) + + def test_trace_metrics_use_only_the_measurement_window(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "continuous", + "csv_text": ( + "arrival_time_ms,input_length\n100,128\n125,128\n160,128\n" + ), + "arrival": { + "kind": "trace", + "warmup_s": 0.01, + "duration_s": 0.04, + }, + } + ) + + result = compare_architectures(config, make_profile(layer_count=1)) + + self.assertEqual(result["afd"]["summary"]["request_count"], 1) + self.assertEqual(result["afd"]["summary"]["offered_qps"], 25) diff --git a/simulator/tests/test_profiles.py b/simulator/tests/test_profiles.py new file mode 100644 index 00000000..68f4f1b3 --- /dev/null +++ b/simulator/tests/test_profiles.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import unittest + +from simulator.profile_builder import ( + DEFAULT_AFD_FFN_TOPOLOGY, + ROUTED_EXPERT_SAMPLE_SHAPES, + ROUTED_EXPERT_SHAPE, + SHARED_EXPERT_SHAPE, + _anchor_grid, + _compose_afd_layers, +) +from simulator.profiles import PROFILE_PHASES, ProfileBundle +from simulator.tests.helpers import make_profile + + +class ProfileTests(unittest.TestCase): + def test_profile_returns_anchor_and_interpolates_both_dimensions(self) -> None: + profile = make_profile(layer_count=1) + exact = profile.duration_ms("afd", 0, "attention_router", 0, 128) + middle = profile.duration_ms("afd", 0, "attention_router", 256, 256) + + self.assertGreater(exact, 0) + self.assertGreater(middle, exact) + + def test_profile_rejects_extrapolation(self) -> None: + profile = make_profile(layer_count=1) + + with self.assertRaisesRegex(ValueError, "outside profile domain"): + profile.duration_ms("afd", 0, "attention_router", 0, 65_536) + + def test_expert_shapes_report_the_actual_interpolation_anchors(self) -> None: + profile = make_profile(layer_count=1) + + samples = profile.expert_shape_samples("afd", 0, "routed_experts", 96) + + self.assertEqual( + samples, + [ + { + "query_tokens": 64, + "weight": 0.5, + "shape": [48, 4_096], + "expert_shapes": [{"shape": [48, 4_096], "count": 1}], + }, + { + "query_tokens": 128, + "weight": 0.5, + "shape": [96, 4_096], + "expert_shapes": [{"shape": [96, 4_096], "count": 1}], + }, + ], + ) + + def test_triangular_grid_covers_full_context_boundary(self) -> None: + max_context, grid = _anchor_grid( + (128, 512, 65_536, 131_072), + (8_192, 32_768, 122_880), + ) + + self.assertEqual(max_context, 131_072) + self.assertIn(0, grid) + self.assertIn(1, grid[8_192]) + self.assertIn(122_880, grid[8_192]) + self.assertEqual(grid[131_071], (1,)) + + def test_triangular_profile_interpolates_near_context_boundary(self) -> None: + points = [] + for prefix, queries in {0: (1, 8, 16), 8: (1, 8), 15: (1,)}.items(): + for query in queries: + layer = {phase: 0.0 for phase in PROFILE_PHASES} + layer["attention_router"] = prefix + 2 * query + points.append( + { + "prefix_tokens": prefix, + "query_tokens": query, + "layers": [layer], + } + ) + profile = ProfileBundle.from_mapping( + { + "schema_version": 1, + "layer_count": 1, + "topologies": { + name: {"max_context_tokens": 16, "points": points} + for name in ("afd", "merged") + }, + } + ) + + value = profile.duration_ms("afd", 0, "attention_router", 4, 12) + + self.assertEqual(value, 28) + + def test_afd_profile_composes_attention_and_single_job_ffn_phases(self) -> None: + attention = [{phase: 1.0 for phase in PROFILE_PHASES}] + ffn = [ + { + **{phase: 2.0 for phase in PROFILE_PHASES}, + ROUTED_EXPERT_SHAPE: [192, 4_096], + ROUTED_EXPERT_SAMPLE_SHAPES: [{"shape": [6, 4_096], "count": 32}], + SHARED_EXPERT_SHAPE: [32, 4_096], + } + ] + + layer = _compose_afd_layers( + attention_layers=attention, + ffn_layers=ffn, + )[0] + + self.assertEqual(layer["attention_router"], 1.0) + self.assertEqual(layer["afd_post"], 1.0) + self.assertEqual(layer["routed_experts"], 2.0) + self.assertEqual(layer["shared_expert"], 2.0) + self.assertEqual(layer[ROUTED_EXPERT_SHAPE], [192, 4_096]) + self.assertEqual( + layer[ROUTED_EXPERT_SAMPLE_SHAPES], + [{"shape": [6, 4_096], "count": 32}], + ) + self.assertEqual(layer[SHARED_EXPERT_SHAPE], [32, 4_096]) + self.assertEqual(DEFAULT_AFD_FFN_TOPOLOGY.dp_size, 8) + self.assertEqual(DEFAULT_AFD_FFN_TOPOLOGY.tp_size, 1) + self.assertFalse(DEFAULT_AFD_FFN_TOPOLOGY.sequence_parallel) diff --git a/simulator/tests/test_server.py b/simulator/tests/test_server.py new file mode 100644 index 00000000..7600be54 --- /dev/null +++ b/simulator/tests/test_server.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import json +import threading +import unittest +import urllib.request + +from simulator.server import SimulatorServer +from simulator.tests.helpers import make_profile + + +class ServerTests(unittest.TestCase): + def test_defaults_page_and_simulation_api(self) -> None: + server = SimulatorServer(("127.0.0.1", 0), make_profile(layer_count=1)) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + host, port = server.server_address + try: + with urllib.request.urlopen( + f"http://{host}:{port}/api/defaults" + ) as response: + defaults = json.load(response) + self.assertEqual(defaults["profile"]["layer_count"], 1) + + with urllib.request.urlopen(f"http://{host}:{port}/") as response: + page = response.read().decode() + self.assertIn('value="vllm_queue_aware"', page) + self.assertIn("TopK 数", page) + self.assertIn("Expert GMM 实采 Shape", page) + self.assertIn("Profile Query", page) + self.assertIn("EP 数", page) + + payload = json.dumps( + { + "mode": "fixed", + "fixed_lengths": [128, 128], + "scheduler": { + "policy": "vllm_queue_aware", + "max_num_batched_tokens": 1024, + }, + } + ).encode() + request = urllib.request.Request( + f"http://{host}:{port}/api/simulate", + data=payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(request) as response: + result = json.load(response) + self.assertIn("afd", result) + self.assertIn("merged", result) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=2) diff --git a/simulator/tests/test_workload.py b/simulator/tests/test_workload.py new file mode 100644 index 00000000..1201bb94 --- /dev/null +++ b/simulator/tests/test_workload.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import unittest + +from simulator.config import SimulationConfig +from simulator.workload import generate_workload, read_csv_requests + + +class WorkloadTests(unittest.TestCase): + def test_csv_length_list_preserves_online_distribution(self) -> None: + rows = read_csv_requests(csv_text="\ufeffinput_length\n512\n8192\n8192\n") + + self.assertEqual([row.input_length for row in rows], [512, 8192, 8192]) + + def test_csv_timestamp_and_cached_prefix_are_replayed(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "continuous", + "csv_text": ( + "request_id,arrival_time_ms,input_length,cached_prefix_tokens\n" + "a,100,1024,512\n" + "b,125,2048,0\n" + ), + "arrival": {"kind": "trace", "duration_s": 1, "warmup_s": 0}, + "scheduler": {"max_num_batched_tokens": 4096}, + "prefix_cache": {"enabled": True}, + } + ) + + workload = generate_workload(config) + + self.assertEqual([item.arrival_ms for item in workload], [0.0, 25.0]) + self.assertEqual(workload[0].cached_prefix_tokens, 512) + self.assertEqual(workload[0].query_tokens, 512) + + def test_csv_trace_is_clipped_to_the_configured_window(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "continuous", + "csv_text": ( + "arrival_time_ms,input_length\n100,128\n125,128\n160,128\n" + ), + "arrival": { + "kind": "trace", + "duration_s": 0.04, + "warmup_s": 0.01, + }, + } + ) + + workload = generate_workload(config) + + self.assertEqual([item.arrival_ms for item in workload], [0.0, 25.0]) + + def test_prefix_cache_sampling_is_deterministic_and_block_aligned(self) -> None: + raw = { + "mode": "fixed", + "fixed_lengths": [1024, 1024, 1024], + "scheduler": {"max_num_batched_tokens": 4096}, + "prefix_cache": { + "enabled": True, + "request_hit_rate": 1.0, + "matched_prefix_ratio": 0.73, + "block_size": 32, + "seed": 7, + }, + } + first = generate_workload(SimulationConfig.from_mapping(raw)) + second = generate_workload(SimulationConfig.from_mapping(raw)) + + self.assertEqual(first, second) + self.assertEqual({item.cached_prefix_tokens for item in first}, {736}) + + def test_full_cache_ratio_keeps_one_aligned_query_block(self) -> None: + config = SimulationConfig.from_mapping( + { + "mode": "fixed", + "fixed_lengths": [1024], + "prefix_cache": { + "enabled": True, + "request_hit_rate": 1, + "matched_prefix_ratio": 1, + "block_size": 32, + }, + } + ) + + request = generate_workload(config)[0] + + self.assertEqual(request.cached_prefix_tokens, 992) + self.assertEqual(request.query_tokens, 32) diff --git a/simulator/web/index.html b/simulator/web/index.html new file mode 100644 index 00000000..2e5d9acc --- /dev/null +++ b/simulator/web/index.html @@ -0,0 +1,652 @@ + + + + + + DSV4-Flash Prefill 仿真器 + + + +
+
+
+

DSV4-Flash Prefill 仿真器

+

相同 16-die 预算下比较 CAMAsync AFD 与 DP4×TP4/EP16 合并部署

+
+
正在读取 profile…
+
+ +
+ + +
+

配置 workload 后运行仿真

页面只负责输入和展示；所有调度、流水、指标均由 Python 后端计算。

+
+
+
AFD 吞吐提升
—
+
AFD / 合并 p99 TTFT
—
+
AFD SLO Goodput 提升
—
+
+
架构吞吐 req/s输入 tok/s计算 tok/sTTFT p50p90p99SLO 达标率SLO goodput
+ +
+

逐层关键路径时间线

显示后端返回的有限事件窗口
+
Attention通信FFN等待/屏障
+
W/S 长按缩放 · A/D 长按平移 · 滚轮缩放 · 拖拽平移 · 双击复位
+ +
+ + +
+
+
+
+ + + + + diff --git a/simulator/workload.py b/simulator/workload.py new file mode 100644 index 00000000..0d3e7a03 --- /dev/null +++ b/simulator/workload.py @@ -0,0 +1,261 @@ +"""Workload generation, CSV replay, and Prefix Cache sampling.""" + +from __future__ import annotations + +import csv +import io +import math +import random +from collections.abc import Iterable +from dataclasses import dataclass +from pathlib import Path + +from simulator.config import PrefixCacheConfig, SimulationConfig + + +@dataclass(frozen=True) +class CsvRequest: + input_length: int + request_id: str | None = None + arrival_time_ms: float | None = None + cached_prefix_tokens: int | None = None + + +@dataclass(frozen=True) +class RequestSpec: + request_id: str + arrival_ms: float + input_tokens: int + cached_prefix_tokens: int + cache_lookup_ms: float + + @property + def query_tokens(self) -> int: + return self.input_tokens - self.cached_prefix_tokens + + +@dataclass +class RuntimeRequest: + spec: RequestSpec + assigned_dp: int + computed_query_tokens: int = 0 + first_scheduled_ms: float | None = None + completion_ms: float | None = None + + @property + def remaining_query_tokens(self) -> int: + return self.spec.query_tokens - self.computed_query_tokens + + @property + def current_prefix_tokens(self) -> int: + return self.spec.cached_prefix_tokens + self.computed_query_tokens + + +def read_csv_requests( + *, + csv_path: str | None = None, + csv_text: str | None = None, +) -> tuple[CsvRequest, ...]: + if bool(csv_path) == bool(csv_text): + raise ValueError("provide exactly one of csv_path or csv_text") + if csv_path: + text = Path(csv_path).read_text(encoding="utf-8-sig") + else: + text = (csv_text or "").lstrip("\ufeff") + reader = csv.DictReader(io.StringIO(text)) + if reader.fieldnames is None or "input_length" not in reader.fieldnames: + raise ValueError("CSV must contain an input_length header") + requests = [] + for row_number, row in enumerate(reader, start=2): + try: + input_length = int(row["input_length"] or "") + request_id = (row.get("request_id") or "").strip() or None + arrival_raw = (row.get("arrival_time_ms") or "").strip() + cached_raw = (row.get("cached_prefix_tokens") or "").strip() + arrival_time_ms = float(arrival_raw) if arrival_raw else None + cached_prefix_tokens = int(cached_raw) if cached_raw else None + except (TypeError, ValueError) as exc: + raise ValueError(f"invalid CSV value on row {row_number}: {exc}") from exc + if input_length <= 0: + raise ValueError(f"input_length must be positive on row {row_number}") + if arrival_time_ms is not None and arrival_time_ms < 0: + raise ValueError(f"arrival_time_ms cannot be negative on row {row_number}") + if ( + cached_prefix_tokens is not None + and not 0 <= cached_prefix_tokens < input_length + ): + raise ValueError( + f"cached_prefix_tokens must be in [0, input_length) on row {row_number}" + ) + requests.append( + CsvRequest( + input_length=input_length, + request_id=request_id, + arrival_time_ms=arrival_time_ms, + cached_prefix_tokens=cached_prefix_tokens, + ) + ) + if not requests: + raise ValueError("CSV contains no requests") + has_arrival = [item.arrival_time_ms is not None for item in requests] + if any(has_arrival) and not all(has_arrival): + raise ValueError("arrival_time_ms must be present on every CSV row or none") + if all(has_arrival): + timestamps = [float(item.arrival_time_ms or 0.0) for item in requests] + if timestamps != sorted(timestamps): + raise ValueError("CSV arrival_time_ms must be non-decreasing") + return tuple(requests) + + +def generate_workload(config: SimulationConfig) -> tuple[RequestSpec, ...]: + csv_requests = None + if config.csv_path or config.csv_text: + csv_requests = read_csv_requests( + csv_path=config.csv_path, + csv_text=config.csv_text, + ) + + if config.mode == "fixed": + raw_requests = ( + csv_requests + if csv_requests is not None + else tuple(CsvRequest(length) for length in config.fixed_lengths) + ) + arrivals = [ + item.arrival_time_ms + if item.arrival_time_ms is not None and config.arrival.kind == "trace" + else 0.0 + for item in raw_requests + ] + return _materialize_requests(raw_requests, arrivals, config.prefix_cache) + + if csv_requests and all(item.arrival_time_ms is not None for item in csv_requests): + if config.arrival.kind != "trace": + raise ValueError( + "CSV arrival_time_ms requires arrival.kind=trace for exact replay" + ) + first_arrival = float(csv_requests[0].arrival_time_ms or 0.0) + arrivals = [ + float(item.arrival_time_ms or 0.0) - first_arrival for item in csv_requests + ] + end_ms = (config.arrival.warmup_s + config.arrival.duration_s) * 1_000.0 + window = [ + (request, arrival) + for request, arrival in zip(csv_requests, arrivals, strict=True) + if arrival < end_ms + ] + raw_requests, arrivals = zip(*window, strict=True) if window else ((), ()) + return _materialize_requests(raw_requests, arrivals, config.prefix_cache) + + arrival_rng = random.Random(config.arrival.seed) + length_rng = random.Random(config.arrival.seed + 1) + end_ms = (config.arrival.warmup_s + config.arrival.duration_s) * 1_000.0 + arrivals = _generate_arrivals( + config.arrival.kind, + config.arrival.qps, + end_ms, + arrival_rng, + ) + if csv_requests: + raw_requests = tuple( + _choose_csv_request( + csv_requests, + index, + config.csv_sampling, + length_rng, + ) + for index in range(len(arrivals)) + ) + else: + raw_requests = tuple( + CsvRequest(_choose_weighted_length(config, length_rng)) for _ in arrivals + ) + return _materialize_requests(raw_requests, arrivals, config.prefix_cache) + + +def _generate_arrivals( + kind: str, + qps: float, + end_ms: float, + rng: random.Random, +) -> list[float]: + if kind == "trace": + raise ValueError("arrival.kind=trace requires CSV arrival_time_ms") + arrivals = [] + current_ms = 0.0 + while current_ms < end_ms: + arrivals.append(current_ms) + if kind == "constant": + current_ms += 1_000.0 / qps + else: + current_ms += rng.expovariate(qps) * 1_000.0 + return arrivals + + +def _choose_csv_request( + requests: tuple[CsvRequest, ...], + index: int, + sampling: str, + rng: random.Random, +) -> CsvRequest: + if sampling == "cycle": + return requests[index % len(requests)] + return requests[rng.randrange(len(requests))] + + +def _choose_weighted_length(config: SimulationConfig, rng: random.Random) -> int: + total = sum(bucket.weight for bucket in config.length_mix) + cursor = rng.random() * total + for bucket in config.length_mix: + cursor -= bucket.weight + if cursor <= 0: + return bucket.tokens + return config.length_mix[-1].tokens + + +def _materialize_requests( + raw_requests: Iterable[CsvRequest], + arrivals: Iterable[float], + prefix_cache: PrefixCacheConfig, +) -> tuple[RequestSpec, ...]: + cache_rng = random.Random(prefix_cache.seed) + result = [] + for index, (raw, arrival_ms) in enumerate(zip(raw_requests, arrivals, strict=True)): + cached_tokens = _cached_tokens(raw, prefix_cache, cache_rng) + request_id = raw.request_id or f"r{index + 1}" + result.append( + RequestSpec( + request_id=request_id, + arrival_ms=float(arrival_ms), + input_tokens=raw.input_length, + cached_prefix_tokens=cached_tokens, + cache_lookup_ms=prefix_cache.lookup_latency_ms(cached_tokens), + ) + ) + return tuple(result) + + +def _cached_tokens( + raw: CsvRequest, + config: PrefixCacheConfig, + rng: random.Random, +) -> int: + if not config.enabled: + return 0 + if raw.cached_prefix_tokens is not None: + return raw.cached_prefix_tokens + if rng.random() >= config.request_hit_rate: + return 0 + candidate = math.floor(raw.input_length * config.matched_prefix_ratio) + candidate = (candidate // config.block_size) * config.block_size + max_cached = ((raw.input_length - 1) // config.block_size) * config.block_size + return min(candidate, max_cached) + + +__all__ = [ + "CsvRequest", + "RequestSpec", + "RuntimeRequest", + "generate_workload", + "read_csv_requests", +]