Added MDP generation to QEff Compile - #930
Conversation
e20d868 to
f393d6e
Compare
|
Tested and working on the following model classes
|
|
@quic-rishinr @mamtsing @ochougul A gentle reminder for review |
There was a problem hiding this comment.
Add warning that ignores mdp_ts_num_partitions whenever seq_len==1
Also, add a warning that ignores when ts_num_devices>1 and seq_len> and mdp_ts_num_partitions>1
|
To be tested on:
|
|
@quic-mohmeh Please let us know once your testing is complete. Also incorporate the review comments. @quic-mamta Can also please review it. |
|
Works on Qwen3-4B(Dense) - without subfunctions |
|
Works on Qwen3-30B-A3B (MOE) model as well - without subfunctions |
|
Update example scripts for Qwen3(with and without VL) and GPTOSS |
d34dbda to
f924a1d
Compare
|
@quic-mohmeh please rebase it on top of release/v1.22.0_tmp branch |
|
Tested with subfunctions on the following model:
|
|
Also verified by Qeff team on following models with subfunctions:
|
| @@ -0,0 +1,263 @@ | |||
| # ----------------------------------------------------------------------------- | |||
There was a problem hiding this comment.
please add unit tests and test for mdp generation. Please make sure the tests are small and exectue under 5 to 10 seconds.
There was a problem hiding this comment.
@quic-rishinr as we have discussed that there is no valid approach to compare compiler MDP dump and QEff MDP dump, all we can ensure is that every node present in compiler MDP dump should be present in QEff MDP dump in correct order. For verifying and testing, we might need to compile no PP and PP QPCs and compare the output. Let me know which direction you want to proceed for testing
There was a problem hiding this comment.
Hi @quic-mohmeh , for unit test case's part, one of the required cases is, add a e2e tiny 2L model output matching verification. This model should be compiled with the manual MDP you have prepared. You can refer and re-use some of the test APIs we already use for output matching against reference. This should test the functionality of the change.
There was a problem hiding this comment.
@quic-mohmeh Add add other unit test cases to flag any edge cases and so on, which causes the functionality to break.
There was a problem hiding this comment.
@quic-mohmeh Have you verified PP+TS combination with this? Or only PP works with this? WHen you say you have verified the above models, have you verified the output sanity too?
As vLLM currently doesn't support PP+TS, so I am unable to verify PP+TS. As for output validation, I have only skim the output, the main purpose of these tests is to check whether the model compiles or not with the MDP generated |
quic-mohmeh
left a comment
quic-mohmeh
left a comment
There was a problem hiding this comment.
Addressed the comments @quic-rishinr
| @@ -0,0 +1,263 @@ | |||
| # ----------------------------------------------------------------------------- | |||
There was a problem hiding this comment.
@quic-rishinr as we have discussed that there is no valid approach to compare compiler MDP dump and QEff MDP dump, all we can ensure is that every node present in compiler MDP dump should be present in QEff MDP dump in correct order. For verifying and testing, we might need to compile no PP and PP QPCs and compare the output. Let me know which direction you want to proceed for testing
|
@quic-rishinr @quic-hemagnih |
4626809 to
e176948
Compare
|
VLMs are not working currently with this MDP partition, both with and without subfunctions(E-P-D). Have to look into compiler codebase |
|
@quic-mohmeh Can you rebase your PR with latest |
|
Working on VLMs - tested on Qwen2.5VL 3B with and without subfunctions(PP4) |
Hi @quic-mohmeh When you say tested on VLMS, have you enabled PP on prefill part of language model? Or have you checked with PP on vision part as well? |
Yes, PP with Prefill, TS on Encode and Decode |
|
Gemma4-26B-A4B Prefill compiled with subfunctions |
67953cc to
96da0b6
Compare
| # nodeList per partition directly from the ONNX graph — no compiler | ||
| # round-trip required. | ||
|
|
||
| num_layers = getattr(self, "num_layers", None) |
There was a problem hiding this comment.
@quic-mohmeh Are you still going with stages in user API? If so, where is it popped before qaic-compile, since it is not a valid compiler option?
4a9e87b to
71d5cb3
Compare
| @@ -0,0 +1,263 @@ | |||
| # ----------------------------------------------------------------------------- | |||
There was a problem hiding this comment.
Hi @quic-mohmeh , for unit test case's part, one of the required cases is, add a e2e tiny 2L model output matching verification. This model should be compiled with the manual MDP you have prepared. You can refer and re-use some of the test APIs we already use for output matching against reference. This should test the functionality of the change.
8273e7f to
e9330a8
Compare
Collection of PRs which include TF 5.5.4, Qwen 3.5, 3.6 Gemma 4 etc. --------- Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com> Signed-off-by: Asmita Goswami <asmigosw@qti.qualcomm.com> Signed-off-by: asmigosw <asmigosw@qti.qualcomm.com> Signed-off-by: Rishin Raj <rishinr@qti.qualcomm.com> Signed-off-by: Vahid Janfaza <vjanfaza@qti.qualcomm.com> Signed-off-by: sanising <sanising@qti.qualcomm.com> Signed-off-by: eplatero <eplatero@qti.qualcomm.com> Signed-off-by: Dipankar Sarkar <dipankar@qti.qualcomm.com> Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com> Signed-off-by: Mohit Soni <mohisoni@qti.qualcomm.com> Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com> Signed-off-by: Tanisha Chawada <tchawada@qti.qualcomm.com> Signed-off-by: Hem Agnihotri <hemagnih@qti.qualcomm.com> Signed-off-by: Abukhoyer Shaik <abukhoye@qti.qualcomm.com> Signed-off-by: Abhishek Kumar Singh <sabhis@qti.qualcomm.com> Signed-off-by: Mamta Singh <mamtsing@qti.qualcomm.com> Signed-off-by: abhishek-singh591 <sabhis@qti.qualcomm.com> Signed-off-by: vtirumal <vtirumal@qti.qualcomm.com> Signed-off-by: quic-xiyushi <xiyushi@qti.qualcomm.com> Co-authored-by: Vinayak Baddi <68580231+vbaddi@users.noreply.github.com> Co-authored-by: Asmita Goswami <asmigosw@qti.qualcomm.com> Co-authored-by: Hem Agnihotri <hemagnih@qti.qualcomm.com> Co-authored-by: vjanfaza <vjanfaza@qti.qualcomm.com> Co-authored-by: Sanidhya Singal <sanising@qti.qualcomm.com> Co-authored-by: Erick Platero <40013722+eplatero97@users.noreply.github.com> Co-authored-by: Dipankar Sarkar <dipankar@qti.qualcomm.com> Co-authored-by: Amit Raj <168538872+quic-amitraj@users.noreply.github.com> Co-authored-by: Mohit Soni <quic_mohisoni@quicinc.com> Co-authored-by: Mohit Soni <mohisoni@qti.qualcomm.com> Co-authored-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com> Co-authored-by: Tanisha Chawada <tchawada@qti.qualcomm.com> Co-authored-by: Abukhoyer Shaik <abukhoye@qti.qualcomm.com> Co-authored-by: Abhishek Kumar Singh <sabhis@qti.qualcomm.com> Co-authored-by: Mamta Singh <mamtsing@qti.qualcomm.com> Co-authored-by: Karthikeya <vtirumal@qti.qualcomm.com> Co-authored-by: quic-xiyushi <xiyushi@qti.qualcomm.com>
## Background PR quic#1046 introduced the optional `kv_cache_prefix` parameter for `QEFFAutoModelForCausalLM`, single-QPC VLMs, and dual-QPC VLMs. When set, it injects an infix token into LLM KV-cache retained-state buffer names: ``` past_key.N_RetainedState → past_key.N_<prefix>_RetainedState ``` This gives vLLM a stable regex handle to select only LLM KV buffers for disaggregated device-to-device transfer, leaving vision buffers (`vision_embeds_RetainedState`, etc.) untouched. The feature worked for non-layerwise paths. This work extends it to the **layerwise E-P-D path** used by large models (e.g. Qwen3.5-122B-A10B). --- ## Bugs Fixed ### Bug 1 — `kv_cache_prefix` dropped in layerwise `compile()` and `export()` **File:** `QEfficient/transformers/models/modeling_auto.py` `_QEffAutoModelForImageTextToTextDualQPC.compile()` and `.export()` both accept `kv_cache_prefix` as a named parameter, but neither forwarded it when taking the `layerwise=True` early-return branch: - `compile()` → `_run_layerwise_compile(...)` — missing - `compile()` → `vision_wrapper.compile(...)` (vision-only branch) — missing - `export()` → `_run_layerwise_export(...)` — missing (explicit named param, not captured by `**kwargs`) **Fix:** Added `kv_cache_prefix=kv_cache_prefix` to all three call sites. --- ### Bug 2 — `_export_layerwise` rebuilt output names from scratch **File:** `QEfficient/base/modeling_qeff.py` `_export_layerwise` constructed its own `output_name` list using hardcoded templates: ```python output_name.append(f"past_key.{layer_idx}_RetainedState") output_name.append(f"past_value.{layer_idx}_RetainedState") ``` It never consulted the `output_names` parameter passed in by the caller, which already carried the prefixed names from `apply_kv_cache_prefix`. The per-window ONNXes were therefore exported with plain names, and the stitched merged ONNX had no prefix — making `kv_cache_prefix` silently a no-op for all layerwise exports. **Fix:** Built a `_caller_retained_map` from the incoming `output_names`: for each KV retained-state name it finds the plain stem (e.g. `past_key.3_vllmKvCache_RetainedState` → plain stem `past_key.3`), then registers the prefixed+suffix-corrected form under both `_RetainedState` and `_InternalRetainedState` keys. When appending each KV buffer to `output_name`, the map is consulted first; the plain default is used as fallback (preserving existing behaviour when no prefix is set). --- ### Bug 3 — KV input names not aligned in `_export_layerwise` **File:** `QEfficient/base/modeling_qeff.py` The regular `_export` path calls `align_kv_input_names_to_retained_outputs` to rename KV input buffers to match the prefixed outputs: ``` past_key.3 → past_key.3_vllmKvCache ``` The AIC compiler pairs an output `X_RetainedState` to the input named `X`; without this rename the compiler cannot establish KV retention for prefixed buffers. `_export_layerwise` skipped this alignment entirely. **Fix:** Added the same `align_kv_input_names_to_retained_outputs` call (with dynamic-axes propagation) immediately before `torch.onnx.export` in `_export_layerwise`. --- ## Files Changed | File | Change | |------|--------| | `QEfficient/transformers/models/modeling_auto.py` | Forward `kv_cache_prefix` in all layerwise `compile()` and `export()` branches | | `QEfficient/base/modeling_qeff.py` | Respect caller's prefixed `output_names` and align input names in `_export_layerwise` | | `examples/image_text_to_text/models/qwen3_5_moe/qwen3_5_122b_epd_disagg.py` | New end-to-end E-P-D disaggregated inference example for Qwen3.5-122B-A10B | --- ## Example Script Notes (`qwen3_5_122b_epd_disagg.py`) - Three separate compile calls: vision encoder, lang prefill, lang decode - `kv_infix` computed once from `KV_CACHE_PREFIX`; `_update_retained_states` uses it to feed the correct buffer names back as inputs for the next step --- ## Runtime Behaviour - **No prefix set:** all paths are no-ops; buffer names unchanged - **Prefix set, non-layerwise:** unchanged from PR quic#1046 - **Prefix set, layerwise:** per-window ONNXes are now exported with prefixed output names and aligned input names; the stitched merged ONNX carries them through to the compiled QPC --------- Signed-off-by: Rishin Raj <rishinr@qti.qualcomm.com>
Minor fix to export 2L onnx for Qwen3_5 Present we are seeing an issue in exporting 2L without layer wise, this fix will resolve that issue Signed-off-by: vtirumal <vtirumal@qti.qualcomm.com>
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>
|
@quic-rishinr As QEff has moved from tmp branch to mdp branch this PR is no longer relevant. Created new PR: #1086 |
8 participants
This PR adds the MDP generation required in case of disaggregated serving for Prefill. This supports both Pipeline Prefill + Tensor Slicing and also supports passing custom cores to the MDP generator