[Attention] Attention head quantization strategy #481
+80
−19
Conversation
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
kylesayrs
changed the base branch from
main
to
kylesayrs/refactor-initialize-tests
kylesayrs
force-pushed
the
kylesayrs/add-attn-head-strat
branch
2 times, most recently
from
bf00a99 to
2ea692d
Compare
kylesayrs
force-pushed
the
kylesayrs/refactor-initialize-tests
branch
from
97a4d16 to
0fdfbd1
Compare
kylesayrs
force-pushed
the
kylesayrs/add-attn-head-strat
branch
from
2ea692d to
326f802
Compare
kylesayrs
force-pushed
the
kylesayrs/refactor-initialize-tests
branch
from
0fdfbd1 to
8973328
Compare
kylesayrs
force-pushed
the
kylesayrs/add-attn-head-strat
branch
2 times, most recently
from
70da261 to
48875e2
Compare
There was a problem hiding this comment.
how are we expecting users to use QuantizationStrategy.ATTN_HEAD in a recipe? If i'm understanding correctly, it would look something like this?
quant_stage:
quant_modifiers:
QuantizationModifier:
config_groups:
group0:
targets: ["re:.*self_attn$"]
weights:
strategy: attn_head
...
group1:
targets: ["re:.*(q|k|v)_proj$"]
weights:
strategy: group
...
|
@brian-dellabetta I’ve decided that giving per-attention strategy its own strategy (rather than reusing group) makes more sense. quant_stage:
quant_modifiers:
QuantizationModifier:
config_groups:
group0:
targets: ["re:.*self_attn$"]
input_activations:
strategy: attn_head
... |
brian-dellabetta
approved these changes
Oct 8, 2025
Signed-off-by: Kyle Sayers <[email protected]>
Signed-off-by: Kyle Sayers <[email protected]>
Signed-off-by: Kyle Sayers <[email protected]>
Signed-off-by: Kyle Sayers <[email protected]>
Signed-off-by: Kyle Sayers <[email protected]>
Signed-off-by: Kyle Sayers <[email protected]>
Signed-off-by: Kyle Sayers <[email protected]>
Signed-off-by: Kyle Sayers <[email protected]>
kylesayrs
force-pushed
the
kylesayrs/add-attn-head-strat
branch
from
48875e2 to
e1ca4fd
Compare
Signed-off-by: Kyle Sayers <[email protected]>
brian-dellabetta
approved these changes
Oct 8, 2025
| raise ValueError("Block quantization cannot be applied to attention") | ||
|
|
||
| if args.strategy == QuantizationStrategy.ATTN_HEAD: | ||
| # (batch_size * seq_len, num_heads, 1, 1, head_dim) |
There was a problem hiding this comment.
why do we want the 1, 1 here?
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.
Purpose
Given an attention state of shape
(batch_size, num_heads, seq_len, head_dim), the head attention strategy will generate scales of shape(num_heads, 1, 1).Prerequisites
Changes
Testing