From 4b2f4961cf1b9b6748a3610bbc2bbc673e60e2a2 Mon Sep 17 00:00:00 2001 From: shsaihdsaiudh <196440533+shsaihdsaiudh@users.noreply.github.com> Date: Sun, 26 Jul 2026 16:14:39 +0800 Subject: [PATCH] feat(issue1): analytical SM/QP budget model reproducing DeepEP get_theoretical_num_sms --- src/code/issue1/README-2026.md | 23 + src/code/issue1/results/sm_qp_budget.csv | 17 + src/code/issue1/results/sm_qp_budget.json | 567 ++++++++++++++++++++++ src/code/issue1/sm_qp_budget.py | 299 ++++++++++++ src/code/issue1/test_sm_qp_budget.py | 36 ++ 5 files changed, 942 insertions(+) create mode 100644 src/code/issue1/README-2026.md create mode 100644 src/code/issue1/results/sm_qp_budget.csv create mode 100644 src/code/issue1/results/sm_qp_budget.json create mode 100644 src/code/issue1/sm_qp_budget.py create mode 100644 src/code/issue1/test_sm_qp_budget.py diff --git a/src/code/issue1/README-2026.md b/src/code/issue1/README-2026.md new file mode 100644 index 0000000..ce36281 --- /dev/null +++ b/src/code/issue1/README-2026.md @@ -0,0 +1,23 @@ +# Issue 1 单机可完成部分：DeepEP SM/QP 预算模型 + +`sm_qp_budget.py` 是 DeepEP V2 `ElasticBuffer.get_theoretical_num_sms()` 与 +`get_theoretical_num_qps()` 的独立复现，保留以下关键逻辑： + +- 用组合数估计 top-k 命中的不同 EP/scale-out rank 数； +- 显式累计 SM read/write、RDMA、NVLink 归一化流量； +- 识别瓶颈链路并由每 SM 读写带宽反推 SM 数； +- 乘 1.25 安全系数、偶数对齐、4 SM 下限和设备上限； +- direct QP=`min(num_sms, 9)`，hybrid QP=`num_sms*16+1`，最后受已分配 QP 限制。 + +对照源码：`deepseek-ai/DeepEP` commit +`dd758caf451848bd150e1046af3d0a73e5fff38d` 的 +`deep_ep/buffers/elastic.py`。 + +```bash +python -m pytest src/code/issue1/test_sm_qp_budget.py +python src/code/issue1/sm_qp_budget.py +``` + +默认输出 1/2/4/8 节点、top-k 1/2/4/8 的 JSON 与 CSV。带宽参数是可替换输入； +最终仍需在目标 DeepEP 多机环境扫描 SM 数，验证预测值能保持实测最优带宽的 95% 以上。 + diff --git a/src/code/issue1/results/sm_qp_budget.csv b/src/code/issue1/results/sm_qp_budget.csv new file mode 100644 index 0000000..0f3fd5f --- /dev/null +++ b/src/code/issue1/results/sm_qp_budget.csv @@ -0,0 +1,17 @@ +nodes,ranks,topk,expected_distinct_ranks,bottleneck,raw_sms,recommended_sms,recommended_qps,sm_saved_vs_fixed_24 +1,8,1,1.0,nvlink,10.285714285714286,14,129,10 +1,8,2,1.8888888888888893,nvlink,10.285714285714286,14,129,10 +1,8,4,3.375462718138551,nvlink,10.285714285714286,14,129,10 +1,8,8,5.432551281938457,nvlink,10.285714285714286,14,129,10 +2,16,1,1.0,rdma,5.0,8,129,16 +2,16,2,1.9523809523809526,rdma,4.915789473684211,8,129,16 +2,16,4,3.720404925587369,rdma,6.003590161090127,8,129,16 +2,16,8,6.750925436277102,rdma,8.76938854208264,12,129,12 +4,32,1,1.0,rdma,3.0,4,65,20 +4,32,2,1.9841269841269842,rdma,2.591591591591592,4,65,20 +4,32,4,3.904761904761905,rdma,2.689954839275614,4,65,20 +4,32,8,7.555555555555557,rdma,3.450977290315941,6,97,18 +8,64,1,1.0,rdma,2.428571428571429,4,65,20 +8,64,2,2.0,rdma,1.9579831932773104,4,65,20 +8,64,4,4.0,rdma,1.8357467379870638,4,65,20 +8,64,8,8.0,rdma,2.036205912506329,4,65,20 diff --git a/src/code/issue1/results/sm_qp_budget.json b/src/code/issue1/results/sm_qp_budget.json new file mode 100644 index 0000000..c8f2a7d --- /dev/null +++ b/src/code/issue1/results/sm_qp_budget.json @@ -0,0 +1,567 @@ +{ + "schema_version": 1, + "source": { + "repository": "deepseek-ai/DeepEP", + "commit": "dd758caf451848bd150e1046af3d0a73e5fff38d", + "path": "deep_ep/buffers/elastic.py" + }, + "configuration": { + "num_experts": 64, + "gpus_per_node": 8, + "mode": "hybrid", + "device_sms": 132, + "prefer_overlap_with_compute": true, + "bandwidths": { + "rdma_gbs": 50.0, + "nvlink_gbs": 450.0, + "sm_read_gbs": 200.0, + "sm_write_gbs": 50.0 + } + }, + "cases": [ + { + "nodes": 1, + "topk": 1, + "topology": { + "num_ranks": 8, + "num_scaleout_ranks": 1, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 1, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 1.0, + "expected_distinct_scaleout_ranks": 0.0, + "normalized_traffic": { + "sm_read": 1.0, + "sm_write": 1.0, + "rdma": 0.0, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.0, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "nvlink", + "read_limited_sms": 2.5714285714285716, + "write_limited_sms": 10.285714285714286, + "raw_sms": 10.285714285714286, + "recommended_sms": 14, + "sm_saved_vs_fixed_24": 10, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 129 + } + }, + { + "nodes": 1, + "topk": 2, + "topology": { + "num_ranks": 8, + "num_scaleout_ranks": 1, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 1, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 1.8888888888888893, + "expected_distinct_scaleout_ranks": 0.0, + "normalized_traffic": { + "sm_read": 0.5294117647058822, + "sm_write": 1.0, + "rdma": 0.0, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.0, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "nvlink", + "read_limited_sms": 1.361344537815126, + "write_limited_sms": 10.285714285714286, + "raw_sms": 10.285714285714286, + "recommended_sms": 14, + "sm_saved_vs_fixed_24": 10, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 129 + } + }, + { + "nodes": 1, + "topk": 4, + "topology": { + "num_ranks": 8, + "num_scaleout_ranks": 1, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 1, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 3.375462718138551, + "expected_distinct_scaleout_ranks": 0.0, + "normalized_traffic": { + "sm_read": 0.29625567914773615, + "sm_write": 1.0, + "rdma": 0.0, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.0, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "nvlink", + "read_limited_sms": 0.7618003178084645, + "write_limited_sms": 10.285714285714286, + "raw_sms": 10.285714285714286, + "recommended_sms": 14, + "sm_saved_vs_fixed_24": 10, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 129 + } + }, + { + "nodes": 1, + "topk": 8, + "topology": { + "num_ranks": 8, + "num_scaleout_ranks": 1, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 1, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 5.432551281938457, + "expected_distinct_scaleout_ranks": 0.0, + "normalized_traffic": { + "sm_read": 0.1840755748270042, + "sm_write": 1.0, + "rdma": 0.0, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.0, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "nvlink", + "read_limited_sms": 0.4733371924122966, + "write_limited_sms": 10.285714285714286, + "raw_sms": 10.285714285714286, + "recommended_sms": 14, + "sm_saved_vs_fixed_24": 10, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 129 + } + }, + { + "nodes": 2, + "topk": 1, + "topology": { + "num_ranks": 16, + "num_scaleout_ranks": 2, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 2, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 1.0, + "expected_distinct_scaleout_ranks": 1.0, + "normalized_traffic": { + "sm_read": 2.0, + "sm_write": 2.5, + "rdma": 0.5, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.01, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 1.0, + "write_limited_sms": 5.0, + "raw_sms": 5.0, + "recommended_sms": 8, + "sm_saved_vs_fixed_24": 16, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 129 + } + }, + { + "nodes": 2, + "topk": 2, + "topology": { + "num_ranks": 16, + "num_scaleout_ranks": 2, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 2, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 1.9523809523809526, + "expected_distinct_scaleout_ranks": 1.507936507936508, + "normalized_traffic": { + "sm_read": 1.2845528455284552, + "sm_write": 1.8983739837398375, + "rdma": 0.3861788617886179, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.007723577235772358, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.8315789473684211, + "write_limited_sms": 4.915789473684211, + "raw_sms": 4.915789473684211, + "recommended_sms": 8, + "sm_saved_vs_fixed_24": 16, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 129 + } + }, + { + "nodes": 2, + "topk": 4, + "topology": { + "num_ranks": 16, + "num_scaleout_ranks": 2, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 2, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 3.720404925587369, + "expected_distinct_scaleout_ranks": 1.8868071818891492, + "normalized_traffic": { + "sm_read": 0.7759389742860818, + "sm_write": 1.5223634603799892, + "rdma": 0.2535755139060924, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.005071510278121848, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.7649979313198181, + "write_limited_sms": 6.003590161090127, + "raw_sms": 6.003590161090127, + "recommended_sms": 8, + "sm_saved_vs_fixed_24": 16, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 129 + } + }, + { + "nodes": 2, + "topk": 8, + "topology": { + "num_ranks": 16, + "num_scaleout_ranks": 2, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 2, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 6.750925436277102, + "expected_distinct_scaleout_ranks": 1.9952472177944165, + "normalized_traffic": { + "sm_read": 0.44367949936152606, + "sm_write": 1.295903669467697, + "rdma": 0.14777582989382898, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.0029555165978765795, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.7505955129473676, + "write_limited_sms": 8.76938854208264, + "raw_sms": 8.76938854208264, + "recommended_sms": 12, + "sm_saved_vs_fixed_24": 12, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 129 + } + }, + { + "nodes": 4, + "topk": 1, + "topology": { + "num_ranks": 32, + "num_scaleout_ranks": 4, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 4, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 1.0, + "expected_distinct_scaleout_ranks": 1.0, + "normalized_traffic": { + "sm_read": 2.0, + "sm_write": 2.25, + "rdma": 0.75, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.015, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.6666666666666667, + "write_limited_sms": 3.0, + "raw_sms": 3.0, + "recommended_sms": 4, + "sm_saved_vs_fixed_24": 20, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 65 + } + }, + { + "nodes": 4, + "topk": 2, + "topology": { + "num_ranks": 32, + "num_scaleout_ranks": 4, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 4, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 1.9841269841269842, + "expected_distinct_scaleout_ranks": 1.7619047619047619, + "normalized_traffic": { + "sm_read": 1.392, + "sm_write": 1.726, + "rdma": 0.666, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.01332, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.5225225225225225, + "write_limited_sms": 2.591591591591592, + "raw_sms": 2.591591591591592, + "recommended_sms": 4, + "sm_saved_vs_fixed_24": 20, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 65 + } + }, + { + "nodes": 4, + "topk": 4, + "topology": { + "num_ranks": 32, + "num_scaleout_ranks": 4, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 4, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 3.904761904761905, + "expected_distinct_scaleout_ranks": 2.775024552391025, + "normalized_traffic": { + "sm_read": 0.9667745804903844, + "sm_write": 1.4337668158543035, + "rdma": 0.533007764636081, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.01066015529272162, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.45345239067505155, + "write_limited_sms": 2.689954839275614, + "raw_sms": 2.689954839275614, + "recommended_sms": 4, + "sm_saved_vs_fixed_24": 20, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 65 + } + }, + { + "nodes": 4, + "topk": 8, + "topology": { + "num_ranks": 32, + "num_scaleout_ranks": 4, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 4, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 7.555555555555557, + "expected_distinct_scaleout_ranks": 3.658983374884153, + "normalized_traffic": { + "sm_read": 0.6166301525581965, + "sm_write": 1.253422244021902, + "rdma": 0.36320790853629453, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.007264158170725891, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.42443331908931853, + "write_limited_sms": 3.450977290315941, + "raw_sms": 3.450977290315941, + "recommended_sms": 6, + "sm_saved_vs_fixed_24": 18, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 97 + } + }, + { + "nodes": 8, + "topk": 1, + "topology": { + "num_ranks": 64, + "num_scaleout_ranks": 8, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 8, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 1.0, + "expected_distinct_scaleout_ranks": 1.0, + "normalized_traffic": { + "sm_read": 2.0, + "sm_write": 2.125, + "rdma": 0.875, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.0175, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.5714285714285715, + "write_limited_sms": 2.428571428571429, + "raw_sms": 2.428571428571429, + "recommended_sms": 4, + "sm_saved_vs_fixed_24": 20, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 65 + } + }, + { + "nodes": 8, + "topk": 2, + "topology": { + "num_ranks": 64, + "num_scaleout_ranks": 8, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 8, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 2.0, + "expected_distinct_scaleout_ranks": 1.8888888888888893, + "normalized_traffic": { + "sm_read": 1.4444444444444446, + "sm_write": 1.6180555555555556, + "rdma": 0.8263888888888891, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.01652777777777778, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.43697478991596633, + "write_limited_sms": 1.9579831932773104, + "raw_sms": 1.9579831932773104, + "recommended_sms": 4, + "sm_saved_vs_fixed_24": 20, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 65 + } + }, + { + "nodes": 8, + "topk": 4, + "topology": { + "num_ranks": 64, + "num_scaleout_ranks": 8, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 8, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 4.0, + "expected_distinct_scaleout_ranks": 3.375462718138551, + "normalized_traffic": { + "sm_read": 1.0938656795346378, + "sm_write": 1.3554832099418297, + "rdma": 0.7383824695928081, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.014767649391856161, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.3703587654707817, + "write_limited_sms": 1.8357467379870638, + "raw_sms": 1.8357467379870638, + "recommended_sms": 4, + "sm_saved_vs_fixed_24": 20, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 65 + } + }, + { + "nodes": 8, + "topk": 8, + "topology": { + "num_ranks": 64, + "num_scaleout_ranks": 8, + "num_scaleup_ranks": 8, + "num_rdma_ranks": 8, + "num_nvlink_ranks": 8 + }, + "estimate": { + "expected_distinct_ranks": 8.0, + "expected_distinct_scaleout_ranks": 5.432551281938457, + "normalized_traffic": { + "sm_read": 0.8040689102423071, + "sm_write": 1.2098836137802884, + "rdma": 0.5941852964620187, + "nvlink": 0.875 + }, + "link_time_units": { + "rdma": 0.011883705929240374, + "nvlink": 0.0019444444444444444 + }, + "bottleneck": "rdma", + "read_limited_sms": 0.3383073070934298, + "write_limited_sms": 2.036205912506329, + "raw_sms": 2.036205912506329, + "recommended_sms": 4, + "sm_saved_vs_fixed_24": 20, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": "Sweep SM counts on the target multi-rank topology and verify the recommendation retains at least 95% of the best measured bandwidth.", + "recommended_qps": 65 + } + } + ] +} diff --git a/src/code/issue1/sm_qp_budget.py b/src/code/issue1/sm_qp_budget.py new file mode 100644 index 0000000..220183a --- /dev/null +++ b/src/code/issue1/sm_qp_budget.py @@ -0,0 +1,299 @@ +#!/usr/bin/env python3 +"""Standalone reproduction of DeepEP V2's analytical SM/QP budget model.""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + + +DEEPEP_SOURCE_COMMIT = "dd758caf451848bd150e1046af3d0a73e5fff38d" +DEEPEP_SOURCE_PATH = "deep_ep/buffers/elastic.py" + + +@dataclass(frozen=True) +class Topology: + num_ranks: int + num_scaleout_ranks: int + num_scaleup_ranks: int + num_rdma_ranks: int + num_nvlink_ranks: int + + def __post_init__(self) -> None: + values = asdict(self) + if any(value <= 0 for value in values.values()): + raise ValueError("all topology sizes must be positive") + if self.num_scaleout_ranks * self.num_scaleup_ranks != self.num_ranks: + raise ValueError("scaleout_ranks * scaleup_ranks must equal num_ranks") + if self.num_nvlink_ranks > self.num_ranks: + raise ValueError("num_nvlink_ranks cannot exceed num_ranks") + + +@dataclass(frozen=True) +class Bandwidths: + rdma_gbs: float + nvlink_gbs: float + sm_read_gbs: float = 200.0 + sm_write_gbs: float = 50.0 + + def __post_init__(self) -> None: + if min(asdict(self).values()) <= 0: + raise ValueError("all bandwidths must be positive") + + +def expected_distinct_groups(num_experts: int, num_topk: int, num_groups: int) -> float: + if num_experts <= 0 or not 1 <= num_topk <= num_experts or num_groups <= 0: + raise ValueError("invalid expert/top-k/group count") + if num_experts % num_groups: + raise ValueError("num_experts must be divisible by num_groups") + misses = math.comb(num_experts - num_experts // num_groups, num_topk) + total = math.comb(num_experts, num_topk) + return num_groups * (1.0 - misses / total) + + +def align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def estimate_sm_budget( + topology: Topology, + bandwidths: Bandwidths, + num_experts: int, + num_topk: int, + device_sms: int, + prefer_overlap_with_compute: bool = True, +) -> dict[str, Any]: + """Reproduce ``ElasticBuffer.get_theoretical_num_sms`` with diagnostics.""" + + if device_sms <= 0: + raise ValueError("device_sms must be positive") + expected_topk = expected_distinct_groups(num_experts, num_topk, topology.num_ranks) + expected_scaleout_topk = ( + expected_distinct_groups(num_experts, num_topk, topology.num_scaleout_ranks) + if topology.num_scaleout_ranks > 1 + else 0.0 + ) + + sm_read = 1.0 / expected_topk + sm_write = 0.0 + rdma_traffic = 0.0 + nvlink_traffic = 0.0 + + if topology.num_scaleout_ranks > 1: + sm_write += 1.0 / expected_topk + sm_write += (1.0 / expected_topk) * ( + expected_scaleout_topk / topology.num_scaleout_ranks + ) + rdma_traffic += (1.0 / expected_topk) * ( + expected_scaleout_topk * (1.0 - 1.0 / topology.num_scaleout_ranks) + ) + sm_read += expected_scaleout_topk / expected_topk + sm_write += 1.0 + nvlink_traffic += 1.0 - 1.0 / topology.num_scaleup_ranks + else: + if topology.num_rdma_ranks > 1: + sm_write += 1.0 / expected_topk + sm_write += topology.num_nvlink_ranks / topology.num_ranks + nvlink_traffic += ( + topology.num_nvlink_ranks + / topology.num_ranks + * (1.0 - 1.0 / topology.num_nvlink_ranks) + ) + rdma_traffic += ( + topology.num_ranks - topology.num_nvlink_ranks + ) / topology.num_ranks + + rdma_time_units = rdma_traffic / bandwidths.rdma_gbs + nvlink_time_units = nvlink_traffic / bandwidths.nvlink_gbs + if topology.num_scaleout_ranks > 1 and rdma_time_units > nvlink_time_units: + bottleneck = "rdma" + bounded_traffic, bounded_gbs = rdma_traffic, bandwidths.rdma_gbs + else: + bottleneck = "nvlink" + bounded_traffic, bounded_gbs = nvlink_traffic, bandwidths.nvlink_gbs + + raw_sms = float(device_sms) + read_limited_sms = 0.0 + write_limited_sms = 0.0 + if bounded_traffic > 0: + read_limited_sms = ( + bounded_gbs / bounded_traffic * sm_read / bandwidths.sm_read_gbs + ) + write_limited_sms = ( + bounded_gbs / bounded_traffic * sm_write / bandwidths.sm_write_gbs + ) + raw_sms = max(read_limited_sms, write_limited_sms) + + recommended_sms = align_up(max(4, math.ceil(raw_sms * 1.25)), 2) + if not prefer_overlap_with_compute: + recommended_sms = max(recommended_sms, 64) + recommended_sms = min(recommended_sms, device_sms) + return { + "expected_distinct_ranks": expected_topk, + "expected_distinct_scaleout_ranks": expected_scaleout_topk, + "normalized_traffic": { + "sm_read": sm_read, + "sm_write": sm_write, + "rdma": rdma_traffic, + "nvlink": nvlink_traffic, + }, + "link_time_units": {"rdma": rdma_time_units, "nvlink": nvlink_time_units}, + "bottleneck": bottleneck, + "read_limited_sms": read_limited_sms, + "write_limited_sms": write_limited_sms, + "raw_sms": raw_sms, + "recommended_sms": recommended_sms, + "sm_saved_vs_fixed_24": 24 - recommended_sms, + "bandwidth_acceptance_floor": 0.95, + "needs_measurement": ( + "Sweep SM counts on the target multi-rank topology and verify the recommendation " + "retains at least 95% of the best measured bandwidth." + ), + } + + +def estimate_qps(num_sms: int, allow_hybrid_mode: bool, num_allocated_qps: int) -> int: + if num_sms <= 0 or num_allocated_qps <= 0: + raise ValueError("SM and allocated-QP counts must be positive") + requested = num_sms * 16 + 1 if allow_hybrid_mode else min(num_sms, 9) + return min(requested, num_allocated_qps) + + +def parse_int_list(value: str) -> list[int]: + result = [int(part) for part in value.split(",")] + if not result or any(item <= 0 for item in result): + raise argparse.ArgumentTypeError("values must be positive comma-separated integers") + return result + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="DeepEP analytical SM/QP budget model") + parser.add_argument("--num-experts", type=int, default=64) + parser.add_argument("--topk-values", type=parse_int_list, default=[1, 2, 4, 8]) + parser.add_argument("--node-counts", type=parse_int_list, default=[1, 2, 4, 8]) + parser.add_argument("--gpus-per-node", type=int, default=8) + parser.add_argument("--mode", choices=("direct", "hybrid"), default="hybrid") + parser.add_argument("--device-sms", type=int, default=132) + parser.add_argument("--rdma-gbs", type=float, default=50.0) + parser.add_argument("--nvlink-gbs", type=float, default=450.0) + parser.add_argument("--sm-read-gbs", type=float, default=200.0) + parser.add_argument("--sm-write-gbs", type=float, default=50.0) + parser.add_argument("--num-allocated-qps", type=int, default=129) + parser.add_argument("--no-overlap", action="store_true") + parser.add_argument("--output-dir", type=Path, default=Path("src/code/issue1/results")) + return parser + + +def build_report(args: argparse.Namespace) -> dict[str, Any]: + bandwidths = Bandwidths( + args.rdma_gbs, args.nvlink_gbs, args.sm_read_gbs, args.sm_write_gbs + ) + cases = [] + for nodes in args.node_counts: + total_ranks = nodes * args.gpus_per_node + if args.num_experts % total_ranks: + raise ValueError( + f"num_experts={args.num_experts} must be divisible by total ranks={total_ranks}" + ) + logical_scaleout = nodes if args.mode == "hybrid" else 1 + logical_scaleup = args.gpus_per_node if args.mode == "hybrid" else total_ranks + topology = Topology( + num_ranks=total_ranks, + num_scaleout_ranks=logical_scaleout, + num_scaleup_ranks=logical_scaleup, + num_rdma_ranks=nodes, + num_nvlink_ranks=args.gpus_per_node, + ) + for topk in args.topk_values: + estimate = estimate_sm_budget( + topology, + bandwidths, + args.num_experts, + topk, + args.device_sms, + not args.no_overlap, + ) + estimate["recommended_qps"] = estimate_qps( + estimate["recommended_sms"], + args.mode == "hybrid", + args.num_allocated_qps, + ) + cases.append( + { + "nodes": nodes, + "topk": topk, + "topology": asdict(topology), + "estimate": estimate, + } + ) + return { + "schema_version": 1, + "source": { + "repository": "deepseek-ai/DeepEP", + "commit": DEEPEP_SOURCE_COMMIT, + "path": DEEPEP_SOURCE_PATH, + }, + "configuration": { + "num_experts": args.num_experts, + "gpus_per_node": args.gpus_per_node, + "mode": args.mode, + "device_sms": args.device_sms, + "prefer_overlap_with_compute": not args.no_overlap, + "bandwidths": asdict(bandwidths), + }, + "cases": cases, + } + + +def write_csv(report: dict[str, Any], path: Path) -> None: + fields = [ + "nodes", + "ranks", + "topk", + "expected_distinct_ranks", + "bottleneck", + "raw_sms", + "recommended_sms", + "recommended_qps", + "sm_saved_vs_fixed_24", + ] + with path.open("w", encoding="utf-8-sig", newline="") as handle: + writer = csv.DictWriter(handle, fields) + writer.writeheader() + for case in report["cases"]: + estimate = case["estimate"] + writer.writerow( + { + "nodes": case["nodes"], + "ranks": case["topology"]["num_ranks"], + "topk": case["topk"], + "expected_distinct_ranks": estimate["expected_distinct_ranks"], + "bottleneck": estimate["bottleneck"], + "raw_sms": estimate["raw_sms"], + "recommended_sms": estimate["recommended_sms"], + "recommended_qps": estimate["recommended_qps"], + "sm_saved_vs_fixed_24": estimate["sm_saved_vs_fixed_24"], + } + ) + + +def main() -> int: + args = build_parser().parse_args() + report = build_report(args) + args.output_dir.mkdir(parents=True, exist_ok=True) + json_path = args.output_dir / "sm_qp_budget.json" + csv_path = args.output_dir / "sm_qp_budget.csv" + json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + write_csv(report, csv_path) + print(f"Wrote {json_path}") + print(f"Wrote {csv_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/code/issue1/test_sm_qp_budget.py b/src/code/issue1/test_sm_qp_budget.py new file mode 100644 index 0000000..1a063f8 --- /dev/null +++ b/src/code/issue1/test_sm_qp_budget.py @@ -0,0 +1,36 @@ +import math + +from sm_qp_budget import ( + Bandwidths, + Topology, + estimate_qps, + estimate_sm_budget, + expected_distinct_groups, +) + + +def test_expected_distinct_groups_matches_occupancy_extremes(): + assert expected_distinct_groups(64, 1, 8) == 1.0 + assert math.isclose(expected_distinct_groups(64, 64, 8), 8.0) + + +def test_qp_budget_matches_deepep_direct_and_hybrid_rules(): + assert estimate_qps(4, allow_hybrid_mode=False, num_allocated_qps=17) == 4 + assert estimate_qps(24, allow_hybrid_mode=False, num_allocated_qps=17) == 9 + assert estimate_qps(4, allow_hybrid_mode=True, num_allocated_qps=129) == 65 + assert estimate_qps(16, allow_hybrid_mode=True, num_allocated_qps=129) == 129 + + +def test_budget_is_even_and_capped_by_device(): + topology = Topology(16, 2, 8, 2, 8) + result = estimate_sm_budget( + topology, + Bandwidths(rdma_gbs=50, nvlink_gbs=450), + num_experts=64, + num_topk=8, + device_sms=132, + ) + assert result["recommended_sms"] % 2 == 0 + assert 4 <= result["recommended_sms"] <= 132 + assert result["bottleneck"] in {"rdma", "nvlink"} +