[JAX] Make dot_product_attention usage configurable - #2549
Open
anko-intel wants to merge 3 commits into
Open
Conversation
Signed-off-by: Andrzej Kotłowski <andrzej.kotlowski@intel.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a configuration knob to control whether quantized JAX attention layers may route through Keras’ fused dot_product_attention path, and wires it through quantization, serialization/deserialization, and the relevant attention layer implementations. The default behavior is now to disable dot_product_attention unless explicitly enabled in the config.
Changes:
- Adds
dot_product_attention_enabletoJaxBaseConfigand propagates it through static/dynamic quantization entrypoints. - Plumbs the flag into quantized layer
prepare(...)APIs and gates fused attention execution in quantizedMultiHeadAttentionandCachedGemma3Attention. - Ensures deserialized quantized models re-apply the flag when preparing layers after load.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
neural_compressor/jax/quantization/saving.py |
Passes dot_product_attention_enable into layer prepare(...) during deserialization. |
neural_compressor/jax/quantization/layers_static.py |
Adds flag to prepare(...) signatures; gates fused attention in static quantized attention layers. |
neural_compressor/jax/quantization/layers_dynamic.py |
Adds flag to prepare(...) signatures; gates fused attention in dynamic quantized attention layers. |
neural_compressor/jax/quantization/config.py |
Introduces dot_product_attention_enable config parameter and serializable state. |
neural_compressor/jax/algorithms/static.py |
Propagates config flag into static layer preparation. |
neural_compressor/jax/algorithms/dynamic.py |
Propagates config flag into dynamic layer preparation. |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Andrzej Kotłowski <andrzej.kotlowski@intel.com>
anko-intel
force-pushed
the
dev/anko/disable_fused_att
branch
from
August 7, 2026 11:15
508aff9 to
cec9aa6
Compare
for more information, see https://pre-commit.ci
wpietka
reviewed
Aug 7, 2026
|
|
||
| self._all_devices_cpu = all(d.platform == "cpu" for d in jax.devices()) | ||
|
|
||
| if self._all_devices_cpu: |
Contributor
There was a problem hiding this comment.
_all_devices_cpu is a little awkward name. I woudn't know what it means without reading definition. Maybe _only_cpu_devices_are_available would be better
| const_weight, | ||
| w_quant_granularity, | ||
| *, | ||
| dot_product_attention_enable=False, |
Contributor
There was a problem hiding this comment.
Why it cannot be a positional parameter?
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
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.
Type of Change
feature
Description
Allows you to control whether to use dot_product_attention from the configuration.
Using dot_product_attention is disabled by default