Skip to content

Add Granite-swa and Granitemoe-swa model support - #47179

Merged
vasqu merged 34 commits into
huggingface:mainfrom
daviswer:granite-swa
Jul 29, 2026
Merged

vasqu merged 34 commits into
huggingface:mainfrom
daviswer:granite-swa

Conversation

@daviswer

@daviswer daviswer commented Jul 8, 2026 •

Copy link
Copy Markdown
Contributor

CI

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

  • I confirm that this is not a pure code agent PR.

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

Who can review?

@ArthurZucker @Cyrilvallez @vasqu

Bharat-Runwal and others added 16 commits May 28, 2026 09:45
…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 vasqu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!)

Comment thread docs/source/en/model_doc/granite_swa.md Resolved
Comment thread docs/source/en/model_doc/granite_swa.md Resolved
Comment thread docs/source/en/model_doc/granite_swa.md Resolved
Comment thread docs/source/en/model_doc/granite_swa.md Resolved
Comment on lines +101 to +125
@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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd rather override those tbh

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
@daviswer

daviswer commented Jul 15, 2026 •

Copy link
Copy Markdown
Contributor Author

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

daviswer added 2 commits July 15, 2026 16:36
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
@github-actions

Copy link
Copy Markdown
Contributor

Thank you for your contribution 🤗!

CI Security Gate — automatic approval blocked

This PR was not automatically approved for CI because the security gate failed.

Possible reasons:

  • The PR touches 50 or more files — only PRs with fewer than 50 changed files are automatically approved
  • A changed file is outside the allowed directories (src/, tests/, docs/, utils/), has a disallowed extension (only .py, .txt, .md permitted outside tests/ and docs/), or is not .md/.yml inside docs/
  • A new high-severity security issue was detected in the changed Python files (Bandit check)

See the workflow run for the exact violations.

A maintainer can review and manually approve CI if a finding is a false positive.

@daviswer

daviswer commented Jul 15, 2026 •

Copy link
Copy Markdown
Contributor Author

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 vasqu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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`).

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh wow so even a diff between different full attention layers? Just fmi

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the intention is more for different SWA layers, but don't quote me on that

Comment on lines +92 to +113
# 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",
}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@daviswer daviswer Jul 17, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

daviswer added 2 commits July 17, 2026 16:43
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
@vasqu

vasqu commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Checking now!

@vasqu vasqu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small comments left but close to merging 🤗

Comment on lines +124 to +126
# 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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: I would move this before the super call if possible

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not really feasible unfortunately - the super call sets / makes consistent the self.rope_parameters["rope_theta"] field

Comment on lines +57 to +64
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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

            # 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:
                continue

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah ok, I wasn't sure if the shared tester harness was a valid target for a new-model PR

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It usually isn't but sometimes we do hit those cases :D

Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
daviswer and others added 2 commits July 28, 2026 21:30
Signed-off-by: Davis Wertheimer <davis.wertheimer@ibm.com>
@github-actions

Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: auto, granite_swa, granitemoe_swa

@vasqu

vasqu commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

run-slow: granite_swa, granitemoe_swa

@vasqu

vasqu commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

@daviswer I updated the values here 5474d5b - not 100% sure if the repetition is expected on the moe swa variation.

@github-actions

Copy link
Copy Markdown
Contributor

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/granite_swa", "models/granitemoe_swa"]
quantizations: []

@github-actions

Copy link
Copy Markdown
Contributor

CI Results

Workflow Run ⚙️

Commit Info

Context Commit Description
RUN 97e985e7 workflow commit (merge commit)
PR 8b242a63 branch commit (from PR)
main 2ef79f87 base commit (on main)

⚠️ No test being reported (jobs are skipped or cancelled)!

@vasqu

vasqu commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

run-slow: granite_swa, granitemoe_swa

@github-actions

Copy link
Copy Markdown
Contributor

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/granite_swa", "models/granitemoe_swa"]
quantizations: []

@github-actions

Copy link
Copy Markdown
Contributor

CI Results

Workflow Run ⚙️

Commit Info

Context Commit Description
RUN 6eb2dc42 workflow commit (merge commit)
PR d4cf007c branch commit (from PR)
main 2ef79f87 base commit (on main)

✅ No failing test specific to this PR 🎉 👏 !

@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

@daviswer

daviswer commented Jul 29, 2026 •

Copy link
Copy Markdown
Contributor Author

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

@vasqu

vasqu commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

I think it's just a subtle difference between A100 and A10 GPUs, it can happen, I'd merge then 🫡

@vasqu
vasqu enabled auto-merge July 29, 2026 21:03
@vasqu
vasqu added this pull request to the merge queue Jul 29, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jul 29, 2026
@vasqu
vasqu added this pull request to the merge queue Jul 29, 2026
Merged via the queue into huggingface:main with commit d24918c Jul 29, 2026
114 checks passed
@daviswer

Copy link
Copy Markdown
Contributor Author

Great, thanks so much for all your help with this @vasqu !

stevhliu pushed a commit to stevhliu/transformers that referenced this pull request Jul 30, 2026
* 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>
Sainava pushed a commit to Sainava/Sai-transformers that referenced this pull request Aug 3, 2026
* 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add SWA+sink option to Granite models

5 participants