From 59799a368793b9f795baf59b067233a65ad8e38e Mon Sep 17 00:00:00 2001 From: Ruiyan Ma <38345787+littlefatfat@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:12:11 +0800 Subject: [PATCH] [Simulator] Add high-fidelity CPU-based inference simulator (#33824) Co-authored-by: zhouhaizhu.zhz Co-authored-by: LinSiyuan814 Co-authored-by: hzh0425 --- .github/workflows/_pr-test-check-changes.yml | 8 + .github/workflows/_pr-test-simulator-cpu.yml | 96 +++ .github/workflows/pr-test-extra.yml | 11 + .gitignore | 1 + benchmark/simulator/bench_runner.py | 168 +++++ benchmark/simulator/bench_serving.py | 280 ++++++++ docs/docs.json | 1 + docs/docs/advanced_features/overview.mdx | 1 + .../advanced_features/sglang_simulator.mdx | 88 +++ tools/sglang-simulator/README.md | 258 +++++++ tools/sglang-simulator/examples/README.md | 50 ++ .../examples/assets/model.pkl | Bin 0 -> 734 bytes .../examples/assets/replay_table.json | 5 + .../examples/assets/tokenizer/tokenizer.json | 59 ++ .../assets/tokenizer/tokenizer_config.json | 6 + .../examples/build_example_assets.py | 55 ++ .../examples/sim_configs/aic_silicon.json | 21 + .../examples/sim_configs/aic_sol.json | 21 + .../examples/sim_configs/ml.json | 22 + .../examples/sim_configs/replay.json | 23 + .../examples/workloads/sharegpt-example.json | 38 + .../workloads/timestamp-trace-example.jsonl | 3 + tools/sglang-simulator/pyproject.toml | 18 + tools/sglang-simulator/setup.py | 21 + .../src/sglang_simulator/__init__.py | 1 + .../src/sglang_simulator/compat.py | 99 +++ .../src/sglang_simulator/dataset/__init__.py | 35 + .../src/sglang_simulator/dataset/autobench.py | 242 +++++++ .../sglang_simulator/dataset/base_dataset.py | 72 ++ .../sglang_simulator/dataset/dataset_args.py | 20 + .../src/sglang_simulator/dataset/random.py | 52 ++ .../src/sglang_simulator/hook/__init__.py | 15 + .../src/sglang_simulator/hook/base_hook.py | 36 + .../sglang_simulator/hook/class_hook_entry.py | 80 +++ .../src/sglang_simulator/hook/utils.py | 8 + .../sglang_simulator/simulation/__init__.py | 0 .../simulation/benchmark/__init__.py | 4 + .../simulation/benchmark/base_runner.py | 18 + .../simulation/benchmark/bench_config.py | 9 + .../simulation/manager/__init__.py | 5 + .../simulation/manager/config.py | 252 +++++++ .../simulation/manager/env.py | 52 ++ .../simulation/manager/state.py | 123 ++++ .../simulation/sglang/__init__.py | 0 .../simulation/sglang/cache_controller.py | 311 +++++++++ .../simulation/sglang/engine.py | 19 + .../simulation/sglang/hicache_storage.py | 155 +++++ .../simulation/sglang/hiradix_cache.py | 25 + .../simulation/sglang/hook_bootstrap.py | 80 +++ .../simulation/sglang/launch_server.py | 109 +++ .../simulation/sglang/mem_cache_allocator.py | 128 ++++ .../simulation/sglang/mem_pool_host.py | 384 +++++++++++ .../simulation/sglang/model_runner.py | 265 +++++++ .../simulation/sglang/req_stats_manager.py | 22 + .../simulation/sglang/scheduler.py | 648 ++++++++++++++++++ .../simulation/sglang/sgl_kernel_hook.py | 15 + .../simulation/sglang/unified_radix_cache.py | 34 + .../simulation/sglang/utils.py | 130 ++++ .../src/sglang_simulator/simulation/types.py | 134 ++++ .../src/sglang_simulator/simulation/utils.py | 239 +++++++ .../src/sglang_simulator/spec/__init__.py | 5 + .../spec/accelerator/__init__.py | 4 + .../sglang_simulator/spec/accelerator/base.py | 100 +++ .../sglang_simulator/spec/accelerator/info.py | 24 + .../src/sglang_simulator/spec/data_type.py | 104 +++ .../sglang_simulator/spec/model/__init__.py | 3 + .../src/sglang_simulator/spec/model/base.py | 44 ++ .../time_predictor/__init__.py | 19 + .../time_predictor/aiconfigurator.py | 291 ++++++++ .../sglang_simulator/time_predictor/base.py | 97 +++ .../src/sglang_simulator/time_predictor/ml.py | 156 +++++ .../sglang_simulator/time_predictor/replay.py | 189 +++++ .../src/sglang_simulator/utils/__init__.py | 3 + .../src/sglang_simulator/utils/json.py | 22 + .../src/sglang_simulator/utils/logger.py | 16 + tools/sglang-simulator/src/usercustomize.py | 15 + .../test/assets/qwen3-8b/config.json | 28 + .../test/test_simulation_cache_hit_ratio.py | 68 ++ .../test/test_simulation_offline_blocking.py | 80 +++ .../test/test_simulation_sglang_runner.py | 118 ++++ .../test/test_simulation_sglang_serving.py | 163 +++++ 81 files changed, 6624 insertions(+) create mode 100644 .github/workflows/_pr-test-simulator-cpu.yml create mode 100644 benchmark/simulator/bench_runner.py create mode 100644 benchmark/simulator/bench_serving.py create mode 100644 docs/docs/advanced_features/sglang_simulator.mdx create mode 100644 tools/sglang-simulator/README.md create mode 100644 tools/sglang-simulator/examples/README.md create mode 100644 tools/sglang-simulator/examples/assets/model.pkl create mode 100644 tools/sglang-simulator/examples/assets/replay_table.json create mode 100644 tools/sglang-simulator/examples/assets/tokenizer/tokenizer.json create mode 100644 tools/sglang-simulator/examples/assets/tokenizer/tokenizer_config.json create mode 100755 tools/sglang-simulator/examples/build_example_assets.py create mode 100644 tools/sglang-simulator/examples/sim_configs/aic_silicon.json create mode 100644 tools/sglang-simulator/examples/sim_configs/aic_sol.json create mode 100644 tools/sglang-simulator/examples/sim_configs/ml.json create mode 100644 tools/sglang-simulator/examples/sim_configs/replay.json create mode 100644 tools/sglang-simulator/examples/workloads/sharegpt-example.json create mode 100644 tools/sglang-simulator/examples/workloads/timestamp-trace-example.jsonl create mode 100644 tools/sglang-simulator/pyproject.toml create mode 100644 tools/sglang-simulator/setup.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/compat.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/dataset/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/dataset/autobench.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/dataset/base_dataset.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/dataset/dataset_args.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/dataset/random.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/hook/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/hook/base_hook.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/hook/class_hook_entry.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/hook/utils.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/base_runner.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/bench_config.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/manager/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/manager/config.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/manager/env.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/manager/state.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/cache_controller.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/engine.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hicache_storage.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hiradix_cache.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hook_bootstrap.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/launch_server.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/mem_cache_allocator.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/mem_pool_host.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/model_runner.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/req_stats_manager.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/scheduler.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/sgl_kernel_hook.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/unified_radix_cache.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/sglang/utils.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/types.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/simulation/utils.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/spec/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/spec/accelerator/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/spec/accelerator/base.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/spec/accelerator/info.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/spec/data_type.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/spec/model/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/spec/model/base.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/time_predictor/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/time_predictor/aiconfigurator.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/time_predictor/base.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/time_predictor/ml.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/time_predictor/replay.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/utils/__init__.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/utils/json.py create mode 100644 tools/sglang-simulator/src/sglang_simulator/utils/logger.py create mode 100644 tools/sglang-simulator/src/usercustomize.py create mode 100644 tools/sglang-simulator/test/assets/qwen3-8b/config.json create mode 100644 tools/sglang-simulator/test/test_simulation_cache_hit_ratio.py create mode 100644 tools/sglang-simulator/test/test_simulation_offline_blocking.py create mode 100644 tools/sglang-simulator/test/test_simulation_sglang_runner.py create mode 100644 tools/sglang-simulator/test/test_simulation_sglang_serving.py diff --git a/.github/workflows/_pr-test-check-changes.yml b/.github/workflows/_pr-test-check-changes.yml index 2dfae1948..660e29f22 100644 --- a/.github/workflows/_pr-test-check-changes.yml +++ b/.github/workflows/_pr-test-check-changes.yml @@ -19,6 +19,8 @@ on: outputs: main_package: value: ${{ jobs.run.outputs.main_package }} + simulator: + value: ${{ jobs.run.outputs.simulator }} sgl_kernel: value: ${{ jobs.run.outputs.sgl_kernel }} jit_kernel: @@ -44,6 +46,7 @@ jobs: runs-on: ubuntu-latest outputs: main_package: ${{ steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }} + simulator: ${{ steps.filter.outputs.simulator || steps.run-mode.outputs.run_all_tests }} sgl_kernel: ${{ steps.filter.outputs.sgl_kernel }} jit_kernel: ${{ steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }} multimodal_gen: ${{ steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} @@ -94,6 +97,11 @@ jobs: - "test/**/!(*.md)" - "rust/**" - "proto/sglang/runtime/v1/sglang.proto" + simulator: + - ".github/workflows/pr-test-extra.yml" + - ".github/workflows/_pr-test-simulator-cpu.yml" + - "benchmark/simulator/**/!(*.md)" + - "tools/sglang-simulator/**/!(*.md)" multimodal_gen: - ".github/workflows/pr-test.yml" - ".github/workflows/pr-test-multimodal-gen.yml" diff --git a/.github/workflows/_pr-test-simulator-cpu.yml b/.github/workflows/_pr-test-simulator-cpu.yml new file mode 100644 index 000000000..a5f050dec --- /dev/null +++ b/.github/workflows/_pr-test-simulator-cpu.yml @@ -0,0 +1,96 @@ +name: PR Test SGLang Simulator (CPU) + +on: + workflow_call: + inputs: + check_changes: + description: 'toJson(needs.check-changes.outputs).' + type: string + required: true + caller_inputs: + description: 'toJson(inputs) from the caller workflow.' + type: string + required: true + rust_ext_artifact: + description: 'Artifact of prebuilt Rust extension modules.' + type: string + default: '' + +env: + SGLANG_IS_IN_CI: true + SKIP_PR_TEST_HEALTH_CHECK: ${{ (fromJson(inputs.caller_inputs).skip_pr_test_health_check || fromJson(inputs.caller_inputs).test_parallel_dispatch || fromJson(inputs.caller_inputs).run_all_tests) && 'true' || 'false' }} + PR_TEST_BYPASS_MAINTENANCE_ON_MAIN: ${{ github.ref == 'refs/heads/main' && 'true' || 'false' }} + USE_VENV: false + +jobs: + run: + name: simulator-test-cpu + if: fromJson(inputs.check_changes).main_package == 'true' || fromJson(inputs.check_changes).simulator == 'true' + runs-on: ubuntu-latest + timeout-minutes: 40 + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ fromJson(inputs.caller_inputs).git_ref || github.sha }} + + - uses: ./.github/actions/check-pr-test-health + + - uses: ./.github/actions/check-maintenance + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - uses: ./.github/actions/download-rust-ext + id: rust_ext + with: + artifact_name: ${{ inputs.rust_ext_artifact }} + + - name: Install protoc + Rust toolchain + if: ${{ steps.rust_ext.outputs.hit != 'true' }} + timeout-minutes: 10 + run: bash scripts/ci/utils/install_rust_protoc.sh + + - name: Rust cache (rust/ workspace) + if: ${{ steps.rust_ext.outputs.hit != 'true' }} + uses: Swatinem/rust-cache@v2 + with: + workspaces: rust + shared-key: "sglang-grpc-cpu" + + - name: Install dependencies + timeout-minutes: 20 + env: + UV_SYSTEM_PYTHON: "1" + run: | + uv pip install -e "python" --index-strategy unsafe-best-match --prerelease allow + uv pip install pytest + uv pip install -e "tools/sglang-simulator[aic]" + + - name: Prebuild HiCache native hash extension + timeout-minutes: 5 + env: + MALLOC_ARENA_MAX: "2" + MAX_JOBS: "1" + TORCH_EXTENSIONS_DIR: ${{ runner.temp }}/torch-extensions + run: | + python3 -c \ + 'from sglang.srt.mem_cache.cpp_utils.native_hash import get_native_hash; get_native_hash([1], None)' + + - name: Run SGLang Simulator compatibility tests + timeout-minutes: 10 + env: + MALLOC_ARENA_MAX: "2" + MAX_JOBS: "1" + PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1" + TORCH_EXTENSIONS_DIR: ${{ runner.temp }}/torch-extensions + run: | + python3 -m pytest -q tools/sglang-simulator/test/test_simulation_sglang_runner.py + python3 -m pytest -q tools/sglang-simulator/test/test_simulation_sglang_serving.py + python3 -m pytest -q tools/sglang-simulator/test/test_simulation_cache_hit_ratio.py + python3 -m pytest -q tools/sglang-simulator/test/test_simulation_offline_blocking.py diff --git a/.github/workflows/pr-test-extra.yml b/.github/workflows/pr-test-extra.yml index 0c5f1115e..3c4cd168a 100644 --- a/.github/workflows/pr-test-extra.yml +++ b/.github/workflows/pr-test-extra.yml @@ -149,6 +149,16 @@ jobs: skip_pr_test_health_check: ${{ inputs.skip_pr_test_health_check == true }} secrets: inherit + simulator-test-cpu: + needs: [check-changes, call-gate, rust-ext-build] + if: ${{ !failure() && !cancelled() && needs.check-changes.result == 'success' && (needs.call-gate.result == 'success' || needs.call-gate.result == 'skipped') }} + uses: ./.github/workflows/_pr-test-simulator-cpu.yml + with: + check_changes: ${{ toJson(needs.check-changes.outputs) }} + caller_inputs: ${{ toJson(inputs) }} + rust_ext_artifact: ${{ needs.rust-ext-build.outputs.artifact_name }} + # No `secrets: inherit`: this hosted CPU job has no secret consumer. + # =============================================== extra-a (1-/2-gpu) =============================================== extra-a-test-1-gpu-small: needs: [check-changes, call-gate, sgl-kernel-build-wheels, rust-ext-build] @@ -249,6 +259,7 @@ jobs: call-gate, sgl-kernel-build-wheels, rust-ext-build, + simulator-test-cpu, extra-a-test-1-gpu-small, extra-a-test-1-gpu-large, extra-a-test-2-gpu-large, diff --git a/.gitignore b/.gitignore index 7c068cb49..581f2c091 100644 --- a/.gitignore +++ b/.gitignore @@ -171,6 +171,7 @@ benchmark/mmlu/data.tar benchmark/llava_bench/images benchmark/llava_bench/mme_pack *.jsonl +!tools/sglang-simulator/examples/replay/trace.jsonl tmp*.txt /tmp/ diff --git a/benchmark/simulator/bench_runner.py b/benchmark/simulator/bench_runner.py new file mode 100644 index 000000000..2186ef3a0 --- /dev/null +++ b/benchmark/simulator/bench_runner.py @@ -0,0 +1,168 @@ +"""In-process benchmark runner for SGLang Simulator.""" + +import asyncio +import atexit +import json +import os +from dataclasses import asdict +from typing import Iterator + +import numpy as np +from sglang_simulator.compat import apply_simulator_server_args +from sglang_simulator.dataset import BaseDataset, GenericRequest +from sglang_simulator.simulation.benchmark import BaseBenchmarkRunner, BenchmarkConfig +from sglang_simulator.utils.logger import get_logger + +SGLANG_SIMULATOR_OUTPUT_DIR = os.getenv( + "SGLANG_SIMULATOR_OUTPUT_DIR", "/tmp/sglang_simulator/output" +) +SIMULATION_METRICS_PATH = f"{SGLANG_SIMULATOR_OUTPUT_DIR}/metrics.json" +os.environ["SGLANG_SIMULATOR_OUTPUT_DIR"] = SGLANG_SIMULATOR_OUTPUT_DIR + +if os.getenv("SGLANG_SIMULATOR_OUTPUT_MODE") is None: + os.environ["SGLANG_SIMULATOR_OUTPUT_MODE"] = "OFFLINE" + +# Import the simulator engine only after configuring its worker environment. +from sglang_simulator.simulation.sglang.engine import ( # noqa: E402 + SGLangSimulationEngine, +) + +# SGLang must be imported after the simulator hooks are installed by engine.py. +from sglang.srt.server_args import ServerArgs # noqa: E402 + +logger = get_logger("sglang_simulator") + + +class SGLangBenchmarkRunner(BaseBenchmarkRunner): + """Run a simulator workload directly through SGLang's in-process Engine.""" + + def __init__(self, server_args: ServerArgs): + # Disable features that are unnecessary for simulation. + server_args_kwargs = asdict(server_args) + apply_simulator_server_args(server_args_kwargs) + self.engine = SGLangSimulationEngine(**server_args_kwargs) + self.server_args = self.engine.server_args + self._shutdown = False + + def flush_cache(self): + self.engine.flush_cache() + + def clear_hicache_storage(self): + self.engine.loop.run_until_complete( + self.engine.tokenizer_manager.clear_hicache_storage() + ) + + def get_request( + self, + dataset: BaseDataset, + ignore_timestamp: bool = False, + request_rate: float = float("inf"), + ) -> Iterator[tuple[GenericRequest, dict]]: + yield_delay = 0 + for req in dataset: + if ignore_timestamp: + created_time = yield_delay + yield_delay += np.random.exponential(1.0 / request_rate) + else: + created_time = req.custom_params.get("created_time", 0) + + simulation_params = { + "total_request": len(dataset), # Include the warmup requests. + "created_time": created_time, + } + + yield (req, simulation_params) + + async def async_benchmark( + self, + benchmark_config: BenchmarkConfig, + dataset: BaseDataset, + ): + await self.engine.tokenizer_manager.start_profile() + + if os.path.exists(SIMULATION_METRICS_PATH): + with open(SIMULATION_METRICS_PATH, "w") as metrics_file: + # Clear data from a previous benchmark in the same process. + pass + + tasks = [] + logger.info(f"Created {len(dataset)} request tasks.") + for req, simulation_params in self.get_request( + dataset, + ignore_timestamp=benchmark_config.ignore_request_timestamp, + request_rate=benchmark_config.request_rate, + ): + task = asyncio.create_task( + self.engine.async_generate( + prompt=req.prompt, + input_ids=req.token_ids, + sampling_params={ + "ignore_eos": True, + "max_new_tokens": req.output_length, + "custom_params": { + # Transfer simulation arguments through sampling params. + "simulation": simulation_params + }, + }, + ) + ) + tasks.append(task) + + _ = await asyncio.gather(*tasks) + + # Trigger the simulator's profile handler to flush final metrics. + await self.engine.tokenizer_manager.start_profile() + + if os.path.exists(SIMULATION_METRICS_PATH): + with open(SIMULATION_METRICS_PATH) as metrics_file: + metrics = json.load(metrics_file) + else: + logger.error( + f"Failed to load metrics from serving backend. The metrics file " + f"should be loaded from {SIMULATION_METRICS_PATH}." + ) + return None + + return metrics + + def benchmark(self, benchmark_config: BenchmarkConfig, dataset: BaseDataset): + return self.engine.loop.run_until_complete( + self.async_benchmark(benchmark_config, dataset) + ) + + def get_iteration_stats(self) -> list[dict]: + data = [] + file_path = f"{SGLANG_SIMULATOR_OUTPUT_DIR}/iteration.jsonl" + if os.path.exists(file_path): + with open(file_path) as stats_file: + line = stats_file.readline() + while line: + data.append(json.loads(line)) + line = stats_file.readline() + else: + logger.error(f"The iteration statistics data({file_path}) does not exist.") + return data + + def get_request_stats(self) -> list[dict]: + data = [] + file_path = f"{SGLANG_SIMULATOR_OUTPUT_DIR}/request.jsonl" + if os.path.exists(file_path): + with open(file_path) as stats_file: + line = stats_file.readline() + while line: + data.append(json.loads(line)) + line = stats_file.readline() + else: + logger.error(f"The request statistics data({file_path}) does not exist.") + return data + + def shutdown(self): + if self._shutdown: + return None + + logger.info("Attempting to shut down the SGLang backend engine.") + try: + return self.engine.shutdown() + finally: + self._shutdown = True + atexit.unregister(self.engine.shutdown) diff --git a/benchmark/simulator/bench_serving.py b/benchmark/simulator/bench_serving.py new file mode 100644 index 000000000..cd5fd0525 --- /dev/null +++ b/benchmark/simulator/bench_serving.py @@ -0,0 +1,280 @@ +"""SGLang serving benchmark adapter for simulator traffic. + +This script deliberately reuses SGLang's benchmark implementation and dataset +loaders. It only owns the simulator-specific parts of the protocol: + +* convert request-rate or trace timestamps into logical arrival timestamps; +* inject the internal ``sampling_params.custom_params.simulation`` metadata; +* avoid client-side pacing in OFFLINE mode; and +* display backend-produced simulator metrics when they are locally available. + +User datasets must not contain simulator metadata. +""" + +import argparse +import contextlib +import json +import os +import re +import sys +from dataclasses import fields +from pathlib import Path +from typing import AsyncGenerator, List, Optional + +import aiohttp +import numpy as np +from sglang_simulator.compat import validate_benchmark_runtime +from sglang_simulator.dataset.autobench import register_autobench_dataset + +register_autobench_dataset() + +from sglang.benchmark import serving +from sglang.benchmark.datasets.common import DatasetRow + +_ORIGINAL_AIOHTTP_REQUEST = None +_ORIGINAL_CALCULATE_METRICS = serving.calculate_metrics +_ORIGINAL_GET_REQUEST = serving.get_request +_ORIGINAL_RUN_BENCHMARK = serving.run_benchmark +_SIMULATOR_MODE = "offline" +_USE_TRACE_TIMESTAMPS = False + + +def _metrics_path() -> Path: + output_dir = Path( + os.getenv("SGLANG_SIMULATOR_OUTPUT_DIR", "/tmp/sglang_simulator/output") + ) + return output_dir / "metrics.json" + + +def _load_backend_metrics() -> Optional[dict]: + metrics_path = _metrics_path() + if not metrics_path.is_file(): + return None + return json.loads(metrics_path.read_text(encoding="utf-8")) + + +class _DurationReplacingStream: + """Keep SGLang's output format but print the simulated duration.""" + + def __init__(self, target): + self.target = target + + def write(self, text): + if "Benchmark duration (s):" in text: + metrics = _load_backend_metrics() + if metrics is not None and "duration" in metrics: + text = "{:<40} {:<10.2f}".format( + "Benchmark duration (s):", metrics["duration"] + ) + return self.target.write(text) + + def flush(self): + return self.target.flush() + + +def _set_simulation_metadata( + request: DatasetRow, *, created_time_ms: float, total_request: int +) -> None: + """Attach transient metadata without replacing dataset-specific parameters.""" + extra_request_body = dict(request.extra_request_body or {}) + extra_request_body["simulation"] = { + "created_time_ms": created_time_ms, + "total_request": total_request, + } + request.extra_request_body = extra_request_body + + +async def simulator_get_request( + input_requests: List[DatasetRow], + request_rate: float, + use_trace_timestamps: bool = False, + slowdown_factor: float = 1.0, +) -> AsyncGenerator[DatasetRow, None]: + """Generate simulator traffic while retaining official BLOCKING pacing.""" + # The benchmark may not forward --use-trace-timestamps to get_request(), + # so preserve the parsed value in this adapter. + use_trace_timestamps = use_trace_timestamps or _USE_TRACE_TIMESTAMPS + if _SIMULATOR_MODE == "blocking": + async for request in _ORIGINAL_GET_REQUEST( + input_requests, + request_rate, + use_trace_timestamps=use_trace_timestamps, + slowdown_factor=slowdown_factor, + ): + yield request + return + + total_request = len(input_requests) + if use_trace_timestamps: + if any(request.timestamp is None for request in input_requests): + raise ValueError( + "--use-trace-timestamps requires every request to have timestamp" + ) + input_requests.sort(key=lambda request: request.timestamp) + trace_start_time_ms = input_requests[0].timestamp if input_requests else 0.0 + for request in input_requests: + created_time_ms = ( + float(request.timestamp) - float(trace_start_time_ms) + ) * slowdown_factor + _set_simulation_metadata( + request, + created_time_ms=created_time_ms, + total_request=total_request, + ) + yield request + return + + created_time_ms = 0.0 + for request in input_requests: + _set_simulation_metadata( + request, + created_time_ms=created_time_ms, + total_request=total_request, + ) + yield request + if request_rate != float("inf"): + created_time_ms += np.random.exponential(1.0 / request_rate) * 1000.0 + + +def install_aiohttp_json_hijack( + *, hijack_url_regex: Optional[str] = r"/generate(?:\?.*)?$" +) -> None: + """Move transient metadata into the already-built sampling parameters.""" + global _ORIGINAL_AIOHTTP_REQUEST + if _ORIGINAL_AIOHTTP_REQUEST is not None: + return + + pattern = re.compile(hijack_url_regex) if hijack_url_regex else None + _ORIGINAL_AIOHTTP_REQUEST = aiohttp.ClientSession._request + + async def patched_request(self, method, url, **kwargs): + if pattern is None or pattern.search(str(url)): + payload = kwargs.get("json") + if isinstance(payload, dict) and "simulation" in payload: + simulation = payload.pop("simulation") + sampling_params = payload.setdefault("sampling_params", {}) + custom_params = sampling_params.setdefault("custom_params", {}) + custom_params["simulation"] = simulation + kwargs["json"] = payload + return await _ORIGINAL_AIOHTTP_REQUEST(self, method, url, **kwargs) + + aiohttp.ClientSession._request = patched_request + + +def simulator_calculate_metrics(*args, **kwargs): + """Use simulator metrics; mark unsupported client-only fields with -1.""" + client_metrics, output_lens = _ORIGINAL_CALCULATE_METRICS(*args, **kwargs) + backend_metrics = _load_backend_metrics() + if backend_metrics is None: + print( + f"Simulator metrics are not available at {_metrics_path()}; " + "showing client-side benchmark metrics." + ) + return client_metrics, output_lens + + metric_names = {field.name for field in fields(serving.BenchmarkMetrics)} + values = {name: backend_metrics.get(name, -1) for name in metric_names} + return serving.BenchmarkMetrics(**values), output_lens + + +def _replace_output_file_duration( + args: argparse.Namespace, simulated_duration: float +) -> None: + output_file = getattr(args, "output_file", None) + if not output_file: + return + path = Path(output_file) + if not path.is_file(): + return + lines = path.read_text(encoding="utf-8").splitlines() + if not lines: + return + last_result = json.loads(lines[-1]) + last_result["duration"] = simulated_duration + lines[-1] = json.dumps(last_result) + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def simulator_run_benchmark(args: argparse.Namespace): + global _USE_TRACE_TIMESTAMPS + if args.backend != "sglang": + raise ValueError( + "benchmark/simulator/bench_serving.py requires --backend sglang" + ) + if args.dataset_name == "mooncake": + raise ValueError( + "Mooncake's multi-round scheduler is not supported by the simulator " + "benchmark adapter" + ) + _USE_TRACE_TIMESTAMPS = getattr(args, "use_trace_timestamps", False) + args.profile = True + with contextlib.redirect_stdout(_DurationReplacingStream(sys.stdout)): + result = _ORIGINAL_RUN_BENCHMARK(args) + + backend_metrics = _load_backend_metrics() + if backend_metrics is not None and "duration" in backend_metrics: + simulated_duration = backend_metrics["duration"] + if isinstance(result, dict): + result["duration"] = simulated_duration + _replace_output_file_duration(args, simulated_duration) + return result + + +def _extract_simulator_args(argv: list[str]) -> tuple[str, list[str]]: + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument( + "--simulator-mode", + choices=("offline", "blocking"), + default="offline", + help=argparse.SUPPRESS, + ) + args, remaining = parser.parse_known_args(argv) + return args.simulator_mode, remaining + + +def _simulator_argument_parser(base_parser): + """Include simulator-owned datasets in SGLang's hard-coded CLI choices.""" + + class SimulatorArgumentParser(base_parser): + def add_argument(self, *name_or_flags, **kwargs): + choices = kwargs.get("choices") + if ( + "--dataset-name" in name_or_flags + and choices is not None + and "autobench" not in choices + ): + kwargs["choices"] = [*choices, "autobench"] + if "--warmup-requests" in name_or_flags: + kwargs["default"] = 0 + return super().add_argument(*name_or_flags, **kwargs) + + return SimulatorArgumentParser + + +def cli_main() -> None: + global _SIMULATOR_MODE + validate_benchmark_runtime() + if any(argument in ("-h", "--help") for argument in sys.argv[1:]): + print( + "SGLang Simulator option: " + "--simulator-mode {offline,blocking} (default: offline)\n" + ) + _SIMULATOR_MODE, remaining = _extract_simulator_args(sys.argv[1:]) + sys.argv = [sys.argv[0], *remaining] + + serving.get_request = simulator_get_request + serving.calculate_metrics = simulator_calculate_metrics + serving.run_benchmark = simulator_run_benchmark + install_aiohttp_json_hijack() + + print(f"SGLang Simulator benchmark mode: {_SIMULATOR_MODE.upper()}") + original_parser = serving.ArgumentParser + serving.ArgumentParser = _simulator_argument_parser(original_parser) + try: + serving.cli_main() + finally: + serving.ArgumentParser = original_parser + + +if __name__ == "__main__": + cli_main() diff --git a/docs/docs.json b/docs/docs.json index ceb41ec22..0ce559710 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -953,6 +953,7 @@ "docs/advanced_features/hicache_storage_runtime_attach_detach" ] }, + "docs/advanced_features/sglang_simulator", "docs/advanced_features/vlm_query", "docs/advanced_features/dp_for_multi_modal_encoder", "docs/advanced_features/cuda_graph_for_multi_modal_encoder", diff --git a/docs/docs/advanced_features/overview.mdx b/docs/docs/advanced_features/overview.mdx index 2a2dbe5f8..a7773d97c 100644 --- a/docs/docs/advanced_features/overview.mdx +++ b/docs/docs/advanced_features/overview.mdx @@ -16,5 +16,6 @@ description: Advanced configuration, optimization, and deployment features for S - [PD Disaggregation](./pd_disaggregation) - [Pipeline Parallelism](./pipeline_parallelism) - [HiCache](./hicache_best_practices) +- [SGLang Simulator](./sglang_simulator) - [Observability](./observability) - [And more…](./server_arguments) diff --git a/docs/docs/advanced_features/sglang_simulator.mdx b/docs/docs/advanced_features/sglang_simulator.mdx new file mode 100644 index 000000000..44567b92f --- /dev/null +++ b/docs/docs/advanced_features/sglang_simulator.mdx @@ -0,0 +1,88 @@ +--- +title: "SGLang Simulator" +metatags: + description: "Run SGLang scheduling and KV-cache simulations without loading model weights or executing GPU kernels." +--- + +SGLang Simulator reuses SGLang's scheduler, request lifecycle, and KV-cache implementation while replacing model forward execution with a latency predictor. Use it to compare scheduling and cache configurations on timestamped or synthetic workloads without loading model weights. + +## Supported scope + +SGLang Simulator tracks the current `main` branch and recent SGLang releases. The current integration is validated with `v0.5.16`, `v0.5.17`, `v0.5.18`, and `main`. + +The initial upstream scope uses one simulated worker with `tp_size=1`, `ep_size=1`, `dp_size=1`, and `pp_size=1`. A simulator configuration can describe a larger target system for latency prediction, but the SGLang runtime process topology remains single-worker. + +The simulator supports: + +- synthetic request rates, ShareGPT workloads, and timestamped Autobench traces; +- OFFLINE logical-time simulation and BLOCKING wall-clock replay; +- AIConfigurator, ML, and replay latency predictors; +- SGLang prefix caching and [HiCache](/docs/advanced_features/hicache); and +- serving-compatible TTFT, TPOT, ITL, throughput, and cache-hit metrics. + +## Install from the SGLang repository + +Use the simulator and SGLang source from the same monorepo checkout: + +```bash +python3 -m pip install -e tools/sglang-simulator +export PYTHONPATH="$PWD/tools/sglang-simulator/src:$PWD/python" +``` + +AIConfigurator is optional. Install the validated extra only when you use an AIConfigurator predictor: + +```bash +python3 -m pip install -e "tools/sglang-simulator[aic]" +``` + +## Start a simulator server + +Choose a fresh output directory for every run. The server owns the simulation mode and writes metrics to this directory. + +```bash +export SGLANG_USE_CPU_ENGINE=1 +export CUDA_VISIBLE_DEVICES="" +export SGLANG_SIMULATOR_OUTPUT_MODE=OFFLINE +export SGLANG_SIMULATOR_OUTPUT_DIR=/tmp/sglang-simulator-quickstart + +python3 -m sglang_simulator.simulation.sglang.launch_server \ + --model-path tools/sglang-simulator/test/assets/qwen3-8b \ + --tokenizer-path tools/sglang-simulator/examples/assets/tokenizer \ + --sim-config-path tools/sglang-simulator/examples/sim_configs/replay.json \ + --port 30000 +``` + +`OFFLINE` advances the simulator's logical clock without sleeping. `BLOCKING` also sleeps for predicted forward and cache-load latency, which is useful when a client must observe simulated wall-clock pacing. + +## Send a workload + +In another terminal, export the same output directory and run the simulator-aware serving benchmark from the repository root: + +```bash +export PYTHONPATH="$PWD/tools/sglang-simulator/src:$PWD/python" +export SGLANG_SIMULATOR_OUTPUT_DIR=/tmp/sglang-simulator-quickstart + +python3 benchmark/simulator/bench_serving.py \ + --simulator-mode offline \ + --backend sglang \ + --base-url http://127.0.0.1:30000 \ + --model tools/sglang-simulator/test/assets/qwen3-8b \ + --tokenizer tools/sglang-simulator/examples/assets/tokenizer \ + --dataset-name sharegpt \ + --dataset-path tools/sglang-simulator/examples/workloads/sharegpt-example.json \ + --sharegpt-output-len 4 \ + --num-prompts 3 \ + --output-file /tmp/sglang-simulator-quickstart/benchmark.json +``` + +The benchmark injects logical arrival metadata into each request and displays the server-side simulator metrics. For timestamped traffic, use the simulator-owned Autobench JSONL format and add `--use-trace-timestamps`. + +## Read the results + +The output directory contains: + +- `metrics.json`: aggregate latency, throughput, and cache metrics; +- `request.jsonl`: per-request timing and cache information; and +- `iteration.jsonl`: scheduler batch composition and predicted iteration latency. + +Use a unique output directory for each run so metrics from separate experiments are not mixed. See the [SGLang Simulator source README](https://github.com/sgl-project/sglang/tree/main/tools/sglang-simulator) for simulator configuration fields, predictor examples, and maintained tests. diff --git a/tools/sglang-simulator/README.md b/tools/sglang-simulator/README.md new file mode 100644 index 000000000..2b453e7b4 --- /dev/null +++ b/tools/sglang-simulator/README.md @@ -0,0 +1,258 @@ +# SGLang Simulator + +SGLang Simulator reuses SGLang's scheduler and cache implementation while +replacing model forward execution with a latency predictor. It supports +timestamped trace replay, synthetic workloads, hierarchical cache simulation, +and serving-compatible metrics without loading model weights. + +See the [SGLang Simulator advanced-feature guide](../../docs/docs/advanced_features/sglang_simulator.mdx) +for the user-facing setup and serving workflow. + +## Compatibility + +SGLang Simulator tracks the current SGLang `main` branch and maintains compatibility +with recent SGLang releases. The current integration is validated with `v0.5.16`, +`v0.5.17`, `v0.5.18`, and `main`. Compatibility code uses API and capability +checks instead of branching on version numbers. + +## Requirements + +- A compatible SGLang checkout. The simulator uses the SGLang source from the + same monorepo checkout. +- A local model directory containing model configuration files. Tokenizer files + are also required unless tokenizer initialization is disabled. +- Predictor data for AIConfigurator, ML, or replay mode. + +Use an official SGLang image matching the checkout when validating GPU and +runtime compatibility. + +## Installation + +From the SGLang repository: + +```bash +pip install -e tools/sglang-simulator +``` + +The simulator does not install or pin a second `sglang` package. Run it from a +checkout whose `python/sglang` package is available on `PYTHONPATH`, or from a +matching official SGLang image. + +AIConfigurator is optional. Install it separately when using the +`aiconfigurator` predictor. The `aic` extra pins AIConfigurator to the exact release +validated with the simulator so upstream API changes cannot silently alter an +installation. In a clean virtual environment, install the extra with: + +```bash +pip install -e "tools/sglang-simulator[aic]" +``` + +In an existing SGLang image, install the same pin without dependency resolution +to avoid replacing its NumPy/CUDA stack: + +```bash +pip install --no-deps "aiconfigurator==0.10.0" +``` + +Upgrade this pin only after rerunning the AIC predictor and compatibility tests. + +## Quick start + +The maintained tests define the supported first-version scope: + +- [`test/test_simulation_sglang_runner.py`](test/test_simulation_sglang_runner.py): + direct Python use of the repository-level + [`SGLangBenchmarkRunner`](../../benchmark/simulator/bench_runner.py); +- [`test/test_simulation_sglang_serving.py`](test/test_simulation_sglang_serving.py): + server plus benchmark-client use through the HTTP serving path with AIC, ML, + and replay predictors and ShareGPT or timestamped traffic; +- [`test/test_simulation_offline_blocking.py`](test/test_simulation_offline_blocking.py): + equivalent logical results in `OFFLINE` and `BLOCKING` modes; +- [`test/test_simulation_cache_hit_ratio.py`](test/test_simulation_cache_hit_ratio.py): + reusable-prefix accounting and cache-tier hit metrics across repeated runs. + +From `tools/sglang-simulator`: + +```bash +python3 -m pytest -q test/test_simulation_sglang_runner.py +python3 -m pytest -q test/test_simulation_sglang_serving.py +``` + +Read these tests as the minimal maintained examples for constructing a dataset, +running a benchmark, starting a simulator server, sending programmatic, ShareGPT, +or timestamped traffic, comparing execution modes, and collecting request, +latency, throughput, and prefix-cache metrics. + +## Serving mode + +Choose a fresh output directory and export it in the server terminal before +starting the server: + +```bash +export SGLANG_USE_CPU_ENGINE=1 +export CUDA_VISIBLE_DEVICES="" +export SGLANG_SIMULATOR_OUTPUT_MODE=OFFLINE +export SIMULATOR_OUTPUT_DIR=/tmp/sglang-simulator-serving-001 +test ! -e "$SIMULATOR_OUTPUT_DIR" +export SGLANG_SIMULATOR_OUTPUT_DIR="$SIMULATOR_OUTPUT_DIR" + +python3 -m sglang_simulator.simulation.sglang.launch_server \ + --model-path /absolute/path/to/model \ + --sim-config-path /absolute/path/to/simulator.json \ + --port 30000 +``` + +In the benchmark terminal, export the same output directory before sending +timestamped traffic with the simulator-aware benchmark adapter: + +```bash +cd /path/to/sglang +export SIMULATOR_OUTPUT_DIR=/tmp/sglang-simulator-serving-001 +export SGLANG_SIMULATOR_OUTPUT_DIR="$SIMULATOR_OUTPUT_DIR" + +python3 benchmark/simulator/bench_serving.py \ + --simulator-mode offline \ + --backend sglang \ + --base-url http://127.0.0.1:30000 \ + --model /absolute/path/to/model \ + --dataset-name autobench \ + --dataset-path /absolute/path/to/trace.jsonl \ + --use-trace-timestamps \ + --num-prompts 100 \ + --profile \ + --output-file "$SIMULATOR_OUTPUT_DIR/benchmark.json" +``` + +The server and benchmark are separate processes, so exporting +`SGLANG_SIMULATOR_OUTPUT_DIR` in the server terminal does not configure the +benchmark terminal. The benchmark adapter reads `metrics.json` from this path +after profiling and uses those server-side logical-time metrics for its serving +table and output file. If the benchmark points at another directory, it may show +unrelated stale metrics or client wall-clock values. Use the same fresh path in +both terminals for every run. + +The simulator always runs the SGLang runtime with `tp_size=ep_size=dp_size=pp_size=1` +and both attention/decode context-parallel sizes set to `1`. Parallel CLI options +accepted by SGLang are therefore ignored by this simulator entry point. This keeps +simulator-only CPU work single-process; it does not change the modeled deployment. +Set the real deployment topology under `scheduler` in `--sim-config-path`. That +topology drives predictor and cache-resource modeling without launching physical +parallel workers. + +Other server options are normal SGLang command-line arguments. For direct Python +integration, see +[`test_simulation_sglang_runner.py`](test/test_simulation_sglang_runner.py); for the +process/HTTP path, see +[`test_simulation_sglang_serving.py`](test/test_simulation_sglang_serving.py). + +## Simulation modes + +| Mode | Behavior | +|---|---| +| `OFFLINE` | Advances the simulator's logical clock without sleeping. | +| `BLOCKING` | Sleeps for predicted forward and visible L2-to-L1 load latency. | + +Use server-side simulator metrics for comparisons. Client wall-clock duration is +not the simulated timeline in `OFFLINE` mode. When using the benchmark adapter, +make sure its `SGLANG_SIMULATOR_OUTPUT_DIR` matches the server's output directory +so the printed table and `benchmark.json` are sourced from the current run's +`metrics.json`. + +## Configuration + +A simulator configuration has three sections: + +```json +{ + "platform": { + "accelerator": {"name": "h20_sxm"}, + "disk_read_bandwidth_gb": 8, + "disk_write_bandwidth_gb": 8, + "memory_read_bandwidth_gb": 64, + "memory_write_bandwidth_gb": 64, + "num_device_per_node": 1 + }, + "predictor": { + "name": "replay", + "database_path": "/absolute/path/to/replay_table.json" + }, + "scheduler": { + "tp_size": 4, + "ep_size": 4, + "dp_size": 1, + "pp_size": 1, + "cp_size": 1, + "cp_style": "none", + "data_type": "BF16", + "kv_cache_data_type": "BF16", + "backend_name": "sglang" + } +} +``` + +- `platform` describes the simulated accelerator and storage bandwidth. +- `predictor` selects forward-latency prediction. +- `scheduler` describes the real target deployment topology and backend metadata. + `tp_size`, `ep_size`, `dp_size`, `pp_size`, and `cp_size` are modeled values; + they do not launch physical workers. For AIConfigurator, `tp_size` is converted + to attention TP after removing modeled DP and CP, while `cp_size` is passed as + AIConfigurator context parallelism. `cp_style` uses the AIConfigurator values + such as `none`, `allgather`, `ulysses`, or `ring`. Decode-only `dcp_size` has no + separate AIConfigurator field and is not modeled yet. + +### Prefix-cache accuracy + +Prefix-cache hit accuracy is highly sensitive to `max_total_tokens`. It controls +the simulated device KV-cache capacity and participates in hierarchical host-cache +sizing, so a mismatch changes eviction timing and device, host, and storage hit +attribution. For deployment-faithful results, copy `max_total_num_tokens=N` from +the real SGLang server startup log and launch the simulator with +`--max-total-tokens N`. Avoid relying on a separately estimated capacity when +comparing the simulator with production traces. + +Supported predictors: + +| Predictor | Purpose | +|---|---| +| `aiconfigurator` | Operator and module performance-database estimation. | +| `ml` | A trained sklearn-compatible 18-feature latency model. | +| `replay` | Exact or nearest-neighbor batch-composition replay. | + +Relative predictor paths are resolved from the simulator configuration location. +Environment variables in paths use `${NAME}` syntax. + +## Workload formats + +The Autobench trace format uses timestamps in milliseconds: + +```json +{"prompt":[1,2,3],"prompt_len":3,"output_len":1,"timestamp":200} +``` + +Random and ShareGPT workloads are also supported by the runner API and serving +benchmark paths. + +## Validation + +Run the CPU compatibility and unit tests from the repository root: + +```bash +pip install -e tools/sglang-simulator +python3 -m pytest -q tools/sglang-simulator/test/test_simulation_sglang_runner.py +python3 -m pytest -q tools/sglang-simulator/test/test_simulation_sglang_serving.py +``` + +Run the two files as separate pytest commands because the runner test installs +process-global simulator hooks and state. + +Run repository checks before submitting: + +```bash +git ls-files -z tools/sglang-simulator | \ + xargs -0 env SKIP=no-commit-to-branch pre-commit run --files +``` + +Runtime changes should also be validated in a matching official SGLang image +with both `OFFLINE` and `BLOCKING` modes. Predictor changes should report +step-level error, and scheduler or cache changes should compare request latency, +throughput, and prefix-cache reuse against measured traces. diff --git a/tools/sglang-simulator/examples/README.md b/tools/sglang-simulator/examples/README.md new file mode 100644 index 000000000..061e3fbb3 --- /dev/null +++ b/tools/sglang-simulator/examples/README.md @@ -0,0 +1,50 @@ +# SGLang Simulator examples + +The example assets are organized by purpose: + +- `sim_configs/`: standalone AIC SOL, AIC SILICON, ML, and replay simulator configs; +- `assets/`: the small illustrative ML model, replay table, and test tokenizer; +- `workloads/`: ShareGPT and timestamped simulator/Autobench workload examples; + +The ML model is an illustrative constant-latency sklearn model, not a calibrated +hardware predictor. Rebuild it and the tokenizer with: + +```bash +python3 examples/build_example_assets.py +``` + +Only load pickle/joblib assets from sources you trust. + +For maintained direct-run and serving examples, see +[`test_simulation_sglang_runner.py`](../test/test_simulation_sglang_runner.py) and +[`test_simulation_sglang_serving.py`](../test/test_simulation_sglang_serving.py). + +Start a server with any example config: + +```bash +python3 -m sglang_simulator.simulation.sglang.launch_server \ + --model-path /path/to/model \ + --sim-config-path examples/sim_configs/aic_sol.json \ + --port 30000 +``` + +Run a ShareGPT workload with at least four output tokens so decode and TPOT are +measured: + +```bash +cd /path/to/sglang +python3 benchmark/simulator/bench_serving.py \ + --simulator-mode=offline \ + --backend=sglang \ + --base-url=http://127.0.0.1:30000 \ + --model=/path/to/model \ + --tokenizer=/path/to/model \ + --dataset-name=sharegpt \ + --dataset-path=examples/workloads/sharegpt-example.json \ + --sharegpt-output-len=4 \ + --num-prompts=3 \ + --profile +``` + +The timestamp trace uses the simulator-owned Autobench JSONL contract. Its +`timestamp` values are request-arrival times in milliseconds. diff --git a/tools/sglang-simulator/examples/assets/model.pkl b/tools/sglang-simulator/examples/assets/model.pkl new file mode 100644 index 0000000000000000000000000000000000000000..c56c7cdd3a7c054a328b52eb8e897335d80b63b7 GIT binary patch literal 734 zcmY*X&1w`u5Y9vwCkrl{M9IMyh?m9Oau9+j2o4!SE}lZuJ=43}+JE+s?g)ZjG(z93 zJVah6Z<2dpt=a52Yae>*=c})(znbq4uUNLVI^W^2%r*s17jC-=3VY)0^AtyOe*57!s00nMRsSISmm> zh|mOzy)@Ls#9(9a=w@lB;~%MZ1Ht!v^mj5?QWR;}e zLmI5AAU)DGZLv6U8u&T29NQGxPxOqdu{dR~(XK;-nfZQaA$ogy48Jdb!8sfj4^s9L z%Cda7#Iw4JtpF?ysg&jmQHol)OwaBoUG*bWikVhi2f~`bJrv`}wwKGh>)V@Cc6M@g zu~t2e$y#gjnL96x#&M1&$CDQjS;{-T-~z=SG8(Em@!~5AG4xRpz6e~`=qS3_a8TSX zleaKSeZU2^pw2|dy}8@-kS@Kq5%ah}$Faw_V None: + model = DummyRegressor(strategy="constant", constant=0.001) + model.fit(np.zeros((1, len(MLTimePredictor.FEATURE_NAMES))), [0.001]) + joblib.dump( + {"model": model, "features": MLTimePredictor.FEATURE_NAMES}, + ASSETS / "model.pkl", + ) + + +def build_tokenizer() -> None: + tokenizer = Tokenizer( + WordLevel( + { + "[UNK]": 0, + "prefix": 1, + "caching": 2, + "latency": 3, + "decode": 4, + "token": 5, + }, + unk_token="[UNK]", + ) + ) + tokenizer.pre_tokenizer = Whitespace() + PreTrainedTokenizerFast( + tokenizer_object=tokenizer, + unk_token="[UNK]", + ).save_pretrained(ASSETS / "tokenizer") + + +def main() -> None: + ASSETS.mkdir(parents=True, exist_ok=True) + build_ml_model() + build_tokenizer() + + +if __name__ == "__main__": + main() diff --git a/tools/sglang-simulator/examples/sim_configs/aic_silicon.json b/tools/sglang-simulator/examples/sim_configs/aic_silicon.json new file mode 100644 index 000000000..69aaba249 --- /dev/null +++ b/tools/sglang-simulator/examples/sim_configs/aic_silicon.json @@ -0,0 +1,21 @@ +{ + "platform": { + "accelerator": {"name": "a100_sxm", "hbm_capacity_gb": 80}, + "disk_read_bandwidth_gb": 8, + "disk_write_bandwidth_gb": 8, + "memory_read_bandwidth_gb": 64, + "memory_write_bandwidth_gb": 64, + "num_device_per_node": 8 + }, + "predictor": { + "name": "aiconfigurator", + "database_mode": "SILICON" + }, + "scheduler": { + "tp_size": 1, + "ep_size": 1, + "dp_size": 1, + "backend_name": "sglang", + "backend_version": "0.5.9" + } +} diff --git a/tools/sglang-simulator/examples/sim_configs/aic_sol.json b/tools/sglang-simulator/examples/sim_configs/aic_sol.json new file mode 100644 index 000000000..1161a6391 --- /dev/null +++ b/tools/sglang-simulator/examples/sim_configs/aic_sol.json @@ -0,0 +1,21 @@ +{ + "platform": { + "accelerator": {"name": "a100_sxm", "hbm_capacity_gb": 80}, + "disk_read_bandwidth_gb": 8, + "disk_write_bandwidth_gb": 8, + "memory_read_bandwidth_gb": 64, + "memory_write_bandwidth_gb": 64, + "num_device_per_node": 8 + }, + "predictor": { + "name": "aiconfigurator", + "database_mode": "SOL" + }, + "scheduler": { + "tp_size": 1, + "ep_size": 1, + "dp_size": 1, + "backend_name": "sglang", + "backend_version": "0.5.9" + } +} diff --git a/tools/sglang-simulator/examples/sim_configs/ml.json b/tools/sglang-simulator/examples/sim_configs/ml.json new file mode 100644 index 000000000..44f42b984 --- /dev/null +++ b/tools/sglang-simulator/examples/sim_configs/ml.json @@ -0,0 +1,22 @@ +{ + "platform": { + "accelerator": {"name": "a100_sxm", "hbm_capacity_gb": 80}, + "disk_read_bandwidth_gb": 8, + "disk_write_bandwidth_gb": 8, + "memory_read_bandwidth_gb": 64, + "memory_write_bandwidth_gb": 64, + "num_device_per_node": 8 + }, + "predictor": { + "name": "ml", + "database_path": "../assets/model.pkl", + "latency_scale": 1.0 + }, + "scheduler": { + "tp_size": 1, + "ep_size": 1, + "dp_size": 1, + "backend_name": "sglang", + "backend_version": "0.5.9" + } +} diff --git a/tools/sglang-simulator/examples/sim_configs/replay.json b/tools/sglang-simulator/examples/sim_configs/replay.json new file mode 100644 index 000000000..58685ce9b --- /dev/null +++ b/tools/sglang-simulator/examples/sim_configs/replay.json @@ -0,0 +1,23 @@ +{ + "platform": { + "accelerator": {"name": "a100_sxm", "hbm_capacity_gb": 80}, + "disk_read_bandwidth_gb": 8, + "disk_write_bandwidth_gb": 8, + "memory_read_bandwidth_gb": 64, + "memory_write_bandwidth_gb": 64, + "num_device_per_node": 8 + }, + "predictor": { + "name": "replay", + "database_path": "../assets/replay_table.json", + "miss_strategy": "knn", + "miss_knn_k": 1 + }, + "scheduler": { + "tp_size": 1, + "ep_size": 1, + "dp_size": 1, + "backend_name": "sglang", + "backend_version": "0.5.9" + } +} diff --git a/tools/sglang-simulator/examples/workloads/sharegpt-example.json b/tools/sglang-simulator/examples/workloads/sharegpt-example.json new file mode 100644 index 000000000..b499ff328 --- /dev/null +++ b/tools/sglang-simulator/examples/workloads/sharegpt-example.json @@ -0,0 +1,38 @@ +[ + { + "conversations": [ + { + "from": "human", + "value": "Explain prefix caching in one concise sentence." + }, + { + "from": "gpt", + "value": "Prefix caching reuses KV states shared by prompt prefixes." + } + ] + }, + { + "conversations": [ + { + "from": "human", + "value": "What does time to first token measure?" + }, + { + "from": "gpt", + "value": "It measures latency from request arrival to the first generated token." + } + ] + }, + { + "conversations": [ + { + "from": "human", + "value": "Why does decode latency matter for serving?" + }, + { + "from": "gpt", + "value": "Decode latency determines the cadence at which later tokens reach the user." + } + ] + } +] diff --git a/tools/sglang-simulator/examples/workloads/timestamp-trace-example.jsonl b/tools/sglang-simulator/examples/workloads/timestamp-trace-example.jsonl new file mode 100644 index 000000000..d7343d511 --- /dev/null +++ b/tools/sglang-simulator/examples/workloads/timestamp-trace-example.jsonl @@ -0,0 +1,3 @@ +{"prompt":[100,101,102,103],"prompt_len":4,"output_len":4,"timestamp":0} +{"prompt":[100,101,102,104],"prompt_len":4,"output_len":4,"timestamp":75} +{"prompt":[200,201,202,203],"prompt_len":4,"output_len":4,"timestamp":250} diff --git a/tools/sglang-simulator/pyproject.toml b/tools/sglang-simulator/pyproject.toml new file mode 100644 index 000000000..949bf3dd7 --- /dev/null +++ b/tools/sglang-simulator/pyproject.toml @@ -0,0 +1,18 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "sglang-simulator" +dynamic = ["version"] +description = "A simulation benchmark tool for sglang" +dependencies = [ + "numpy", + "scikit-learn", + "joblib", +] + +[project.optional-dependencies] +aic = [ + "aiconfigurator==0.10.0" +] diff --git a/tools/sglang-simulator/setup.py b/tools/sglang-simulator/setup.py new file mode 100644 index 000000000..bc7ecf81a --- /dev/null +++ b/tools/sglang-simulator/setup.py @@ -0,0 +1,21 @@ +from setuptools import find_packages, setup + + +def get_version(): + version = "0.1.0" + with open("src/sglang_simulator/__init__.py") as f: + for line in f: + if line.startswith("__version__"): + version = line.split("=")[1].strip(' \n"') + return version + + +setup( + name="sglang-simulator", + version=get_version(), + url="https://github.com/sgl-project/sglang.git", + description="A High-Fidelity LLM inference simulator for SGLang", + packages=find_packages(where="src"), + package_dir={"": "src"}, + py_modules=["usercustomize"], +) diff --git a/tools/sglang-simulator/src/sglang_simulator/__init__.py b/tools/sglang-simulator/src/sglang_simulator/__init__.py new file mode 100644 index 000000000..3dc1f76bc --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/__init__.py @@ -0,0 +1 @@ +__version__ = "0.1.0" diff --git a/tools/sglang-simulator/src/sglang_simulator/compat.py b/tools/sglang-simulator/src/sglang_simulator/compat.py new file mode 100644 index 000000000..a344b07d5 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/compat.py @@ -0,0 +1,99 @@ +"""Early compatibility checks for the SGLang surfaces used by the simulator.""" + +import inspect +from importlib import metadata + +SIMULATOR_SERVER_ARG_OVERRIDES = { + # The simulator models the target deployment topology separately through + # sim_config.scheduler. Keep the SGLang runtime single-process so host-side + # simulator work is not multiplied by the modeled parallel world size. + "tp_size": 1, + "ep_size": 1, + "dp_size": 1, + "pp_size": 1, + "attn_cp_size": 1, + "dcp_size": 1, + "disable_overlap_schedule": True, + "disable_cuda_graph": True, + "attention_backend": "torch_native", + "prefill_attention_backend": "torch_native", + "decode_attention_backend": "torch_native", +} + + +class SGLangCompatibilityError(RuntimeError): + pass + + +def _sglang_version() -> str: + try: + return metadata.version("sglang") + except metadata.PackageNotFoundError: + return "source-checkout" + + +def _require_parameters(function, required: set[str], surface: str) -> None: + parameters = set(inspect.signature(function).parameters) + missing = required - parameters + if missing: + raise SGLangCompatibilityError( + f"SGLang {_sglang_version()} is missing {surface} parameters: " + f"{', '.join(sorted(missing))}. The simulator must be adapted to " + "this SGLang revision before it can run." + ) + + +def apply_simulator_server_args(target) -> None: + """Apply simulator-owned values before constructing the final ServerArgs.""" + if isinstance(target, dict): + target.update(SIMULATOR_SERVER_ARG_OVERRIDES) + return + + for name, value in SIMULATOR_SERVER_ARG_OVERRIDES.items(): + setattr(target, name, value) + + +def validate_simulator_server_args(server_args) -> None: + """Fail early if a process bypassed a simulator-owned launch entry point.""" + mismatches = [ + f"{name}={getattr(server_args, name, None)!r} (expected {expected!r})" + for name, expected in SIMULATOR_SERVER_ARG_OVERRIDES.items() + if getattr(server_args, name, None) != expected + ] + if mismatches: + raise SGLangCompatibilityError( + "SGLang Simulator server arguments were not prepared by a supported " + f"entry point: {', '.join(mismatches)}" + ) + + +def validate_launch_runtime() -> None: + from sglang.srt.entrypoints.http_server import launch_server + + _require_parameters( + launch_server, + {"server_args", "run_scheduler_process_func", "run_detokenizer_process_func"}, + "launch_server", + ) + + +def validate_benchmark_runtime() -> None: + from sglang.benchmark import serving + + missing = [ + name + for name in ( + "BenchmarkMetrics", + "calculate_metrics", + "cli_main", + "get_request", + "run_benchmark", + ) + if not hasattr(serving, name) + ] + if missing: + raise SGLangCompatibilityError( + f"SGLang {_sglang_version()} is missing benchmark surfaces: " + f"{', '.join(missing)}. The simulator must be adapted to this " + "SGLang revision before it can run." + ) diff --git a/tools/sglang-simulator/src/sglang_simulator/dataset/__init__.py b/tools/sglang-simulator/src/sglang_simulator/dataset/__init__.py new file mode 100644 index 000000000..2c93e1666 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/dataset/__init__.py @@ -0,0 +1,35 @@ +from sglang_simulator.dataset.base_dataset import ( + BaseDataset, + GenericRequest, + SimpleDataset, +) +from sglang_simulator.dataset.dataset_args import DatasetArgs +from sglang_simulator.dataset.random import RandomDataset, RandomIDsDataset +from transformers import PreTrainedTokenizer + +dataset_registry: dict[str, BaseDataset] = { + "random": RandomDataset, + "random_ids": RandomIDsDataset, +} + + +def get_dataset( + dataset_args: DatasetArgs, tokenizer: PreTrainedTokenizer | None = None +) -> BaseDataset: + if dataset_args.name not in dataset_registry: + raise ValueError(f"unknown dataset name: {dataset_args.name}") + + dataset: BaseDataset = dataset_registry[dataset_args.name]( + args=dataset_args, tokenizer=tokenizer + ) + + return dataset + + +__all__ = ( + "DatasetArgs", + "BaseDataset", + "SimpleDataset", + "GenericRequest", + "get_dataset", +) diff --git a/tools/sglang-simulator/src/sglang_simulator/dataset/autobench.py b/tools/sglang-simulator/src/sglang_simulator/dataset/autobench.py new file mode 100644 index 000000000..ef81dd5e1 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/dataset/autobench.py @@ -0,0 +1,242 @@ +"""Simulator-owned loader for timestamped Autobench JSONL traces. + +The trace format is a public SGLang Simulator input contract. Keep its parser +here instead of importing SGLang's benchmark-internal Autobench module, which +may be moved or removed independently of the simulator. +""" + +import json +from argparse import Namespace +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +import numpy as np +from transformers import PreTrainedTokenizerBase + +from sglang.benchmark.datasets.common import BaseDataset, DatasetRow + +_RESERVED_FIELDS = { + "prompt", + "messages", + "prompt_origin", + "output_len", + "max_tokens", + "max_completion_tokens", + "completion_tokens", + "prompt_len", + "text_prompt_len", + "vision_prompt_len", + "image_data", + "timestamp", + "routing_key", + "metadata", + "extra_request_body", + "param_send", +} + + +def _load_json_if_needed(value: Any) -> Any: + if not isinstance(value, str): + return value + value = value.strip() + if not value or value[0] not in "[{": + return value + try: + return json.loads(value) + except json.JSONDecodeError: + return value + + +def _normalize_messages(messages: Any) -> Optional[list[dict[str, Any]]]: + messages = _load_json_if_needed(messages) + if not isinstance(messages, list) or not messages: + return None + if not all(isinstance(message, dict) for message in messages): + return None + + normalized = [] + for message in messages: + if "role" not in message or message.get("content") is None: + return None + normalized.append({"role": message["role"], "content": message["content"]}) + return normalized + + +def _normalize_prompt(row: dict[str, Any]) -> tuple[Any, str]: + for key in ("messages", "prompt_origin"): + normalized = _normalize_messages(row.get(key)) + if normalized is not None: + return normalized, "messages" + + prompt = _load_json_if_needed(row.get("prompt")) + if isinstance(prompt, list) and prompt: + if isinstance(prompt[0], dict): + normalized = _normalize_messages(prompt) + if normalized is not None: + return normalized, "messages" + if all(isinstance(item, int) for item in prompt): + return prompt, "token_ids" + if all(isinstance(item, str) for item in prompt): + return prompt, "multi_turn" + if all( + isinstance(turn, list) + and turn + and all( + isinstance(message, dict) and "role" in message and "content" in message + for message in turn + ) + for turn in prompt + ): + return prompt, "multi_turn" + if isinstance(prompt, str) and prompt: + return prompt, "prompt" + + if isinstance(row.get("content"), list): + turns = [str(item) for item in row["content"]] + if len(turns) % 2 == 0: + turns = turns[:-1] + messages = [] + if row.get("system"): + messages.append({"role": "system", "content": str(row["system"])}) + messages.extend( + { + "role": "user" if index % 2 == 0 else "assistant", + "content": turn, + } + for index, turn in enumerate(turns) + ) + if messages: + return messages, "messages" + + raise ValueError("Unsupported Autobench row: missing prompt/messages") + + +def _prompt_lengths( + row: dict[str, Any], + prompt: Any, + prompt_kind: str, + tokenizer: Optional[PreTrainedTokenizerBase], +) -> tuple[int, int, int]: + if row.get("prompt_len") is not None: + prompt_len = int(row["prompt_len"]) + return ( + prompt_len, + int(row.get("text_prompt_len", prompt_len)), + int(row.get("vision_prompt_len", 0)), + ) + if prompt_kind == "token_ids": + return len(prompt), len(prompt), 0 + if tokenizer is None: + raise ValueError("Autobench rows without prompt_len require a tokenizer") + if prompt_kind == "messages": + prompt_len = len( + tokenizer.apply_chat_template( + prompt, tokenize=True, add_generation_prompt=True + ) + ) + return prompt_len, prompt_len, 0 + if prompt_kind == "prompt": + prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False)) + return prompt_len, prompt_len, 0 + return 0, 0, 0 + + +def _extra_request_body(row: dict[str, Any]) -> dict[str, Any]: + extra = {} + param_send = _load_json_if_needed(row.get("param_send")) + if isinstance(param_send, dict): + extra.update(param_send) + extra.update( + {key: value for key, value in row.items() if key not in _RESERVED_FIELDS} + ) + explicit = _load_json_if_needed(row.get("extra_request_body")) + if isinstance(explicit, dict): + extra.update(explicit) + return extra + + +def sample_autobench_requests( + dataset_path: str, + num_requests: int, + tokenizer: Optional[PreTrainedTokenizerBase], + fixed_output_len: Optional[int] = None, +) -> list[DatasetRow]: + dataset = [] + with Path(dataset_path).open(encoding="utf-8") as file: + for line_number, line in enumerate(file, start=1): + if num_requests > 0 and len(dataset) >= num_requests: + break + if not line.strip(): + continue + try: + row = json.loads(line) + prompt, prompt_kind = _normalize_prompt(row) + prompt_len, text_prompt_len, vision_prompt_len = _prompt_lengths( + row, prompt, prompt_kind, tokenizer + ) + except (TypeError, ValueError, json.JSONDecodeError) as error: + raise ValueError( + f"Invalid Autobench row {line_number} in {dataset_path}: {error}" + ) from error + + output_len = fixed_output_len + for key in ( + "output_len", + "max_tokens", + "max_completion_tokens", + "completion_tokens", + ): + output_len = output_len or row.get(key) + dataset.append( + DatasetRow( + prompt=prompt, + prompt_len=prompt_len, + output_len=int(output_len or 256), + text_prompt_len=text_prompt_len, + vision_prompt_len=vision_prompt_len, + image_data=row.get("image_data"), + timestamp=row.get("timestamp"), + routing_key=row.get("routing_key"), + extra_request_body=_extra_request_body(row), + ) + ) + + print(f"Loaded {len(dataset)} Autobench requests") + print(f"#Input tokens: {np.sum([row.prompt_len for row in dataset])}") + print(f"#Output tokens: {np.sum([row.output_len for row in dataset])}") + return dataset + + +@dataclass +class AutoBenchmarkDataset(BaseDataset): + dataset_path: str + num_requests: int + fixed_output_len: Optional[int] + + @classmethod + def from_args(cls, args: Namespace) -> "AutoBenchmarkDataset": + return cls( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + fixed_output_len=getattr(args, "sharegpt_output_len", None), + ) + + def load( + self, + tokenizer: PreTrainedTokenizerBase, + model_id: Optional[str] = None, + ) -> list[DatasetRow]: + return sample_autobench_requests( + self.dataset_path, + self.num_requests, + tokenizer, + self.fixed_output_len, + ) + + +def register_autobench_dataset() -> None: + """Register the simulator trace contract with SGLang's serving benchmark.""" + from sglang.benchmark import datasets + + datasets.DATASET_MAPPING["autobench"] = AutoBenchmarkDataset diff --git a/tools/sglang-simulator/src/sglang_simulator/dataset/base_dataset.py b/tools/sglang-simulator/src/sglang_simulator/dataset/base_dataset.py new file mode 100644 index 000000000..1cc3ad317 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/dataset/base_dataset.py @@ -0,0 +1,72 @@ +from dataclasses import dataclass, field +from typing import Optional, overload + +from sglang_simulator.dataset.dataset_args import DatasetArgs +from transformers import PreTrainedTokenizerBase + + +@dataclass(slots=True) +class GenericRequest: + prompt: Optional[str] = None + token_ids: Optional[list[int]] = None + input_length: int = -1 + output_length: int = -1 + custom_params: dict = field(default_factory=dict) + + def __post_init__(self): + if self.prompt is None and self.token_ids is None: + raise ValueError("Invalid Request") + + +class BaseDataset: + def __init__(self, tokenizer: PreTrainedTokenizerBase, args: DatasetArgs): + self.tokenizer: PreTrainedTokenizerBase = tokenizer + self.args = args + self._name = "" + + @overload + def __getitem__(self, index: int) -> GenericRequest: ... + + @overload + def __getitem__(self, index: slice) -> list[GenericRequest]: ... + + def __getitem__(self, index): + """Get item(s) by index or slice. Delegates to _get_single_item for single items.""" + if isinstance(index, slice): + start, stop, step = index.indices(len(self)) + return [self[i] for i in range(start, stop, step)] + if index >= len(self): + raise IndexError + return self._get_single_item(index) + + def _get_single_item(self, index: int) -> GenericRequest: + raise NotImplementedError + + def __len__(self) -> int: + raise NotImplementedError + + @property + def name(self): + if self._name: + return self._name + else: + return self.__class__.__name__ + + +class SimpleDataset(BaseDataset): + def __init__( + self, tokenizer=None, args=None, reqs: list[GenericRequest] | None = None + ): + super().__init__(tokenizer, args) + self.data: list[GenericRequest] = [] + if reqs is not None: + self.data.extend(reqs) + + def add_request(self, req: GenericRequest): + self.data.append(req) + + def _get_single_item(self, index: int) -> GenericRequest: + return self.data[index] + + def __len__(self): + return len(self.data) diff --git a/tools/sglang-simulator/src/sglang_simulator/dataset/dataset_args.py b/tools/sglang-simulator/src/sglang_simulator/dataset/dataset_args.py new file mode 100644 index 000000000..302d8c0af --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/dataset/dataset_args.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass + + +@dataclass +class DatasetArgs: + name: str = "" + filepath: str = "" + num_prompts: int = -1 + min_input_len: int = -1 + max_input_len: int = -1 + min_output_len: int = -1 + max_output_len: int = -1 + + @property + def mean_input_length(self): + return (self.min_input_len + self.max_input_len) // 2 + + @property + def mean_output_length(self): + return (self.min_output_len + self.max_output_len) // 2 diff --git a/tools/sglang-simulator/src/sglang_simulator/dataset/random.py b/tools/sglang-simulator/src/sglang_simulator/dataset/random.py new file mode 100644 index 000000000..154e1b091 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/dataset/random.py @@ -0,0 +1,52 @@ +from random import randint + +from sglang_simulator.dataset.base_dataset import BaseDataset, GenericRequest + + +class RandomIDsDataset(BaseDataset): + def __init__(self, tokenizer, args): + super().__init__(tokenizer, args) + self.cached: list[GenericRequest] = [] + self._name = "random_ids" + + def __len__(self): + return self.args.num_prompts + + def _get_single_item(self, index: int) -> GenericRequest: + if index < len(self.cached): + return self.cached[index] + min_id, max_id = ( + int(self.tokenizer.vocab_size * 0.25), + int(self.tokenizer.vocab_size * 0.75), + ) + + input_len = randint(self.args.min_input_len, self.args.max_input_len) + input_ids = [randint(min_id, max_id) for _ in range(input_len)] + + req = GenericRequest( + token_ids=input_ids, + input_length=input_len, + output_length=randint(self.args.min_output_len, self.args.max_output_len), + ) + self.cached.append(req) + return req + + +class RandomDataset(RandomIDsDataset): + def __init__(self, tokenizer, args): + super().__init__(tokenizer, args) + self.cached: list[GenericRequest] = [] + self._name = "random" + + def __len__(self): + return self.args.num_prompts + + def _get_single_item(self, index: int) -> GenericRequest: + if index < len(self.cached): + return self.cached[index] + req = super()._get_single_item(index) + if req.token_ids is not None: + req.prompt = self.tokenizer.decode(req.token_ids, skip_special_tokens=True) + req.token_ids = None + self.cached.append(req) + return req diff --git a/tools/sglang-simulator/src/sglang_simulator/hook/__init__.py b/tools/sglang-simulator/src/sglang_simulator/hook/__init__.py new file mode 100644 index 000000000..afb88b1df --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/hook/__init__.py @@ -0,0 +1,15 @@ +from sglang_simulator.hook.base_hook import BaseHook +from sglang_simulator.hook.class_hook_entry import ( + install_class_hooks, + is_class_hook_matched, + remove_class_hooks, + validate_required_class_hooks, +) + +__all__ = ( + install_class_hooks, + is_class_hook_matched, + remove_class_hooks, + validate_required_class_hooks, + BaseHook, +) diff --git a/tools/sglang-simulator/src/sglang_simulator/hook/base_hook.py b/tools/sglang-simulator/src/sglang_simulator/hook/base_hook.py new file mode 100644 index 000000000..dd8f1eb9d --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/hook/base_hook.py @@ -0,0 +1,36 @@ +from typing import List, Optional, Union + +from sglang_simulator.utils import get_logger + +logger = get_logger("sgl_simulator") + + +class BaseHook: + HOOK_CLASS_NAME: Optional[str] = None + HOOK_MODULE_NAME: Optional[str] = None + REGEX: bool = False + REQUIRED: bool = True + + def __init__(self): + pass + + @classmethod + def hook(cls, target) -> None: + """ + Return a new target or simply modify the target reference. + """ + raise NotImplementedError + + +def _register_hooks(HOOKS: List[BaseHook], hooks: Union[List[BaseHook], BaseHook]): + if isinstance(hooks, list): + for hook in hooks: + if not issubclass(hook, BaseHook): + raise TypeError("The hook should inherit from BaseHook.") + HOOKS.append(hook) + elif isinstance(hooks, type) and issubclass(hooks, BaseHook): + HOOKS.append(hooks) + else: + raise TypeError( + "The type of registered hook should be a list of BaseHook or a single BaseHook." + ) diff --git a/tools/sglang-simulator/src/sglang_simulator/hook/class_hook_entry.py b/tools/sglang-simulator/src/sglang_simulator/hook/class_hook_entry.py new file mode 100644 index 000000000..5f94981f9 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/hook/class_hook_entry.py @@ -0,0 +1,80 @@ +import builtins +import re +from types import FunctionType +from typing import List, Union + +from sglang_simulator.hook.base_hook import BaseHook, _register_hooks +from sglang_simulator.utils import get_logger + +logger = get_logger("sgl_simulator") + + +CLASS_HOOKS: List[BaseHook] = [] +_MATCHED_CLASS_HOOKS = set() + +_builtins_build_class_ = builtins.__build_class__ + + +def _custom_build_class_(func, name: str, *bases, **kwargs): + for hook in CLASS_HOOKS: + if ( + hook.REGEX + and hook.HOOK_CLASS_NAME + and re.search(hook.HOOK_CLASS_NAME, name) + ) or name == hook.HOOK_CLASS_NAME: + module_name = None + if isinstance(func, FunctionType): + module_name = getattr(func, "__globals__", {}).get("__name__", "") + if ( + hook.REGEX and re.search(hook.HOOK_MODULE_NAME, module_name) + ) or module_name == hook.HOOK_MODULE_NAME: + logger.debug( + f"Hooking Class: {hook.__name__} into {module_name}|{name}" + + ( + "(Regex is enabled, which might cause unexpected behavior.)" + if hook.REGEX + else "" + ) + ) + target_class = _builtins_build_class_(func, name, *bases, **kwargs) + hook.hook(target_class) + _MATCHED_CLASS_HOOKS.add(hook) + return target_class + + return _builtins_build_class_(func, name, *bases, **kwargs) + + +def install_class_hooks(hooks: Union[List[BaseHook], BaseHook]): + _register_hooks(CLASS_HOOKS, hooks) + builtins.__build_class__ = _custom_build_class_ + + +def is_class_hook_matched(hook: BaseHook) -> bool: + return hook in _MATCHED_CLASS_HOOKS + + +def validate_required_class_hooks() -> None: + unmatched = [ + hook + for hook in CLASS_HOOKS + if hook.REQUIRED and hook not in _MATCHED_CLASS_HOOKS + ] + if not unmatched: + return + + hook_names = ", ".join( + f"{hook.__name__} ({hook.HOOK_MODULE_NAME}.{hook.HOOK_CLASS_NAME})" + for hook in unmatched + ) + raise RuntimeError( + "Required SGLang Simulator hooks did not match imported SGLang classes: " + f"{hook_names}. The simulator must be adapted to this SGLang revision." + ) + + +def remove_class_hooks(): + # Clear the registered hooks and reset the build class function. + # Note: The classes that have been hooked will not be reset. + CLASS_HOOKS.clear() + _MATCHED_CLASS_HOOKS.clear() + builtins.__build_class__ = _builtins_build_class_ diff --git a/tools/sglang-simulator/src/sglang_simulator/hook/utils.py b/tools/sglang-simulator/src/sglang_simulator/hook/utils.py new file mode 100644 index 000000000..a85433d2f --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/hook/utils.py @@ -0,0 +1,8 @@ +def get_obj_from_args(type_name: str, *args, **kwargs): + for obj in args: + if type_name == f"{type(obj).__module__}.{type(obj).__name__}": + return obj + for obj in kwargs.values(): + if type_name == f"{type(obj).__module__}.{type(obj).__name__}": + return obj + return None diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/__init__.py b/tools/sglang-simulator/src/sglang_simulator/simulation/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/__init__.py b/tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/__init__.py new file mode 100644 index 000000000..28c5b5bb8 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/__init__.py @@ -0,0 +1,4 @@ +from sglang_simulator.simulation.benchmark.base_runner import BaseBenchmarkRunner +from sglang_simulator.simulation.benchmark.bench_config import BenchmarkConfig + +__all__ = ["BaseBenchmarkRunner", "BenchmarkConfig"] diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/base_runner.py b/tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/base_runner.py new file mode 100644 index 000000000..3a14b156a --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/base_runner.py @@ -0,0 +1,18 @@ +from abc import ABC, abstractmethod + + +class BaseBenchmarkRunner(ABC): + def __init__(self): + pass + + @abstractmethod + def benchmark(self) -> dict: + pass + + @abstractmethod + def flush_cache(self): + pass + + @abstractmethod + def shutdown(self): + pass diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/bench_config.py b/tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/bench_config.py new file mode 100644 index 000000000..a8d6a5ec4 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/benchmark/bench_config.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class BenchmarkConfig: + request_rate: float = float("inf") + max_concurrency: Optional[int] = None + ignore_request_timestamp: bool = False diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/manager/__init__.py b/tools/sglang-simulator/src/sglang_simulator/simulation/manager/__init__.py new file mode 100644 index 000000000..fee67d151 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/manager/__init__.py @@ -0,0 +1,5 @@ +from sglang_simulator.simulation.manager.config import ConfigManager +from sglang_simulator.simulation.manager.env import Envs +from sglang_simulator.simulation.manager.state import StateManager + +__all__ = ["StateManager", "Envs", "ConfigManager"] diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/manager/config.py b/tools/sglang-simulator/src/sglang_simulator/simulation/manager/config.py new file mode 100644 index 000000000..f25e39150 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/manager/config.py @@ -0,0 +1,252 @@ +import json +from pathlib import Path +from typing import Optional + +from sglang_simulator.simulation.manager.env import Envs +from sglang_simulator.simulation.types import PlatformConfig, SchedulerConfig +from sglang_simulator.simulation.utils import ( + calc_kv_cache_cell_elems, + calc_kv_cache_per_layer_elems, +) +from sglang_simulator.spec import AcceleratorInfo, DataType, ModelInfo +from sglang_simulator.time_predictor import ( + AIConfiguratorTimePredictor, + InferTimePredictor, +) +from sglang_simulator.utils import get_logger + +logger = get_logger() + + +class ConfigManager: + """Centralized configuration manager with caching.""" + + _model_info: Optional[ModelInfo] = None + _platform_config: Optional[PlatformConfig] = None + _scheduler_config: Optional[SchedulerConfig] = None + _raw_config: Optional[dict] = None + + @classmethod + def _get_raw_config(cls) -> dict: + if cls._raw_config is None: + with open(Envs.config_path()) as f: + cls._raw_config = json.load(f) + return cls._raw_config + + @classmethod + def resolve_config_relative_path(cls, path: str | None) -> str | None: + """Resolve predictor assets without depending on the process cwd.""" + if not path or Path(path).is_absolute(): + return path + + cwd_candidate = Path(path) + if cwd_candidate.exists(): + return str(cwd_candidate.resolve()) + + config_path = Path(Envs.config_path()).resolve() + for parent in config_path.parents: + candidate = parent / path + if candidate.exists(): + return str(candidate) + + # Keep the original value so predictor-specific errors remain clear. + return path + + @classmethod + def reset_config_cache(cls): + cls._raw_config = None + cls._model_info = None + cls._platform_config = None + cls._scheduler_config = None + + @classmethod + def set_model_info(cls, model: ModelInfo): + cls._model_info = model + + @classmethod + def get_model_info(cls) -> ModelInfo | None: + return cls._model_info + + @classmethod + def get_accelerator_info(cls) -> AcceleratorInfo: + config = cls._get_raw_config() + platform_config = config.get("platform", {}) + acc_info = platform_config.get("accelerator", {}) + hw = AcceleratorInfo.find_by_hw_name(acc_info.get("name")) + if hw is None: + logger.debug( + f"Failed to initialize device info with {acc_info.get('name')}. All available devices are: {AcceleratorInfo.list_all_hws().keys()}" + ) + hw = AcceleratorInfo( + name=acc_info.get("name"), + vendor=acc_info.get("vendor"), + hbm_bandwidth_gb=acc_info.get("hbm_bandwidth_gb"), + hbm_capacity_gb=acc_info.get("hbm_capacity_gb"), + inter_node_bandwidth_gb=acc_info.get("inter_node_bandwidth_gb"), + intra_node_bandwidth_gb=acc_info.get("intra_node_bandwidth_gb"), + tflops=acc_info.get("tflops"), + ) + else: + logger.info(f"Device info initialized: {hw}") + return hw + + @classmethod + def get_platform_config(cls) -> PlatformConfig: + if cls._platform_config is None: + hw = cls.get_accelerator_info() + config = cls._get_raw_config() + platform_config = config.get("platform", {}) + cls._platform_config = PlatformConfig( + device=hw, + disk_read_bandwidth_gb=platform_config.get("disk_read_bandwidth_gb"), + disk_write_bandwidth_gb=platform_config.get("disk_write_bandwidth_gb"), + memory_read_bandwidth_gb=platform_config.get( + "memory_read_bandwidth_gb" + ), + memory_write_bandwidth_gb=platform_config.get( + "memory_write_bandwidth_gb" + ), + num_device_per_node=platform_config.get("num_device_per_node"), + ) + + logger.info( + f"Platform configuration initialized successfully. {cls._platform_config}" + ) + + return cls._platform_config + + @classmethod + def set_scheduler_config(cls, config: SchedulerConfig): + # The configuration from the external config file has higher priority. + external_config = cls._get_raw_config().get("scheduler", {}) + for field_name in [ + "tp_size", + "dp_size", + "ep_size", + "pp_size", + "cp_size", + "cp_style", + "backend_name", + "backend_version", + "kv_bytes_per_token_per_gpu", + "hicache_ratio", + "moe_quant_mode_override", + "fmha_quant_mode_override", + "comm_quant_mode_override", + ]: + field_value = external_config.get(field_name) + if field_value is not None: + setattr(config, field_name, field_value) + + for field_name in ["data_type", "kv_cache_data_type"]: + field_value = external_config.get(field_name) + if field_value is not None: + setattr(config, field_name, DataType(field_value)) + + cls._scheduler_config = config + + @classmethod + def get_kv_cache_bytes(cls) -> int: + model = cls._model_info + scheduler_config = cls._scheduler_config + return ( + calc_kv_cache_cell_elems( + model, scheduler_config.tp_size, scheduler_config.pp_size + ) + * scheduler_config.kv_cache_data_type.bytes + ) + + @classmethod + def get_kv_cache_bytes_per_layer(cls) -> int: + model = cls._model_info + scheduler_config = cls._scheduler_config + return ( + calc_kv_cache_per_layer_elems( + model, scheduler_config.tp_size, scheduler_config.pp_size + ) + * scheduler_config.kv_cache_data_type.bytes + ) + + @classmethod + def get_scheduler_config(cls): + return cls._scheduler_config + + @classmethod + def _parse_server_args(cls, server_args: dict, backend: str) -> SchedulerConfig: + if backend == "sglang": + return SchedulerConfig( + tp_size=server_args.get("tp_size", 1), + ep_size=server_args.get("ep_size", 1), + dp_size=server_args.get("dp_size", 1), + pp_size=server_args.get("pp_size", 1), + cp_size=server_args.get("attn_cp_size", 1), + cp_style=server_args.get("cp_style", "none"), + mem_fraction_static=server_args.get("mem_fraction_static"), + backend_name="sglang", + ) + else: + raise RuntimeError(f"Unsupported backend[{backend}] server args parser.") + + @classmethod + def get_inference_time_predictor( + cls, model: ModelInfo, hw: AcceleratorInfo, sched_config: SchedulerConfig + ) -> InferTimePredictor: + config = cls._get_raw_config() + predictor_config = config.get("predictor", {}) + if predictor_config.get("name") == "aiconfigurator": + database_mode = predictor_config.get("database_mode", "SILICON") + prefill_scale_factor = predictor_config.get("prefill_scale_factor", 1) + decode_scale_factor = predictor_config.get("decode_scale_factor", 1) + prefill_min_latency = predictor_config.get("prefill_min_latency", 0) + workload_distribution = predictor_config.get( + "workload_distribution", "balanced" + ) + enable_oom_check = predictor_config.get("enable_oom_check", False) + + return AIConfiguratorTimePredictor( + model, + hw=hw, + config=sched_config, + database_path=cls.resolve_config_relative_path( + predictor_config.get("database_path") + ), + database_mode=database_mode, + prefill_scale_factor=prefill_scale_factor, + decode_scale_factor=decode_scale_factor, + prefill_min_latency=prefill_min_latency, + workload_distribution=workload_distribution, + enable_oom_check=enable_oom_check, + ) + elif predictor_config.get("name") == "ml": + from sglang_simulator.time_predictor.ml import MLTimePredictor + + return MLTimePredictor( + model, + hw=hw, + config=sched_config, + database_path=cls.resolve_config_relative_path( + predictor_config.get("database_path") + ), + latency_scale=predictor_config.get("latency_scale", 1.0), + ) + elif predictor_config.get("name") == "replay": + from sglang_simulator.time_predictor.replay import ReplayTimePredictor + + return ReplayTimePredictor( + model, + hw=hw, + config=sched_config, + database_path=cls.resolve_config_relative_path( + predictor_config.get("database_path") + ), + miss_fallback_seconds=predictor_config.get( + "miss_fallback_seconds", 0.0 + ), + miss_strategy=predictor_config.get("miss_strategy", "zero"), + miss_knn_k=predictor_config.get("miss_knn_k", 3), + ) + else: + raise ValueError( + f"Unknown predictor name: {predictor_config.get('name')}. " + f"Supported: aiconfigurator, ml, replay" + ) diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/manager/env.py b/tools/sglang-simulator/src/sglang_simulator/simulation/manager/env.py new file mode 100644 index 000000000..172f2a191 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/manager/env.py @@ -0,0 +1,52 @@ +import os + +from sglang_simulator.utils.logger import get_logger + +logger = get_logger("sgl_simulator") + + +class Envs: + @classmethod + def config_path(cls) -> str: + SGLANG_SIMULATOR_CONFIG_PATH = os.getenv("SGLANG_SIMULATOR_CONFIG_PATH") + if not SGLANG_SIMULATOR_CONFIG_PATH or not os.path.exists( + SGLANG_SIMULATOR_CONFIG_PATH + ): + raise RuntimeError( + f"The mock configuration path is not set or does not exist({SGLANG_SIMULATOR_CONFIG_PATH}). Please set it using the system variable SGLANG_SIMULATOR_CONFIG_PATH" + ) + return SGLANG_SIMULATOR_CONFIG_PATH + + @classmethod + def output_dir(cls) -> str: + SGLANG_SIMULATOR_OUTPUT_DIR = os.getenv( + "SGLANG_SIMULATOR_OUTPUT_DIR", "/tmp/sglang_simulator/output/" + ) + SGLANG_SIMULATOR_OUTPUT_DIR = os.path.realpath(SGLANG_SIMULATOR_OUTPUT_DIR) + if os.path.exists(SGLANG_SIMULATOR_OUTPUT_DIR) and os.path.isfile( + SGLANG_SIMULATOR_OUTPUT_DIR + ): + logger.error( + f"The metrics output path, {SGLANG_SIMULATOR_OUTPUT_DIR}, exists and is a file." + ) + raise RuntimeError( + f"{SGLANG_SIMULATOR_OUTPUT_DIR} exists but is not a directory." + ) + os.makedirs(SGLANG_SIMULATOR_OUTPUT_DIR, exist_ok=True) + return SGLANG_SIMULATOR_OUTPUT_DIR + + @classmethod + def hicache_storage_keys_path(cls) -> str: + SGLANG_SIMULATOR_HICACHE_STORAGE_KEYS_PATH = os.getenv( + "SGLANG_SIMULATOR_HICACHE_STORAGE_KEYS_PATH", + "/tmp/sglang_simulator/hicache_storage_keys.txt", + ) + return SGLANG_SIMULATOR_HICACHE_STORAGE_KEYS_PATH + + @classmethod + def simulation_mode(cls) -> str: + SGLANG_SIMULATOR_OUTPUT_MODE = os.getenv( + "SGLANG_SIMULATOR_OUTPUT_MODE", "OFFLINE" + ).upper() + assert SGLANG_SIMULATOR_OUTPUT_MODE in ("BLOCKING", "OFFLINE") + return SGLANG_SIMULATOR_OUTPUT_MODE diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/manager/state.py b/tools/sglang-simulator/src/sglang_simulator/simulation/manager/state.py new file mode 100644 index 000000000..d90d18d19 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/manager/state.py @@ -0,0 +1,123 @@ +class StateManager: + _iteration: int = 0 + _global_clock: float = 0 + _last_inference_dur: float = 0 + _current_inference_dur: float = 0 + _hicache_l2_load_dur: float = 0 + _hicache_l2_backup_dur: float = 0 + _hicache_l2_load_call_count: int = 0 + _hicache_l2_load_segment_count: int = 0 + _hicache_l2_load_units: int = 0 + _hicache_l2_load_bytes: float = 0 + _last_real_time_ts: float = 0 + _last_flush_time_ts: float = 0 + + @classmethod + def reset(cls): + cls._iteration = 0 + cls._global_clock = 0 + cls._last_inference_dur = 0 + cls._current_inference_dur = 0 + cls._hicache_l2_backup_dur = 0 + cls._hicache_l2_load_dur = 0 + cls._hicache_l2_load_call_count = 0 + cls._hicache_l2_load_segment_count = 0 + cls._hicache_l2_load_units = 0 + cls._hicache_l2_load_bytes = 0 + cls._last_real_time_ts = 0 + + @classmethod + def inc_iteration(cls) -> None: + cls._iteration += 1 + + @classmethod + def get_iteration(cls) -> int: + return cls._iteration + + @classmethod + def inc_hicache_l2_load_dur(cls, dur: float) -> None: + cls._hicache_l2_load_dur += dur + + @classmethod + def inc_hicache_l2_load_stats( + cls, + call_count: int = 0, + segment_count: int = 0, + units: int = 0, + bytes_: float = 0, + ) -> None: + cls._hicache_l2_load_call_count += call_count + cls._hicache_l2_load_segment_count += segment_count + cls._hicache_l2_load_units += units + cls._hicache_l2_load_bytes += bytes_ + + @classmethod + def inc_hicache_l2_backup_dur(cls, dur: float) -> None: + cls._hicache_l2_backup_dur += dur + + @classmethod + def pop_hicache_l2_load_dur(cls) -> float: + dur = cls._hicache_l2_load_dur + cls._hicache_l2_load_dur = 0 + return dur + + @classmethod + def pop_hicache_l2_load_stats(cls) -> dict: + stats = { + "h2d_load_call_count": cls._hicache_l2_load_call_count, + "h2d_load_segment_count": cls._hicache_l2_load_segment_count, + "h2d_load_units": cls._hicache_l2_load_units, + "h2d_load_bytes": cls._hicache_l2_load_bytes, + } + cls._hicache_l2_load_call_count = 0 + cls._hicache_l2_load_segment_count = 0 + cls._hicache_l2_load_units = 0 + cls._hicache_l2_load_bytes = 0 + return stats + + @classmethod + def pop_hicache_l2_backup_dur(cls) -> float: + dur = cls._hicache_l2_backup_dur + cls._hicache_l2_backup_dur = 0 + return dur + + @classmethod + def get_global_clock(cls) -> float: + return cls._global_clock + + @classmethod + def step_global_clock(cls, dur: float) -> None: + cls._global_clock += dur + + @classmethod + def set_global_clock(cls, clock: float) -> None: + cls._global_clock = clock + + @classmethod + def set_current_inference_dur(cls, dur: float) -> None: + cls._last_inference_dur = cls._current_inference_dur + cls._current_inference_dur = dur + + @classmethod + def get_last_inference_dur(cls) -> float: + return cls._last_inference_dur + + @classmethod + def get_current_inference_dur(cls) -> float: + return cls._current_inference_dur + + @classmethod + def set_last_real_time_ts(cls, ts): + cls._last_real_time_ts = ts + + @classmethod + def get_last_real_time_ts(cls): + return cls._last_real_time_ts + + @classmethod + def set_last_flush_time_ts(cls, ts: float): + cls._last_flush_time_ts = ts + + @classmethod + def get_last_flush_time_ts(cls) -> float: + return cls._last_flush_time_ts diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/__init__.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/cache_controller.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/cache_controller.py new file mode 100644 index 000000000..8421caf10 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/cache_controller.py @@ -0,0 +1,311 @@ +from queue import Empty, Queue +from typing import Optional + +from sglang_simulator.hook import BaseHook +from sglang_simulator.simulation.manager import ConfigManager, StateManager +from sglang_simulator.simulation.sglang.req_stats_manager import request_stats_manager + + +class C_HiCacheController(BaseHook): + HOOK_CLASS_NAME = "HiCacheController" + HOOK_MODULE_NAME = "sglang.srt.managers.cache_controller" + REQUIRED = False + + KV_CACHE_BYTES: Optional[int] = None + DISK_READ_BANDWIDTH_BYTES: Optional[float] = None + DISK_WRITE_BANDWIDTH_BYTES: Optional[float] = None + + @staticmethod + def calc_prefetch_pages( + required_pages: int, page_size_byte: int, max_dur: float, bandwidth: float + ) -> tuple[float, float]: + _prefetch_dur = required_pages * page_size_byte / bandwidth + if _prefetch_dur > max_dur: + _completed_pages = max(max_dur * bandwidth / page_size_byte, 1) + return _completed_pages, max_dur + else: + return required_pages, _prefetch_dur + + @classmethod + def hook(cls, target): + + original_terminate_prefetch = target.terminate_prefetch + original_storage_hit_query = target._storage_hit_query + original_init = target.__init__ + original_append_host_mem_release = target.append_host_mem_release + + def wrapped_init(self, *args, **kwargs): + self.sim_prefetch_buffer = Queue() + result = original_init(self, *args, **kwargs) + # The real IO thread normally creates this queue. The simulator + # replaces that thread, so initialize the handoff queue here. + if hasattr(self, "prefetch_hit_queue"): + self.prefetch_buffer = Queue() + return result + + def wrapped_append_host_mem_release(self, host_indices): + # A terminated prefetch may not have allocated host memory yet. + if host_indices is None: + return + return original_append_host_mem_release(self, host_indices) + + def override_backup_thread_func(self, *args, **kwargs): + # Async thread: perform no action + # The action will be performed by `handle_backup_operation` + pass + + def override_prefetch_thread_func(self, *args, **kwargs): + # Async thread: perform no action + # The action will be performed by `handle_prefetch_operation` + pass + + def handle_backup_operation(self): + if not self.enable_storage: + return + while True: + try: + operation = self.backup_queue.get(block=False) + if operation is None: + return + + if not self.backup_skip: + self._page_backup(operation) + # TODO: Track the backup operation according to the global clock + self.ack_backup_queue.put(operation) + + except Empty: + return + + def handle_prefetch_operation(self): + if not self.enable_storage: + return + + if C_HiCacheController.KV_CACHE_BYTES is None: + C_HiCacheController.KV_CACHE_BYTES = ConfigManager.get_kv_cache_bytes() + if C_HiCacheController.DISK_READ_BANDWIDTH_BYTES is None: + C_HiCacheController.DISK_READ_BANDWIDTH_BYTES = ( + ConfigManager.get_platform_config().disk_read_bandwidth + ) + + # TODO: Overlap schedule + remain_dur = StateManager.get_current_inference_dur() + + # Process all operations in the prefetch_queue: place those meeting + # the prefetch criteria into the sim_prefetch_buffer, and release the + # remaining operations along with any excess memory they have allocated. + while not self.prefetch_queue.empty(): + try: + operation = self.prefetch_queue.get(block=False) + if operation is None: + break + + # Ignore terminated operation + if operation._terminated_flag: + if hasattr(self, "prefetch_revoke_queue"): + self.prefetch_revoke_queue.put(operation.request_id) + else: + self.append_host_mem_release(operation.host_indices) + continue + + hash_value, storage_hit_count = self._storage_hit_query(operation) + # not to prefetch if not enough benefits + if ( + self.prefetch_threshold is not None + and storage_hit_count < self.prefetch_threshold + ): + if hasattr(self, "prefetch_revoke_queue"): + self.prefetch_revoke_queue.put(operation.request_id) + continue + operation.mark_terminate() + self.append_host_mem_release(operation.host_indices) + continue + else: + operation.hash_value = hash_value[ + : (storage_hit_count // self.page_size) + ] + if hasattr(self, "prefetch_hit_queue"): + # Allocate only the storage-hit range on the scheduler. + operation.storage_hit_count = storage_hit_count + self.prefetch_hit_queue.put(operation) + continue + + storage_hit_count = ( + storage_hit_count // self.page_size * self.page_size + ) + # free the pre-allocated memory for pages that are not hit + self.append_host_mem_release( + operation.host_indices[storage_hit_count:] + ) + operation.host_indices = operation.host_indices[ + :storage_hit_count + ] + self.sim_prefetch_buffer.put(operation) + except Empty: + break + + # handle operation which not yet fully prefetched + chunked_prefetch_operation = getattr( + self, "chunked_prefetch_operation", None + ) + if chunked_prefetch_operation is not None: + operation = chunked_prefetch_operation["operation"] + if operation._terminated_flag: + setattr(self, "chunked_prefetch_operation", None) + self.append_host_mem_release( + operation.host_indices[int(operation.completed_tokens) :] + ) + else: + storage_hit_count = chunked_prefetch_operation["storage_hit_count"] + completed_tokens, prefetch_dur = ( + C_HiCacheController.calc_prefetch_pages( + (storage_hit_count - operation.completed_tokens), + C_HiCacheController.KV_CACHE_BYTES, + remain_dur, + C_HiCacheController.DISK_READ_BANDWIDTH_BYTES, + ) + ) + if ( + completed_tokens + < storage_hit_count - operation.completed_tokens + ): + operation.completed_tokens += completed_tokens + remain_dur = 0 + else: + operation.completed_tokens = int(storage_hit_count) + operation.mark_terminate() + remain_dur -= prefetch_dur + setattr(self, "chunked_prefetch_operation", None) + # Release host memory after current operation is finished + self.append_host_mem_release( + operation.host_indices[storage_hit_count:] + ) + + # Feed operations whose host pages were allocated by the scheduler + # into the virtual-time transfer loop. + prefetch_buffer = getattr(self, "prefetch_buffer", None) + if prefetch_buffer is not None: + while not prefetch_buffer.empty(): + try: + operation = prefetch_buffer.get(block=False) + if operation is not None: + self.sim_prefetch_buffer.put(operation) + except Empty: + break + + # handle operation in sim_prefetch_buffer + while remain_dur > 0: + try: + operation = self.sim_prefetch_buffer.get(block=False) + if operation is None: + return + + # Ignore terminated operation + if operation._terminated_flag: + self.append_host_mem_release( + operation.host_indices[int(operation.completed_tokens) :] + ) + continue + + storage_hit_count = len(operation.host_indices) + completed_tokens, prefetch_dur = ( + C_HiCacheController.calc_prefetch_pages( + storage_hit_count, + C_HiCacheController.KV_CACHE_BYTES, + remain_dur, + C_HiCacheController.DISK_READ_BANDWIDTH_BYTES, + ) + ) + if completed_tokens < storage_hit_count: + # Continue to prefetch data next time. + operation.completed_tokens = completed_tokens + setattr( + self, + "chunked_prefetch_operation", + { + "operation": operation, + "storage_hit_count": storage_hit_count, + }, + ) + remain_dur = 0 + else: + operation.completed_tokens = int( + storage_hit_count // self.page_size * self.page_size + ) + # TODO: Track the prefetch operation according to the global clock + operation.mark_terminate() + remain_dur -= prefetch_dur + + except Empty: + return + + def override_generic_page_set( + self, hash_values, host_indices, extra_info=None + ) -> bool: + host_pool = getattr(self, "storage_host_pool", self.mem_pool_host) + # Always pass extra_info to storage_backend. + data = [ + host_pool.get_data_page(host_indices[i * self.page_size]) + for i in range(len(hash_values)) + ] + return self.storage_backend.batch_set(hash_values, data, extra_info) + + def wrapped_terminate_prefetch(self, operator): + result = original_terminate_prefetch(self, operator) + # This value may be a float if prefetch progress is interrupted by HiRadixCache.check_prefetch_progress. + result = (int(result[0]), result[1]) + # operation.completed_tokens, operation.hash_value = result + req_stats = request_stats_manager.get_req_stats(operator.request_id) + req_stats.final_storage_hit_len = result[0] + return result + + def wrapped_storage_hit_query(self, operator): + result = original_storage_hit_query(self, operator) + # hash_value, storage_query_count = result + req_stats = request_stats_manager.get_req_stats(operator.request_id) + req_stats.recv_storage_hit_len = result[1] + return result + + target.__init__ = wrapped_init + target.prefetch_thread_func = override_prefetch_thread_func + target.backup_thread_func = override_backup_thread_func + target.handle_backup_operation = handle_backup_operation + target.handle_prefetch_operation = handle_prefetch_operation + target.append_host_mem_release = wrapped_append_host_mem_release + target._generic_page_set = override_generic_page_set + target.terminate_prefetch = wrapped_terminate_prefetch + target.storage_hit_query = wrapped_storage_hit_query + if hasattr(target, "_storage_hit_query"): + target._storage_hit_query = wrapped_storage_hit_query + + +class C_HybridCacheController(BaseHook): + """Adapt UnifiedRadixCache's controller without duplicating legacy logic.""" + + HOOK_CLASS_NAME = "HybridCacheController" + HOOK_MODULE_NAME = "sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller" + REQUIRED = False + + @classmethod + def hook(cls, target): + # HybridCacheController inherits the deterministic thread replacements and + # handle_* methods installed on HiCacheController. Its own initialization + # creates Unified's control queues after the base initializer returns. + original_init = target.__init__ + original_storage_hit_query = target._storage_hit_query + + def wrapped_init(self, *args, **kwargs): + result = original_init(self, *args, **kwargs) + if hasattr(self, "prefetch_hit_queue") and not hasattr( + self, "prefetch_buffer" + ): + self.prefetch_buffer = Queue() + return result + + def wrapped_storage_hit_query(self, operator): + result = original_storage_hit_query(self, operator) + req_stats = request_stats_manager.get_req_stats(operator.request_id) + req_stats.recv_storage_hit_len = result[1] + return result + + target.__init__ = wrapped_init + target._storage_hit_query = wrapped_storage_hit_query diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/engine.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/engine.py new file mode 100644 index 000000000..0ed499562 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/engine.py @@ -0,0 +1,19 @@ +"""SGLang engine entry point with simulator-aware worker processes.""" + +from sglang_simulator.simulation.sglang.hook_bootstrap import ( + install_simulator_hooks, + run_simulator_detokenizer_process, + run_simulator_scheduler_process, +) + +install_simulator_hooks() + +# Install hooks before importing Engine so its worker entry points are patched. +from sglang.srt.entrypoints.engine import Engine # noqa: E402 + + +class SGLangSimulationEngine(Engine): + """Engine whose spawned workers install SGLang Simulator hooks explicitly.""" + + run_scheduler_process_func = staticmethod(run_simulator_scheduler_process) + run_detokenizer_process_func = staticmethod(run_simulator_detokenizer_process) diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hicache_storage.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hicache_storage.py new file mode 100644 index 000000000..dcb2e8921 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hicache_storage.py @@ -0,0 +1,155 @@ +import os +from typing import Any, List, Optional + +from sglang_simulator.hook import BaseHook +from sglang_simulator.simulation.manager.env import Envs +from sglang_simulator.utils.logger import get_logger + +logger = get_logger("sglang-simulator") + + +class C_StorageBackendFactory(BaseHook): + HOOK_CLASS_NAME = "StorageBackendFactory" + HOOK_MODULE_NAME = "sglang.srt.mem_cache.storage.backend_factory" + REQUIRED = False + + @classmethod + def hook(cls, target): + def override_create_backend(cls, *args, **kwargs): + logger.info("Creating hijacked cache storage backend.") + return MockHiCacheStorage() + + target.create_backend = override_create_backend + + +class MockHiCacheStorage: + def __init__(self, *args, **kwargs): + + self.storage: set = set() + self.storage_file_path: str = Envs.hicache_storage_keys_path() + os.makedirs(os.path.dirname(self.storage_file_path), exist_ok=True) + + if os.path.exists(self.storage_file_path): + with open(self.storage_file_path) as f: + line = f.readline() + while line: + self.storage.add(line.strip()) + line = f.readline() + + self.registered_pools = {} + + def register_mem_pool_host(self, mem_pool_host): + pass + + def register_mem_host_pool_v2(self, host_pool, host_pool_name): + """Register one pool from UnifiedRadixCache's multi-pool HiCache stack.""" + self.registered_pools[host_pool_name] = host_pool + + @staticmethod + def _pool_storage_key(key: str, pool_name) -> str: + name = str(pool_name) + return key if name == "kv" else f"{key}.{name}" + + def set( + self, + key: str, + value: Optional[Any] = None, + target_location: Optional[Any] = None, + target_sizes: Optional[Any] = None, + ) -> bool: + if self.exists(key): + return True + self.storage.add(key) + with open(self.storage_file_path, "a+") as f: + f.write(key + "\n") + return True + + def batch_set( + self, + keys: List[str], + values: Optional[Any] = None, + extra_info=None, # HiCacheStorageExtraInfo + target_locations: Optional[Any] = None, + target_sizes: Optional[Any] = None, + ) -> bool: + + for key, value in zip(keys, values): + if not self.set(key, value): + return False + return True + + def exists(self, key: str) -> bool: + return key in self.storage + + def batch_exists(self, keys: List[str], extra_info) -> int: + for i in range(len(keys)): + if not self.exists(keys[i]): + return i + return len(keys) + + def batch_exists_v2(self, keys, pool_transfers=None, extra_info=None): + """Return Unified HiCache's per-pool longest-prefix result.""" + from sglang.srt.mem_cache.hicache_storage import PoolTransferResult + + kv_hit_pages = self.batch_exists(keys, extra_info) + extra_pool_hit_pages = {} + final_pages = kv_hit_pages + for transfer in pool_transfers or []: + + def has_component(page_idx): + return self.exists( + self._pool_storage_key(keys[page_idx], transfer.name) + ) + + hit_policy = getattr(transfer.hit_policy, "value", transfer.hit_policy) + if hit_policy == "all_pages": + boundary = next( + (i for i in range(kv_hit_pages) if not has_component(i)), + kv_hit_pages, + ) + else: + trailing = max(1, len(transfer.keys) if transfer.keys else 1) + boundary = 0 + for prefix_len in range(kv_hit_pages, 0, -1): + if all( + has_component(i) + for i in range(max(0, prefix_len - trailing), prefix_len) + ): + boundary = prefix_len + break + extra_pool_hit_pages[transfer.name] = boundary + final_pages = min(final_pages, boundary) + + return PoolTransferResult( + kv_hit_pages=final_pages, + extra_pool_hit_pages=extra_pool_hit_pages, + ) + + def batch_get_v2(self, transfers, extra_info=None): + """Simulate loading every available pool page into registered host pools.""" + results = {} + for transfer in transfers: + keys = transfer.keys or [] + results[transfer.name] = [ + self.exists(self._pool_storage_key(key, transfer.name)) for key in keys + ] + return results + + def batch_set_v2(self, transfers, extra_info=None): + """Persist Unified HiCache component keys without materializing payloads.""" + results = {} + for transfer in transfers: + keys = transfer.keys or [] + pool_results = [] + for key in keys: + pool_results.append( + self.set(self._pool_storage_key(key, transfer.name)) + ) + results[transfer.name] = pool_results + return results + + def clear(self) -> bool: + self.storage.clear() + with open(self.storage_file_path, "w"): + pass + return True diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hiradix_cache.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hiradix_cache.py new file mode 100644 index 000000000..0536a12a3 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hiradix_cache.py @@ -0,0 +1,25 @@ +from sglang_simulator.hook import BaseHook + + +class C_HiRadixCacheHook(BaseHook): + HOOK_CLASS_NAME = "HiRadixCache" + HOOK_MODULE_NAME = "sglang.srt.mem_cache.hiradix_cache" + REQUIRED = False + + @classmethod + def hook(cls, target): + original_check_hicache_events = target.check_hicache_events + + def wrapped_check_hicache_events(self, *args, **kwargs): + # The async thread for prefetching and backup in `HiCacheController` has been deprecated. + # So we have to handle the backup or prefetch operation manually. + self.cache_controller.handle_backup_operation() + self.cache_controller.handle_prefetch_operation() + result = original_check_hicache_events(self, *args, **kwargs) + # Host pages are allocated after the storage query. Run the + # simulated transfer only after that allocation step. + if hasattr(self.cache_controller, "prefetch_hit_queue"): + self.cache_controller.handle_prefetch_operation() + return result + + target.check_hicache_events = wrapped_check_hicache_events diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hook_bootstrap.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hook_bootstrap.py new file mode 100644 index 000000000..dd7725a6e --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/hook_bootstrap.py @@ -0,0 +1,80 @@ +"""Install SGLang Simulator hooks in the parent and spawned SGLang worker processes.""" + +import os + +# Spawned interpreters inherit this marker before usercustomize runs. +os.environ["SGLANG_SIMULATOR_BOOTSTRAP"] = "1" + +import sglang_simulator.hook as sglang_simulator_hook +from sglang_simulator.simulation.sglang import ( + cache_controller, + hicache_storage, + hiradix_cache, + mem_cache_allocator, + mem_pool_host, + model_runner, + scheduler, + sgl_kernel_hook, + unified_radix_cache, +) + +# A spawned worker imports this module while unpickling its target. ModelConfig +# can import GPU kernels while later arguments are still being unpickled, before +# the target wrapper executes, so the loader stub must already be present here. +sgl_kernel_hook.install_load_utils_stub() + +_HOOKS_INSTALLED = False + + +def install_simulator_hooks() -> None: + """Install hooks once in the current Python interpreter.""" + global _HOOKS_INSTALLED + if _HOOKS_INSTALLED: + return + + # The package __init__ loads GPU ops before a child-module import hook can + # run reliably under spawn. Seed the loader module before importing SGLang. + sgl_kernel_hook.install_load_utils_stub() + + sglang_simulator_hook.install_class_hooks( + [ + scheduler.C_SchedulerHook, + scheduler.C_SglangPrefillAdderHook, + scheduler.C_SchedulerRequestReceiver, + model_runner.C_ModelRunnerHook, + model_runner.C_KVCacheConfiguratorHook, + hicache_storage.C_StorageBackendFactory, + cache_controller.C_HiCacheController, + cache_controller.C_HybridCacheController, + hiradix_cache.C_HiRadixCacheHook, + unified_radix_cache.C_UnifiedRadixCacheHook, + mem_cache_allocator.C_PagedTokenToKVPoolAllocatorHook, + mem_pool_host.C_MHATokenToKVPoolHostHook, + mem_pool_host.C_HostKVCacheHook, + mem_pool_host.C_PackedSingleKVPoolHook, + mem_pool_host.C_GenericHostKVCacheSubclassHook, + ] + ) + _HOOKS_INSTALLED = True + + +def run_simulator_scheduler_process(*args, **kwargs): + """Spawn-safe scheduler entry point which installs SGLang Simulator before SGLang imports.""" + install_simulator_hooks() + + # Spawned workers do not inherit parent-process monkey patches, so import + # the scheduler only after installing hooks in this process. + from sglang.srt.managers.scheduler import run_scheduler_process + + return run_scheduler_process(*args, **kwargs) + + +def run_simulator_detokenizer_process(*args, **kwargs): + """Spawn-safe detokenizer entry point which installs SGLang Simulator before imports.""" + install_simulator_hooks() + + # Install hooks before transitive schedule_batch and memory_pool imports + # so this CPU-only process does not load real GPU kernels. + from sglang.srt.managers.detokenizer_manager import run_detokenizer_process + + return run_detokenizer_process(*args, **kwargs) diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/launch_server.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/launch_server.py new file mode 100644 index 000000000..ec8d5fbe3 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/launch_server.py @@ -0,0 +1,109 @@ +import argparse +import dataclasses +import os +import sys +from typing import Optional + +from sglang_simulator.compat import ( + apply_simulator_server_args, + validate_launch_runtime, +) +from sglang_simulator.simulation.sglang.hook_bootstrap import ( + install_simulator_hooks, + run_simulator_detokenizer_process, + run_simulator_scheduler_process, +) +from sglang_simulator.utils import get_logger + +install_simulator_hooks() + + +logger = get_logger("sgl_simulator") + + +@dataclasses.dataclass +class SimulationArgs: + sim_config_path: Optional[str] = None + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument( + "--sim-config-path", + type=str, + default=None, + help="Path to simulation JSON config (same as SGLANG_SIMULATOR_CONFIG_PATH).", + ) + + @classmethod + def from_cli_args(cls, ns: argparse.Namespace) -> "SimulationArgs": + return SimulationArgs(sim_config_path=ns.sim_config_path) + + +def _has_cli_option(argv: list[str], option: str) -> bool: + return any(arg == option or arg.startswith(f"{option}=") for arg in argv) + + +def apply_simulator_defaults(raw_args: argparse.Namespace, argv: list[str]) -> None: + """Avoid real model execution while preserving explicit SGLang options.""" + if not _has_cli_option(argv, "--load-format"): + raw_args.load_format = "dummy" + + if os.getenv("SGLANG_USE_CPU_ENGINE") != "1": + return + + if not _has_cli_option(argv, "--device"): + raw_args.device = "cpu" + if not _has_cli_option(argv, "--attention-backend"): + raw_args.attention_backend = "torch_native" + if not _has_cli_option(argv, "--sampling-backend"): + raw_args.sampling_backend = "pytorch" + if not ( + _has_cli_option(argv, "--cuda-graph-backend-decode") + or _has_cli_option(argv, "--cuda-graph-backend-prefill") + or _has_cli_option(argv, "--disable-cuda-graph") + ): + raw_args.disable_cuda_graph = True + + # CPU-only model validation may still query CUDA capability while + # constructing ServerArgs, before the simulator runner is spawned. + import torch + + torch.cuda.get_device_capability = lambda *_args, **_kwargs: (10, 0) + + +if __name__ == "__main__": + validate_launch_runtime() + + from sglang.srt.entrypoints.http_server import launch_server + from sglang.srt.server_args import ServerArgs + from sglang.srt.utils import kill_process_tree + + parser = argparse.ArgumentParser() + + g = parser.add_argument_group("sglang") + ServerArgs.add_cli_args(g) + + g = parser.add_argument_group("simulation") + SimulationArgs.add_cli_args(g) + + argv = sys.argv[1:] + raw_args = parser.parse_args(argv) + apply_simulator_defaults(raw_args, argv) + apply_simulator_server_args(raw_args) + server_args = ServerArgs.from_cli_args(raw_args) + simulation_args = SimulationArgs.from_cli_args(raw_args) + + config_path = os.getenv("SGLANG_SIMULATOR_CONFIG_PATH") + if config_path and os.path.exists(config_path): + logger.info(f"Using config from {config_path}") + elif simulation_args.sim_config_path: + os.environ["SGLANG_SIMULATOR_CONFIG_PATH"] = simulation_args.sim_config_path + + try: + launch_server( + server_args, + run_scheduler_process_func=run_simulator_scheduler_process, + run_detokenizer_process_func=run_simulator_detokenizer_process, + ) + finally: + kill_process_tree(os.getpid(), include_parent=False) diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/mem_cache_allocator.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/mem_cache_allocator.py new file mode 100644 index 000000000..748dc73eb --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/mem_cache_allocator.py @@ -0,0 +1,128 @@ +import types + +import torch +from sglang_simulator.hook import BaseHook + + +def _alloc_extend_cpu( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + num_new_pages: int = None, +): + """CPU implementation using SGLang's native paged-allocation helper.""" + from sglang.srt.mem_cache.allocator import alloc_extend_naive + from sglang.srt.utils import get_num_new_pages + + if num_new_pages is None: + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, + page_size=self.page_size, + prefix_lens=prefix_lens_cpu, + ) + if self.need_sort and num_new_pages > len(self.free_pages): + self.merge_and_sort_free() + if num_new_pages > len(self.free_pages): + return None + + out_indices = torch.empty( + (extend_num_tokens,), + dtype=self.free_pages.dtype, + device=self.device, + ) + alloc_extend_naive( + prefix_lens, + seq_lens, + last_loc, + self.free_pages, + out_indices, + self.page_size, + self.device, + ) + self.free_pages = self.free_pages[num_new_pages:] + return out_indices + + +def _alloc_decode_cpu( + self, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, +): + """CPU decode allocation through the allocator's public method contract.""" + from sglang.srt.utils import get_num_new_pages + + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, + page_size=self.page_size, + decode=True, + ) + if self.need_sort and num_new_pages > len(self.free_pages): + self.merge_and_sort_free() + if num_new_pages > len(self.free_pages): + return None + + out_indices = (last_loc + 1).to(dtype=self.free_pages.dtype) + need_new_page = seq_lens % self.page_size == 1 + if num_new_pages: + out_indices = out_indices.clone() + out_indices[need_new_page] = self.free_pages[:num_new_pages] * self.page_size + + self.free_pages = self.free_pages[num_new_pages:] + return out_indices + + +def alloc_extend_cpu(*args, **kwargs): + """Compatibility entry plus the native allocator-method implementation.""" + if args and isinstance(args[0], torch.Tensor): + from sglang.srt.mem_cache.allocator import alloc_extend_naive + + prefix_lens, seq_lens, last_loc, free_pages, out_indices = args[:5] + alloc_extend_naive( + prefix_lens, + seq_lens, + last_loc, + free_pages, + out_indices, + kwargs["page_size"], + prefix_lens.device, + ) + return None + return _alloc_extend_cpu(*args, **kwargs) + + +def alloc_decode_cpu(*args, **kwargs): + """Compatibility entry plus the native allocator-method implementation.""" + if args and isinstance(args[0], torch.Tensor): + seq_lens, last_loc, free_pages, out_indices = args[:4] + page_size = kwargs["page_size"] + need_new_page = seq_lens % page_size == 1 + result = last_loc + 1 + result[need_new_page] = ( + free_pages[: int(need_new_page.sum().item())] * page_size + ) + out_indices.copy_(result) + return None + return _alloc_decode_cpu(*args, **kwargs) + + +class C_PagedTokenToKVPoolAllocatorHook(BaseHook): + HOOK_CLASS_NAME = "PagedTokenToKVPoolAllocator" + HOOK_MODULE_NAME = r"^sglang\.srt\.mem_cache\.allocator(?:\.paged)?$" + REGEX = True + + @classmethod + def hook(cls, target): + original_init = target.__init__ + + def wrapped_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + if self.device == "cpu": + self.alloc_extend = types.MethodType(_alloc_extend_cpu, self) + self.alloc_decode = types.MethodType(_alloc_decode_cpu, self) + + target.__init__ = wrapped_init diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/mem_pool_host.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/mem_pool_host.py new file mode 100644 index 000000000..86e0ea0a8 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/mem_pool_host.py @@ -0,0 +1,384 @@ +from abc import ABC, abstractmethod +from enum import Enum +from functools import lru_cache + +import numpy as np +import torch +from sglang_simulator.hook import BaseHook +from sglang_simulator.simulation.manager import ConfigManager, StateManager +from sglang_simulator.utils import get_logger + +logger = get_logger() + + +class TransportDirection(Enum): + H2D = "H2D" + D2H = "D2H" + + +class HicacheTransportEstimator(ABC): + def __init__( + self, + memory_read_bandwidth_bytes: float, + memory_write_bandwidth_bytes: float, + ): + self.memory_read_bandwidth_bytes = memory_read_bandwidth_bytes + self.memory_write_bandwidth_bytes = memory_write_bandwidth_bytes + + @abstractmethod + def estimate_bandwidth( + self, size_bytes: np.ndarray, direction: TransportDirection + ) -> np.ndarray: + raise NotImplementedError + + +class HicacheTransportOverheadEstimator(HicacheTransportEstimator): + """Bandwidth model with a fixed launch overhead and 85% efficiency.""" + + def estimate_bandwidth( + self, size_bytes: np.ndarray, direction: TransportDirection + ) -> np.ndarray: + if direction is TransportDirection.H2D: + overhead_s = 6.67e-6 + bandwidth = self.memory_read_bandwidth_bytes * 0.85 + else: + overhead_s = 4e-6 + bandwidth = self.memory_write_bandwidth_bytes * 0.85 + return size_bytes * bandwidth / (overhead_s * bandwidth + size_bytes) + + +def compute_contiguous_index_lengths( + host_indices: torch.Tensor, + device_indices: torch.Tensor, +) -> np.ndarray: + if len(host_indices) != len(device_indices): + raise ValueError("Host and device cache index lists must have the same length.") + if len(host_indices) == 0: + return np.empty(0, dtype=np.float64) + + host = np.asarray(host_indices.cpu(), dtype=np.int64) + device = np.asarray(device_indices.cpu(), dtype=np.int64) + contiguous = (np.diff(host) == 1) & (np.diff(device) == 1) + cuts = np.flatnonzero(~contiguous) + 1 + starts = np.r_[0, cuts] + ends = np.r_[cuts, len(host_indices)] + return (ends - starts).astype(np.float64) + + +def allocate_meta_tensor( + dims, + dtype: torch.dtype, + device: str, + pin_memory: bool, + allocator=None, + registration_granularity_bytes=None, +) -> torch.Tensor: + """Allocate metadata-only host cache payload for simulation.""" + return torch.empty(dims, dtype=dtype, device="meta") + + +def _install_meta_allocators() -> None: + modules = [] + try: + from sglang.srt.mem_cache import memory_pool_host + + modules.append(memory_pool_host) + except ImportError: + pass + try: + from sglang.srt.mem_cache.pool_host import common + + modules.append(common) + except ImportError: + pass + + for module in modules: + allocators = getattr(module, "ALLOC_MEMORY_FUNCS", None) + if allocators is None: + continue + allocators.default_factory = lambda: allocate_meta_tensor + for key in list(allocators): + allocators[key] = allocate_meta_tensor + + +_SIMULATED_AVAILABLE_HOST_MEMORY_BYTES = 1 << 60 + + +class _PsutilProxy: + def __init__(self, psutil_module): + self._psutil_module = psutil_module + + def virtual_memory(self): + snapshot = self._psutil_module.virtual_memory() + return snapshot._replace( + available=max( + snapshot.available, + _SIMULATED_AVAILABLE_HOST_MEMORY_BYTES, + ) + ) + + def __getattr__(self, name): + return getattr(self._psutil_module, name) + + +def _call_with_meta_host_memory(original_init, self, *args, **kwargs): + """Bypass physical host-payload checks while meta allocation is active.""" + init_globals = getattr(original_init, "__globals__", None) + psutil_module = init_globals.get("psutil") if init_globals is not None else None + if psutil_module is None: + return original_init(self, *args, **kwargs) + + proxy = _PsutilProxy(psutil_module) + init_globals["psutil"] = proxy + try: + return original_init(self, *args, **kwargs) + finally: + if init_globals.get("psutil") is proxy: + init_globals["psutil"] = psutil_module + + +@lru_cache(maxsize=256) +def get_refined_cache_size_per_token(host_pool) -> float: + internal_size = float(host_pool.get_size_per_token()) + scheduler_config = ConfigManager.get_scheduler_config() + if scheduler_config is None or scheduler_config.kv_cache_data_type is None: + logger.warning( + "Scheduler KV-cache dtype is unavailable; using %s's native " + "size-per-token value.", + host_pool.__class__.__name__, + ) + return internal_size + + internal_dtype = host_pool.dtype + dtype_factor = scheduler_config.kv_cache_data_type.bytes / internal_dtype.itemsize + return internal_size * dtype_factor + + +_DSV4_TRANSFER_SIZE_MULTIPLIERS = { + "swa": 130, + "deepseek_v4_c4": 65, + "deepseek_v4_c4_indexer": 132, + "deepseek_v4_c128": 3, + "deepseek_v4_c4_state": 256, + "deepseek_v4_c128_state": 256, + "deepseek_v4_indexer_state": 128, + "deepseek_v4_c4_indexer_state": 128, +} + +_DSV4_PAGED_POOL_NAMES = { + "swa", + "deepseek_v4_c4", + "deepseek_v4_c4_indexer", + "deepseek_v4_c128", +} + + +def _dsv4_transfer_size_multiplier(host_pool) -> int | None: + return _DSV4_TRANSFER_SIZE_MULTIPLIERS.get(str(getattr(host_pool, "pool_name", ""))) + + +def get_transfer_size_per_unit(host_pool, *, all_layers: bool) -> float: + """Return calibrated bytes moved for one transfer unit. + + DSv4's paged and state pools expose physical page-row geometry through Unified + HiCache. Preserve the 0714 estimator's calibrated logical-byte multipliers while + keeping transfer dispatch on the current Unified pool interfaces. + """ + size = get_refined_cache_size_per_token(host_pool) + dsv4_multiplier = _dsv4_transfer_size_multiplier(host_pool) + if dsv4_multiplier is not None: + return size * dsv4_multiplier + + layer_num = max(int(getattr(host_pool, "layer_num", 1)), 1) + per_layer_size = size / layer_num + return per_layer_size * layer_num if all_layers else per_layer_size + + +def _transport_estimator() -> HicacheTransportEstimator: + platform = ConfigManager.get_platform_config() + return HicacheTransportOverheadEstimator( + memory_read_bandwidth_bytes=platform.memory_read_bandwidth, + memory_write_bandwidth_bytes=platform.memory_write_bandwidth, + ) + + +def _normalize_transfer_indices(self, host_indices, device_indices): + if host_indices is None or device_indices is None: + return None, None + if hasattr(self, "_to_page_indices"): + host_indices = self._to_page_indices(host_indices) + device_indices = self._to_page_indices(device_indices) + return host_indices, device_indices + + +def _transfer_segment_lengths( + self, host_indices, device_indices, *, count_logical_tokens: bool = False +) -> np.ndarray: + if host_indices is None or device_indices is None: + return np.empty(0, dtype=np.float64) + + original_unit_count = len(host_indices) + host_indices, device_indices = _normalize_transfer_indices( + self, host_indices, device_indices + ) + lengths = compute_contiguous_index_lengths(host_indices, device_indices) + if ( + len(lengths) + and count_logical_tokens + and str(getattr(self, "pool_name", "")) in _DSV4_PAGED_POOL_NAMES + ): + # The 0714 DSv4 paged-pool H2D estimator counted logical token slots + # while using page-row contiguity to determine transfer segments. + lengths[-1] += original_unit_count - len(host_indices) + return lengths + + +def _sim_load_to_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend, + *, + is_draft: bool = False, +) -> None: + segment_lengths = _transfer_segment_lengths( + self, host_indices, device_indices, count_logical_tokens=True + ) + if not len(segment_lengths): + return + + size_bytes = segment_lengths * get_transfer_size_per_unit(self, all_layers=False) + StateManager.inc_hicache_l2_load_stats( + call_count=1, + segment_count=len(size_bytes), + units=int(np.sum(segment_lengths)), + bytes_=float(np.sum(size_bytes)), + ) + bandwidth = _transport_estimator().estimate_bandwidth( + size_bytes, TransportDirection.H2D + ) + StateManager.inc_hicache_l2_load_dur(float(np.sum(size_bytes / bandwidth))) + + +def _sim_backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend +) -> None: + segment_lengths = _transfer_segment_lengths(self, host_indices, device_indices) + if not len(segment_lengths): + return + + size_bytes = segment_lengths * get_transfer_size_per_unit(self, all_layers=True) + bandwidth = _transport_estimator().estimate_bandwidth( + size_bytes, TransportDirection.D2H + ) + StateManager.inc_hicache_l2_backup_dur(float(np.sum(size_bytes / bandwidth))) + + +def _sim_get_data_page(self, index, flat: bool = True) -> torch.Tensor: + return torch.ones(size=(1, 1)) * index + + +def _sim_set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None: + return None + + +def _install_transport_methods(target) -> None: + original_init = target.__init__ + + def wrapped_init(self, *args, **kwargs): + _install_meta_allocators() + if "pin_memory" in kwargs: + kwargs["pin_memory"] = False + return _call_with_meta_host_memory(original_init, self, *args, **kwargs) + + target.__init__ = wrapped_init + target.load_to_device_per_layer = _sim_load_to_device_per_layer + target.backup_from_device_all_layer = _sim_backup_from_device_all_layer + target.get_data_page = _sim_get_data_page + target.set_from_flat_data_page = _sim_set_from_flat_data_page + + +class C_MHATokenToKVPoolHostHook(BaseHook): + HOOK_CLASS_NAME = "MHATokenToKVPoolHost" + HOOK_MODULE_NAME = r"^sglang\.srt\.mem_cache\.(memory_pool_host|pool_host\.mha)$" + REGEX = True + REQUIRED = False + + @classmethod + def hook(cls, target): + _install_transport_methods(target) + + +class C_HostKVCacheHook(BaseHook): + HOOK_CLASS_NAME = "HostKVCache" + HOOK_MODULE_NAME = r"^sglang\.srt\.mem_cache\.(memory_pool_host|pool_host\.base)$" + REGEX = True + REQUIRED = False + + @classmethod + def hook(cls, target): + original_init = target.__init__ + + def wrapped_init(self, *args, **kwargs): + _install_meta_allocators() + if "pin_memory" in kwargs: + kwargs["pin_memory"] = False + elif len(args) > 5: + args = list(args) + args[5] = False + return _call_with_meta_host_memory(original_init, self, *args, **kwargs) + + target.__init__ = wrapped_init + + +class C_PackedSingleKVPoolHook(BaseHook): + """Allocate byte-packed single KV pools from their runtime geometry.""" + + HOOK_CLASS_NAME = r".*SingleKVPool$" + HOOK_MODULE_NAME = r"^sglang\.srt\.mem_cache\..+$" + REGEX = True + REQUIRED = False + + @classmethod + def hook(cls, target): + original_create_buffer = target.create_buffer + + def wrapped_create_buffer(self, *, num_pages: int): + if self.store_dtype != torch.uint8 or not hasattr( + self, "get_bytes_per_token" + ): + return original_create_buffer(self, num_pages=num_pages) + + try: + return original_create_buffer(self, num_pages=num_pages) + except AssertionError: + # Some packed pools validate production-only geometry. Simulation + # still needs a correctly sized byte buffer for dummy model configs. + pass + + bytes_per_token = self.get_bytes_per_token() + self.kv_cache_total_dim = bytes_per_token + bytes_per_page = self.page_size * bytes_per_token + self.bytes_per_page_padded = (bytes_per_page + 575) // 576 * 576 + return torch.zeros( + num_pages, + self.bytes_per_page_padded, + dtype=self.store_dtype, + device=self.device, + ) + + target.create_buffer = wrapped_create_buffer + + +class C_GenericHostKVCacheSubclassHook(BaseHook): + HOOK_CLASS_NAME = r".*(?:PoolHost|HostPool)$" + HOOK_MODULE_NAME = r"^sglang\.srt\.mem_cache\.(memory_pool_host|pool_host\..+)$" + REGEX = True + REQUIRED = False + + @classmethod + def hook(cls, target): + if any(base.__name__ == "HostKVCache" for base in target.__mro__[1:]): + _install_transport_methods(target) diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/model_runner.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/model_runner.py new file mode 100644 index 000000000..56fa992e6 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/model_runner.py @@ -0,0 +1,265 @@ +import torch +from sglang_simulator.hook import BaseHook +from sglang_simulator.simulation.manager import ConfigManager +from sglang_simulator.simulation.sglang.utils import ( + resolve_model_info, + resolve_scheduler_config, +) +from sglang_simulator.simulation.utils import profile_device_available_bytes + + +class _MockModel(torch.nn.Module): + """Minimal model surface needed by SGLang's native runner initialization.""" + + def forward(self, *args, **kwargs): + return None + + +class _MockModelLoader: + """Minimal loader state for upstream resident-weight accounting.""" + + preloaded_weights_bytes = 0 + + +def _make_mock_model_loader(model_runner_type): + if hasattr(model_runner_type, "preloaded_weights_bytes"): + return _MockModelLoader() + return None + + +def _resolve_kv_page_size(configurator): + return ( + getattr(configurator, "page_size", None) + or getattr(configurator.server_args, "page_size", None) + or 1 + ) + + +class C_ModelRunnerHook(BaseHook): + HOOK_CLASS_NAME = "ModelRunner" + HOOK_MODULE_NAME = "sglang.srt.model_executor.model_runner" + + @classmethod + def hook(cls, target): + def override_load_model(self): + from sglang.srt.model_executor.model_runner import ( + resolve_sliding_window_size, + ) + + self.model = _MockModel() + self.dtype = self.model_config.dtype + self.sliding_window_size = resolve_sliding_window_size( + self.model, self.model_config + ) + self.prefill_aware_swa = False + self.weight_load_mem_usage = 0 + self.load_config = None + self.loader = _make_mock_model_loader(type(self)) + + if ConfigManager.get_model_info() is None: + ConfigManager.set_model_info(resolve_model_info(self.model_config)) + + def wrapped_forward(self, *args, **kwargs): + batch = args[0] + from sglang.srt.layers.logits_processor import LogitsProcessorOutput + from sglang.srt.model_executor.model_runner import ModelRunnerOutput + + output = LogitsProcessorOutput( + next_token_logits=torch.empty( + size=(batch.batch_size, self.model_config.vocab_size), + device=self.device, + ) + ) + return ModelRunnerOutput( + logits_output=output, + can_run_graph=False, + expert_distribution_metrics=None, + ) + + def wrapped_sample(self, *args, **kwargs): + logits = args[0] + return torch.ones( + size=(logits.next_token_logits.shape[0],), + device=self.device, + dtype=torch.int64, + ) + + def wrapped_compute_logprobs_only(*args, **kwargs): + return None + + def wrapped_init_attention_backends(self): + try: + from sglang.srt.model_executor.model_runner_components.attention_backend_setup import ( + resolve_attention_backend_strs, + ) + except ImportError: + default_backend = self.server_args.attention_backend + self.prefill_attention_backend_str = ( + self.server_args.prefill_attention_backend or default_backend + ) + self.decode_attention_backend_str = ( + self.server_args.decode_attention_backend or default_backend + ) + else: + resolved = resolve_attention_backend_strs(model_runner=self) + self.prefill_attention_backend_str = resolved.prefill + self.decode_attention_backend_str = resolved.decode + + self.attn_backend = None + self.decode_attn_backend = None + self.decode_attn_backend_group = None + + def wrapped_init_cuda_graphs(self, capture_decode_cuda_graph=True): + self.graph_mem_usage = 0 + self.cuda_graph_runner = None + self.eager_runner = None + self.prefill_cuda_graph_runner = None + self.decode_cuda_graph_runner = None + + # Keep SGLang's native initialize() and alloc_memory_pool() lifecycle. + # Only operations that require model weights or GPU kernels are mocked. + target.load_model = override_load_model + target.forward = wrapped_forward + target.sample = wrapped_sample + target.compute_logprobs_only = wrapped_compute_logprobs_only + target.init_attention_backends = wrapped_init_attention_backends + target.init_cuda_graphs = wrapped_init_cuda_graphs + + +class C_KVCacheConfiguratorHook(BaseHook): + HOOK_CLASS_NAME = "KVCacheConfigurator" + HOOK_MODULE_NAME = "sglang.srt.mem_cache.kv_cache_configurator" + + @classmethod + def hook(cls, target): + original_configure = target.configure + original_init_pools = target._init_pools + supports_cpu_fp8_quant_method = hasattr(target, "_build_mha_quant_method") + + def wrapped_configure(self, *args, **kwargs): + if not ( + supports_cpu_fp8_quant_method + and getattr(self, "device", None) == "cpu" + and getattr(self, "kv_cache_dtype", None) == torch.float8_e4m3fn + ): + return original_configure(self, *args, **kwargs) + + # Newer SGLang runtimes validate CPU FP8 KV-cache support and select + # an AMX-only quant method. The simulator executes scheduler state on + # CPU while modeling the target accelerator's FP8 cache. Suppress the + # physical-CPU predicate only while native compact pools are built; + # all other platform capabilities and the logical dtype stay intact. + from sglang.srt.mem_cache.kv_cache_configurator import current_platform + + original_is_cpu = current_platform.is_cpu + current_platform.is_cpu = lambda: False + try: + return original_configure(self, *args, **kwargs) + finally: + current_platform.is_cpu = original_is_cpu + + def override_profile_available_bytes(self, pre_model_load_memory): + if self.server_args.max_total_tokens is not None: + from sglang.srt.model_executor.pool_configurator import ( + create_memory_pool_configurator, + ) + + configurator = create_memory_pool_configurator(self) + target_tokens = self.server_args.max_total_tokens + + page_size = _resolve_kv_page_size(self) + + def resolved_tokens(budget_bytes): + try: + config = configurator.calculate_pool_sizes( + budget_bytes, page_size + ) + except RuntimeError: + return 0 + return config.max_total_num_tokens + + lower, upper = 0, 1 + while resolved_tokens(upper) < target_tokens: + lower, upper = upper, upper * 2 + while lower + 1 < upper: + middle = (lower + upper) // 2 + if resolved_tokens(middle) < target_tokens: + lower = middle + else: + upper = middle + return upper + + model = ConfigManager.get_model_info() + if model is None: + model = resolve_model_info(self.model_config) + ConfigManager.set_model_info(model) + hardware = ConfigManager.get_accelerator_info() + scheduler_config = resolve_scheduler_config( + server_args=self.server_args, + model_config=self.model_config, + ) + if hardware is None or scheduler_config is None: + raise RuntimeError( + "Simulator model, accelerator, and scheduler configuration " + "must be resolved before KV-cache pool sizing." + ) + + available_bytes = profile_device_available_bytes( + model=model, + device=hardware, + scheduler_config=scheduler_config, + ) + if self.mambaish_config is not None: + rest_memory_gb = self._handle_max_mamba_cache( + available_bytes / (1 << 30) + ) + available_bytes = int(rest_memory_gb * (1 << 30)) + return available_bytes + + def wrapped_init_pools(self, *args, **kwargs): + # Pool payload is never read during simulation. Preserve the native + # pool classes and allocator wiring, but allocate minimal payload + # dimensions and restore their logical metadata afterwards. + compact_attrs = ( + "qk_nope_head_dim", + "qk_rope_head_dim", + "index_head_dim", + "kv_lora_rank", + "head_dim", + "v_head_dim", + "linear_value_head_dim", + "linear_key_head_dim", + "linear_conv_kernel_dim", + ) + original_attrs = { + name: getattr(self.model_config, name) + for name in compact_attrs + if hasattr(self.model_config, name) + } + try: + for name in original_attrs: + setattr(self.model_config, name, 1) + pools = original_init_pools(self, *args, **kwargs) + finally: + for name, value in original_attrs.items(): + setattr(self.model_config, name, value) + + token_pool = pools.token_to_kv_pool + for name, value in original_attrs.items(): + if hasattr(token_pool, name): + setattr(token_pool, name, value) + + if ( + hasattr(token_pool, "kv_cache_dim") + and token_pool.kv_cache_dim == 2 + and "kv_lora_rank" in original_attrs + and "qk_rope_head_dim" in original_attrs + ): + token_pool.kv_cache_dim = ( + original_attrs["kv_lora_rank"] + original_attrs["qk_rope_head_dim"] + ) + return pools + + target.configure = wrapped_configure + target._profile_available_bytes = override_profile_available_bytes + target._init_pools = wrapped_init_pools diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/req_stats_manager.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/req_stats_manager.py new file mode 100644 index 000000000..2b6031d8d --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/req_stats_manager.py @@ -0,0 +1,22 @@ +from sglang_simulator.simulation.types import RequestStats + + +class RequestStatsManager: + """Shared request statistics manager for `Scheduler` and `HicacheController`.""" + + def __init__(self): + self.stats: dict[str, RequestStats] = {} + + def get_req_stats(self, rid: str) -> RequestStats: + if rid not in self.stats: + self.stats[rid] = RequestStats(rid=rid) + return self.stats[rid] + + def get_all_req_stats(self) -> list[RequestStats]: + return list(self.stats.values()) + + def reset(self): + self.stats.clear() + + +request_stats_manager = RequestStatsManager() diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/scheduler.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/scheduler.py new file mode 100644 index 000000000..cd9ae04b6 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/scheduler.py @@ -0,0 +1,648 @@ +import heapq +import importlib +import json +import os +import time +from dataclasses import asdict +from typing import Any + +from sglang_simulator.compat import validate_simulator_server_args +from sglang_simulator.hook import ( + BaseHook, + is_class_hook_matched, + validate_required_class_hooks, +) +from sglang_simulator.hook.utils import get_obj_from_args +from sglang_simulator.simulation.manager import ConfigManager, Envs, StateManager +from sglang_simulator.simulation.sglang.req_stats_manager import request_stats_manager +from sglang_simulator.simulation.sglang.utils import ( + resolve_model_info, + resolve_scheduler_config, +) +from sglang_simulator.simulation.types import ( + RequestStats, + SimulationMode, +) +from sglang_simulator.simulation.utils import ( + calc_iteration_metrics, + calc_metrics, +) +from sglang_simulator.time_predictor import InferTimePredictor +from sglang_simulator.time_predictor import ScheduleBatch as SimulationScheduleBatch +from sglang_simulator.time_predictor import ScheduleRequest +from sglang_simulator.utils import get_logger +from sglang_simulator.utils.json import CustomJsonEncoder + +logger = get_logger("sgl_simulator") + + +def simulation_mode_log_message(mode: SimulationMode) -> str: + return f"SGLang Simulator simulation mode: {mode.value}" + + +def effective_l2_load_delay( + load_duration: float, + last_inference_duration: float, + overlap_schedule: bool, +) -> float: + if overlap_schedule: + return max(load_duration - last_inference_duration, 0.0) + return max(load_duration, 0.0) + + +def block_on_l2_load(mode: SimulationMode, delay: float) -> float: + """Sleep for visible L2 load time and return actual blocked wall time.""" + if mode != SimulationMode.BLOCKING or delay <= 0: + return 0.0 + start = time.perf_counter() + time.sleep(delay) + return time.perf_counter() - start + + +class C_SglangPrefillAdderHook(BaseHook): + HOOK_CLASS_NAME = "PrefillAdder" + HOOK_MODULE_NAME = "sglang.srt.managers.schedule_policy" + + @classmethod + def hook(cls, target): + original_add_one_req = target.add_one_req + + def wrapped_add_one_req(self, *args, **kwargs): + req = get_obj_from_args( + "sglang.srt.managers.schedule_batch.Req", + *args, + **kwargs, + ) + req_infos = request_stats_manager.get_req_stats(req.rid) + req_infos.before_adder_device_hit_len = len(req.prefix_indices) + req_infos.final_host_hit_len = req.host_hit_length + + return original_add_one_req(self, *args, **kwargs) + + target.add_one_req = wrapped_add_one_req + + +class ReqDispatcher: + _instance = None + _initialized = False + + def __new__(cls, mode): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self, mode: SimulationMode): + if self.__class__._initialized: + return + + self.mode = mode + # If the simulation mode is `BLOCKING`, all requests are released immediately. + # If the simulation mode is `OFFLINE`, only control requests, such as `flush_cache` + # and `server_info`, are released immediately. + self.immediate_release_requests = [] + self.future_queue: list[ + tuple[float, int, Any] + ] = [] # tuple(created time, salt, request) + self.offline_recv_all_requests = False + self.profile_active = False + + @staticmethod + def simulation_created_time_s(simulation_args: dict) -> float: + if "created_time_ms" in simulation_args: + return simulation_args["created_time_ms"] / 1000.0 + return simulation_args["created_time"] + + def has_next(self) -> bool: + return len(self.future_queue) > 0 + + def next_req_from_future_ts(self) -> float: + return self.future_queue[0][0] + + def reset(self) -> None: + self.immediate_release_requests.clear() + self.future_queue.clear() + self.offline_recv_all_requests = False + + def add(self, reqs: list): + if self.mode == SimulationMode.BLOCKING: + self.immediate_release_requests.extend(reqs) + elif self.mode == SimulationMode.OFFLINE: + if self.offline_recv_all_requests: + self.immediate_release_requests.extend(reqs) + return + + gen_requests = [] + time.sleep(0.05) # waiting requests + + for req in reqs: + if req.__class__.__name__ == "TokenizedGenerateReqInput": + gen_requests.append(req) + else: + # Such as: /profile_start, /flush_cache, etc. + self.immediate_release_requests.append(req) + + # Add requests to future queue + for req in gen_requests: + sim_params = None + if req.sampling_params.custom_params is not None: + sim_params = req.sampling_params.custom_params.get("simulation") + if sim_params is None: + # There are some warm-up requests when starting the server without --skip-server-warmup. + self.immediate_release_requests.append(req) + logger.warning( + "Failed to extract the simulation parameters required for simulation from the request. Ignore this warning if the request is a warm-up request." + ) + continue + if sim_params.get("queue_start"): + logger.debug( + "Add request to waiting queue with custom queue start timestamp." + ) + + self.future_queue.append( + ( + sim_params.get("queue_start") + or self.simulation_created_time_s(sim_params), + time.time_ns(), # The request is not comparable, so add the salt to avoid comparison. + req, + ) + ) + + if len(self.future_queue) != 0: + _, _, gen_req = self.future_queue[-1] + total_request = gen_req.sampling_params.custom_params["simulation"][ + "total_request" + ] + + if len(self.future_queue) == total_request: + self.offline_recv_all_requests = True + heapq.heapify(self.future_queue) + logger.info("All requests received. Starting simulation now.") + else: + logger.info( + f"Offline simulation mode enabled. {total_request} requests expected in total. Received {len(self.future_queue)} requests so far." + ) + + def dispatch(self) -> list: + recv_reqs = [] + + recv_reqs.extend(self.immediate_release_requests) + self.immediate_release_requests.clear() + + if self.mode == SimulationMode.OFFLINE and self.offline_recv_all_requests: + # Process the arrived requests only after all requests have been added to the future queue + current_timestamp = StateManager.get_global_clock() + while len(self.future_queue) > 0: + enqueue_time, _, req = self.future_queue[0] + if enqueue_time > current_timestamp: + break + recv_reqs.append(req) + heapq.heappop(self.future_queue) + + now = time.time() + for req in recv_reqs: + if req.__class__.__name__ in [ + "BatchTokenizedGenerateReqInput", + "TokenizedGenerateReqInput", + ]: + simulation_args = None + if req.sampling_params.custom_params is not None: + simulation_args = req.sampling_params.custom_params.get( + "simulation" + ) + # The warm-up request might not include any simulation arguments. + if simulation_args is None: + if self.mode != SimulationMode.BLOCKING or not self.profile_active: + continue + simulation_args = {} + req_stats = request_stats_manager.get_req_stats(req.rid) + req_stats.rid = req.rid + req_stats.input_length = len(req.input_ids) + req_stats.output_length = req.sampling_params.max_new_tokens + + if self.mode == SimulationMode.BLOCKING: + req_stats.created_time = simulation_args.get( + "server_created_time", now + ) + req_stats.last_event_time = req_stats.created_time + req_stats.queue_start = now + elif self.mode == SimulationMode.OFFLINE: + req_stats.created_time = self.simulation_created_time_s( + simulation_args + ) + req_stats.last_event_time = req_stats.created_time + # Align with the real queue start timestamp if queue_start is not None. For debugging only. + queue_start = simulation_args.get("queue_start") + if queue_start is not None: + StateManager.set_global_clock(queue_start) + req_stats.queue_start = StateManager.get_global_clock() + + if recv_reqs and StateManager.get_last_real_time_ts() == 0: + StateManager.set_last_real_time_ts(time.time()) + StateManager.set_global_clock( + now if self.mode == SimulationMode.BLOCKING else 0 + ) + + return recv_reqs + + +class C_SchedulerRequestReceiver(BaseHook): + HOOK_CLASS_NAME = "SchedulerRequestReceiver" + HOOK_MODULE_NAME = "sglang.srt.managers.scheduler_components.request_receiver" + + # Older SGLang versions receive requests directly on Scheduler; that path is + # patched by C_SchedulerHook instead. + REQUIRED = False + + REQ_DISPATCHER: ReqDispatcher = ReqDispatcher( + SimulationMode(Envs.simulation_mode()) + ) + + @classmethod + def hook(cls, target): + original_recv_requests = target.recv_requests + + def wrapped_recv_requests(self, *args, **kwargs): + recv_reqs = original_recv_requests(self, *args, **kwargs) + C_SchedulerRequestReceiver.REQ_DISPATCHER.add(recv_reqs) + return C_SchedulerRequestReceiver.REQ_DISPATCHER.dispatch() + + target.recv_requests = wrapped_recv_requests + + +class C_SchedulerHook(BaseHook): + HOOK_CLASS_NAME = "Scheduler" + HOOK_MODULE_NAME = "sglang.srt.managers.scheduler" + + INFERENCE_PREDICTOR: InferTimePredictor = None + + ITERATION_STATS: list[dict] = [] + TOTAL_PREDICTOR_TIME_COST = 0 + GET_NEW_BATCH_PREFILL_TIME_COST = 0 + + SIMULATION_BATCH: SimulationScheduleBatch = None + OVERLAP_SCHEDULE: bool = False + SIM_MODE = SimulationMode(Envs.simulation_mode()) + # Shared singleton instance with `C_SchedulerRequestReceiver.REQ_DISPATCHER`. + REQ_DISPATCHER = ReqDispatcher(SIM_MODE) + + @classmethod + def hook(cls, target): + original_init = target.__init__ + original_recv_requests = getattr(target, "recv_requests", None) + original_prefetch_kvcache = target._prefetch_kvcache + original_get_new_batch_prefill = target.get_new_batch_prefill + original_run_batch = target.run_batch + original_process_batch_result = target.process_batch_result + original_event_loop_normal = target.event_loop_normal + original_init_request_dispatcher = target.init_request_dispatcher + + def override_event_loop_overlap(self, *args, **kwargs): + # To reduce the complexity of the simulation, the overlapping schedule is not needed. + return original_event_loop_normal(self, *args, **kwargs) + + def wrapped_init(self, *args, **kwargs): + logger.info(simulation_mode_log_message(C_SchedulerHook.SIM_MODE)) + # Supported entry points prepare the final config before publication. + server_args = get_obj_from_args( + "sglang.srt.server_args.ServerArgs", *args, **kwargs + ) + validate_simulator_server_args(server_args) + C_SchedulerHook.OVERLAP_SCHEDULE = not getattr( + server_args, "disable_overlap_schedule", False + ) + logger.debug( + f"Overlap schedule simulation mode: {C_SchedulerHook.OVERLAP_SCHEDULE}." + ) + original_init(self, *args, **kwargs) + validate_required_class_hooks() + if original_recv_requests is None and not is_class_hook_matched( + C_SchedulerRequestReceiver + ): + raise RuntimeError( + "SGLang Simulator could not hook a request receiver. The " + "simulator must be adapted to this SGLang revision." + ) + + try: + if ConfigManager.get_model_info() is None: + model = resolve_model_info(self.model_config) + ConfigManager.set_model_info(model) + + model = ConfigManager.get_model_info() + + hw = ConfigManager.get_accelerator_info() + + if ConfigManager.get_scheduler_config() is None: + sched_config = resolve_scheduler_config( + server_args=self.server_args, + model_config=self.model_config, + ) + ConfigManager.set_scheduler_config(sched_config) + sched_config = ConfigManager.get_scheduler_config() + + C_SchedulerHook.INFERENCE_PREDICTOR = ( + ConfigManager.get_inference_time_predictor(model, hw, sched_config) + ) + except Exception as e: + logger.error( + f"Failed to initialize inference time predictor. Error: {e}" + ) + raise e + + def wrapped_recv_requests(self, *args, **kwargs) -> list: + recv_reqs = original_recv_requests(self, *args, **kwargs) + C_SchedulerHook.REQ_DISPATCHER.add(recv_reqs) + return C_SchedulerHook.REQ_DISPATCHER.dispatch() + + def wrapped_get_new_batch_prefill(self, *args, **kwargs): + start = time.perf_counter() + result = original_get_new_batch_prefill(self, *args, **kwargs) + C_SchedulerHook.GET_NEW_BATCH_PREFILL_TIME_COST = ( + time.perf_counter() - start + ) + + # Accept both a plan wrapper and a direct batch return value. + new_batch = getattr(result, "batch_to_run", result) + + # A plan reports the running batch before self.running_batch is updated. + running_batch = getattr(result, "running_batch", self.running_batch) + + now = time.time() + if new_batch is not None: + for req in new_batch.reqs: + req_stats = request_stats_manager.get_req_stats(req.rid) + req_stats.final_device_hit_len = req.cached_tokens + if req_stats.queue_end == -1: + if C_SchedulerHook.SIM_MODE == SimulationMode.BLOCKING: + req_stats.queue_end = now + else: + req_stats.queue_end = StateManager.get_global_clock() + else: + # Chunked request + pass + elif len(running_batch.reqs) == 0 and len(self.waiting_queue) > 0: + # Prefetching + StateManager.step_global_clock(0.005) + StateManager.set_current_inference_dur(0.005) + else: + # Idle stage, there are some requests pendding in the future queue. + if C_SchedulerHook.SIM_MODE == SimulationMode.OFFLINE and ( + C_SchedulerHook.REQ_DISPATCHER.has_next() + and len(running_batch.reqs) == 0 + ): + next_created_time = ( + C_SchedulerHook.REQ_DISPATCHER.next_req_from_future_ts() + ) + StateManager.set_global_clock(next_created_time + 1e-6) + logger.debug( + f"Get new batch prefill: global iteration={StateManager.get_iteration()}, " + f"new batch={new_batch.batch_size() if new_batch is not None else 0}, " + f"waiting queue={len(self.waiting_queue)}" + ) + + return result + + def wrapped_prefetch_kvcache(self, *args, **kwargs): + original_prefetch_kvcache(self, *args, **kwargs) + + req = get_obj_from_args( + "sglang.srt.managers.schedule_batch.Req", + *args, + **kwargs, + ) + req_stats = request_stats_manager.get_req_stats(req.rid) + req_stats.recv_device_hit_len = len(req.prefix_indices) + req_stats.recv_host_hit_len = req.host_hit_length + + def wrapped_run_batch(self, *args, **kwargs): + ret = original_run_batch(self, *args, **kwargs) + + batch = get_obj_from_args( + "sglang.srt.managers.schedule_batch.ScheduleBatch", *args, **kwargs + ) + + if ret.__class__.__name__ == "GenerationBatchResult": + simulation_batch = SimulationScheduleBatch(reqs=[]) + if batch.forward_mode.is_extend(): + for req in batch.reqs: + extend_length = getattr(req, "extend_input_len", None) + if extend_length is None: + # The range API represents extend tokens as a half-open interval. + extend_length = req.extend_range.length + simulation_batch.reqs.append( + ScheduleRequest( + extend_length=extend_length, + past_kv_length=len(req.prefix_indices) + + len(req.output_ids), + ) + ) + elif batch.forward_mode.is_decode(): + for req in batch.reqs: + simulation_batch.reqs.append( + ScheduleRequest( + extend_length=1, + past_kv_length=len(req.prefix_indices) + + len(req.output_ids), + ) + ) + + if not simulation_batch.is_empty(): + StateManager.inc_iteration() + pred_start = time.perf_counter() + predicted_latency = ( + C_SchedulerHook.INFERENCE_PREDICTOR.predict_infer_time( + simulation_batch + ) + ) + # Accumulate predictor execution time for performance analysis. + C_SchedulerHook.TOTAL_PREDICTOR_TIME_COST += ( + time.perf_counter() - pred_start + ) + predicted_latency = float(predicted_latency) + + forward_latency = 0 + if C_SchedulerHook.SIM_MODE == SimulationMode.BLOCKING: + time.sleep(abs(predicted_latency)) + now = time.time() + forward_latency = now - StateManager.get_last_real_time_ts() + StateManager.set_last_real_time_ts(now) + else: + forward_latency = predicted_latency + + StateManager.set_current_inference_dur(forward_latency) + + C_SchedulerHook.SIMULATION_BATCH = simulation_batch + + return ret + + def wrapped_process_batch_result(self, *args, **kwargs): + process_batch_result_start = time.perf_counter() + ret = original_process_batch_result(self, *args, **kwargs) + process_batch_result_end = time.perf_counter() + + batch = get_obj_from_args( + "sglang.srt.managers.schedule_batch.ScheduleBatch", *args, **kwargs + ) + if batch is not None: + if len(batch.reqs) == 0: + return ret + + hicache_l2_load_dur = StateManager.pop_hicache_l2_load_dur() + hicache_l2_load_stats = StateManager.pop_hicache_l2_load_stats() + hicache_l2_backup_dur = StateManager.pop_hicache_l2_backup_dur() + current_inference_dur = StateManager.get_current_inference_dur() + visible_l2_load_dur = effective_l2_load_delay( + hicache_l2_load_dur, + StateManager.get_last_inference_dur(), + C_SchedulerHook.OVERLAP_SCHEDULE, + ) + blocked_l2_wall_dur = block_on_l2_load( + C_SchedulerHook.SIM_MODE, + visible_l2_load_dur, + ) + + StateManager.step_global_clock(visible_l2_load_dur) + StateManager.step_global_clock(current_inference_dur) + # Step CPU overhead BEFORE recording latencies, + # so current iter's CPU time is reflected in current iter's TTFT. + now = time.time() + cpu_overhead = max( + now - StateManager.get_last_real_time_ts() - blocked_l2_wall_dur, + 0.0, + ) + StateManager.step_global_clock(cpu_overhead) + StateManager.set_last_real_time_ts(now) + + request_response_time = StateManager.get_global_clock() + # Request statistics + for req in batch.reqs: + if len(req.output_ids) != 0: # not chunked + req_stats = request_stats_manager.get_req_stats(req.rid) + req_stats.gen_token_latencies.append( + request_response_time + - req_stats.last_event_time # queue duration + ) + req_stats.last_event_time = request_response_time + else: + # Chunked request: nothing to do + pass + # Iteration statistics + C_SchedulerHook.ITERATION_STATS.append( + { + "requests": C_SchedulerHook.SIMULATION_BATCH.request_info(), + "forward_latency": current_inference_dur, + "l2_load_latency": hicache_l2_load_dur, + "l2_blocking_wall_latency": blocked_l2_wall_dur, + **hicache_l2_load_stats, + "l2_backup_latency": hicache_l2_backup_dur, + "preprocess_latency": C_SchedulerHook.GET_NEW_BATCH_PREFILL_TIME_COST, + "postprocess_latency": process_batch_result_end + - process_batch_result_start, + "cpu_overhead": cpu_overhead, + } + ) + else: + now = time.time() + StateManager.step_global_clock( + now - StateManager.get_last_real_time_ts() + ) + StateManager.set_last_real_time_ts(now) + + return ret + + def override_profile(req, *args, **kwargs): + is_start_profile = req.req_type.name == "START_PROFILE" + stats: list[RequestStats] = [] + for item in request_stats_manager.get_all_req_stats(): + if item.rid is not None and item.input_length > 0: + stats.append(item) + + stats = sorted(stats, key=lambda req: req.created_time) + + output_dir = Envs.output_dir() + os.makedirs(output_dir, exist_ok=True) + + if len(stats) > 0: + min_created_time = stats[0].created_time + # Align timestamps + for item in stats: + item.created_time -= min_created_time + item.queue_start -= min_created_time + item.queue_end -= min_created_time + item.last_event_time -= min_created_time + + metrics = calc_metrics(stats) + metrics["time_cost"] = ( + time.time() - StateManager.get_last_flush_time_ts() + ) + metrics["predictor_time_cost"] = ( + C_SchedulerHook.TOTAL_PREDICTOR_TIME_COST + ) + metrics.update( + calc_iteration_metrics(C_SchedulerHook.ITERATION_STATS, metrics) + ) + metrics.update(C_SchedulerHook.INFERENCE_PREDICTOR.get_metrics()) + + try: + with open(f"{output_dir}/metrics.json", "w") as f: + f.write(json.dumps(metrics, cls=CustomJsonEncoder) + "\n") + + with open(f"{output_dir}/iteration.jsonl", "w") as f: + for item in C_SchedulerHook.ITERATION_STATS: + f.write(json.dumps(item) + "\n") + + with open(f"{output_dir}/request.jsonl", "w") as f: + for item in stats: + f.write(json.dumps(asdict(item)) + "\n") + + logger.info(f"Simulation results saved to {output_dir}.") + + except Exception as e: + logger.error(f"Failed to dump results. Error: {e}") + else: + logger.warning("No request statistics available.") + + StateManager.reset() + StateManager.set_last_flush_time_ts(time.time()) + request_stats_manager.reset() + C_SchedulerHook.ITERATION_STATS.clear() + C_SchedulerHook.TOTAL_PREDICTOR_TIME_COST = 0 + C_SchedulerHook.REQ_DISPATCHER.reset() + C_SchedulerHook.REQ_DISPATCHER.profile_active = is_start_profile + C_SchedulerHook.INFERENCE_PREDICTOR.reset_metrics() + + ProfileReqOutput = getattr( + importlib.import_module("sglang.srt.managers.io_struct"), + "ProfileReqOutput", + ) + result = { + "total_request": len(stats), + "output_directory": output_dir, + } + + return ProfileReqOutput( + success=True, + message=json.dumps(result), + ) + + def wrapped_init_request_dispatcher(self, *args, **kwargs): + ret = original_init_request_dispatcher(self, *args, **kwargs) + + _request_dispatcher = getattr(self, "_request_dispatcher", None) + + if _request_dispatcher is not None: + for ty in _request_dispatcher._mapping.keys(): + if ty.__name__ == "ProfileReq": + _request_dispatcher._mapping[ty] = override_profile + return ret + + target.event_loop_overlap = override_event_loop_overlap + target.__init__ = wrapped_init + target.get_new_batch_prefill = wrapped_get_new_batch_prefill + target.run_batch = wrapped_run_batch + target.process_batch_result = wrapped_process_batch_result + target._prefetch_kvcache = wrapped_prefetch_kvcache + target.init_request_dispatcher = wrapped_init_request_dispatcher + + if original_recv_requests: + target.recv_requests = wrapped_recv_requests diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/sgl_kernel_hook.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/sgl_kernel_hook.py new file mode 100644 index 000000000..4b13dc4aa --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/sgl_kernel_hook.py @@ -0,0 +1,15 @@ +import sys +import types + + +def install_load_utils_stub() -> None: + """Install the kernel loader stub before importing the sgl_kernel package.""" + module_name = "sgl_kernel.load_utils" + module = sys.modules.get(module_name) + if module is None: + module = types.ModuleType(module_name) + module.__package__ = "sgl_kernel" + sys.modules[module_name] = module + + module._load_architecture_specific_ops = lambda *args, **kwargs: None + module._preload_cuda_library = lambda *args, **kwargs: None diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/unified_radix_cache.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/unified_radix_cache.py new file mode 100644 index 000000000..3e7827be8 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/unified_radix_cache.py @@ -0,0 +1,34 @@ +from sglang_simulator.hook import BaseHook + + +class C_UnifiedRadixCacheHook(BaseHook): + """Drive Unified HiCache storage work from the simulator's logical clock.""" + + HOOK_CLASS_NAME = "UnifiedRadixCache" + HOOK_MODULE_NAME = "sglang.srt.mem_cache.unified_radix_cache" + REQUIRED = False + + @classmethod + def hook(cls, target): + original_check_hicache_events = target.check_hicache_events + + def handle_pending_operations(controller): + if controller is None: + return + backup_handler = getattr(controller, "handle_backup_operation", None) + prefetch_handler = getattr(controller, "handle_prefetch_operation", None) + if backup_handler is not None: + backup_handler() + if prefetch_handler is not None: + prefetch_handler() + + def wrapped_check_hicache_events(self, *args, **kwargs): + controller = getattr(self, "cache_controller", None) + handle_pending_operations(controller) + result = original_check_hicache_events(self, *args, **kwargs) + # Unified allocates host pages while draining its scheduler-side + # control queues. Process those newly admitted reads immediately. + handle_pending_operations(controller) + return result + + target.check_hicache_events = wrapped_check_hicache_events diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/utils.py b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/utils.py new file mode 100644 index 000000000..4e070080c --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/sglang/utils.py @@ -0,0 +1,130 @@ +import typing + +from sglang_simulator.simulation.types import SchedulerConfig +from sglang_simulator.spec import DataType, ModelInfo + +if typing.TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.server_args import ServerArgs + + +def _resolve_model_config(server_args: "ServerArgs", model_config=None): + if model_config is not None: + return model_config + + get_model_config = getattr(server_args, "get_model_config", None) + if callable(get_model_config): + return get_model_config() + + return server_args.model_config + + +def _resolved_server_args(server_args: "ServerArgs") -> dict: + """Return effective ServerArgs values when the runtime exposes them.""" + resolved_dict = getattr(server_args, "resolved_dict", None) + if not callable(resolved_dict): + return {} + + try: + values = resolved_dict() + except (AttributeError, RuntimeError, TypeError): + return {} + + return values if isinstance(values, dict) else {} + + +def resolve_scheduler_config( + server_args: "ServerArgs", + model_config: typing.Optional["ModelConfig"] = None, +) -> SchedulerConfig: + from sglang.version import __version__ + + resolved = _resolved_server_args(server_args) + + def get_arg(name: str, default=None): + value = resolved.get(name) + if value is None: + value = getattr(server_args, name, None) + return default if value is None else value + + dtype = get_arg("dtype", "auto") + if dtype == "auto": + model_config = _resolve_model_config(server_args, model_config) + dtype = str(model_config.dtype).strip("torch.") + data_type = DataType.from_torch_dtype(dtype) + return SchedulerConfig( + data_type=data_type, + kv_cache_data_type=DataType.from_torch_dtype(get_arg("kv_cache_dtype")) + or data_type, + mem_fraction_static=get_arg("mem_fraction_static"), + max_total_tokens=get_arg("max_total_tokens"), + tp_size=get_arg("tp_size"), + ep_size=get_arg("ep_size"), + dp_size=get_arg("dp_size"), + pp_size=get_arg("pp_size"), + cp_size=get_arg("attn_cp_size", 1), + cp_style=get_arg("cp_style", "none"), + page_size=get_arg("page_size"), + swa_full_tokens_ratio=get_arg("swa_full_tokens_ratio"), + kv_bytes_per_token_per_gpu=get_arg("kv_bytes_per_token_per_gpu"), + hicache_ratio=get_arg("hicache_ratio"), + enable_hierarchical_cache=get_arg("enable_hierarchical_cache"), + backend_name="sglang", + backend_version=__version__, + ) + + +def resolve_model_info(model_config: "ModelConfig") -> ModelInfo: + from sglang.srt.configs.model_config import AttentionArch + + torch_dtype = str(model_config.dtype).strip("torch.") + if model_config.attention_arch == AttentionArch.MHA: + return ModelInfo( + hf_config=model_config.hf_text_config, + model_path=model_config.model_path, + attention_arch="MHA", + context_len=model_config.context_len, + hidden_size=model_config.hidden_size, + head_dim=model_config.head_dim, + num_attention_heads=model_config.num_attention_heads, + num_hidden_layers=model_config.num_hidden_layers, + num_key_value_heads=model_config.num_key_value_heads, + v_head_dim=model_config.v_head_dim, + vocab_size=model_config.vocab_size, + # DSv4-style models (e.g. DSv4-Pro) report attention_arch=MHA because + # sglang routes them through a custom `attention_backend='dsv4'`, not + # MLA. But they still carry compress_ratios + indexer + SWA fields on + # ModelConfig, and is_dsv4() needs them to take the right calculator + # branch. getattr makes this a no-op for true MHA models. + compression_ratios=getattr(model_config, "compress_ratios", None), + indexer_head_dim=getattr(model_config, "index_head_dim", None), + window_size=getattr(model_config, "window_size", None), + qk_nope_head_dim=getattr(model_config, "qk_nope_head_dim", None), + qk_rope_head_dim=getattr(model_config, "qk_rope_head_dim", None), + torch_dtype=torch_dtype, + ) + elif model_config.attention_arch == AttentionArch.MLA: + return ModelInfo( + hf_config=model_config.hf_text_config, + model_path=model_config.model_path, + attention_arch="MLA", + context_len=model_config.context_len, + hidden_size=model_config.hidden_size, + head_dim=model_config.head_dim, + num_attention_heads=model_config.num_attention_heads, + num_hidden_layers=model_config.num_hidden_layers, + num_key_value_heads=model_config.num_key_value_heads, + v_head_dim=model_config.v_head_dim, + vocab_size=model_config.vocab_size, + qk_rope_head_dim=model_config.qk_rope_head_dim, + qk_nope_head_dim=model_config.qk_nope_head_dim, + kv_lora_rank=model_config.kv_lora_rank, + compression_ratios=getattr(model_config, "compress_ratios", None), + indexer_head_dim=getattr(model_config, "index_head_dim", None), + window_size=getattr(model_config, "window_size", None), + torch_dtype=torch_dtype, + ) + else: + raise ValueError( + f"The attention type of `{model_config.attention_arch}` is not supported now." + ) diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/types.py b/tools/sglang-simulator/src/sglang_simulator/simulation/types.py new file mode 100644 index 000000000..73a4c140a --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/types.py @@ -0,0 +1,134 @@ +from dataclasses import dataclass, field +from enum import Enum +from typing import Optional, Union + +from sglang_simulator.spec import AcceleratorInfo, DataType + + +@dataclass +class SchedulerConfig: + data_type: Optional[DataType] = ( + None # Data type for model weights and activations. If none is set, it will be automatically detected. + ) + kv_cache_data_type: Optional[DataType] = None + # AIC adapter overrides — bypass MAP_DTYPE_TO_* lookup when set. + # Pass aiconfigurator MoEQuantMode/FMHAQuantMode/CommQuantMode enum name as string + # (e.g. 'w4a8_mxfp4_mxfp8' for DSv4-Pro on Blackwell). + moe_quant_mode_override: Optional[str] = None + fmha_quant_mode_override: Optional[str] = None + comm_quant_mode_override: Optional[str] = None + mem_fraction_static: Optional[float] = None + max_total_tokens: Optional[int] = None + + tp_size: int = 1 + ep_size: int = 1 + dp_size: int = 1 + pp_size: int = 1 + cp_size: int = 1 + cp_style: str = "none" + + # DSv4 KV cache calculator inputs (sourced from server_args) + page_size: Optional[int] = None + swa_full_tokens_ratio: Optional[float] = None + + # Optional explicit override of per-GPU KV bytes/token, sourced from + # sglang server startup log: "KV Cache is allocated. #tokens: N, KV size: G GB" + # kv_bytes_per_token_per_gpu = G * 1024**3 / N + # When set, takes priority over the model-derived calculator path. + # Useful for models where sglang doesn't expose its KV calculator output + # (e.g. GlmMoeDsa) and we want metrics to match the live sglang server. + kv_bytes_per_token_per_gpu: Optional[float] = None + + # L2 host KV pool sizing: host_pool_tokens = hicache_ratio * max_total_tokens. + hicache_ratio: Optional[float] = None + enable_hierarchical_cache: Optional[bool] = None + + # framework backend + backend_name: str = "sglang" + backend_version: Optional[str] = None + + @property + def attn_tp_size(self) -> int: + divisor = self.dp_size * self.cp_size + if self.tp_size % divisor != 0: + raise ValueError( + "tp_size must be divisible by dp_size * cp_size: " + f"{self.tp_size} % ({self.dp_size} * {self.cp_size}) != 0" + ) + return self.tp_size // divisor + + @property + def attn_dp_size(self) -> int: + return self.dp_size + + @property + def moe_tp_size(self) -> int: + if self.tp_size % self.ep_size != 0: + raise ValueError( + "tp_size must be divisible by ep_size: " + f"{self.tp_size} % {self.ep_size} != 0" + ) + return self.tp_size // self.ep_size + + @property + def moe_ep_size(self) -> int: + return self.ep_size + + +class SimulationMode(Enum): + BLOCKING = "BLOCKING" + OFFLINE = "OFFLINE" + + +@dataclass(slots=True) +class RequestStats: + rid: str = "" + last_event_time: float = 0.0 + input_length: int = 1 + output_length: int = 1 + + # Prefix cache stats + recv_device_hit_len: int = 0 + # Device hit length before `get_new_batch_prefill`. + # It may decrease if queued requests trigger KV eviction. + before_adder_device_hit_len: int = 0 + final_device_hit_len: int = 0 + recv_host_hit_len: int = 0 # Host hit length before prefetch + final_host_hit_len: int = 0 # Host hit length after prefetch + recv_storage_hit_len: int = 0 # Storage hit length at prefetch enqueue + final_storage_hit_len: int = 0 # Storage hit length at prefetch end + + queue_start: float = -1 + queue_end: float = -1 + created_time: float = -1 + gen_token_latencies: list[float] = field(default_factory=list) + + def is_complete(self) -> bool: + return True + + +def _bandwidth_property(gb_attr: str): + def getter(self): + gb_value = getattr(self, gb_attr) + return gb_value * 1e9 if gb_value else None + + return property(getter) + + +@dataclass +class PlatformConfig: + device: Union[AcceleratorInfo, str] + # Storage configuration for hierarchical cache management. + disk_capacity_gb: Optional[float] = None + disk_read_bandwidth_gb: Optional[float] = None + disk_write_bandwidth_gb: Optional[float] = None + memory_capacity_gb: Optional[float] = None + memory_read_bandwidth_gb: Optional[float] = None + memory_write_bandwidth_gb: Optional[float] = None + num_device_per_node: int = 8 + + # Bandwidth properties (in bytes, converted from GB) + disk_read_bandwidth = _bandwidth_property("disk_read_bandwidth_gb") + disk_write_bandwidth = _bandwidth_property("disk_write_bandwidth_gb") + memory_read_bandwidth = _bandwidth_property("memory_read_bandwidth_gb") + memory_write_bandwidth = _bandwidth_property("memory_write_bandwidth_gb") diff --git a/tools/sglang-simulator/src/sglang_simulator/simulation/utils.py b/tools/sglang-simulator/src/sglang_simulator/simulation/utils.py new file mode 100644 index 000000000..2de2b5f2c --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/simulation/utils.py @@ -0,0 +1,239 @@ +import numpy as np +from sglang_simulator.simulation.types import RequestStats, SchedulerConfig +from sglang_simulator.spec.accelerator import AcceleratorInfo +from sglang_simulator.spec.model import ModelInfo +from sglang_simulator.time_predictor.aiconfigurator import get_perf_model + + +def calc_kv_cache_cell_elems(model_info: ModelInfo, tp_size: int, pp_size: int) -> int: + num_layers = model_info.num_hidden_layers // pp_size + if model_info.is_mla(): + return (model_info.kv_lora_rank + model_info.qk_rope_head_dim) * num_layers + else: + num_kv_heads = max(model_info.num_key_value_heads // tp_size, 1) + return num_kv_heads * model_info.head_dim * num_layers * 2 + + +def calc_kv_cache_per_layer_elems( + model_info: ModelInfo, tp_size: int, pp_size: int +) -> int: + if model_info.is_mla(): + return model_info.kv_lora_rank + model_info.qk_rope_head_dim + else: + num_kv_heads = max(model_info.num_key_value_heads // tp_size, 1) + return num_kv_heads * model_info.head_dim * 2 + + +def profile_device_available_bytes( + model: ModelInfo, device: AcceleratorInfo, scheduler_config: SchedulerConfig +) -> int: + """Return the simulated per-GPU byte budget available to KV-cache pools.""" + # Simulation capacity must come from the declared target accelerator. Do + # not fall back to the local CUDA device: doing so would make an identical + # simulation config host-dependent and could silently simulate the wrong + # hardware. + if device.hbm_capacity_gb is None: + raise ValueError( + "Cannot estimate max_total_num_tokens: the simulated accelerator " + f"{device.name!r} has no hbm_capacity_gb. Add the accelerator to " + "the simulator hardware registry, provide hbm_capacity_gb in the " + "simulation config, or set max_total_tokens explicitly. The " + "simulator never falls back to the local GPU memory capacity." + ) + + perf_model = get_perf_model(scheduler_config, model) + weights = 0 + for op in perf_model.context_ops: + weights += op.get_weights() + # Count weights on a single GPU + weights /= perf_model.config.pp_size + framework_reserved_mem_gb = 1.4 + rest_memory = ( + scheduler_config.mem_fraction_static * device.hbm_capacity_gb + - framework_reserved_mem_gb + ) * (1 << 30) - weights + return int(rest_memory) + + +def calc_input_token_metrics( + total_input: int, + total_reused_tokens: int, + total_dur_s: float, +) -> dict: + """Compute model-independent new-input token count and throughput.""" + dur_s = max(total_dur_s, 1e-9) + + total_new_input_tokens = total_input - total_reused_tokens + new_input_write_thr_tokens = total_new_input_tokens / dur_s + + return { + "total_new_input": total_new_input_tokens, + "new_input_write_throughput_tokens_per_s": new_input_write_thr_tokens, + } + + +def calc_iteration_metrics( + iteration_stats: list[dict], request_metrics: dict | None = None +) -> dict: + """Aggregate per-iteration simulator latency into result metrics.""" + iterations = len(iteration_stats) + forward_s = sum( + float(item.get("forward_latency", 0) or 0) for item in iteration_stats + ) + l2_load_s = sum( + float(item.get("l2_load_latency", 0) or 0) for item in iteration_stats + ) + cpu_s = sum(float(item.get("cpu_overhead", 0) or 0) for item in iteration_stats) + total_s = forward_s + l2_load_s + cpu_s + avg_iter_latency_ms = total_s / iterations * 1000 if iterations else 0 + metrics = { + "iterations": iterations, + "avg_iter_latency_ms": avg_iter_latency_ms, + } + + if request_metrics: + mean_ttft_ms = request_metrics.get("mean_ttft_ms") + mean_queue_ms = request_metrics.get("mean_queue_ms") + if mean_ttft_ms is not None and mean_queue_ms is not None: + mean_exec_ms = mean_ttft_ms - mean_queue_ms + metrics["mean_exec_ms"] = mean_exec_ms + metrics["avg_iters_per_req"] = ( + mean_exec_ms / avg_iter_latency_ms if avg_iter_latency_ms else None + ) + + return metrics + + +def calc_metrics(requests: list[RequestStats]) -> dict: + ttfts = [] + tpots = [] + itls = [] + e2e_latencies = [] + total_dur_s = 1e-9 + total_input = 0 + total_output = 0 + completed = 0 + total_reused_tokens = 0 + total_device_hit_tokens = 0 + total_host_hit_tokens = 0 + total_storage_hit_tokens = 0 + queue_durs = [] + dispatch_wait_durs = [] + arrival_to_prefill_durs = [] + output_token_timestamps = [] + concurrency_events = [] + for req in requests: + if not req.is_complete(): + continue + completed += 1 + ttfts.append(req.gen_token_latencies[0]) + # Queue latency is the time spent in SGLang's waiting queue before + # the request's first prefill admission. + queue_durs.append(req.queue_end - req.queue_start) + dispatch_wait_durs.append(req.queue_start - req.created_time) + arrival_to_prefill_durs.append(req.queue_end - req.created_time) + if len(req.gen_token_latencies) > 1: + # output length > 1 + tpots.append(np.mean(req.gen_token_latencies[1:])) + itls.extend(req.gen_token_latencies[1:]) + e2e_latencies.append(sum(req.gen_token_latencies)) + token_timestamp = req.created_time + for token_latency in req.gen_token_latencies: + token_timestamp += token_latency + output_token_timestamps.append(token_timestamp) + concurrency_events.append((req.created_time, 1)) + concurrency_events.append((req.last_event_time, -1)) + total_dur_s = max(total_dur_s, req.last_event_time) + total_input += req.input_length + total_output += req.output_length + total_reused_tokens += req.final_device_hit_len + total_device_hit_tokens += req.final_device_hit_len - req.final_host_hit_len + total_host_hit_tokens += req.final_host_hit_len - req.final_storage_hit_len + total_storage_hit_tokens += req.final_storage_hit_len + + input_token_metrics = calc_input_token_metrics( + total_input=total_input, + total_reused_tokens=total_reused_tokens, + total_dur_s=total_dur_s, + ) + + max_output_tokens_per_s = 0.0 + if output_token_timestamps: + first_created_time = min( + req.created_time for req in requests if req.is_complete() + ) + num_buckets = int(max(output_token_timestamps) - first_created_time) + 1 + output_tokens_per_s = np.zeros(max(num_buckets, 1)) + for timestamp in output_token_timestamps: + bucket = int(timestamp - first_created_time) + output_tokens_per_s[bucket] += 1 + max_output_tokens_per_s = float(np.max(output_tokens_per_s)) + + max_concurrent_requests = 0 + current_concurrent_requests = 0 + # Treat request intervals as [created_time, last_event_time): requests that + # finish exactly when another arrives are not simultaneously active. + for _, delta in sorted(concurrency_events, key=lambda event: (event[0], event[1])): + current_concurrent_requests += delta + max_concurrent_requests = max( + max_concurrent_requests, current_concurrent_requests + ) + + return { + "num_requests": len(requests), + "completed": completed, + "total_input": total_input, + "total_output": total_output, + "duration": total_dur_s, + "request_throughput": completed / total_dur_s, + "input_throughput": total_input / total_dur_s, + "output_throughput": total_output / total_dur_s, + "total_throughput": (total_input + total_output) / total_dur_s, + "prefix_cache_reused_ratio": ( + 0 if total_input == 0 else total_reused_tokens / total_input + ), + "kv_cache_storage_hit_ratio": ( + 0 if total_input == 0 else total_storage_hit_tokens / total_input + ), + "kv_cache_host_hit_ratio": ( + 0 if total_input == 0 else total_host_hit_tokens / total_input + ), + "kv_cache_device_hit_ratio": ( + 0 if total_input == 0 else total_device_hit_tokens / total_input + ), + **input_token_metrics, + "mean_ttft_ms": np.mean(ttfts or 0) * 1000, + "median_ttft_ms": np.median(ttfts or 0) * 1000, + "std_ttft_ms": np.std(ttfts or 0) * 1000, + "p90_ttft_ms": np.percentile(ttfts or 0, 90) * 1000, + "p95_ttft_ms": np.percentile(ttfts or 0, 95) * 1000, + "p99_ttft_ms": np.percentile(ttfts or 0, 99) * 1000, + "mean_queue_ms": max(np.mean(queue_durs or 0), 0.0) * 1000, + "mean_dispatch_wait_ms": (max(np.mean(dispatch_wait_durs or 0), 0.0) * 1000), + "mean_arrival_to_prefill_ms": ( + max(np.mean(arrival_to_prefill_durs or 0), 0.0) * 1000 + ), + "mean_tpot_ms": np.mean(tpots or 0) * 1000, + "median_tpot_ms": np.median(tpots or 0) * 1000, + "std_tpot_ms": np.std(tpots or 0) * 1000, + "p90_tpot_ms": np.percentile(tpots or 0, 90) * 1000, + "p95_tpot_ms": np.percentile(tpots or 0, 95) * 1000, + "p99_tpot_ms": np.percentile(tpots or 0, 99) * 1000, + "mean_itl_ms": np.mean(itls or 0) * 1000, + "median_itl_ms": np.median(itls or 0) * 1000, + "std_itl_ms": np.std(itls or 0) * 1000, + "p90_itl_ms": np.percentile(itls or 0, 90) * 1000, + "p95_itl_ms": np.percentile(itls or 0, 95) * 1000, + "p99_itl_ms": np.percentile(itls or 0, 99) * 1000, + "max_itl_ms": np.max(itls or 0) * 1000, + "mean_e2e_latency_ms": np.mean(e2e_latencies or 0) * 1000, + "median_e2e_latency_ms": np.median(e2e_latencies or 0) * 1000, + "std_e2e_latency_ms": np.std(e2e_latencies or 0) * 1000, + "p90_e2e_latency_ms": np.percentile(e2e_latencies or 0, 90) * 1000, + "p95_e2e_latency_ms": np.percentile(e2e_latencies or 0, 95) * 1000, + "p99_e2e_latency_ms": np.percentile(e2e_latencies or 0, 99) * 1000, + "concurrency": np.sum(e2e_latencies or 0) / total_dur_s, + "max_output_tokens_per_s": max_output_tokens_per_s, + "max_concurrent_requests": max_concurrent_requests, + "time_cost": -1, # Updated by external benchmark caller + } diff --git a/tools/sglang-simulator/src/sglang_simulator/spec/__init__.py b/tools/sglang-simulator/src/sglang_simulator/spec/__init__.py new file mode 100644 index 000000000..2acaa63a5 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/spec/__init__.py @@ -0,0 +1,5 @@ +from sglang_simulator.spec.accelerator import AcceleratorInfo +from sglang_simulator.spec.data_type import DataType +from sglang_simulator.spec.model import ModelInfo + +__all__ = ["AcceleratorInfo", "ModelInfo", "DataType"] diff --git a/tools/sglang-simulator/src/sglang_simulator/spec/accelerator/__init__.py b/tools/sglang-simulator/src/sglang_simulator/spec/accelerator/__init__.py new file mode 100644 index 000000000..f9ca1b4c5 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/spec/accelerator/__init__.py @@ -0,0 +1,4 @@ +from sglang_simulator.spec.accelerator.base import AcceleratorInfo +from sglang_simulator.spec.accelerator.info import NVIDIA + +__all__ = ["AcceleratorInfo", "NVIDIA"] diff --git a/tools/sglang-simulator/src/sglang_simulator/spec/accelerator/base.py b/tools/sglang-simulator/src/sglang_simulator/spec/accelerator/base.py new file mode 100644 index 000000000..8d5522144 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/spec/accelerator/base.py @@ -0,0 +1,100 @@ +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Dict, Optional, Union + +from sglang_simulator.spec.data_type import DataType +from sglang_simulator.utils import get_logger + +_all_accs_: Dict[str, "AcceleratorInfo"] = {} +_acc_alias: Dict[str, str] = {} + +logger = get_logger("sgl_simulator") + + +@dataclass +class AcceleratorInfo: + name: str + vendor: str + hbm_capacity_gb: int + hbm_bandwidth_gb: int + intra_node_bandwidth_gb: Optional[int] = None # scale up + inter_node_bandwidth_gb: int = 64 # scale out + device_alias: list = field(default_factory=list) + tflops: dict = field(default_factory=dict) + ref: str = "" + + @classmethod + def from_dict(cls, config: Dict, save_to_registry: bool = False): + acc = cls(**config) + if save_to_registry: + if acc.name in _acc_alias: + logger.error(f"{acc.name} is already in registry") + _all_accs_[acc.name.upper()] = acc + for alias in acc.device_alias: + if alias in _acc_alias: + logger.warning(f"Device alias [{alias}] is already in registry.") + else: + _acc_alias[alias] = acc.name.upper() + return acc + + def flops(self, datatype: Union[str, DataType] = DataType.FP16): + if isinstance(datatype, DataType): + datatype = datatype.value + return self.tflops.get(datatype, 1) * 1e12 + + def tensor_flops(self, datatype: Union[str, DataType] = DataType.FP16_TENSOR): + if isinstance(datatype, DataType): + datatype = datatype.value + if not datatype.endswith(DataType.tensor_suffix()): + datatype += DataType.tensor_suffix() + tflops = self.tflops.get(datatype, None) + return None if tflops is None else tflops * 1e12 + + @property + def hbm_io_bw(self): + return self.hbm_bandwidth_gb * 1e9 + + @property + def hbm_bytes(self): + return self.hbm_capacity_gb * 1e9 + + @property + def intra_node_bw(self) -> Optional[float]: + if self.intra_node_bandwidth_gb is None: + return None + return self.intra_node_bandwidth_gb * 1e9 + + @property + def inter_node_bw(self): + return self.inter_node_bandwidth_gb * 1e9 + + @staticmethod + def find_by_hw_name(hw_name: str) -> Union[None, "AcceleratorInfo"]: + if hw_name in _acc_alias: + hw = _all_accs_.get(_acc_alias[hw_name], None) + if hw is not None: + hw = deepcopy(hw) + hw.name = hw_name + return hw + else: + return _all_accs_.get(hw_name.upper(), None) + + @staticmethod + def list_all_hws() -> Dict[str, "AcceleratorInfo"]: + return _all_accs_ + + @classmethod + def from_config(cls, config: Dict): + hw_info = cls.find_by_hw_name(config["name"]) + return cls(**config) if hw_info is None else hw_info + + def __eq__(self, value): + if isinstance(value, str): + value = self.find_by_hw_name(value) + + if isinstance(value, AcceleratorInfo): + return _acc_alias.get(value.name, value.name.upper()) == _acc_alias.get( + self.name, self.name.upper() + ) + + return False diff --git a/tools/sglang-simulator/src/sglang_simulator/spec/accelerator/info.py b/tools/sglang-simulator/src/sglang_simulator/spec/accelerator/info.py new file mode 100644 index 000000000..f6582546b --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/spec/accelerator/info.py @@ -0,0 +1,24 @@ +from sglang_simulator.spec.accelerator.base import AcceleratorInfo + + +class NVIDIA: + NVIDIA_H20 = AcceleratorInfo.from_dict( + config={ + "name": "NVIDIA H20", + "device_alias": ["H20", "h20_sxm"], + "tflops": { + "FP8_TENSOR": 296, + "INT8_TENSOR": 296, + "FP16_TENSOR": 148, + "BF16_TENSOR": 148, + "FP32": 74, + }, + "hbm_capacity_gb": 96, + "hbm_bandwidth_gb": 4022, + "inter_node_bandwidth_gb": 64, + "intra_node_bandwidth_gb": 450, + "vendor": "NVIDIA", + "ref": "https://viperatech.com/product/nvidia-hgx-h20", + }, + save_to_registry=True, + ) diff --git a/tools/sglang-simulator/src/sglang_simulator/spec/data_type.py b/tools/sglang-simulator/src/sglang_simulator/spec/data_type.py new file mode 100644 index 000000000..b3547dbae --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/spec/data_type.py @@ -0,0 +1,104 @@ +from enum import Enum, unique +from typing import Dict, Optional + +_BYTES_MAP: dict["DataType", float] = {} +_ALIAS_MAP: Dict[str, str] = {} +_TORCH_DTYPE_TO_DATA_TYPE: Dict[str, "DataType"] = {} + + +@unique +class DataType(Enum): + INT4 = "INT4" + INT8 = "INT8" + INT16 = "INT16" + INT32 = "INT32" + INT64 = "INT64" + FP4 = "FP4" + FP8 = "FP8" + FP16 = "FP16" + BF16 = "BF16" + TF32 = "TF32" + FP32 = "FP32" + FP64 = "FP64" + # tensor + INT4_TENSOR = "INT4_TENSOR" + INT8_TENSOR = "INT8_TENSOR" + INT16_TENSOR = "INT16_TENSOR" + INT32_TENSOR = "INT32_TENSOR" + INT64_TENSOR = "INT64_TENSOR" + FP4_TENSOR = "FP4_TENSOR" + FP8_TENSOR = "FP8_TENSOR" + FP16_TENSOR = "FP16_TENSOR" + BF16_TENSOR = "BF16_TENSOR" + TF32_TENSOR = "TF32_TENSOR" + FP32_TENSOR = "FP32_TENSOR" + FP64_TENSOR = "FP64_TENSOR" + + # FIXME: This map will be added as a enum member. + + @property + def bytes(self) -> float: + return _BYTES_MAP.get(self, 1) + + @classmethod + def tensor_suffix(cls) -> str: + return "_TENSOR" + + @classmethod + def alias(cls): + return _ALIAS_MAP + + @classmethod + def from_torch_dtype(cls, dtype: str) -> Optional["DataType"]: + return _TORCH_DTYPE_TO_DATA_TYPE.get(dtype.lower()) + + +_BYTES_MAP.update( + { + DataType.INT4: 0.5, + DataType.INT8: 1, + DataType.INT16: 2, + DataType.INT32: 4, + DataType.INT64: 8, + DataType.FP4: 0.5, + DataType.FP8: 1, + DataType.FP16: 2, + DataType.BF16: 2, + DataType.TF32: 4, + DataType.FP32: 4, + DataType.FP64: 8, + DataType.INT4_TENSOR: 0.5, + DataType.INT8_TENSOR: 1, + DataType.INT16_TENSOR: 2, + DataType.INT32_TENSOR: 4, + DataType.INT64_TENSOR: 8, + DataType.FP4_TENSOR: 0.5, + DataType.FP8_TENSOR: 1, + DataType.FP16_TENSOR: 2, + DataType.BF16_TENSOR: 2, + DataType.TF32_TENSOR: 4, + DataType.FP32_TENSOR: 4, + DataType.FP64_TENSOR: 8, + } +) + +_ALIAS_MAP.update( + { + "int8": "INT8", + "float8": "FP8", + "float16": "FP16", + "float32": "FP32", + "bfloat16": "BF16", + } +) + +_TORCH_DTYPE_TO_DATA_TYPE.update( + { + "fp8": DataType.FP8, + "int8": DataType.INT8, + "float8": DataType.FP8, + "float16": DataType.FP16, + "float32": DataType.FP32, + "bfloat16": DataType.BF16, + } +) diff --git a/tools/sglang-simulator/src/sglang_simulator/spec/model/__init__.py b/tools/sglang-simulator/src/sglang_simulator/spec/model/__init__.py new file mode 100644 index 000000000..cfc17afcb --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/spec/model/__init__.py @@ -0,0 +1,3 @@ +from sglang_simulator.spec.model.base import ModelInfo + +__all__ = ["ModelInfo"] diff --git a/tools/sglang-simulator/src/sglang_simulator/spec/model/base.py b/tools/sglang-simulator/src/sglang_simulator/spec/model/base.py new file mode 100644 index 000000000..17b2beb10 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/spec/model/base.py @@ -0,0 +1,44 @@ +from dataclasses import dataclass +from typing import Optional + +from sglang_simulator.utils import get_logger + +logger = get_logger("sgl_simulator") + + +@dataclass +class ModelInfo: + hf_config: Optional[dict] = None + model_path: Optional[str] = None + + attention_arch: Optional[str] = None # MLA | MHA + context_len: Optional[int] = None + hidden_size: Optional[int] = None + head_dim: Optional[int] = None + num_attention_heads: Optional[int] = None + num_hidden_layers: Optional[int] = None + num_key_value_heads: Optional[int] = None + v_head_dim: Optional[int] = None + vocab_size: Optional[int] = None + + kv_lora_rank: Optional[int] = None + qk_rope_head_dim: Optional[int] = None + qk_nope_head_dim: Optional[int] = None + + # DSv4-specific (DSv4-Pro: per-layer compression ratios + sparse indexer + SWA) + compression_ratios: Optional[list] = None # per-layer: 4 or 128 + indexer_head_dim: Optional[int] = None + window_size: Optional[int] = None + + torch_dtype: Optional[str] = None + + # deepseek v4 model config + qk_nope_head_dim: Optional[int] = None + qk_rope_head_dim: Optional[int] = None + indexer_head_dim: Optional[int] = None + + def is_mla(self) -> bool: + return self.attention_arch == "MLA" + + def is_dsv4(self) -> bool: + return self.compression_ratios is not None diff --git a/tools/sglang-simulator/src/sglang_simulator/time_predictor/__init__.py b/tools/sglang-simulator/src/sglang_simulator/time_predictor/__init__.py new file mode 100644 index 000000000..4b5215e19 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/time_predictor/__init__.py @@ -0,0 +1,19 @@ +from sglang_simulator.time_predictor.aiconfigurator import ( + AIConfiguratorTimePredictor, +) +from sglang_simulator.time_predictor.base import ( + InferTimePredictor, + ScheduleBatch, + ScheduleRequest, +) +from sglang_simulator.time_predictor.ml import MLTimePredictor +from sglang_simulator.time_predictor.replay import ReplayTimePredictor + +__all__ = ( + ScheduleRequest, + ScheduleBatch, + InferTimePredictor, + AIConfiguratorTimePredictor, + MLTimePredictor, + ReplayTimePredictor, +) diff --git a/tools/sglang-simulator/src/sglang_simulator/time_predictor/aiconfigurator.py b/tools/sglang-simulator/src/sglang_simulator/time_predictor/aiconfigurator.py new file mode 100644 index 000000000..12d202357 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/time_predictor/aiconfigurator.py @@ -0,0 +1,291 @@ +from typing import Optional + +import numpy as np +from aiconfigurator.sdk import models +from aiconfigurator.sdk.backends.factory import get_backend +from aiconfigurator.sdk.common import ( + CommQuantMode, + DatabaseMode, + FMHAQuantMode, + GEMMQuantMode, + KVCacheQuantMode, + MoEQuantMode, +) +from aiconfigurator.sdk.config import ModelConfig, RuntimeConfig +from aiconfigurator.sdk.inference_session import InferenceSession +from aiconfigurator.sdk.perf_database import get_database, get_systems_paths +from sglang_simulator.simulation.types import ( + SchedulerConfig, +) +from sglang_simulator.spec.accelerator import AcceleratorInfo +from sglang_simulator.spec.data_type import DataType +from sglang_simulator.spec.model import ModelInfo +from sglang_simulator.time_predictor.base import ( + InferTimePredictor, + ScheduleBatch, + ScheduleRequest, +) +from sglang_simulator.utils import get_logger + +# Map the common data types to AIConfigurator data types. +MAP_DTYPE_TO_GEMMQuantMode = { + DataType.FP16: GEMMQuantMode.bfloat16, + DataType.BF16: GEMMQuantMode.bfloat16, + DataType.FP8: GEMMQuantMode.fp8_block, + DataType.INT8: GEMMQuantMode.int8_wo, + DataType.FP4: GEMMQuantMode.nvfp4, + DataType.INT4: GEMMQuantMode.int4_wo, + DataType.FP16_TENSOR: GEMMQuantMode.bfloat16, + DataType.BF16_TENSOR: GEMMQuantMode.bfloat16, + DataType.FP8_TENSOR: GEMMQuantMode.fp8, + DataType.INT8_TENSOR: GEMMQuantMode.int8_wo, + DataType.FP4_TENSOR: GEMMQuantMode.nvfp4, + DataType.INT4_TENSOR: GEMMQuantMode.int4_wo, +} + +MAP_DTYPE_TO_KVCacheQuantMode = { + DataType.FP16: KVCacheQuantMode.bfloat16, + DataType.BF16: KVCacheQuantMode.bfloat16, + DataType.FP8: KVCacheQuantMode.fp8, + DataType.INT8: KVCacheQuantMode.int8, +} + +MAP_DTYPE_TO_FMHAQuantMode = { + DataType.FP16: FMHAQuantMode.bfloat16, + DataType.BF16: FMHAQuantMode.bfloat16, + DataType.FP8: FMHAQuantMode.fp8, +} + +MAP_DTYPE_TO_MoEQuantMode = { + DataType.FP16: MoEQuantMode.bfloat16, + DataType.BF16: MoEQuantMode.bfloat16, + DataType.FP8: MoEQuantMode.fp8_block, + DataType.INT8: MoEQuantMode.fp8, + DataType.FP4: MoEQuantMode.nvfp4, + DataType.INT4: MoEQuantMode.int4_wo, +} + +MAP_DTYPE_TO_CommQuantMode = { + DataType.FP16: CommQuantMode.half, + DataType.BF16: CommQuantMode.half, + DataType.FP8: CommQuantMode.fp8, + DataType.INT8: CommQuantMode.int8, +} + + +logger = get_logger("sgl_simulator") + + +def _resolve_comm_quant_mode(sched_config: SchedulerConfig) -> CommQuantMode: + if sched_config.comm_quant_mode_override: + return getattr(CommQuantMode, sched_config.comm_quant_mode_override) + + if sched_config.data_type is None: + return CommQuantMode.half + + try: + return MAP_DTYPE_TO_CommQuantMode[sched_config.data_type] + except KeyError: + raise ValueError( + "AIConfigurator has no communication quantization mapping for " + f"model data type {sched_config.data_type.value}. Set " + "comm_quant_mode_override explicitly to half, int8, or fp8." + ) from None + + +def get_perf_model( + sched_config: SchedulerConfig, + model: ModelInfo, + workload_distribution: str = "balanced", +) -> models.BaseModel: + model_config = ModelConfig( + pp_size=sched_config.pp_size, + tp_size=sched_config.attn_tp_size, + moe_tp_size=sched_config.moe_tp_size, + moe_ep_size=sched_config.moe_ep_size, + attention_dp_size=sched_config.attn_dp_size, + cp_size=sched_config.cp_size, + cp_style=sched_config.cp_style, + gemm_quant_mode=MAP_DTYPE_TO_GEMMQuantMode.get( + sched_config.data_type, GEMMQuantMode.bfloat16 + ), + moe_quant_mode=( + getattr(MoEQuantMode, sched_config.moe_quant_mode_override) + if sched_config.moe_quant_mode_override + else MAP_DTYPE_TO_MoEQuantMode.get( + sched_config.data_type, MoEQuantMode.bfloat16 + ) + ), + kvcache_quant_mode=MAP_DTYPE_TO_KVCacheQuantMode.get( + sched_config.kv_cache_data_type, KVCacheQuantMode.bfloat16 + ), + fmha_quant_mode=( + getattr(FMHAQuantMode, sched_config.fmha_quant_mode_override) + if sched_config.fmha_quant_mode_override + else MAP_DTYPE_TO_FMHAQuantMode.get( + sched_config.kv_cache_data_type, FMHAQuantMode.bfloat16 + ) + ), + comm_quant_mode=_resolve_comm_quant_mode(sched_config), + workload_distribution=workload_distribution, + ) + + logger.info(f"Model config for AIConfigurator: {model_config}") + + return models.get_model( + model_path=model.model_path, + model_config=model_config, + backend_name=sched_config.backend_name, + ) + + +class AIConfiguratorTimePredictor(InferTimePredictor): + def __init__( + self, + model: ModelInfo, + hw: AcceleratorInfo, + config: SchedulerConfig, + database_path: Optional[str] = None, + database_mode: DatabaseMode | str = DatabaseMode.SILICON, + prefill_scale_factor: float = 1, + decode_scale_factor: float = 1, + prefill_min_latency: float = 0, + workload_distribution: str = "balanced", + enable_oom_check: bool = False, + ): + super().__init__(model, hw, config) + + self.prefill_scale_factor = prefill_scale_factor + self.decode_scale_factor = decode_scale_factor + self.prefill_min_latency = prefill_min_latency + if isinstance(database_mode, str): + database_mode = self._get_database_mode(database_mode) + + database = get_database( + system=hw.name, + backend=config.backend_name, + version=config.backend_version, + systems_paths=( + [database_path] if database_path is not None else get_systems_paths() + ), + ) + + if database is None: + raise ValueError("Failed to initialize the database.") + + database.set_default_database_mode(database_mode) + logger.info(f"AIC Database mode: {database_mode}") + + self._session = InferenceSession( + model=get_perf_model(config, model, workload_distribution), + backend=get_backend(self.config.backend_name), + database=database, + ) + + self.enable_oom_check = enable_oom_check + self._is_oom = False + + def _get_database_mode(self, mode: str) -> DatabaseMode: + return { + "SILICON": DatabaseMode.SILICON, + "HYBRID": DatabaseMode.HYBRID, + "EMPIRICAL": DatabaseMode.EMPIRICAL, + "SOL": DatabaseMode.SOL, + "SOL_FULL": DatabaseMode.SOL_FULL, + }.get(mode.upper(), DatabaseMode.SILICON) + + def ctx_attn_flops_ratio_with_avg(self, reqs: list[ScheduleRequest]) -> float: + if len(reqs) == 1: + return 1.0 + mean_past = np.mean([req.past_kv_length for req in reqs]) + mean_input = np.mean([req.extend_length for req in reqs]) + avg_flops = (mean_past + mean_past + mean_input) * mean_input / 2 * len(reqs) + + actual_flops = 0 + for req in reqs: + actual_flops += ( + (req.past_kv_length + req.past_kv_length + req.extend_length) + * req.extend_length + / 2 + ) + + return actual_flops / avg_flops + + def predict_infer_latency_dict(self, batch: ScheduleBatch) -> dict: + # Returns latency details for debugging operators. + if batch.is_decode(): + # Decode: output sequence length (osl) = 2, input sequence length (isl) = mean(past_kv_length) + isl = int(np.mean([req.past_kv_length for req in batch.reqs])) + runtime_config = RuntimeConfig(batch_size=batch.batch_size, isl=isl, osl=2) + if self.enable_oom_check: + summary = self._session.run_static(runtime_config, mode="static_gen") + latency_dict = summary.get_generation_latency_dict() + else: + # faster path + results = self._session._backend._run_static_breakdown( + self._session._model, + self._session._database, + runtime_config, + mode="static_gen", + ) + latency_dict = results[2] + else: + # Prefill: output sequence length (osl) = 1, input sequence length (isl) = mean(past_kv + input), prefix = mean(past_kv) + mean_past = np.mean([req.past_kv_length for req in batch.reqs]) + mean_input = np.mean([req.extend_length for req in batch.reqs]) + isl = int(mean_past + mean_input) + prefix = int(mean_past) + runtime_config = RuntimeConfig( + batch_size=batch.batch_size, isl=isl, prefix=prefix, osl=1 + ) + + seq_imbalance_correction_scale = self.ctx_attn_flops_ratio_with_avg( + batch.reqs + ) + if seq_imbalance_correction_scale >= 0.4: + runtime_config = RuntimeConfig( + batch_size=batch.batch_size, + isl=isl, + prefix=prefix, + osl=1, + seq_imbalance_correction_scale=seq_imbalance_correction_scale, + ) + else: + runtime_config = RuntimeConfig( + batch_size=batch.batch_size, isl=isl, prefix=prefix, osl=1 + ) + + if self.enable_oom_check: + summary = self._session.run_static(runtime_config, mode="static_ctx") + latency_dict = summary.get_context_latency_dict() + else: + # faster path + results = self._session._backend._run_static_breakdown( + self._session._model, + self._session._database, + runtime_config, + mode="static_ctx", + ) + latency_dict = results[0] + return latency_dict + + def predict_infer_time(self, batch: ScheduleBatch) -> float: + latency_dict = self.predict_infer_latency_dict(batch) + infer_time = sum(latency_dict.values()) + + if self._is_oom: + logger.warning("Out of memory detected during estimation.") + infer_time = -infer_time + if batch.is_decode(): + infer_time *= self.decode_scale_factor + else: + infer_time *= self.prefill_scale_factor + + if not batch.is_decode(): + infer_time = ( + max(infer_time, self.prefill_min_latency) + if infer_time > 0 + else infer_time + ) + + return infer_time / 1e3 diff --git a/tools/sglang-simulator/src/sglang_simulator/time_predictor/base.py b/tools/sglang-simulator/src/sglang_simulator/time_predictor/base.py new file mode 100644 index 000000000..2be0b6986 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/time_predictor/base.py @@ -0,0 +1,97 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass, field + +from sglang_simulator.simulation.types import SchedulerConfig +from sglang_simulator.spec.accelerator import AcceleratorInfo +from sglang_simulator.spec.model import ModelInfo +from sglang_simulator.utils import get_logger + +logger = get_logger("sgl_simulator") + + +@dataclass +class ScheduleRequest: + extend_length: int = 0 + past_kv_length: int = 0 + + +@dataclass +class ScheduleBatch: + reqs: list[ScheduleRequest] = field(default_factory=list) + + def __repr__(self) -> str: + return f"batch_size={len(self.reqs)},reqs={[(req.extend_length, req.past_kv_length) for req in self.reqs]}" + + def __eq__(self, batch: "ScheduleBatch"): + if self.batch_size != batch.batch_size: + return False + + req1, req2 = [], [] + for idx in range(self.batch_size): + req1.append((self.reqs[idx].extend_length, self.reqs[idx].past_kv_length)) + req2.append((batch.reqs[idx].extend_length, batch.reqs[idx].past_kv_length)) + + return sorted(req1) == sorted(req2) + + def request_info(self) -> list[list[int, int]]: + # The request information organized in the format `(input_len, past_kv_len)` + return [[req.extend_length, req.past_kv_length] for req in self.reqs] + + @property + def num_context_tokens(self) -> int: + return sum(req.extend_length for req in self.reqs) + + @property + def total_past_kv_length(self) -> int: + return sum(req.past_kv_length for req in self.reqs) + + @property + def batch_size(self) -> int: + return len(self.reqs) + + def is_empty(self) -> bool: + return len(self.reqs) == 0 + + def is_prefill(self) -> bool: + return not self.is_decode() + + def is_decode(self) -> bool: + for req in self.reqs: + if req.extend_length > 1: + return False + return True + + @property + def num_ctx_requests(self) -> int: + return self.batch_size if self.is_prefill() else 0 + + @property + def num_gen_requests(self) -> int: + return self.batch_size if self.is_decode() else 0 + + +class InferTimePredictor(ABC): + def __init__( + self, + model: ModelInfo, + hw: AcceleratorInfo, + config: SchedulerConfig, + *args, + **kwargs, + ): + self.model: ModelInfo = model + self.hw: AcceleratorInfo = hw + self.config: SchedulerConfig = config + + @abstractmethod + def predict_infer_time(self, batch: ScheduleBatch) -> float: + # Return the inference time in seconds. Return a negative value if an exception occurs (e.g., out of memory). + pass + + def get_metrics(self) -> dict: + """Return predictor-specific metrics for the current profile interval.""" + return {} + + def reset_metrics(self) -> None: + """Reset predictor-specific metrics after a profile flush.""" + return None diff --git a/tools/sglang-simulator/src/sglang_simulator/time_predictor/ml.py b/tools/sglang-simulator/src/sglang_simulator/time_predictor/ml.py new file mode 100644 index 000000000..d49569556 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/time_predictor/ml.py @@ -0,0 +1,156 @@ +"""ML-trained per-iter latency predictor. + +Loads a joblib pickle of a sklearn-compatible regressor and predicts forward latency +from batch composition features. Train one with `train_latency_model.py`. + +sim_config.json usage: + "predictor": { + "name": "ml", + "database_path": "/path/to/latency_model.pkl" + } +""" + +import math +import os + +import joblib +from sglang_simulator.simulation.types import SchedulerConfig +from sglang_simulator.spec.accelerator import AcceleratorInfo +from sglang_simulator.spec.model import ModelInfo +from sglang_simulator.time_predictor.base import InferTimePredictor, ScheduleBatch +from sglang_simulator.utils import get_logger + +logger = get_logger("sgl_simulator") + + +class MLTimePredictor(InferTimePredictor): + """Per-iter latency predictor backed by an offline-trained sklearn regressor. + + Features (18 dim) extracted from ScheduleBatch: + batch_size, sum/max/min(extend), sum/max/min(past), + sum(extend*past), sum(extend^2), sum(past^2), + sum_attn_flops (= sum(e*(p+e/2))), + sum(extend × max_past), log1p(sum_past), log1p(sum_attn_flops), + batch_size × sum_extend, max_past - min_past, + is_decode, is_prefill + """ + + # This ordered list is the ABI between offline training and simulation. + # The concrete regressor algorithm is intentionally unrestricted as long + # as it exposes sklearn-compatible predict([[18 features]]) -> [seconds]. + + FEATURE_NAMES = [ + "batch_size", + "sum_extend", + "max_extend", + "min_extend", + "sum_past", + "max_past", + "min_past", + "sum_extend_x_past", + "sum_extend_squared", + "sum_past_squared", + "sum_attn_flops", + "sum_extend_x_max_past", + "log1p_sum_past", + "log1p_sum_attn_flops", + "batch_size_x_sum_extend", + "max_past_minus_min_past", + "is_decode", + "is_prefill", + ] + + def __init__( + self, + model: ModelInfo, + hw: AcceleratorInfo, + config: SchedulerConfig, + database_path: str, + latency_scale: float = 1.0, + **kwargs, + ): + super().__init__(model, hw, config) + database_path = os.path.expandvars(os.path.expanduser(database_path)) + if not database_path or not os.path.exists(database_path): + raise FileNotFoundError( + f"MLTimePredictor database_path not found: {database_path}. " + "Train one with `train_latency_model.py` first." + ) + + bundle = joblib.load(database_path) + if ( + not isinstance(bundle, dict) + or "model" not in bundle + or "features" not in bundle + ): + raise ValueError( + "MLTimePredictor requires a joblib bundle containing both " + "'model' and ordered 'features' metadata" + ) + self._model = bundle["model"] + saved_features = list(bundle["features"]) + + if saved_features != self.FEATURE_NAMES: + raise ValueError( + "MLTimePredictor feature contract mismatch: " + f"saved={saved_features}, expected={self.FEATURE_NAMES}. " + "Retrain or export the model with the exact 18-feature ABI." + ) + if not callable(getattr(self._model, "predict", None)): + raise TypeError( + "MLTimePredictor model must expose a callable predict() method" + ) + + self._features = saved_features + self._call_count = 0 + self._latency_scale = float(latency_scale) + logger.info( + "MLTimePredictor loaded from %s (model=%s, n_features=%d, latency_scale=%.4f)", + database_path, + type(self._model).__name__, + len(self._features), + self._latency_scale, + ) + + def predict_infer_time(self, batch: ScheduleBatch) -> float: + if batch.is_empty(): + return 0.0 + + exts = [req.extend_length for req in batch.reqs] + pasts = [req.past_kv_length for req in batch.reqs] + + bs = len(exts) + sum_e = sum(exts) + sum_p = sum(pasts) + sum_ep = sum(e * p for e, p in zip(exts, pasts)) + sum_e2 = sum(e * e for e in exts) + sum_p2 = sum(p * p for p in pasts) + sum_attn = sum(e * (p + e / 2) for e, p in zip(exts, pasts)) + max_e = max(exts) + max_p = max(pasts) + min_e = min(exts) + min_p = min(pasts) + + feats = [ + bs, + sum_e, + max_e, + min_e, + sum_p, + max_p, + min_p, + sum_ep, + sum_e2, + sum_p2, + sum_attn, + sum_e * max_p, + math.log1p(sum_p), + math.log1p(sum_attn), + bs * sum_e, + max_p - min_p, + int(all(e == 1 for e in exts)), + int(any(e > 1 for e in exts)), + ] + + self._call_count += 1 + return float(self._model.predict([feats])[0]) * self._latency_scale diff --git a/tools/sglang-simulator/src/sglang_simulator/time_predictor/replay.py b/tools/sglang-simulator/src/sglang_simulator/time_predictor/replay.py new file mode 100644 index 000000000..ca613869f --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/time_predictor/replay.py @@ -0,0 +1,189 @@ +"""Oracle lookup predictor — replays real GPU iter_latency from a pre-built table. + +Replay is a diagnostic predictor for separating latency-prediction error from +scheduler, cache, and simulator behavior. A replay table maps a JSON-encoded, +sorted list of ``[extend_input_length, prefix_length]`` pairs to measured +iteration latency in seconds. + +Example simulator configuration: + "predictor": { + "name": "replay", + "database_path": "/path/to/replay_table.json", + "miss_strategy": "knn", # "zero" (legacy) or "knn" (interpolated) + "miss_knn_k": 3, # KNN k for "knn" strategy + "miss_fallback_seconds": 0.0 # used only when strategy=="zero" + } +""" + +import json +import math +import os + +from sglang_simulator.simulation.types import SchedulerConfig +from sglang_simulator.spec.accelerator import AcceleratorInfo +from sglang_simulator.spec.model import ModelInfo +from sglang_simulator.time_predictor.base import InferTimePredictor, ScheduleBatch +from sglang_simulator.utils import get_logger + +logger = get_logger("sgl_simulator") + + +def _decode_key(key: str): + """Parse a sorted-tuple lookup key back to list of (extend, past) pairs.""" + return [tuple(pair) for pair in json.loads(key)] + + +def _shape_feat(extends, pasts): + """3-D feature for KNN: (batch_size, sum_extend, sum_past). + + Low-dim and aligned with the dominant latency drivers; avoids the curse of + dimensionality on the typically small (~1-3K) entries per replay table. + """ + return (len(extends), sum(extends), sum(pasts)) + + +class ReplayTimePredictor(InferTimePredictor): + """Oracle lookup predictor: returns real GPU iter_latency for matching batch compositions. + + Compositions are matched exactly by sorted (extend_len, past_kv_len) tuples. + On lookup miss the behavior is controlled by `miss_strategy`: + + - "zero" (default, legacy): returns `miss_fallback_seconds` (default 0.0). + Use for "is the gap NOT from the predictor?" diagnostic. Only meaningful + when miss rate is small AND you accept the bias of dropping miss work. + + - "knn": KNN-interpolated latency from the k batches in the table with + the closest (batch_size, sum_extend, sum_past) shape (per-dim z-scored + Euclidean distance). Use when miss rate matters or you want a logically + complete oracle. Typical k=3. + + For workloads where sim batch composition diverges from real (e.g., bursty + max-tps), the miss rate can be high — always check the post-run hit ratio. + """ + + def __init__( + self, + model: ModelInfo, + hw: AcceleratorInfo, + config: SchedulerConfig, + database_path: str, + miss_fallback_seconds: float = 0.0, + miss_strategy: str = "zero", + miss_knn_k: int = 3, + **kwargs, + ): + super().__init__(model, hw, config) + if not database_path or not os.path.exists(database_path): + raise FileNotFoundError( + f"ReplayTimePredictor database_path not found: {database_path}" + ) + with open(database_path) as f: + self._table = json.load(f) + if miss_strategy not in ("zero", "knn"): + raise ValueError( + f"miss_strategy must be 'zero' or 'knn', got {miss_strategy!r}" + ) + self._miss_strategy = miss_strategy + self._miss_fallback = float(miss_fallback_seconds) + self._miss_knn_k = int(miss_knn_k) + self._hits = 0 + self._misses = 0 + + if self._miss_strategy == "knn": + self._prep_knn_index() + logger.info( + "ReplayTimePredictor loaded %d unique compositions from %s " + "(miss_strategy=knn, k=%d)", + len(self._table), + database_path, + self._miss_knn_k, + ) + else: + self._knn_feats = None + logger.info( + "ReplayTimePredictor loaded %d unique compositions from %s " + "(miss_strategy=zero, fallback=%.4fs)", + len(self._table), + database_path, + self._miss_fallback, + ) + + def _prep_knn_index(self): + """Build per-feature mean/std + cached (feat, lat) arrays for KNN fallback. + + Uses plain Python (no numpy / sklearn dependency from the predictor side) + — table size is ~1K-3K so this is fine. Distance is z-scored Euclidean + across (batch_size, sum_extend, sum_past). + """ + feats = [] + lats = [] + for key, lat in self._table.items(): + extends_pasts = _decode_key(key) + extends = [e for e, _ in extends_pasts] + pasts = [p for _, p in extends_pasts] + feats.append(_shape_feat(extends, pasts)) + lats.append(float(lat)) + n = len(feats) + # per-dim mean/std for z-scoring + means = [sum(f[d] for f in feats) / n for d in range(3)] + var = [sum((f[d] - means[d]) ** 2 for f in feats) / n for d in range(3)] + stds = [math.sqrt(v) if v > 1e-12 else 1.0 for v in var] + self._knn_feats = feats + self._knn_lats = lats + self._knn_mean = means + self._knn_std = stds + + def _knn_predict(self, query_feat): + """k nearest neighbors over z-scored 3-D shape feature, simple mean.""" + qz = tuple( + (query_feat[d] - self._knn_mean[d]) / self._knn_std[d] for d in range(3) + ) + # compute squared distance to every table entry; pick smallest k + # n ≤ ~3K so O(n) per query is fine; sim only calls on misses + dists = [] + for i, f in enumerate(self._knn_feats): + fz = ( + (f[0] - self._knn_mean[0]) / self._knn_std[0], + (f[1] - self._knn_mean[1]) / self._knn_std[1], + (f[2] - self._knn_mean[2]) / self._knn_std[2], + ) + d2 = (fz[0] - qz[0]) ** 2 + (fz[1] - qz[1]) ** 2 + (fz[2] - qz[2]) ** 2 + dists.append((d2, i)) + dists.sort(key=lambda t: t[0]) + k = min(self._miss_knn_k, len(dists)) + return sum(self._knn_lats[i] for _, i in dists[:k]) / k + + def predict_infer_time(self, batch: ScheduleBatch) -> float: + if batch.is_empty(): + return 0.0 + extends = [req.extend_length for req in batch.reqs] + pasts = [req.past_kv_length for req in batch.reqs] + key = json.dumps( + sorted([req.extend_length, req.past_kv_length] for req in batch.reqs) + ) + v = self._table.get(key) + if v is None: + self._misses += 1 + if self._miss_strategy == "knn": + return self._knn_predict(_shape_feat(extends, pasts)) + return self._miss_fallback + self._hits += 1 + return float(v) + + def get_metrics(self) -> dict: + total = self._hits + self._misses + return { + "replay_exact_match_steps": self._hits, + "replay_miss_steps": self._misses, + "replay_zero_fallback_steps": ( + self._misses if self._miss_strategy == "zero" else 0 + ), + "replay_knn_fallback_steps": ( + self._misses if self._miss_strategy == "knn" else 0 + ), + "replay_fallback_rate": self._misses / total if total else 0.0, + } + + def reset_metrics(self) -> None: + self._hits = 0 + self._misses = 0 diff --git a/tools/sglang-simulator/src/sglang_simulator/utils/__init__.py b/tools/sglang-simulator/src/sglang_simulator/utils/__init__.py new file mode 100644 index 000000000..9ce512c9d --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/utils/__init__.py @@ -0,0 +1,3 @@ +from sglang_simulator.utils.logger import get_logger + +__all__ = ["get_logger"] diff --git a/tools/sglang-simulator/src/sglang_simulator/utils/json.py b/tools/sglang-simulator/src/sglang_simulator/utils/json.py new file mode 100644 index 000000000..3910442f0 --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/utils/json.py @@ -0,0 +1,22 @@ +import json +from dataclasses import asdict, is_dataclass +from enum import Enum + +import numpy as np + + +class CustomJsonEncoder(json.JSONEncoder): + def default(self, obj): + # Enum + if isinstance(obj, Enum): + return obj.value + # Dataclass + if is_dataclass(obj): + return asdict(obj) + # Numpy + if isinstance(obj, (np.int32, np.int64, np.float32, np.float64)): + return int(obj) if isinstance(obj, (np.int32, np.int64)) else float(obj) + if isinstance(obj, np.ndarray): + return obj.tolist() + # Other + return super().default(obj) diff --git a/tools/sglang-simulator/src/sglang_simulator/utils/logger.py b/tools/sglang-simulator/src/sglang_simulator/utils/logger.py new file mode 100644 index 000000000..722bd248a --- /dev/null +++ b/tools/sglang-simulator/src/sglang_simulator/utils/logger.py @@ -0,0 +1,16 @@ +import logging + + +def get_logger(name: str = "sglang_simulator") -> logging.Logger: + logger = logging.getLogger(name) + if not logger.handlers: + logger.setLevel(logging.INFO) + handler = logging.StreamHandler() + formatter = logging.Formatter( + fmt="%(asctime)s %(levelname)s [%(name)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + handler.setFormatter(formatter) + logger.addHandler(handler) + logger.propagate = False + return logger diff --git a/tools/sglang-simulator/src/usercustomize.py b/tools/sglang-simulator/src/usercustomize.py new file mode 100644 index 000000000..450c21622 --- /dev/null +++ b/tools/sglang-simulator/src/usercustomize.py @@ -0,0 +1,15 @@ +"""Early CPU-simulation compatibility for spawned SGLang workers.""" + +import os + +if ( + os.environ.get("SGLANG_SIMULATOR_BOOTSTRAP") == "1" + and os.environ.get("SGLANG_USE_CPU_ENGINE") == "1" +): + import torch + + # Some model-specific import-time checks probe the target GPU even though + # SGLang Simulator never executes a real model forward. Spawned workers reach those + # imports before the simulator target wrapper can run. CPU simulation is + # explicit, so physical GPU visibility must not affect this shim. + torch.cuda.get_device_capability = lambda *_args, **_kwargs: (10, 0) diff --git a/tools/sglang-simulator/test/assets/qwen3-8b/config.json b/tools/sglang-simulator/test/assets/qwen3-8b/config.json new file mode 100644 index 000000000..50965de9e --- /dev/null +++ b/tools/sglang-simulator/test/assets/qwen3-8b/config.json @@ -0,0 +1,28 @@ +{ + "architectures": ["Qwen3ForCausalLM"], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 12288, + "max_position_embeddings": 40960, + "max_window_layers": 36, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 36, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 1000000, + "sliding_window": null, + "tie_word_embeddings": false, + "torch_dtype": "bfloat16", + "transformers_version": "5.12.1", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936 +} diff --git a/tools/sglang-simulator/test/test_simulation_cache_hit_ratio.py b/tools/sglang-simulator/test/test_simulation_cache_hit_ratio.py new file mode 100644 index 000000000..bbafadfe2 --- /dev/null +++ b/tools/sglang-simulator/test/test_simulation_cache_hit_ratio.py @@ -0,0 +1,68 @@ +import json + +import pytest +from sglang_simulator.simulation.benchmark import BenchmarkConfig +from test_simulation_sglang_runner import make_fixed_dataset, make_sglang_runner +from test_simulation_sglang_serving import ( + SIM_CONFIGS, + SGLangServingRunner, + assert_decode_metrics, +) + + +def test_in_process_runner_reports_each_cache_tier(tmp_path): + runner = make_sglang_runner(tmp_path) + benchmark_config = BenchmarkConfig(request_rate=10, ignore_request_timestamp=False) + cached_ds = make_fixed_dataset(1000, 8) + evict_l1_ds = make_fixed_dataset(2000, 10) + evict_l2_ds = make_fixed_dataset(3000, 20) + + try: + metrics = runner.benchmark(benchmark_config, dataset=cached_ds) + assert metrics["completed"] == len(cached_ds) + assert metrics["prefix_cache_reused_ratio"] == 0 + + metrics = runner.benchmark(benchmark_config, dataset=cached_ds) + assert metrics["kv_cache_device_hit_ratio"] > 0.95 + + runner.benchmark(benchmark_config, dataset=evict_l1_ds) + metrics = runner.benchmark(benchmark_config, dataset=cached_ds) + assert metrics["kv_cache_host_hit_ratio"] > 0.95 + + runner.benchmark(benchmark_config, dataset=evict_l2_ds) + metrics = runner.benchmark(benchmark_config, dataset=cached_ds) + assert metrics["kv_cache_storage_hit_ratio"] > 0.95 + finally: + runner.shutdown() + + +def test_second_replay_benchmark_hits_all_reusable_prefix_tokens(tmp_path, monkeypatch): + # This test validates cache reuse across consecutive benchmark runs. + monkeypatch.setenv("SGLANG_IS_IN_CI", "false") + + runner = SGLangServingRunner(SIM_CONFIGS["replay"], tmp_path) + try: + first_metrics = runner.benchmark(tmp_path / "benchmark-first.json") + second_metrics = runner.benchmark(tmp_path / "benchmark-second.json") + finally: + runner.shutdown() + + assert_decode_metrics(first_metrics) + assert_decode_metrics(second_metrics) + + assert second_metrics["total_input"] == 24 + assert second_metrics["total_new_input"] == 3 + assert second_metrics["prefix_cache_reused_ratio"] == pytest.approx(0.875) + assert second_metrics["kv_cache_device_hit_ratio"] == pytest.approx(0.875) + assert second_metrics["kv_cache_host_hit_ratio"] == 0 + assert second_metrics["kv_cache_storage_hit_ratio"] == 0 + + requests = [ + json.loads(line) + for line in (runner.output_dir / "request.jsonl") + .read_text(encoding="utf-8") + .splitlines() + ] + assert len(requests) == 3 + assert all(request["input_length"] == 8 for request in requests) + assert all(request["final_device_hit_len"] == 7 for request in requests) diff --git a/tools/sglang-simulator/test/test_simulation_offline_blocking.py b/tools/sglang-simulator/test/test_simulation_offline_blocking.py new file mode 100644 index 000000000..dbb89faef --- /dev/null +++ b/tools/sglang-simulator/test/test_simulation_offline_blocking.py @@ -0,0 +1,80 @@ +import json + +import pytest +from test_simulation_sglang_serving import ( + SIM_CONFIGS, + SGLangServingRunner, + assert_decode_metrics, +) + +REQUEST_RATE = 1 +SEED = 123 +RELATIVE_TOLERANCES = { + "duration": 0.01, + "request_throughput": 0.01, + "input_throughput": 0.01, + "output_throughput": 0.01, + "mean_e2e_latency_ms": 0.10, + "mean_ttft_ms": 0.10, + "mean_tpot_ms": 0.10, + "mean_itl_ms": 0.10, +} + + +def _relative_error(actual, expected): + return abs(actual - expected) / abs(expected) + + +def _run_mode(mode, tmp_path): + case_dir = tmp_path / mode + case_dir.mkdir() + runner = SGLangServingRunner(SIM_CONFIGS["aic_sol"], case_dir, mode=mode) + try: + metrics = runner.benchmark( + case_dir / "benchmark.json", request_rate=REQUEST_RATE, seed=SEED + ) + finally: + runner.shutdown() + + requests = [ + json.loads(line) + for line in (runner.output_dir / "request.jsonl") + .read_text(encoding="utf-8") + .splitlines() + ] + requests.sort(key=lambda request: request["created_time"]) + return metrics, requests + + +def test_request_rate_offline_matches_blocking(tmp_path): + offline_metrics, offline_requests = _run_mode("offline", tmp_path) + blocking_metrics, blocking_requests = _run_mode("blocking", tmp_path) + + for metrics in (offline_metrics, blocking_metrics): + assert_decode_metrics(metrics) + + assert len(offline_requests) == len(blocking_requests) == 3 + + offline_arrivals = [request["created_time"] for request in offline_requests] + blocking_arrivals = [request["created_time"] for request in blocking_requests] + assert offline_arrivals[1] > 0.5 + assert blocking_arrivals[1] > 0.5 + assert offline_arrivals == pytest.approx(blocking_arrivals, abs=0.02) + assert ( + offline_metrics["max_concurrent_requests"] + == blocking_metrics["max_concurrent_requests"] + == 1 + ) + + for key in ("completed", "total_input", "total_output"): + assert offline_metrics[key] == blocking_metrics[key] + + for key, tolerance in RELATIVE_TOLERANCES.items(): + error = _relative_error(offline_metrics[key], blocking_metrics[key]) + assert error <= tolerance, ( + key, + offline_metrics[key], + blocking_metrics[key], + error, + tolerance, + ) diff --git a/tools/sglang-simulator/test/test_simulation_sglang_runner.py b/tools/sglang-simulator/test/test_simulation_sglang_runner.py new file mode 100644 index 000000000..b0d4d60ff --- /dev/null +++ b/tools/sglang-simulator/test/test_simulation_sglang_runner.py @@ -0,0 +1,118 @@ +import atexit +import json +import os +import sys +from pathlib import Path +from unittest.mock import patch + +from sglang_simulator.dataset import GenericRequest, SimpleDataset +from sglang_simulator.simulation.benchmark import BenchmarkConfig + +ASSETS = Path(__file__).parent / "assets" +SGLANG_ROOT = Path(__file__).parents[3] +if str(SGLANG_ROOT) not in sys.path: + sys.path.insert(0, str(SGLANG_ROOT)) +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") + + +def make_fixed_dataset( + start_token: int, + count: int, + *, + input_length: int = 1025, + output_length: int = 1, +) -> SimpleDataset: + return SimpleDataset( + reqs=[ + GenericRequest( + token_ids=[start_token + i] * input_length, + input_length=input_length, + output_length=output_length, + custom_params={"created_time": i / 10}, + ) + for i in range(count) + ] + ) + + +def _write_sim_config(tmp_path: Path) -> Path: + table_path = tmp_path / "replay.json" + table_path.write_text( + json.dumps({"[[1, 1024]]": 0.001, "[[1025, 0]]": 0.01}), + encoding="utf-8", + ) + config = { + "platform": { + "accelerator": {"name": "a100_sxm", "hbm_capacity_gb": 80}, + "disk_read_bandwidth_gb": 8, + "disk_write_bandwidth_gb": 8, + "memory_read_bandwidth_gb": 64, + "memory_write_bandwidth_gb": 64, + "num_device_per_node": 8, + }, + "predictor": { + "name": "replay", + "database_path": str(table_path), + "miss_strategy": "knn", + "miss_knn_k": 1, + }, + "scheduler": {"tp_size": 1, "ep_size": 1, "dp_size": 1}, + } + config_path = tmp_path / "sim_config.json" + config_path.write_text(json.dumps(config), encoding="utf-8") + return config_path + + +def make_sglang_runner(tmp_path: Path): + os.environ["SGLANG_SIMULATOR_CONFIG_PATH"] = str(_write_sim_config(tmp_path)) + + from benchmark.simulator.bench_runner import SGLangBenchmarkRunner + from sglang.srt.server_args import ServerArgs + + runner = SGLangBenchmarkRunner( + server_args=ServerArgs( + model_path=str(ASSETS / "qwen3-8b"), + load_format="dummy", + device="cpu", + enable_hierarchical_cache=True, + hicache_ratio=2, + hicache_storage_backend="file", + hicache_storage_prefetch_policy="wait_complete", + max_total_tokens=10 * 1024, + page_size=256, + skip_tokenizer_init=True, + ) + ) + runner.clear_hicache_storage() + return runner + + +def test_benchmark_sglang_runs_paged_decode(tmp_path): + runner = make_sglang_runner(tmp_path) + dataset = make_fixed_dataset( + 1000, + 2, + input_length=1024, + output_length=2, + ) + + try: + metrics = runner.benchmark( + BenchmarkConfig(request_rate=10, ignore_request_timestamp=False), + dataset=dataset, + ) + request_stats = runner.get_request_stats() + finally: + with patch.object(atexit, "unregister", wraps=atexit.unregister) as unregister: + runner.shutdown() + runner.shutdown() + + unregister.assert_called_once_with(runner.engine.shutdown) + + assert metrics["completed"] == len(dataset) + assert metrics["total_input"] == 2 * 1024 + assert metrics["total_output"] == 2 * 2 + assert metrics["mean_tpot_ms"] > 0 + assert all( + idx == 0 or req["created_time"] > 0 for idx, req in enumerate(request_stats) + ) diff --git a/tools/sglang-simulator/test/test_simulation_sglang_serving.py b/tools/sglang-simulator/test/test_simulation_sglang_serving.py new file mode 100644 index 000000000..c7ab574f3 --- /dev/null +++ b/tools/sglang-simulator/test/test_simulation_sglang_serving.py @@ -0,0 +1,163 @@ +import json +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path + +import pytest +import requests + +ASSETS = Path(__file__).parent / "assets" +SGLANG_ROOT = Path(__file__).parents[3] +BENCH_SERVING = SGLANG_ROOT / "benchmark" / "simulator" / "bench_serving.py" +EXAMPLES = Path(__file__).parent.parent / "examples" +SIM_CONFIGS = { + "aic_sol": EXAMPLES / "sim_configs" / "aic_sol.json", + "aic_silicon": EXAMPLES / "sim_configs" / "aic_silicon.json", + "ml": EXAMPLES / "sim_configs" / "ml.json", + "replay": EXAMPLES / "sim_configs" / "replay.json", +} + + +class SGLangServingRunner: + def __init__(self, config_path: Path, tmp_path: Path, mode: str = "offline"): + self.mode = mode + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + self.port = sock.getsockname()[1] + + self.output_dir = tmp_path / "output" + env = os.environ.copy() + env.update( + CUDA_VISIBLE_DEVICES="", + SGLANG_USE_CPU_ENGINE="1", + SGLANG_SIMULATOR_CONFIG_PATH=str(config_path), + SGLANG_SIMULATOR_OUTPUT_MODE=mode.upper(), + SGLANG_SIMULATOR_OUTPUT_DIR=str(self.output_dir), + ) + cmd = [ + sys.executable, + "-m", + "sglang_simulator.simulation.sglang.launch_server", + "--model-path", + str(ASSETS / "qwen3-8b"), + "--sim-config-path", + str(config_path), + "--port", + str(self.port), + "--tokenizer-path", + str(EXAMPLES / "assets" / "tokenizer"), + "--max-total-tokens", + "8192", + "--max-running-requests", + "8", + "--disable-overlap-schedule", + ] + self.server_proc = subprocess.Popen(cmd, env=env, preexec_fn=os.setsid) + for _ in range(120): + if self.server_proc.poll() is not None: + raise RuntimeError("SGLang Simulator server exited during startup") + try: + if requests.get(self.base_url, timeout=1).status_code < 500: + return + except requests.RequestException: + pass + time.sleep(1) + self.shutdown() + raise RuntimeError("SGLang Simulator server did not become ready") + + @property + def base_url(self) -> str: + return f"http://127.0.0.1:{self.port}" + + def benchmark( + self, + output_file: Path, + workload: str = "sharegpt", + request_rate=None, + seed=42, + ) -> dict: + cmd = [ + sys.executable, + str(BENCH_SERVING), + f"--simulator-mode={self.mode}", + "--backend=sglang", + f"--base-url={self.base_url}", + f"--model={ASSETS / 'qwen3-8b'}", + f"--tokenizer={EXAMPLES / 'assets' / 'tokenizer'}", + "--num-prompts=3", + "--disable-tqdm", + "--profile", + f"--output-file={output_file}", + ] + if request_rate is not None: + cmd.extend([f"--request-rate={request_rate}", f"--seed={seed}"]) + + if workload == "sharegpt": + cmd.extend( + [ + "--dataset-name=sharegpt", + f"--dataset-path={EXAMPLES / 'workloads' / 'sharegpt-example.json'}", + "--sharegpt-output-len=4", + ] + ) + else: + assert workload == "timestamp_trace" + cmd.extend( + [ + "--dataset-name=autobench", + f"--dataset-path={EXAMPLES / 'workloads' / 'timestamp-trace-example.jsonl'}", + "--use-trace-timestamps", + ] + ) + + subprocess.run(cmd, check=True) + assert output_file.is_file() + return json.loads( + (self.output_dir / "metrics.json").read_text(encoding="utf-8") + ) + + def shutdown(self): + if self.server_proc.poll() is not None: + return + os.killpg(self.server_proc.pid, signal.SIGTERM) + try: + self.server_proc.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(self.server_proc.pid, signal.SIGKILL) + self.server_proc.wait() + + +def assert_decode_metrics(metrics): + assert metrics["completed"] == 3 + assert metrics["total_output"] == 12 + assert metrics["mean_ttft_ms"] >= 0 + assert metrics["mean_tpot_ms"] > 0 + assert metrics["mean_itl_ms"] > 0 + assert metrics["input_throughput"] > 0 + + +@pytest.mark.parametrize("config_name", SIM_CONFIGS) +def test_benchmark(config_name, tmp_path): + runner = SGLangServingRunner(SIM_CONFIGS[config_name], tmp_path) + try: + metrics = runner.benchmark(tmp_path / "benchmark.json") + finally: + runner.shutdown() + + assert_decode_metrics(metrics) + + +def test_timestamp_trace(tmp_path): + runner = SGLangServingRunner(SIM_CONFIGS["replay"], tmp_path) + try: + metrics = runner.benchmark( + tmp_path / "benchmark.json", workload="timestamp_trace" + ) + finally: + runner.shutdown() + + assert_decode_metrics(metrics)