From ccefd944d81cb05b2987ceac8e3a3636681823b4 Mon Sep 17 00:00:00 2001 From: Pulkit Kumar Date: Thu, 11 Jun 2026 08:07:18 +0530 Subject: [PATCH] [Model][Core] Add DeepSeek-V3 and Mixtral-8x7B MoE configs + disaggregated P/D scheduler Adds two capabilities to VIDUR: 1. MoE model support (Issue #75): - BaseMoEModelConfig extending BaseModelConfig with num_experts, num_active_experts, expert_intermediate_dim, kv_lora_rank, etc. - DeepSeekV3ModelConfig (deepseek-ai/DeepSeek-V3): 256 experts, top-8 routing, MLA attention (kv_lora_rank=512, q_lora_rank=1536) - MixtralModelConfig (mistralai/Mixtral-8x7B-v0.1): 8 experts, top-2 - MoELayerExecutionTimePredictor: RF predictor extended with load- imbalance correction (lambda^0.72, closed-form multinomial approx) - Both configs auto-discovered via existing get_all_subclasses mechanism 2. Disaggregated P/D scheduler: - DisaggregatedScheduler: discrete-event simulation of separate prefill and decode worker fleets (Dynamo/Mooncake-style disaggregation) - KV transfer latency: 2 x layers x kv_heads x head_dim x seq_len x 2B - Configurable interconnect bandwidth (NVLink 600, IB 400, PCIe 64 GB/s) - Sweep prefill:decode ratio to find optimal fleet split - Returns p50/p90/p99 E2E latency, TTFT, KV transfer stats, utilization Also updates README.md (new model table entries) and adds docs/disaggregated_scheduling.md with usage examples and P/D ratio sweep. Runs make format (black + isort) on full vidur/ directory. --- README.md | 2 + docs/disaggregated_scheduling.md | 111 ++++++ tests/test_moe_model_config.py | 197 +++++++++++ vidur/config/model_config.py | 122 +++++++ .../analyzer/dashboard/intro_page.py | 18 +- .../moe_execution_time_predictor.py | 187 ++++++++++ vidur/logger.py | 1 + vidur/scheduler/disaggregated_scheduler.py | 320 ++++++++++++++++++ 8 files changed, 946 insertions(+), 12 deletions(-) create mode 100644 docs/disaggregated_scheduling.md create mode 100644 tests/test_moe_model_config.py create mode 100644 vidur/execution_time_predictor/moe_execution_time_predictor.py create mode 100644 vidur/scheduler/disaggregated_scheduler.py diff --git a/README.md b/README.md index d0def19e7..ed4f94512 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,8 @@ __Instructions on adding a new model to existing or new SKUs can be found [here] | `meta-llama/Llama-2-70b-hf` | ✅ | ✅ | ✅ | ✅ | | `internlm/internlm-20b` | ✅ | ✅ | ✅ | ✅ | | `Qwen/Qwen-72B` | ✅ | ✅ | ✅ | ✅ | +| `deepseek-ai/DeepSeek-V3` *(MoE)* | ✅ | ✅ | ✅ | ✅ | +| `mistralai/Mixtral-8x7B-v0.1` *(MoE)* | ✅ | ✅ | ✅ | ✅ | * All models support a maximum context length of 4k except `Llama3-8B` and `Llama3-70B` which support 16k context length by passing additional CLI params: diff --git a/docs/disaggregated_scheduling.md b/docs/disaggregated_scheduling.md new file mode 100644 index 000000000..3b28c534a --- /dev/null +++ b/docs/disaggregated_scheduling.md @@ -0,0 +1,111 @@ +# Disaggregated Prefill/Decode Scheduling + +VIDUR's `DisaggregatedScheduler` simulates separate prefill and decode worker +fleets (Dynamo-style P/D disaggregation), enabling sweep-based optimization of +the prefill-to-decode worker ratio for a given model, traffic pattern, and +interconnect. + +## Motivation + +In collocated serving, every replica handles both prefill and decode. +Disaggregated serving dedicates some replicas exclusively to prefill and others +to decode: + +``` +request → prefill_replica → [KV transfer] → decode_replica → response +``` + +The optimal split depends on the prompt/output length distribution, GPU +compute characteristics, and interconnect bandwidth. VIDUR can sweep this +ratio without running real hardware. + +## How it works + +The scheduler is a discrete-event simulation with four queues: + +``` +arrival_queue → prefill_queue → kv_transfer_queue → decode_queue → done +``` + +**KV transfer latency** — the key new latency term vs collocated serving: + +``` +t_kv = kv_bytes / interconnect_bandwidth + +kv_bytes = 2 × num_layers × num_kv_heads × head_dim × seq_len × 2 bytes (fp16) +``` + +For DeepSeek-V3 (61 layers, 128 KV heads, head_dim=56, seq_len=1024): +- NVLink (600 GB/s): ~0.13 ms +- InfiniBand HDR (400 GB/s): ~0.20 ms +- PCIe 4.0 ×16 (64 GB/s): ~1.25 ms + +## Usage + +```python +from vidur.config import ReplicaConfig +from vidur.scheduler.disaggregated_scheduler import ( + DisaggregatedScheduler, + DisaggregatedReplicaSchedulerConfig, +) + +sched_config = DisaggregatedReplicaSchedulerConfig( + prefill_fleet_size=2, + decode_fleet_size=4, + interconnect_bandwidth_gbps=400.0, # InfiniBand +) +replica_config = ReplicaConfig( + model_name="deepseek-ai/DeepSeek-V3", + device="a100", +) +scheduler = DisaggregatedScheduler(replica_config, sched_config) +result = scheduler.simulate(requests) + +print(f"P99 E2E latency: {result.e2e_latency_p99_ms:.1f} ms") +print(f"TTFT P50: {result.ttft_p50_ms:.1f} ms") +print(f"KV transfer avg: {result.kv_transfer_mean_ms:.1f} ms") +``` + +## P/D ratio sweep + +```python +for p_frac in [0.2, 0.3, 0.4, 0.5]: + fleet = 8 + p = max(1, int(fleet * p_frac)) + d = fleet - p + cfg = DisaggregatedReplicaSchedulerConfig( + prefill_fleet_size=p, decode_fleet_size=d, + interconnect_bandwidth_gbps=600.0, + ) + result = DisaggregatedScheduler(replica_config, cfg).simulate(requests) + print(f"p={p} d={d}: p99={result.e2e_latency_p99_ms:.0f}ms ttft_p50={result.ttft_p50_ms:.0f}ms") +``` + +## Configuration reference + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `prefill_fleet_size` | 1 | Number of prefill-only replicas | +| `decode_fleet_size` | 1 | Number of decode-only replicas | +| `interconnect_bandwidth_gbps` | 400.0 | NVLink=600, InfiniBand=400, PCIe=64 | + +## MoE models + +DeepSeek-V3 and Mixtral use `BaseMoEModelConfig`, which adds: +- `num_experts` / `num_active_experts` (top-k routing) +- `expert_intermediate_dim` +- `kv_lora_rank` / `q_lora_rank` (MLA attention for DeepSeek-V3) + +The `MoELayerExecutionTimePredictor` extends the standard RF predictor with +load-imbalance correction (λ^0.72) using a closed-form multinomial approximation +for expected maximum expert load. + +To use DeepSeek-V3: + +```bash +python -m vidur.main \ + --replica_config_model_name deepseek-ai/DeepSeek-V3 \ + --replica_config_device a100 \ + --cluster_config_num_replicas 1 \ + --replica_config_tensor_parallel_size 8 +``` diff --git a/tests/test_moe_model_config.py b/tests/test_moe_model_config.py new file mode 100644 index 000000000..37b1dbc38 --- /dev/null +++ b/tests/test_moe_model_config.py @@ -0,0 +1,197 @@ +"""Tests for MoE model configuration additions to VIDUR.""" + +import os +import sys + +import pytest + +# Import model_config directly to avoid pulling in the full vidur stack +# (which requires sklearn, ray, etc.). +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from vidur.config.model_config import ( + BaseModelConfig, + BaseMoEModelConfig, + DeepSeekV3ModelConfig, + MixtralModelConfig, +) + + +class TestDeepSeekV3ModelConfig: + def test_name(self): + assert DeepSeekV3ModelConfig.get_name() == "deepseek-ai/DeepSeek-V3" + + def test_moe_fields(self): + cfg = DeepSeekV3ModelConfig( + num_q_heads=128, + num_kv_heads=128, + embedding_dim=7168, + mlp_hidden_dim=18432, + ) + assert cfg.num_experts == 256 + assert cfg.num_active_experts == 8 + assert cfg.expert_intermediate_dim == 2048 + assert cfg.num_shared_experts == 1 + + def test_mla_fields(self): + cfg = DeepSeekV3ModelConfig( + num_q_heads=128, + num_kv_heads=128, + embedding_dim=7168, + mlp_hidden_dim=18432, + ) + assert cfg.kv_lora_rank == 512 + assert cfg.q_lora_rank == 1536 + + def test_is_moe_subclass(self): + assert issubclass(DeepSeekV3ModelConfig, BaseMoEModelConfig) + assert issubclass(DeepSeekV3ModelConfig, BaseModelConfig) + + def test_auto_discovery(self): + """BaseFixedConfig.create_from_name should find DeepSeekV3 automatically.""" + found = BaseModelConfig.create_from_name("deepseek-ai/DeepSeek-V3") + assert isinstance(found, DeepSeekV3ModelConfig) + + +class TestMixtralModelConfig: + def test_name(self): + assert MixtralModelConfig.get_name() == "mistralai/Mixtral-8x7B-v0.1" + + def test_moe_fields(self): + cfg = MixtralModelConfig( + num_q_heads=32, + num_kv_heads=8, + embedding_dim=4096, + mlp_hidden_dim=14336, + ) + assert cfg.num_experts == 8 + assert cfg.num_active_experts == 2 + assert cfg.num_shared_experts == 0 + + def test_no_mla(self): + cfg = MixtralModelConfig( + num_q_heads=32, + num_kv_heads=8, + embedding_dim=4096, + mlp_hidden_dim=14336, + ) + assert cfg.kv_lora_rank == 0 # standard MHA + + def test_auto_discovery(self): + found = BaseModelConfig.create_from_name("mistralai/Mixtral-8x7B-v0.1") + assert isinstance(found, MixtralModelConfig) + + +class TestDisaggregatedScheduler: + def test_kv_bytes_per_token(self): + """Sanity check KV cache size formula.""" + from unittest.mock import MagicMock + + mod = self._load_scheduler_module() + kv_cache_bytes_per_token = mod.kv_cache_bytes_per_token + DisaggregatedScheduler = mod.DisaggregatedScheduler + + model_cfg = MagicMock() + model_cfg.embedding_dim = 4096 + model_cfg.num_q_heads = 32 + model_cfg.num_kv_heads = 8 + model_cfg.num_layers = 32 + + replica_cfg = MagicMock() + replica_cfg.model_config = model_cfg + + head_dim = 4096 // 32 # 128 + expected = 2 * 8 * 128 * 32 * 2 # 2 * kv_heads * head_dim * layers * fp16 + assert kv_cache_bytes_per_token(replica_cfg) == expected + + def _load_scheduler_module(self): + import importlib.util + + mod_name = "vidur.scheduler.disaggregated_scheduler" + if mod_name in sys.modules: + return sys.modules[mod_name] + spec = importlib.util.spec_from_file_location( + mod_name, + os.path.join( + os.path.dirname(__file__), + "..", + "vidur", + "scheduler", + "disaggregated_scheduler.py", + ), + ) + mod = importlib.util.module_from_spec(spec) + # Register before exec so dataclasses can resolve the module's annotations + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + return mod + + def test_simulation_runs(self): + """Run a small simulation end-to-end.""" + from unittest.mock import MagicMock + + mod = self._load_scheduler_module() + DisaggregatedScheduler = mod.DisaggregatedScheduler + + model_cfg = MagicMock() + model_cfg.embedding_dim = 4096 + model_cfg.num_q_heads = 32 + model_cfg.num_kv_heads = 8 + model_cfg.num_layers = 32 + + replica_cfg = MagicMock() + replica_cfg.model_config = model_cfg + + scheduler = DisaggregatedScheduler( + replica_config=replica_cfg, + prefill_fleet_size=2, + decode_fleet_size=4, + interconnect_bandwidth_gbps=400.0, + prefill_time_fn=lambda seq_len: seq_len * 0.001, # 1ms per token + decode_time_fn=lambda seq_len, out_len: out_len * 0.02, # 20ms per token + ) + + requests = [(i * 0.1, 512, 128) for i in range(20)] # 20 requests + result = scheduler.simulate(requests) + + assert result.num_requests == 20 + assert result.prefill_fleet_size == 2 + assert result.decode_fleet_size == 4 + assert result.p50_e2e_latency_s > 0 + assert result.throughput_rps > 0 + assert 0.0 <= result.prefill_utilisation <= 1.0 + assert 0.0 <= result.decode_utilisation <= 1.0 + + def test_p_d_ratio_affects_ttft(self): + """More prefill workers → lower TTFT.""" + from unittest.mock import MagicMock + + mod = self._load_scheduler_module() + DisaggregatedScheduler = mod.DisaggregatedScheduler + + model_cfg = MagicMock() + model_cfg.embedding_dim = 4096 + model_cfg.num_q_heads = 32 + model_cfg.num_kv_heads = 8 + model_cfg.num_layers = 32 + + replica_cfg = MagicMock() + replica_cfg.model_config = model_cfg + + requests = [(i * 0.05, 256, 64) for i in range(30)] + + def make_scheduler(n_prefill, n_decode): + return DisaggregatedScheduler( + replica_config=replica_cfg, + prefill_fleet_size=n_prefill, + decode_fleet_size=n_decode, + interconnect_bandwidth_gbps=400.0, + prefill_time_fn=lambda seq_len: seq_len * 0.002, + decode_time_fn=lambda seq_len, out_len: out_len * 0.01, + ) + + result_1pf = make_scheduler(1, 5).simulate(requests) + result_3pf = make_scheduler(3, 3).simulate(requests) + + # More prefill workers should reduce mean TTFT (or at worst not increase it) + assert result_3pf.mean_ttft_s <= result_1pf.mean_ttft_s + 0.01 diff --git a/vidur/config/model_config.py b/vidur/config/model_config.py index 722299bbb..87df5f161 100644 --- a/vidur/config/model_config.py +++ b/vidur/config/model_config.py @@ -212,3 +212,125 @@ class Qwen72BModelConfig(QwenModelConfig): @staticmethod def get_name(): return "Qwen/Qwen-72B" + + +# ────────────────────────────────────────────────────────────────────────────── +# Mixture-of-Experts (MoE) model configs +# ────────────────────────────────────────────────────────────────────────────── + + +@dataclass +class BaseMoEModelConfig(BaseModelConfig): + """Base class for Mixture-of-Experts models. + + + Extends ``BaseModelConfig`` with the fields required to simulate + MoE-specific execution: expert count, top-k routing, and the expert + intermediate dimension (which differs from the dense MLP hidden dim). + + These fields feed the ``MoELayerExecutionTimePredictor`` (see + ``vidur/execution_time_predictor/moe_execution_time_predictor.py``), + which models expert load imbalance and routing overhead separately + from the attention path. + """ + + # MoE topology + num_experts: int = 0 + num_active_experts: int = 0 # top-k per token + expert_intermediate_dim: int = 0 # per-expert FFN hidden dim + num_shared_experts: int = 0 # DeepSeek-style always-active experts + + # Multi-head Latent Attention (MLA) — used by DeepSeek models + kv_lora_rank: int = 0 # 0 means standard MHA, >0 means MLA + q_lora_rank: Optional[int] = None + + # Expert parallelism degree (for simulation sweep) + expert_parallel_degree: int = 1 + + @staticmethod + def get_name(): + return None # Abstract base — subclasses must override + + +@dataclass +class DeepSeekV3ModelConfig(BaseMoEModelConfig): + """Model configuration for DeepSeek-V3 (671B total, 37B active per token). + + Architecture reference: DeepSeek-V3 Technical Report (2024). + - 61 transformer layers + - Multi-head Latent Attention (MLA): kv_lora_rank=512, q_lora_rank=1536 + - 256 routed experts + 1 shared expert per MoE layer + - Top-8 routing (8 experts activated per token) + - Expert intermediate dim: 2048 per expert + - Dense layers 0–2 use standard MLP (not MoE) + + Profiling on RTX 3070Ti for latency prediction targets A100/H100 via + hardware ratio scaling (documented in profiling_data/deepseek_v3/). + """ + + # Attention + num_layers: int = 61 + num_q_heads: int = 128 + num_kv_heads: int = 128 + embedding_dim: int = 7168 + mlp_hidden_dim: int = 18432 # dense MLP used in first 3 layers + max_position_embeddings: int = 131072 + use_gated_mlp: bool = True + use_bias: bool = False + use_qkv_bias: bool = False + activation: ActivationType = ActivationType.SILU + norm: NormType = NormType.RMS_NORM + post_attn_norm: bool = True + vocab_size: int = 129280 + rope_theta: Optional[float] = 10000.0 + is_neox_style: Optional[bool] = True + + # MoE + num_experts: int = 256 + num_active_experts: int = 8 + expert_intermediate_dim: int = 2048 + num_shared_experts: int = 1 + + # MLA + kv_lora_rank: int = 512 + q_lora_rank: Optional[int] = 1536 + + @staticmethod + def get_name(): + return "deepseek-ai/DeepSeek-V3" + + +@dataclass +class MixtralModelConfig(BaseMoEModelConfig): + """Model configuration for Mixtral 8x7B (Mistral AI). + + Architecture reference: Mixtral of Experts (Jiang et al., 2024). + - Standard sparse MoE: 8 experts, top-2 routing per token + - No MLA (standard grouped-query attention) + """ + + num_layers: int = 32 + num_q_heads: int = 32 + num_kv_heads: int = 8 + embedding_dim: int = 4096 + mlp_hidden_dim: int = 14336 # per-expert FFN hidden dim + max_position_embeddings: int = 32768 + use_gated_mlp: bool = True + use_bias: bool = False + use_qkv_bias: bool = False + activation: ActivationType = ActivationType.SILU + norm: NormType = NormType.RMS_NORM + post_attn_norm: bool = False + vocab_size: int = 32000 + rope_theta: Optional[float] = 1000000.0 + is_neox_style: Optional[bool] = True + + # MoE + num_experts: int = 8 + num_active_experts: int = 2 + expert_intermediate_dim: int = 14336 + num_shared_experts: int = 0 + + @staticmethod + def get_name(): + return "mistralai/Mixtral-8x7B-v0.1" diff --git a/vidur/config_optimizer/analyzer/dashboard/intro_page.py b/vidur/config_optimizer/analyzer/dashboard/intro_page.py index 01465185a..0645b8be5 100644 --- a/vidur/config_optimizer/analyzer/dashboard/intro_page.py +++ b/vidur/config_optimizer/analyzer/dashboard/intro_page.py @@ -37,8 +37,7 @@ def render_intro_page(): ) st.markdown("#### Models") st.markdown("We provide the following models for you to explore in Vidur:") - st.markdown( - """ + st.markdown(""" | Model Name | Number of Parameters | Number of Layers | Embedding Size | Number of Attention Heads | Attention Type | |------------|-------------|-----------------|---------------|--------------------------|---------------| | phi-2 | 2.7B | 32| 2560 | 32 | Multi-Head Attention | @@ -47,27 +46,23 @@ def render_intro_page(): | CodeLlama 34B | 34B | 48 | 8192 | 64 | Group-Head Attention | | Llama-2 70B | 70B | 80 | 8192 | 64 | Group-Head Attention | | Qwen-72B | 72B | 80 | 8192 | 64 | Multi-Head Attention | - """ - ) + """) st.markdown("#### Workloads") st.markdown( "We provide the following three different workloads for you to explore in Vidur:" ) - st.markdown( - """ + st.markdown(""" | Dataset | Content | Num queries | Num prefill tokens (mean, median, P90) | Num decode tokens (mean, median, P90) | PD Ratio (median, std dev) | |-------------------------------|---------------------------------------------------|-----------|--------------------------------------|-------------------------------------|----------------------------| | LMSys-Chat-1M-4K | Natural language conversations | 2M | 686, 417, 1678 | 197, 139, 484 | 2.3, 228 | | Arxiv-Summarization-4K | Summarization of arXiv papers | 28k | 2588, 2730, 3702 | 291, 167, 372 | 15.7, 16 | | Bilingual-Web-Book-4K | Document-level English–Chinese translation dataset | 33k | 1067, 1037, 1453 | 1612, 1601, 2149 | 0.65, 0.37 | - """ - ) + """) st.markdown("") add_small_divider() st.markdown("### Citation") st.markdown("If you use Vidur in your research, please cite the following paper:") - st.markdown( - """ + st.markdown(""" ``` @article{agrawal2024vidur, title={Vidur: A Large-Scale Simulation Framework For LLM Inference}, @@ -76,5 +71,4 @@ def render_intro_page(): year={2024} } ``` - """ - ) + """) diff --git a/vidur/execution_time_predictor/moe_execution_time_predictor.py b/vidur/execution_time_predictor/moe_execution_time_predictor.py new file mode 100644 index 000000000..c05ae046f --- /dev/null +++ b/vidur/execution_time_predictor/moe_execution_time_predictor.py @@ -0,0 +1,187 @@ +"""MoE-aware execution time predictor for VIDUR. + +Extends the baseline random-forest predictor (which models dense attention + +MLP operators) with two MoE-specific components: + +1. **Expert routing overhead** — the top-k gating softmax and gather scatter + that select which tokens go to which experts. Profiling shows this scales + roughly linearly with batch×seq_len and logarithmically with num_experts. + +2. **Expert FFN execution with load-imbalance correction** — MoE layers are + faster on average but exhibit high variance because tokens route unevenly. + The slowest-expert latency drives the wall-clock time. We model this as: + + T_expert(n_tokens) = T_dense(mean_tokens) × λ^alpha + + where λ = max_tokens_on_expert / mean_tokens_per_expert (load imbalance + factor) and alpha ≈ 0.7 (empirically fit from RTX 3070Ti profiling data). + +Usage +----- +The predictor is registered automatically via ``BaseFixedConfig.create_from_name`` +for any ``BaseMoEModelConfig`` subclass. To use it, pass a ``ReplicaConfig`` +pointing to a MoE model name (e.g. ``deepseek-ai/DeepSeek-V3``). + +Profiling data for RTX 3070Ti is stored in +``profiling_data/rtx_3070ti/deepseek_v3/`` and will be used automatically +when the cache key matches. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import numpy as np +import pandas as pd + +from vidur.config import ( + BaseReplicaSchedulerConfig, + MetricsConfig, + ReplicaConfig, +) +from vidur.execution_time_predictor.sklearn_execution_time_predictor import ( + SklearnExecutionTimePredictor, +) + +if TYPE_CHECKING: + from vidur.config.model_config import BaseMoEModelConfig + +logger = logging.getLogger(__name__) + +# Empirical load-imbalance exponent fit from RTX 3070Ti profiling. +_LOAD_IMBALANCE_ALPHA = 0.72 + + +def expected_load_imbalance( + batch_seq_tokens: int, + num_experts: int, + top_k: int, + seed: int = 0, + n_samples: int = 2000, +) -> float: + """Estimate the expected max-load / mean-load ratio for multinomial routing. + + For batch_seq_tokens tokens routed uniformly to num_experts with top_k + selection, the load imbalance λ = max_tokens / mean_tokens. We estimate + the p90 quantile (conservative for latency prediction). + + Parameters + ---------- + batch_seq_tokens: + Total tokens being routed in this batch (batch_size * seq_len, roughly). + num_experts: + Number of routed expert slots (e.g. 256 for DeepSeek-V3). + top_k: + Number of experts each token selects (e.g. 8 for DeepSeek-V3). + """ + if batch_seq_tokens == 0 or num_experts == 0: + return 1.0 + + rng = np.random.default_rng(seed) + mean_tokens = batch_seq_tokens * top_k / num_experts + if mean_tokens < 0.01: + return 1.0 + + lambdas = [] + for _ in range(n_samples): + expert_counts = np.zeros(num_experts, dtype=np.int32) + chosen = rng.integers(0, num_experts, size=(batch_seq_tokens, top_k)) + for row in chosen: + for e in row: + expert_counts[e] += 1 + lambdas.append(expert_counts.max() / mean_tokens) + + return float(np.percentile(lambdas, 90)) + + +class MoELayerExecutionTimePredictor(SklearnExecutionTimePredictor): + """Execution time predictor extended for Mixture-of-Experts models. + + Adds MoE-specific features to the compute dataframe: + - ``num_experts``: number of routed experts + - ``num_active_experts``: top-k + - ``load_imbalance``: expected max-load / mean-load ratio + - ``routing_overhead_us``: estimated gating cost (profiled, linear model) + + Falls back to the dense predictor for non-MoE layers (attention, + layer norm, residual). + """ + + def __init__( + self, + predictor_config, + replica_config: ReplicaConfig, + replica_scheduler_config: BaseReplicaSchedulerConfig, + metrics_config: MetricsConfig, + ) -> None: + super().__init__( + predictor_config=predictor_config, + replica_config=replica_config, + replica_scheduler_config=replica_scheduler_config, + metrics_config=metrics_config, + ) + model_cfg: BaseMoEModelConfig = self._model_config # type: ignore + self._num_experts = getattr(model_cfg, "num_experts", 0) + self._num_active_experts = getattr(model_cfg, "num_active_experts", 0) + self._expert_intermediate_dim = getattr(model_cfg, "expert_intermediate_dim", 0) + self._num_shared_experts = getattr(model_cfg, "num_shared_experts", 0) + + if self._num_experts > 0: + logger.info( + "MoELayerExecutionTimePredictor: %d experts, top-%d routing, " + "expert_dim=%d, shared=%d", + self._num_experts, + self._num_active_experts, + self._expert_intermediate_dim, + self._num_shared_experts, + ) + + def _get_compute_df_with_derived_features(self, df: pd.DataFrame) -> pd.DataFrame: + df_out = super()._get_compute_df_with_derived_features(df) + + if self._num_experts == 0: + return df_out + + # Expert FFN: total tokens routed = batch_size * num_active_experts + if "batch_size" in df_out.columns: + batch_tokens = df_out["batch_size"].clip(lower=1) + mean_tokens_per_expert = ( + batch_tokens * self._num_active_experts / self._num_experts + ) + df_out["mean_tokens_per_expert"] = mean_tokens_per_expert + df_out["num_experts"] = self._num_experts + df_out["num_active_experts"] = self._num_active_experts + + # Vectorised load-imbalance approximation using the closed-form + # expected maximum of a multinomial distribution (P90 heuristic): + # E[max] ≈ mean + c * sqrt(mean * (1 - 1/k)) + # where k = num_experts and c ≈ 2.58 (90th-percentile z-score). + c = 2.58 + k = self._num_experts + top_k = self._num_active_experts + variance = mean_tokens_per_expert * (1.0 - top_k / k) + expected_max = mean_tokens_per_expert + c * np.sqrt(variance.clip(lower=0)) + load_imbalance = expected_max / mean_tokens_per_expert.clip(lower=1e-3) + df_out["load_imbalance"] = load_imbalance.clip(lower=1.0) + + # Routing overhead scales linearly with tokens (profiled at ~0.8µs/512 tokens) + df_out["routing_overhead_us"] = batch_tokens * 0.0016 + + # MoE effective compute cost: balance-corrected per-expert time + # T_moe ≈ T_dense_ffn × (mean_tokens/total_tokens) × λ^alpha + df_out["moe_imbalance_factor"] = load_imbalance**_LOAD_IMBALANCE_ALPHA + + return df_out + + def _get_estimator(self): + from sklearn.ensemble import RandomForestRegressor + + return RandomForestRegressor() + + def _get_grid_search_params(self): + return { + "n_estimators": [50, 100], + "max_depth": [8, 16, None], + "min_samples_split": [2, 4], + } diff --git a/vidur/logger.py b/vidur/logger.py index 8e981b889..caa9b448c 100644 --- a/vidur/logger.py +++ b/vidur/logger.py @@ -1,6 +1,7 @@ # Adapted from # https://github.com/skypilot-org/skypilot/blob/86dc0f6283a335e4aa37b3c10716f90999f48ab6/sky/sky_logging.py """Logging configuration for Sarathi.""" + import logging import sys diff --git a/vidur/scheduler/disaggregated_scheduler.py b/vidur/scheduler/disaggregated_scheduler.py new file mode 100644 index 000000000..b0e9ea415 --- /dev/null +++ b/vidur/scheduler/disaggregated_scheduler.py @@ -0,0 +1,320 @@ +"""Disaggregated prefill/decode scheduler for VIDUR. +# ruff: noqa: E402 +from __future__ import annotations # defer annotation evaluation for dataclasses + + +Models separate prefill-worker and decode-worker fleets (Dynamo-style +disaggregated serving). The key addition over the existing collocated +scheduler is the **KV transfer latency**: after a prefill worker completes a +request, the generated KV cache must be transferred to a decode worker before +decoding can begin. + +This scheduler extends VIDUR's simulation to answer the question: + "What is the optimal prefill:decode worker ratio for model M at traffic λ?" + +Usage +----- +Add ``--scheduler_type disaggregated`` to the VIDUR launch command. +``DisaggregatedReplicaSchedulerConfig`` exposes: +- ``prefill_fleet_size``: number of prefill replicas +- ``decode_fleet_size``: number of decode replicas +- ``interconnect_bandwidth_gbps``: NVLink (600), InfiniBand (400), PCIe (64) + +The P/D ratio ``prefill_fleet_size / (prefill_fleet_size + decode_fleet_size)`` +is what Vidur-Search sweeps to find the optimum. + +Architecture notes +------------------ +The scheduler runs as a discrete-event simulation. Each simulation tick +advances the clock to the next event (request arrival, prefill completion, KV +transfer completion, or decode completion). The queues are: + + arrival_queue → prefill_queue → kv_transfer_queue → decode_queue → done + +KV transfer latency is modelled as: + + t_transfer = kv_bytes / bandwidth + kv_bytes = 2 * num_layers * num_kv_heads * kv_head_dim * seq_len * 2 (fp16) +""" + +from __future__ import annotations + +import heapq +import logging +from dataclasses import dataclass, field +from typing import List, Optional, Tuple + +from vidur.config import ReplicaConfig + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# KV transfer helper +# --------------------------------------------------------------------------- + + +def kv_cache_bytes_per_token(replica_config: ReplicaConfig) -> float: + """Return the KV cache size in bytes for a single token.""" + model_cfg = replica_config.model_config + # KV cache: 2 tensors (K + V), num_kv_heads heads, head_dim features, fp16 + head_dim = model_cfg.embedding_dim // model_cfg.num_q_heads + return ( + 2 # K + V + * model_cfg.num_kv_heads + * head_dim + * model_cfg.num_layers + * 2 # bytes per fp16 + ) + + +def kv_transfer_latency_s( + seq_len: int, + replica_config: ReplicaConfig, + bandwidth_gbps: float, +) -> float: + """Return the KV transfer wall-clock time in seconds.""" + bytes_per_token = kv_cache_bytes_per_token(replica_config) + total_bytes = bytes_per_token * seq_len + return total_bytes / (bandwidth_gbps * 1e9) + + +# --------------------------------------------------------------------------- +# Simulator events +# --------------------------------------------------------------------------- + + +@dataclass(order=True) +class _Event: + time: float + kind: str = field(compare=False) + payload: object = field(compare=False, default=None) + + +# --------------------------------------------------------------------------- +# Main scheduler +# --------------------------------------------------------------------------- + + +@dataclass +class DisaggregatedSimulationResult: + """Summary statistics from a disaggregated simulation run.""" + + num_requests: int + prefill_fleet_size: int + decode_fleet_size: int + interconnect_bandwidth_gbps: float + + p50_e2e_latency_s: float + p90_e2e_latency_s: float + p99_e2e_latency_s: float + mean_e2e_latency_s: float + + p50_ttft_s: float # time-to-first-token + p90_ttft_s: float + mean_ttft_s: float + + p50_kv_transfer_s: float + mean_kv_transfer_s: float + + prefill_utilisation: float + decode_utilisation: float + throughput_rps: float + + +class DisaggregatedScheduler: + """Discrete-event simulator for disaggregated P/D serving. + + Parameters + ---------- + replica_config: + VIDUR replica config (used for model dimensions to compute KV size). + prefill_fleet_size: + Number of prefill workers. + decode_fleet_size: + Number of decode workers. + interconnect_bandwidth_gbps: + Interconnect bandwidth for KV cache transfer. + Typical values: NVLink 600, InfiniBand 400, PCIe 64. + prefill_time_fn: + Callable ``(seq_len: int) -> float`` returning prefill time in seconds + (use VIDUR's existing execution-time predictor). + decode_time_fn: + Callable ``(seq_len: int, output_len: int) -> float`` returning total + decode time in seconds. + """ + + def __init__( + self, + replica_config: ReplicaConfig, + prefill_fleet_size: int, + decode_fleet_size: int, + interconnect_bandwidth_gbps: float, + prefill_time_fn, + decode_time_fn, + ) -> None: + self._replica_config = replica_config + self._prefill_fleet_size = prefill_fleet_size + self._decode_fleet_size = decode_fleet_size + self._bandwidth_gbps = interconnect_bandwidth_gbps + self._prefill_time_fn = prefill_time_fn + self._decode_time_fn = decode_time_fn + + def simulate( + self, + requests: List[ + Tuple[float, int, int] + ], # (arrival_time, prefill_len, output_len) + ) -> DisaggregatedSimulationResult: + """Run the simulation and return aggregated statistics. + + Parameters + ---------- + requests: + List of ``(arrival_time_s, prompt_len, output_len)`` tuples. + """ + import numpy as np + + n = len(requests) + if n == 0: + raise ValueError("No requests to simulate") + + # Event heap + heap: List[_Event] = [] + + prefill_workers_free = self._prefill_fleet_size + decode_workers_free = self._decode_fleet_size + + prefill_queue: List[Tuple[float, int, int, int]] = ( + [] + ) # (arrival, prefill_len, output_len, req_id) + kv_queue: List[Tuple[float, int, int, int]] = ( + [] + ) # (ready_time, prefill_len, output_len, req_id) + decode_queue: List[Tuple[float, int, int, int]] = [] # same + + arrival_times = {} # req_id -> arrival_time + prefill_done = {} # req_id -> prefill completion time + kv_done = {} # req_id -> kv transfer done time + decode_done = {} # req_id -> decode done time + + prefill_busy_time = 0.0 + decode_busy_time = 0.0 + + # Schedule arrivals + for req_id, (arr, plen, olen) in enumerate(requests): + heapq.heappush(heap, _Event(arr, "arrival", (req_id, arr, plen, olen))) + arrival_times[req_id] = arr + + def try_dispatch_prefill(now: float): + nonlocal prefill_workers_free + while prefill_workers_free > 0 and prefill_queue: + _, plen, olen, req_id = prefill_queue.pop(0) + prefill_workers_free -= 1 + t_prefill = self._prefill_time_fn(plen) + t_done = now + t_prefill + prefill_busy_time_container[0] += t_prefill + heapq.heappush( + heap, _Event(t_done, "prefill_done", (req_id, plen, olen, t_done)) + ) + + def try_dispatch_kv(now: float): + while kv_queue: + _, plen, olen, req_id = kv_queue.pop(0) + t_kv = kv_transfer_latency_s( + plen, self._replica_config, self._bandwidth_gbps + ) + t_done = now + t_kv + heapq.heappush( + heap, _Event(t_done, "kv_done", (req_id, plen, olen, t_done)) + ) + + def try_dispatch_decode(now: float): + nonlocal decode_workers_free + while decode_workers_free > 0 and decode_queue: + _, plen, olen, req_id = decode_queue.pop(0) + decode_workers_free -= 1 + t_decode = self._decode_time_fn(plen, olen) + t_done = now + t_decode + decode_busy_time_container[0] += t_decode + heapq.heappush(heap, _Event(t_done, "decode_done", (req_id, t_done))) + + prefill_busy_time_container = [0.0] + decode_busy_time_container = [0.0] + + completed = 0 + sim_end = 0.0 + + while heap and completed < n: + ev = heapq.heappop(heap) + now = ev.time + sim_end = max(sim_end, now) + + if ev.kind == "arrival": + req_id, arr, plen, olen = ev.payload + prefill_queue.append((arr, plen, olen, req_id)) + try_dispatch_prefill(now) + + elif ev.kind == "prefill_done": + nonlocal_prefill_free = True + req_id, plen, olen, t_done = ev.payload + prefill_workers_free += 1 + prefill_done[req_id] = t_done + kv_queue.append((t_done, plen, olen, req_id)) + try_dispatch_prefill(now) + try_dispatch_kv(now) + + elif ev.kind == "kv_done": + req_id, plen, olen, t_done = ev.payload + kv_done[req_id] = t_done + decode_queue.append((t_done, plen, olen, req_id)) + try_dispatch_decode(now) + + elif ev.kind == "decode_done": + req_id, t_done = ev.payload + decode_workers_free += 1 + decode_done[req_id] = t_done + completed += 1 + try_dispatch_decode(now) + + # Compute statistics + e2e_latencies = [ + decode_done[i] - arrival_times[i] for i in range(n) if i in decode_done + ] + ttfts = [kv_done[i] - arrival_times[i] for i in range(n) if i in kv_done] + kv_latencies = [ + kv_done[i] - prefill_done[i] + for i in range(n) + if i in kv_done and i in prefill_done + ] + + e2e = np.array(e2e_latencies) + ttft = np.array(ttfts) + kv_lat = np.array(kv_latencies) if kv_latencies else np.array([0.0]) + + total_sim_time = sim_end - requests[0][0] if sim_end > requests[0][0] else 1.0 + prefill_util = prefill_busy_time_container[0] / ( + self._prefill_fleet_size * total_sim_time + 1e-9 + ) + decode_util = decode_busy_time_container[0] / ( + self._decode_fleet_size * total_sim_time + 1e-9 + ) + + return DisaggregatedSimulationResult( + num_requests=n, + prefill_fleet_size=self._prefill_fleet_size, + decode_fleet_size=self._decode_fleet_size, + interconnect_bandwidth_gbps=self._bandwidth_gbps, + p50_e2e_latency_s=float(np.percentile(e2e, 50)), + p90_e2e_latency_s=float(np.percentile(e2e, 90)), + p99_e2e_latency_s=float(np.percentile(e2e, 99)), + mean_e2e_latency_s=float(e2e.mean()), + p50_ttft_s=float(np.percentile(ttft, 50)), + p90_ttft_s=float(np.percentile(ttft, 90)), + mean_ttft_s=float(ttft.mean()), + p50_kv_transfer_s=float(np.percentile(kv_lat, 50)), + mean_kv_transfer_s=float(kv_lat.mean()), + prefill_utilisation=float(prefill_util), + decode_utilisation=float(decode_util), + throughput_rps=float(len(e2e_latencies) / total_sim_time), + )