[Model][Core] Add DeepSeek-V3 and Mixtral-8x7B MoE configs + disaggregated P/D scheduler - #80
Open
Pulkit Kumar (buddywhitman) wants to merge 1 commit into
Conversation
…gated P/D scheduler Adds two capabilities to VIDUR: 1. MoE model support (Issue microsoft#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.
Author
|
@microsoft-github-policy-service agree |
Copilot started reviewing on behalf of
Jayashree Mohan (jayashreemohan29)
September 4, 2026 05:23
View session
There was a problem hiding this comment.
🟡 Changes recommended
Several additions are inconsistent or incomplete (notably MoE predictor wiring/feature computation and disaggregated scheduler/docs API mismatches), which can lead to incorrect behavior or confusing public documentation.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds new model configuration primitives and simulation components to extend VIDUR toward Mixture-of-Experts (MoE) modeling and disaggregated prefill/decode scheduling, along with documentation and tests.
Changes:
- Introduces
BaseMoEModelConfigplus concrete DeepSeek-V3 and Mixtral MoE model configs. - Adds a MoE-aware execution-time predictor and a new disaggregated prefill/decode discrete-event scheduler.
- Updates docs/README and adds tests covering the new configs and scheduler helpers.
File summaries
| File | Description |
|---|---|
vidur/scheduler/disaggregated_scheduler.py |
New discrete-event scheduler for separated prefill/decode fleets with KV-transfer modeling and summary stats. |
vidur/execution_time_predictor/moe_execution_time_predictor.py |
New MoE-oriented predictor intended to add routing/load-imbalance features. |
vidur/config/model_config.py |
Adds MoE config base class and DeepSeek-V3 / Mixtral configs. |
docs/disaggregated_scheduling.md |
New user doc for disaggregated scheduling and KV-transfer derivation + examples. |
tests/test_moe_model_config.py |
New tests for MoE config discovery and basic disaggregated scheduler behavior. |
README.md |
Adds DeepSeek-V3 and Mixtral to the supported models list. |
vidur/config_optimizer/analyzer/dashboard/intro_page.py |
Markdown formatting cleanup in the dashboard intro page. |
vidur/logger.py |
Minor formatting change (blank line). |
Review details
Suppressed comments (1)
vidur/execution_time_predictor/moe_execution_time_predictor.py:150
- This block checks for a
batch_sizecolumn, but VIDUR’s compute models are trained/predicted overnum_tokens(seeSklearnExecutionTimePredictor._train_compute_models). As a result, the MoE derived features never get added to the compute dataframe. Also, the variance term should use(1 - 1/k)for multinomial/binomial routing with p=1/k;(1 - top_k/k)underestimates variance when top_k>1.
# 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
- Files reviewed: 8/8 changed files
- Comments generated: 9
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+182
to
+187
| # Event heap | ||
| heap: List[_Event] = [] | ||
|
|
||
| prefill_workers_free = self._prefill_fleet_size | ||
| decode_workers_free = self._decode_fleet_size | ||
|
|
| 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] |
| 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 |
Comment on lines
+46
to
+66
| 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") |
Comment on lines
+19
to
+23
| 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``). |
Comment on lines
+1
to
+3
| """Disaggregated prefill/decode scheduler for VIDUR. | ||
| # ruff: noqa: E402 | ||
| from __future__ import annotations # defer annotation evaluation for dataclasses |
Comment on lines
+15
to
+21
| 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) |
Comment on lines
+201
to
+202
| prefill_busy_time = 0.0 | ||
| decode_busy_time = 0.0 |
Comment on lines
+259
to
+260
| nonlocal_prefill_free = True | ||
| req_id, plen, olen, t_done = ev.payload |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds two new capabilities:
MoE model support (related to Issue Support simulation of MoE model and Expert Parallelism #75):
Disaggregated prefill/decode scheduler:
FIX #75
Summary
1. MoE model configs
BaseMoEModelConfigextendingBaseModelConfigwithnum_experts,num_active_experts,expert_intermediate_dim,kv_lora_rank,q_lora_rank,expert_parallel_degree.Two concrete configs auto-discovered via the existing
get_all_subclassesmechanism:DeepSeekV3ModelConfig(deepseek-ai/DeepSeek-V3): 256 experts, top-8, MLA attentionMixtralModelConfig(mistralai/Mixtral-8x7B-v0.1): 8 experts, top-2MoELayerExecutionTimePredictorextends the RF predictor with load-imbalance features (closed-form multinomial approximation, λ^0.72 latency correction).2. Disaggregated prefill/decode scheduler
DisaggregatedSchedulersimulates separate prefill and decode worker fleets.2 × layers × kv_heads × head_dim × seq_len × 2B (fp16).Tests
Documentation
docs/disaggregated_scheduling.md: motivation, architecture diagram, KV transfer latency derivation with a worked DeepSeek-V3 example (NVLink ~0.13ms, IB ~0.20ms, PCIe ~1.25ms at seq_len=1024), usage example, and a P/D ratio sweep example.README.md"Supported Models" table updated with the two new MoE entries.Testing / Code quality
make format(isort --profile black+blackon the fullvidur/tree).python -m vidur.main -hlisting the new model names.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 8runs end-to-end.No breaking changes to existing models or configs — both new model configs and the disaggregated scheduler are additive and opt-in.
PR Checklist (Click to Expand)
Thank you for your contribution to Vidur! Before submitting the pull request, please ensure the PR meets the following criteria. This helps Vidur maintain the code quality and improve the efficiency of the review process.
PR Title and Classification
Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:
[Bugfix]for bug fixes.[CI/Build]for build or continuous integration improvements.[Doc]for documentation fixes and improvements.[Model]for adding a new model or improving an existing model. Model name should appear in the title.[Profiling]For changes on the profiling module.[Core]for changes in the core simulator logic[Misc]for PRs that do not fit the above categories. Please use this sparingly.Note: If the PR spans more than one category, please include all relevant prefixes.
Code Quality
The PR need to meet the following code quality standards:
make formatto format your code.docs/source/if the PR modifies the user-facing behaviors of Vidur. It helps user understand and utilize the new features or changes.Notes for Large Changes
Please keep the changes as concise as possible. For major architectural changes (>500 LOC), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with
rfc-requiredand might not go through the PR.Thank You
Finally, thank you for taking the time to read these guidelines and for your interest in contributing to Vidur. Your contributions make Vidur a great tool for everyone!