Repository navigation
Add Granite-swa and Granitemoe-swa model support - #47179
Conversation
…le attention sinks - New model type 'granite_swa' with per-layer sliding window attention - LSE-based sink mechanism: sink_scale = sigmoid(lse - sinks) applied post-attention - Supports MuP parameters (embedding_multiplier, residual_multiplier, logits_scaling, attention_multiplier) - Per-layer attention pattern: configurable full_attention vs sliding_attention layers - Includes standalone modeling, configuration, modular source, and auto registration
- Uses flash_attn_interface._flash_attn_forward for FA3 path (returns LSE natively) - Falls back to eager attention with manual LSE computation - Adds _supports_flash_attn and _compatible_flash_implementations flags - Verified: both eager and FA3 produce correct output
- SWA layers store only last (sliding_window - 1) tokens in KV cache - Full attention layers keep complete KV history - ~98% KV memory savings on SWA layers for long sequences - Verified: identical results with and without sliding window cache
Make modular_granite_swa.py the single source of truth and regenerate the modeling/config: - Define GraniteSWAConfig in the modular file (inheriting GraniteConfig) so the converter no longer wipes the config, and rope_parameters standardizes correctly. - Apply the learnable attention sink through the standard `s_aux` attention dispatch (as in GPT-OSS) instead of the bespoke `_lse`/`_fa3_lse` path. Eager keeps the original sigmoid(logsumexp - sink) math; FlashAttention-3 and 4 now work via the shared interface (`s_aux`/`learnable_sink`). - Disable SDPA and FlexAttention (sink not expressible there); supported backends are eager, flash_attention_3, flash_attention_4. - Add tests/models/granite_swa (CausalLMModelTest framework) and the model doc + _toctree entry. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add GraniteMoeSWA, sublassing GraniteMoeShared, with SWA sliding window + sink added. - MoeSWAConfig inherits MoeSharedConfig, including default no shared experts (`shared_intermediate_size=0`). - Attention uses SWAAttention and thus supports Flash3+4 and eager. No SDPA or Flex. - MoeSWAModel handles per-layer causal masks during forward. All other layers inherit unchanged. - Add tests/models/granitemoe_swa (CausalLMModelTest framework), model doc and _toctree entry, and register model in auto mappings. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
vasqu
left a comment
There was a problem hiding this comment.
Looks good for the most part, added some comments what we could simplify imo and parts we should change (esp the tests I feel like we should enable more re TP/EP!)
| @unittest.skip("TP test mixin forces attn_implementation='sdpa', unsupported by GraniteMoeSWA (sink).") | ||
| def test_tp_forward(self): | ||
| pass | ||
|
|
||
| @unittest.skip("TP test mixin forces attn_implementation='sdpa', unsupported by GraniteMoeSWA (sink).") | ||
| def test_tp_backward(self): | ||
| pass | ||
|
|
||
| @unittest.skip("TP test mixin forces attn_implementation='sdpa', unsupported by GraniteMoeSWA (sink).") | ||
| def test_tp_generation(self): | ||
| pass | ||
|
|
||
| @unittest.skip("TP test mixin forces attn_implementation='sdpa', unsupported by GraniteMoeSWA (sink).") | ||
| def test_tp_generation_quantized(self): | ||
| pass | ||
|
|
||
| # Skipped for consistency with the prior Granite MoE models (granitemoe, granitemoeshared), which | ||
| # do not exercise the expert-parallel test path. | ||
| @unittest.skip("EP test path not exercised by the prior Granite MoE models; skipped for consistency.") | ||
| def test_ep_forward(self): | ||
| pass | ||
|
|
||
| @unittest.skip("EP test path not exercised by the prior Granite MoE models; skipped for consistency.") | ||
| def test_ep_backward(self): | ||
| pass |
There was a problem hiding this comment.
I'd rather override those tbh
There was a problem hiding this comment.
I added tp plans so the TP tests no longer skip (and run correctly). I left the EP test exclusion in granitemoe_swa for now. LMK if you had something different in mind (not entirely sure what you meant by "override" here, sorry!)
There was a problem hiding this comment.
I originally meant that if it's something specific to the model but it should run, it makes more sense to override the tests than to fully skip them
But ig TP has resolved itself :D for EP I wouldn't skip, even if it goes for prior models, we ought to skip as little as possible. Possibly we can unskip for the other granite moes as well
There was a problem hiding this comment.
Sure, I'll try and get those running
-Fix header years -TP and integration tests -Mark flex support -Rope/Nope based on position embed passing -Top level model defs via modular subclassing Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
|
Thanks @vasqu for the comments! I've made changes incorporating your feedback, up to 1) keeping the custom forward implementation for eager, 2) still skipping the EP tests, not entirely sure what you had in mind there? and 3) need to find out if tying rope/nope to layer type is future-safe on our end |
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
|
Thank you for your contribution 🤗! CI Security Gate — automatic approval blockedThis PR was not automatically approved for CI because the security gate failed. Possible reasons:
See the workflow run for the exact violations. A maintainer can review and manually approve CI if a finding is a false positive. |
|
I updated the moe code so it runs against the EP tests (though isn't quite modular anymore). It also turns out that we want to support different rope thetas for different layers, regardless of attention type, not just a binary rope/nope decision, so the old 1/0 list has been replaced with a list of per-layer frequencies (with 0 for Nope, and defaulting to the global rope_theta param when None or unspecified). Current model correctness is unaffected. I think this addresses all of your first-round comments at this point @vasqu |
vasqu
left a comment
There was a problem hiding this comment.
Looks very ready, my comments are mostly small but you uncovered that we missed EP on the granite moe model 😢 can we fix this along this PR or in a separate one? We really shouldnt need to override in this implementation for the moe portion
| Per-layer RoPE base (`theta`) frequency. `0` sets NoPE (no positional embedding) for | ||
| that layer. Overrides global `rope_parameters["rope_theta"]`, which is only used when | ||
| this list is not provided or specified (`layer_rope_theta = None`). | ||
|
|
There was a problem hiding this comment.
Oh wow so even a diff between different full attention layers? Just fmi
There was a problem hiding this comment.
I think the intention is more for different SWA layers, but don't quote me on that
| # Attention shards like Granite (+ per-head `sinks` colwise to track the head-sharding); the | ||
| # routed experts shard tensor-parallel (packed gate/up colwise, down rowwise, `moe_tp_experts`) | ||
| # with the router replicated. The optional shared expert (`shared_mlp`, off by default) is left | ||
| # replicated -- it is small and its full output sums consistently with the all-reduced MoE output. | ||
| base_model_tp_plan = { | ||
| "layers.*.self_attn.q_proj": "colwise", | ||
| "layers.*.self_attn.k_proj": "colwise", | ||
| "layers.*.self_attn.v_proj": "colwise", | ||
| "layers.*.self_attn.o_proj": "rowwise", | ||
| "layers.*.self_attn.sinks": "colwise", | ||
| "layers.*.block_sparse_moe.experts.gate_up_proj": "packed_colwise", | ||
| "layers.*.block_sparse_moe.experts.down_proj": "rowwise", | ||
| "layers.*.block_sparse_moe.experts": "moe_tp_experts", | ||
| } | ||
| # Expert-parallel plan: shard the routed experts across ranks (each rank owns a slice of the | ||
| # experts) with the router driving the dispatch. The optional shared expert is left replicated. | ||
| base_model_ep_plan = { | ||
| "layers.*.block_sparse_moe.router": "ep_router", | ||
| "layers.*.block_sparse_moe.experts.gate_up_proj": "grouped_gemm", | ||
| "layers.*.block_sparse_moe.experts.down_proj": "grouped_gemm", | ||
| "layers.*.block_sparse_moe.experts": "moe_tp_experts", | ||
| } |
There was a problem hiding this comment.
I just checked granitemoe and it seems to completely miss the tp/ep plans, I think we can at least inherit the ep plan then
There was a problem hiding this comment.
See below - it appears the older Granite Moe models (#33207 #35894 #37658 ) predate the tp_plan/ep_plan logic and tests (#42697 #43722 ). So I'd prefer to update the moe router implementation for the new models, for correctness in the current env, but keep the older models grandfathered in, in their current form, if possible
| self.assertIsNotNone(output) | ||
|
|
||
| @unittest.skip("GraniteSWA does not support FlashAttention-2 (only FA3/FA4).") | ||
| def test_flash_attn_2_equivalence(self): |
There was a problem hiding this comment.
Uh oh this should be automatically disabled based on compatible fa; the other tests are skipped properly?
So something along
# Custom kernel which needs the mask interface to be properly usable on these models
if not model_class._supports_attention_backend and not attn_implementation.startswith("flash_attention"):
continue
There was a problem hiding this comment.
It looks like there's two different flash 2 tests: ModelTesterMixin.test_flash_attn_2_inference_equivalence from test_modeling_common.py, and CausalLMModelTest.test_flash_attn_2_equivalence from causal_lm_tester.py. The former does automatically disable based on supported flash attentions, but the latter only skips when flash is not supported at all, so we (and gpt-oss) skip that one manually. In our case, the test does actually pass because FA2 gets auto-corrected to FA3 during model construction, but now the test isn't really testing FA2 equivalence anymore. I can keep the skip with an updated comment or just let it run, lmk what you think would be better
There was a problem hiding this comment.
I'd rather skip properly as in the mixin, does was not supposed to fallback for these tests otherwise the meaning of these tests get messy
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
|
Checking now! |
vasqu
left a comment
There was a problem hiding this comment.
Small comments left but close to merging 🤗
| # Per-layer RoPE base theta (0 => NoPE). Default: global rope_theta | ||
| if self.layer_rope_theta is None: | ||
| self.layer_rope_theta = [self.rope_parameters["rope_theta"]] * self.num_hidden_layers |
There was a problem hiding this comment.
Nit: I would move this before the super call if possible
There was a problem hiding this comment.
Not really feasible unfortunately - the super call sets / makes consistent the self.rope_parameters["rope_theta"] field
| embedding_multiplier (`float`, *optional*, defaults to 1.0): | ||
| embedding multiplier | ||
| logits_scaling (`float`, *optional*, defaults to 1.0): | ||
| divisor for output logits | ||
| residual_multiplier (`float`, *optional*, defaults to 1.0): | ||
| residual multiplier | ||
| attention_multiplier (`float`, *optional*, defaults to 1.0): | ||
| attention multiplier |
There was a problem hiding this comment.
A lot of multipliers that are actually not doing anything (based on the default values). Are they really necessary for this architecture?
- Could we remove
- Or at least shorten where we don't need them
If it's for modular it's fine as well. Just for my clarification
There was a problem hiding this comment.
The multiplier fields are necessary because we publish models with non-default values, but the docstring entries aren't needed here and I'll remove them
| # Mirror CausalLMModelTest's Flash2 check; otherwise FA2 auto-corrects to FA3 and no longer tests FA2. | ||
| if "flash_attention_2" not in (self.model_tester.base_model_class._compatible_flash_implementations or []): | ||
| self.skipTest("flash_attention_2 not in _compatible_flash_implementations (FA3/FA4-only)") | ||
| super().test_flash_attn_2_equivalence() |
There was a problem hiding this comment.
Ah no I actually meant to add this to the causallm tester. This should never happen for any of these models (with the same skip reason as the base mixin)
There was a problem hiding this comment.
# Some models only support a sub set of all FA implementations
valid_fa_implementations = model._compatible_flash_implementations
if valid_fa_implementations is not None and attn_implementation not in valid_fa_implementations:
continueThere was a problem hiding this comment.
ah ok, I wasn't sure if the shared tester harness was a valid target for a new-model PR
There was a problem hiding this comment.
It usually isn't but sometimes we do hit those cases :D
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, granite_swa, granitemoe_swa |
|
run-slow: granite_swa, granitemoe_swa |
|
This comment contains models: ["models/granite_swa", "models/granitemoe_swa"] |
|
run-slow: granite_swa, granitemoe_swa |
|
This comment contains models: ["models/granite_swa", "models/granitemoe_swa"] |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
Thanks @vasqu , the repetition is a bit surprising but it's not a correctness issue, the reference model is just somewhat experimental. FWIW the old output should still be correct for both cuda 8 and 9, but if adding that back in means rerunning the slow tests again then feel free to skip |
|
I think it's just a subtle difference between A100 and A10 GPUs, it can happen, I'd merge then 🫡 |
|
Great, thanks so much for all your help with this @vasqu ! |
* Add GraniteSWA model: Granite with Sliding Window Attention + learnable attention sinks - New model type 'granite_swa' with per-layer sliding window attention - LSE-based sink mechanism: sink_scale = sigmoid(lse - sinks) applied post-attention - Supports MuP parameters (embedding_multiplier, residual_multiplier, logits_scaling, attention_multiplier) - Per-layer attention pattern: configurable full_attention vs sliding_attention layers - Includes standalone modeling, configuration, modular source, and auto registration * Add FA3 attention dispatch with native LSE for sink scaling - Uses flash_attn_interface._flash_attn_forward for FA3 path (returns LSE natively) - Falls back to eager attention with manual LSE computation - Adds _supports_flash_attn and _compatible_flash_implementations flags - Verified: both eager and FA3 produce correct output * Enable DynamicSlidingWindowCache for memory-efficient KV cache - SWA layers store only last (sliding_window - 1) tokens in KV cache - Full attention layers keep complete KV history - ~98% KV memory savings on SWA layers for long sequences - Verified: identical results with and without sliding window cache * Finalize GraniteSWA: modular impl, FA4 support, formatting, tests, docs Make modular_granite_swa.py the single source of truth and regenerate the modeling/config: - Define GraniteSWAConfig in the modular file (inheriting GraniteConfig) so the converter no longer wipes the config, and rope_parameters standardizes correctly. - Apply the learnable attention sink through the standard `s_aux` attention dispatch (as in GPT-OSS) instead of the bespoke `_lse`/`_fa3_lse` path. Eager keeps the original sigmoid(logsumexp - sink) math; FlashAttention-3 and 4 now work via the shared interface (`s_aux`/`learnable_sink`). - Disable SDPA and FlexAttention (sink not expressible there); supported backends are eager, flash_attention_3, flash_attention_4. - Add tests/models/granite_swa (CausalLMModelTest framework) and the model doc + _toctree entry. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Add GraniteMoeSWA (MoeShared+SWA) Add GraniteMoeSWA, sublassing GraniteMoeShared, with SWA sliding window + sink added. - MoeSWAConfig inherits MoeSharedConfig, including default no shared experts (`shared_intermediate_size=0`). - Attention uses SWAAttention and thus supports Flash3+4 and eager. No SDPA or Flex. - MoeSWAModel handles per-layer causal masks during forward. All other layers inherit unchanged. - Add tests/models/granitemoe_swa (CausalLMModelTest framework), model doc and _toctree entry, and register model in auto mappings. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Update model refs, add per-layer rope/nope support a la Llama4 Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Small formatting/date fixes Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Test comments Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * automodel import style fix * Regenerate modeling_ files post-resync to main Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Fix upload date in model cards Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Additional TP / historic test exclusions post-sync to main Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Docs cleanup, one final round on submit date Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Fix grante_swa init and imports Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Addressing review comments: -Fix header years -TP and integration tests -Mark flex support -Rope/Nope based on position embed passing -Top level model defs via modular subclassing Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Replace binary rope/nope support with dynamic rope freq support Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Reorder router outputs to support native EP (and tests) Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Clean flex flags, flash correction tests, rope constructors Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Correct Flash2 test skip check Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Clean docstrings/comments, fix flash2 causal lm test Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Modelcard dates Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Fix rope typing and slow test shape checking Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Update integration test expectations by device Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Regenerate modeling files to match inherited modular upstream changes Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * adjust to ci * fixup repo * try * kernels sync --------- Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> Co-authored-by: Bharat-Runwal <bharatrunwal@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: vasqu <antonprogamer@gmail.com> Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>
* Add GraniteSWA model: Granite with Sliding Window Attention + learnable attention sinks - New model type 'granite_swa' with per-layer sliding window attention - LSE-based sink mechanism: sink_scale = sigmoid(lse - sinks) applied post-attention - Supports MuP parameters (embedding_multiplier, residual_multiplier, logits_scaling, attention_multiplier) - Per-layer attention pattern: configurable full_attention vs sliding_attention layers - Includes standalone modeling, configuration, modular source, and auto registration * Add FA3 attention dispatch with native LSE for sink scaling - Uses flash_attn_interface._flash_attn_forward for FA3 path (returns LSE natively) - Falls back to eager attention with manual LSE computation - Adds _supports_flash_attn and _compatible_flash_implementations flags - Verified: both eager and FA3 produce correct output * Enable DynamicSlidingWindowCache for memory-efficient KV cache - SWA layers store only last (sliding_window - 1) tokens in KV cache - Full attention layers keep complete KV history - ~98% KV memory savings on SWA layers for long sequences - Verified: identical results with and without sliding window cache * Finalize GraniteSWA: modular impl, FA4 support, formatting, tests, docs Make modular_granite_swa.py the single source of truth and regenerate the modeling/config: - Define GraniteSWAConfig in the modular file (inheriting GraniteConfig) so the converter no longer wipes the config, and rope_parameters standardizes correctly. - Apply the learnable attention sink through the standard `s_aux` attention dispatch (as in GPT-OSS) instead of the bespoke `_lse`/`_fa3_lse` path. Eager keeps the original sigmoid(logsumexp - sink) math; FlashAttention-3 and 4 now work via the shared interface (`s_aux`/`learnable_sink`). - Disable SDPA and FlexAttention (sink not expressible there); supported backends are eager, flash_attention_3, flash_attention_4. - Add tests/models/granite_swa (CausalLMModelTest framework) and the model doc + _toctree entry. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Add GraniteMoeSWA (MoeShared+SWA) Add GraniteMoeSWA, sublassing GraniteMoeShared, with SWA sliding window + sink added. - MoeSWAConfig inherits MoeSharedConfig, including default no shared experts (`shared_intermediate_size=0`). - Attention uses SWAAttention and thus supports Flash3+4 and eager. No SDPA or Flex. - MoeSWAModel handles per-layer causal masks during forward. All other layers inherit unchanged. - Add tests/models/granitemoe_swa (CausalLMModelTest framework), model doc and _toctree entry, and register model in auto mappings. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Update model refs, add per-layer rope/nope support a la Llama4 Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Small formatting/date fixes Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Test comments Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * automodel import style fix * Regenerate modeling_ files post-resync to main Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Fix upload date in model cards Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Additional TP / historic test exclusions post-sync to main Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Docs cleanup, one final round on submit date Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Fix grante_swa init and imports Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Addressing review comments: -Fix header years -TP and integration tests -Mark flex support -Rope/Nope based on position embed passing -Top level model defs via modular subclassing Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Replace binary rope/nope support with dynamic rope freq support Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Reorder router outputs to support native EP (and tests) Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Clean flex flags, flash correction tests, rope constructors Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Correct Flash2 test skip check Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Clean docstrings/comments, fix flash2 causal lm test Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Modelcard dates Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Fix rope typing and slow test shape checking Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Update integration test expectations by device Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * Regenerate modeling files to match inherited modular upstream changes Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> * adjust to ci * fixup repo * try * kernels sync --------- Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com> Co-authored-by: Bharat-Runwal <bharatrunwal@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: vasqu <antonprogamer@gmail.com> Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>
What does this PR do?
Adds two new models for upcoming Granite model releases, combining existing Granite/Granitemoe architectures respectively with per-layer sliding window attention and learnable per-head attention sinks. Supports shared experts, flash3+4 attention, and per-layer Rope/Nope toggling as well, following existing patterns (granitemoe, gpt-oss, llama4 respectively). SDPA and FlexAttention not currently supported.
Fixes #46645.
Code Agent Policy
Claude Code assisted in implementing these code changes; I have reviewed every line and run the HF formatting and correctness tests (SDPA, Flash2, and TP are skipped, as SDPA and Flash2 do not support sink tokens, and TP specifies SDPA).
Before submitting
Pull Request checks?
to it if that's the case.
(See Add SWA+sink option to Granite models #46645)
Who can review?
@ArthurZucker @Cyrilvallez @vasqu