Repeatkv transform#1037
Open
quic-dhirajku wants to merge 8 commits into
Open
Conversation
…VLMs. Based on PR quic#625. Addressed most of the comments made on the previous PR. Repeat check is done on a subset of models during CI, primarily due to difference in configs of such models. Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
…ng with changes made for the new transforms. TODO: Check for the ONNX directory path name being different. Check if the list of classes for mapping covers all the models that we support. Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
…oder Wrappers were added to string mapping list to enable dummy model export for CI. Changes were made to prevent multiple application of ReplicateKVTransform if done in either Encoder or Decoder Wrapper already. Modeling files updated to access config in EncoderWrapper as well. Infra added for causalLM and VLM checks for repeatKV setup CI tests. CausalLM script APIRunner instantiation moved to allow updated input shapes to be made. Similarly commented export in VLM script since compile will call it with updated changes already. TODO: Confirm the changes that were made for DeepSeekV3 model for RepeatKV, currently they were removed for a generic approach. Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
Made changes to allow generic name based transformation of heads (num_attention_heads, n_heads, n_head etc). Minor edits and utils created for this task. Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
Edited the changes as suggested by quic-mamta. Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
This reverts commit b40a34d.
Enabled method to calculate best possible repeat_kv count based on model and num devices. Added repeat_kv method for AWQ quantized models. Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
quic-dhirajku
force-pushed
the
repeatkv_transform
branch
from
cb92eaf to
12ca65e
Compare
Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
quic-dhirajku
force-pushed
the
repeatkv_transform
branch
2 times, most recently
from
fe75c92 to
008adc7
Compare
Closed
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
2 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 improves repeat_kv to be more reliable and portable across model/config variants and quantized backends. It adds shared config-key resolution (for common head aliases), centralized repeat-value calculation from (num_devices, model_config), and consistent propagation of num_kv_heads_repeat into qaic_config and hashing so export/compile state matches the effective runtime settings. It also hardens transform behavior: MLA models are explicitly skipped with a warning (instead of attempting unsafe replication), and ReplicateKVHeadTransform.apply() now reports transformed status only when replication is actually applied.
The main functional upgrade is quantized RepeatKV support. _duplicate_weights_for_linear_layer was updated to correctly handle packed layouts (instead of relying on generic reshape assumptions), including AWQ/GPTQ packed tensors and a full QuantLinearORT path (dequantize → replicate → resize buffers → repack). This directly fixes the observed failures in RepeatKV quantized tests (shape mismatch and MatMulNBits zero-point dtype issues), and the previously failing AWQ TinyLlama RepeatKV test now passes.