Created ReplicateKVHeadTransform to integrate KV-heads replication module within Qefficient library.#625
Open
quic-dhirajku wants to merge 6 commits intoquic:mainfrom
Open
Conversation
quic-dhirajku
requested review from
ochougul,
quic-amitraj,
quic-hemagnih and
quic-rishinr
as code owners
quic-dhirajku
force-pushed
the
replicate_kv_heads_transform
branch
from
1dfdea6 to
3c90390
Compare
Contributor
|
@ochougul @quic-amitraj please review |
quic-dhirajku
force-pushed
the
replicate_kv_heads_transform
branch
from
8c4a1fc to
e502542
Compare
…dule within Qefficient library. The Transform enables KV-head replication for CausalLMs and VLMs as well. The feature is enabled by passing n_kv_head_repeat parameter during initialization of the QEff wrapper class for the corresponding model. n_kv_head_repeat param acts as the multiplier for the number of repeats to be done to original count of KV heads. This operation also causes the config and the hash params of the respective model to update the num_key_value_heads parameter and add a paramter orig_kv_heads to it; It allows us to export the same model with different number of kv_heads without causing a hash conflict. Also added tests for both CausalLMs and VLMs with this functionality to compare outputs of Pytorch HF model and the AIC model. Two new optional paramters n_kv_head_repeat and test_kv_replicate are added for testing purpose. Setting test_kv_replicate to True performs a KV-head replication of every model such that the number of KV-heads and attention heads becomes equal. This was done to ensure tests don't fail due to misalignment issues when we simply repeat num_key_value_heads twice and thus cause a divisibility error on hum_heads. Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
… Doing so would prevent any issues during Transforms when we don't wish to apply it. Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
…orm. Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
quic-dhirajku
force-pushed
the
replicate_kv_heads_transform
branch
from
870cc8d to
08032e1
Compare
vbaddi
requested changes
Jan 21, 2026
…changes to repeat Bias factor appropriately on quantized layers. Signed-off-by: Dhiraj Kumar Sah <dhirajku@qti.qualcomm.com>
ochougul
requested changes
Feb 3, 2026
Contributor
ochougul
left a comment
ochougul
left a comment
There was a problem hiding this comment.
Write a test that makes sure onnx hash is different when different number of kv heads are passed.
Comment on lines
+2498
to
+2499
| # InternVL causes an error if we pass the num_kv_heads_repeat parameter | ||
| num_kv_heads_repeat = kwargs.pop("num_kv_heads_repeat", 1) |
Comment on lines
+2398
to
+2400
| self.model, replicate_kv_transformed = ReplicateKVHeadTransform.apply(self.model, **kwargs) | ||
| if replicate_kv_transformed: | ||
| self.hash_params["config"] = model.config.to_diff_dict() |
Contributor
There was a problem hiding this comment.
better add it to _pytorch_transforms if we are always going to call it.
Comment on lines
+893
to
+895
| if replicate_kv_transformed: | ||
| self.lang_model.hash_params["config"] = model.config.to_diff_dict() | ||
| self.vision_model.hash_params["config"] = model.config.to_diff_dict() |
Contributor
There was a problem hiding this comment.
don't we already dump config somewhere? in _generate_export_hash?
You can just always add repeat_kv_heads value to self.hash_params which will be 1 if nothing is passed.
| } | ||
|
|
||
|
|
||
| class ReplicateKVHeadTransform: |
Contributor
There was a problem hiding this comment.
Make this inherit ModuleMutatorTransform
You may need to implement mutate method which is similar to apply here
Contributor
|
@quic-dhirajku Please take this PR post 595 |
quic-mamta
reviewed
Mar 17, 2026
| layer.bias.data = torch.repeat_interleave( | ||
| layer.bias.data.view(orig_kv_heads, head_dim), repeat, 0 | ||
| ).view(new_kv_heads * head_dim) | ||
| if layer.bias is not None: |
Contributor
There was a problem hiding this comment.
lines 782-785 are repeated here, please remove
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
7 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.
The Transform enables KV-head replication for CausalLMs and VLMs as well.
The feature is enabled by passing n_kv_head_repeat parameter during initialization of the QEff wrapper class for the corresponding model.
n_kv_head_repeatparam acts as the multiplier for the number of repeats to be done to original count of KV heads. This operation also causes the config and the hash params of the respective model to update the num_key_value_heads parameter and add a paramter orig_kv_heads to it; It allows us to export the same model with different number of kv_heads without causing a hash conflict.Added tests for both CausalLMs and VLMs with this functionality to compare outputs of Pytorch HF model and the AIC model. Two new optional paramters
n_kv_head_repeatandtest_kv_replicateare added for testing purpose. Settingtest_kv_replicateto True performs a KV-head replication of every model such that the number of KV-heads and attention heads becomes equal. This was done to ensure tests don't fail due to misalignment issues when we simply repeat num_key_value_heads twice and thus cause a divisibility error on hum_heads.