Added MDP generation to QEff Compile - #1086
Merged
Merged
Conversation
ochougul
reviewed
Jun 16, 2026
ochougul
requested changes
Jun 16, 2026
ochougul
left a comment
ochougul
left a comment
Contributor
There was a problem hiding this comment.
Please document what is tested with this here
Contributor
|
Also can you check, if this info is getting included in hash/ dumped qconfigs for compile? |
Contributor
Author
Both |
Contributor
Author
|
The unit tests are failing due to unable to connect to HF. It is not related to changes done in PR: |
Contributor
Author
|
Working for(with
with
|
quic-akuruvil
left a comment
quic-akuruvil
left a comment
Contributor
There was a problem hiding this comment.
@quic-mohmeh Have you verified thischanges with layerwise export also?
quic-akuruvil
requested changes
Jun 19, 2026
quic-akuruvil
left a comment
quic-akuruvil
left a comment
Contributor
There was a problem hiding this comment.
Please add unit-tests for this change, as mentioned in PR #930
quic-mohmeh
force-pushed
the
mdp
branch
2 times, most recently
from
June 19, 2026 09:51
2b1f322 to
d28db2f
Compare
quic-hemagnih
requested changes
Jun 19, 2026
quic-hemagnih
left a comment
quic-hemagnih
left a comment
Contributor
There was a problem hiding this comment.
Can you please do the following, post to that I will merge it.
- Add unit test cases.
- Can you also add the results for all the models which you have verified with your PR.
Contributor
Author
Done, add the unit test for both Causal LM and MoE VLMs
|
Contributor
Author
Not yet, haven't got the time |
Contributor
Author
Unit Tests AddedTwo suites of CPU-only unit tests are added in 1.
|
| Test | What it checks |
|---|---|
test_8_layers_3_partitions_counts |
8 layers / 3 partitions produces counts [3, 3, 2] |
test_10_layers_4_partitions_counts |
10 layers / 4 partitions produces counts [2, 3, 3, 2] |
test_max_minus_min_le_1_for_8_3 |
Max − min partition size ≤ 1 (balanced invariant) for 8/3 |
test_max_minus_min_le_1_for_10_4 |
Max − min partition size ≤ 1 (balanced invariant) for 10/4 |
test_last_partition_receives_base_allocation_8_3 |
Last partition always gets only floor(N/P) layers — never a remainder — for 8/3 |
test_last_partition_receives_base_allocation_10_4 |
Same invariant for 10/4 |
test_remainder_goes_to_middle_partitions_for_10_4 |
For 10/4, both remainder layers land on middle partitions (indices 1 and 2); partition 0 stays at base |
test_evenly_divisible_has_zero_spread_12_4 |
When evenly divisible, all four partitions get exactly 3 layers each |
test_two_partition_remainder_goes_to_first_5_2 |
Edge case: 5/2 with no middle zone — remainder goes to the first partition, giving [3, 2] |
test_bounds_length_equals_num_partitions_minus_one |
The returned bounds list always has length num_partitions − 1 across all cases |
test_total_layers_preserved |
Sum of per-partition counts always equals num_layers — no layers lost or duplicated |
2. TestMdpCompileIntegration — end-to-end compile artifact validation (@pytest.mark.mdp @pytest.mark.cpu_only)
Uses a synthetic ONNX graph (no Hugging Face Hub download) whose nodes carry h.N/attn and
h.N/mlp names to exercise the full ONNX MDP strategy code path.
subprocess.run is monkeypatched so no real qaic-compile binary is needed.
The fixture builds a 2-layer GPT-2-shaped graph in a tmp_path, runs QEFFAutoModelForCausalLM._compile
with mdp_ts_num_devices=4, mdp_num_partitions=2, mdp_strategy="onnx", and then asserts on the
produced artifacts.
| Test | What it checks |
|---|---|
test_mdp_disagg_json_is_created |
mdp_disagg_4d_2p.json is written to the compile directory |
test_mdp_disagg_json_has_two_partitions |
The JSON contains exactly 2 partition entries |
test_mdp_disagg_partition0_devices_are_0_1 |
Partition 0 is assigned device IDs [0, 1] |
test_mdp_disagg_partition1_devices_are_2_3 |
Partition 1 is assigned device IDs [2, 3] |
test_hashed_compile_params_command_references_mdp_json |
The compile command recorded in hashed_compile_params.json contains -mdp-load-partition-config= pointing at the disagg JSON |
test_hashed_compile_params_contains_mdp_num_partitions |
hashed_compile_params.json records mdp_num_partitions=2 for cache-key stability |
test_hashed_compile_params_contains_mdp_strategy_onnx |
hashed_compile_params.json records mdp_strategy="onnx" |
test_hashed_compile_params_contains_mdp_ts_json |
The full generated MDP partition dict is embedded in hashed_compile_params.json under mdp_ts_json |
test_qconfig_compiler_config_contains_mdp_num_partitions |
qconfig.json compiler_config section contains mdp_num_partitions=2 |
test_qconfig_compiler_config_contains_mdp_strategy_onnx |
qconfig.json compiler_config section contains mdp_strategy="onnx" |
3. On-device compile-only integration tests (require @pytest.mark.on_qaic)
Two new parametrized tests exercise the full export → compile pipeline on real hardware in
compile-only mode (no inference run), keeping CI time low.
test_few_causal_lm_onnx_mdp_compile_only (tests/transformers/models/causal_lm_models/test_causal_lm_models.py)
- Parametrized over all causal LM models in the test config ×
use_onnx_subfunctions=[False, True]. - Calls
check_causal_lm_pytorch_vs_kv_vs_ort_vs_ai100withcompile_only=True,
mdp_num_partitions=2,mdp_strategy="onnx"to verify the compiler accepts the MDP partition
config for every supported causal LM.
test_few_image_text_to_text_onnx_mdp_compile_only (tests/transformers/models/image_text_to_text/test_image_text_to_text_models.py)
- Scoped to MoE VLMs only (
test_mm_moe_models— models wheremodel_typecontains"moe"). - Parametrized over
kv_offload=[True, False]; dual-QPC models are skipped unlesskv_offload=True. - Runs with
use_onnx_subfunctions=True,mdp_num_partitions=2,mdp_strategy="onnx",
compile_only=True. For non-InternVL / non-Molmo models,skip_vision=Trueis also set so only
the text backbone is compiled with MDP.
Contributor
Author
|
@quic-hemagnih The PR is good to merge |
quic-amitraj
requested changes
Jun 24, 2026
quic-amitraj
requested changes
Jun 24, 2026
quic-amitraj
left a comment
quic-amitraj
left a comment
Contributor
There was a problem hiding this comment.
The PR adds a significant new API surface (mdp_num_partitions, mdp_strategy, mdp_compiler_dump_path) but ships no example showing how to use it. please consider adding one example script to demonstrate its use.
Contributor
Author
|
Addressed the latest feedback and pushed updates. Added the Qwen3-VL compile-only MDP example script, documented it in the examples READMEs, fixed Ruff formatting, and CI checks are green. @quic-amitraj @quic-hemagnih PR is good to merge |
Contributor
|
Others LGTM @quic-rishinr |
Contributor
Author
|
Addressed the latest review feedback and pushed 39af48f. |
Contributor
|
This PR doesnt handle Qwen 3.5 122B and 397B. A separate PR will be raised for it. |
Contributor
|
the mdp_ts_num_devices is sent as 4 in case of TS4PP6 case, but this value is never used in the code, as it's overridden by num_devices while calling self._compile causing same argument passed twice, once via passed argument and once via kwargs. Just tried to run the example script errors out with following |
Add QEfficient/compile/mdp_generator.py implementing two strategies for
generating MDP partition configs for pipeline-parallel (PP) disaggregated
prefill serving:
MdpStrategy.ONNX (default)
Single-pass ONNX topsort walk; assigns each node to a partition by
transformer layer index (layers.N / h.N / layer_N// naming patterns).
Non-layer nodes (embeddings, lm_head) follow the nearest layer.
Inlined local function call-sites (CtxScatterCB / CtxGatherCB) are
expanded inline so nodeList order matches the ONNX topsort. Known
custom ops (CustomRMSNorm, CastToUInt4) are included by call-site name.
Produces a ~19 MB JSON superset of the compiler IR.
MdpStrategy.INTERSECTION
Runs the ONNX strategy first (superset), then filters each partition's
nodeList to only names present in a prior qaic-compile
-mdp-dump-partition-config dump. Result is compact (~1-2 MB) and
contains only exact-match Glow IR names for doManualPartitioning().
Requires a prior compile run to produce the dump path.
Integrate into QEfficient/base/modeling_qeff.py:
- New compile() parameter mdp_num_partitions (default 1, existing
template/TS behaviour unchanged).
- When mdp_num_partitions > 1, auto-generate the MDP JSON from the
exported ONNX and pass it to qaic-compile via mdp_load_partition_config.
- Strategy and optional compiler dump path controllable via
mdp_strategy / mdp_compiler_dump_path in compiler_options.
- MdpStrategy enum exported for type-safe caller use.
- Compile hash includes mdp_num_partitions to avoid stale QPC reuse.
Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com>
Previously, generate_disagg_mdp_partition_config distributed remainder
layers (num_layers % num_partitions) entirely into the last partition
via a min(..., num_partitions - 1) clamp. This left the last stage
holding postprocessing (lm_head, final norm) plus all the overflow,
creating a consistent pipeline bottleneck.
Replace the flat floor-division heuristic with a balanced breakpoint
scheme implemented in the new _compute_partition_breakpoints helper:
- Last partition always receives the base allocation (base = num_layers
// num_partitions), keeping postprocessing light.
- Remainder layers fill the middle stages first (indices 1 ..
num_partitions-2), distributed centre-outward within the middle zone
so the load stays as symmetric as possible.
- The first partition (embedding preprocessing) absorbs overflow only
when the remainder exceeds the number of available middle stages.
- Worst case (remainder == num_partitions - 1): every stage except the
last receives one extra layer.
Example: 8 layers / 3 stages
before: [2, 2, 4]
after: [3, 3, 2]
Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com>
Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com>
Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com>
Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com>
Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com>
Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com>
Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com>
Contributor
|
There are two pending issues with the MDP changes in this PR:
|
ochougul
pushed a commit
that referenced
this pull request
Jun 26, 2026
This PR adds MDP generation required for disaggregated serving for Prefill. Supports both Pipeline Prefill + Tensor Slicing and passing custom cores to the MDP generator. Also adds support for VLMs, compiler 'stages' option, and layerwise export. --------- Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com> Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
quic-rishinr
pushed a commit
that referenced
this pull request
Jun 26, 2026
This PR adds MDP generation required for disaggregated serving for Prefill. Supports both Pipeline Prefill + Tensor Slicing and passing custom cores to the MDP generator. Also adds support for VLMs, compiler 'stages' option, and layerwise export. --------- Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com> Signed-off-by: ochougul <ochougul@qti.qualcomm.com> (cherry picked from commit 5ca7a31) Signed-off-by: Rishin Raj <rishinr@qti.qualcomm.com>
quic-rishinr
added a commit
that referenced
this pull request
Jun 29, 2026
This PR adds MDP generation required for disaggregated serving for Prefill. Supports both Pipeline Prefill + Tensor Slicing and passing custom cores to the MDP generator. Also adds support for VLMs, compiler 'stages' option, and layerwise export. --------- Signed-off-by: Mohit Mehta <mohmeh@qti.qualcomm.com> Signed-off-by: ochougul <ochougul@qti.qualcomm.com> Signed-off-by: Rishin Raj <rishinr@qti.qualcomm.com> Co-authored-by: quic-mohmeh <mohmeh@qti.qualcomm.com> Co-authored-by: Rishin Raj <rishinr@qti.qualcomm.com>
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
6 participants
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.
This PR adds MDP generation required for disaggregated serving for Prefill. Supports both Pipeline Prefill + Tensor Slicing and passing custom cores to the MDP generator. Also adds support for VLMs, compiler 'stages' option, and layerwise export.