Add GraniteSWA (granite_swa) with a shared SinkAttention component - #590
Open
justinchuby wants to merge 5 commits into
Open
Add GraniteSWA (granite_swa) with a shared SinkAttention component#590justinchuby wants to merge 5 commits into
justinchuby wants to merge 5 commits into
Conversation
…onent GraniteSWA is Granite plus three architectural changes, all of which this change implements faithfully against upstream `GraniteSWAForCausalLM`: * Mixed attention spans driven by `config.layer_types` (full vs a `sliding_window`-wide local window), dispatched per layer from two precomputed float additive biases. * Learnable per-head attention sinks. Upstream's eager kernel rescales the attention output by `sigmoid(logsumexp(scores) - sink)`; that is algebraically the same as keeping the sink as one extra column in the softmax denominator, which is how GPT-OSS already expresses it in mobius. The GPT-OSS implementation is therefore promoted to a reusable `SinkAttention` component (subclassing `Attention` so the projection, Q/K-norm and RoPE setup is shared) and GPT-OSS now uses it. A unit test runs the exported graph in ORT against the upstream sigmoid-LSE formula. * Per-layer RoPE base frequency, with `0` meaning NoPE. One rotary module is built per distinct non-zero theta, mirroring `GraniteSWAModel.rotary_embs`. Because the sink lives inside the softmax, no fused `Attention` or `GroupQueryAttention` op can be used; the score matrix is built explicitly and the bias carries causality, the window and padding. Reduced-precision builds upcast the sink softmax to float32, matching the upstream forced-fp32 path, while float32 builds keep an unchanged (Cast-free) graph. `GraniteSwaConfig` re-applies HuggingFace's `__post_init__` defaults (every fourth layer full attention, per-layer theta defaulting to the global one) and re-asserts RoPE, so a raw `config.json` mapping builds the same architecture as a materialised HF config object. Also pins the revision and the eager attention kernel through the causal-LM golden reference loader: the sink is not SDPA-expressible, so the reference must never silently drift onto a kernel that drops it. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchu@microsoft.com>
…rity tests The golden references were produced by an independently invoked HuggingFace pipeline (transformers 5.15.0 GraniteSWAForCausalLM, float32, eager attention, revision af1e3227100b61088eead48389ab5409b5d0e39c), not by the implementation under test: L4 prefill top-1 = 330, top-2 = 2355 L5 greedy, 20/20 tokens, decoding to: ' "The Last of the Mohicans" by James Fenimore Cooper. It is a poem about the' Both e2e_golden_test levels pass against the float32 CPU export. The new integration tests prompt with more than sliding_window (128) tokens on purpose. Below that length a windowed layer and a full-attention layer see exactly the same keys, so a window bug would pass unnoticed; the decode test then takes a cached step whose query lies outside the window of the oldest cached keys. Both pin the checkpoint revision and assert that the HuggingFace reference really resolved to the eager kernel. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchu@microsoft.com>
…sink softmax The decode parity test cross-fed HuggingFace's KV cache into a plain DynamicCache, which crashed inside HuggingFace's own mask builder. The root cause is a representation difference, not a numerical one: HuggingFace's DynamicSlidingWindowLayer physically drops keys once they leave the window (prefill of 20 tokens with sliding_window=8 returns 7 cached keys for sliding layers and 20 for full-attention layers), whereas the exported graph keeps the full cache and masks the same keys out through the sliding attention bias. Both are correct and neither cache can be fed to the other, so each side now runs prefill and decode against its own cache and only the decode logits are compared. The test asserts the trimming difference explicitly so the reason for the split is not lost. Also adds graph assertions for the float16/bfloat16 sink softmax: upstream forces the logsumexp/sigmoid and the softmax to float32 regardless of compute dtype, which here means an upcast before the softmax and a downcast after the sink column is dropped. float32 builds are asserted to stay Cast-free. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchu@microsoft.com>
Performance Comparison
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Comment on lines
+215
to
+217
| logits, present_key_values = super().forward( | ||
| op, input_ids, attention_mask, position_ids, past_key_values | ||
| ) |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds Mobius support for HuggingFace granite_swa (GraniteSWAForCausalLM) by introducing a reusable SinkAttention component (shared with GPT-OSS), wiring it through decoder-layer construction, and adding config extraction + parity/golden coverage for ibm-granite/granite-swash-2b.
Changes:
- Add
SinkAttentionas anAttentionsubclass and allowcreate_decoder_layer/DecoderLayerto swap attention implementations viaattention_class. - Introduce
GraniteSwaConfig+GraniteSwaCausalLMModel/backbone with mixed sliding/full attention, per-layer RoPE theta (including NoPE), and sink logits. - Add/extend synthetic parity, integration parity, and L4/L5 golden test cases for
granite_swa, plus golden-generation plumbing (revision + eager attention pinning).
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/synthetic_parity_test.py | Adds eager-pinned HF factory for granite_swa and tiny-config override for in-range BOS/EOS IDs. |
| tests/integration_test.py | Adds revision-pinned real-weight GraniteSWA prefill+decode parity tests against HF eager attention. |
| tests/_test_configs.py | Registers a tiny representative granite_swa config for graph-build and validation suites. |
| testdata/golden/causal-lm/granite-swash-2b.json | Adds L4 prefill golden reference for ibm-granite/granite-swash-2b. |
| testdata/golden/causal-lm/granite-swash-2b_generation.json | Adds L5 greedy generation golden reference for ibm-granite/granite-swash-2b. |
| testdata/cases/causal-lm/granite-swash-2b.yaml | Adds L4+L5 golden test case metadata (model_id/type/revision, generation settings). |
| src/mobius/models/granite_swa.py | Implements GraniteSWA model/backbone: mixed attention spans, per-layer RoPE theta (NoPE via 0), sink attention, Granite multipliers. |
| src/mobius/models/granite_swa_test.py | Adds unit tests for GraniteSWA config extraction, RoPE dispatch, sink algebra, and attention-op constraints. |
| src/mobius/models/gptoss.py | Refactors GPT-OSS sink attention to reuse shared SinkAttention. |
| src/mobius/models/init.py | Exports GraniteSwaCausalLMModel. |
| src/mobius/components/_decoder.py | Adds attention_class override to DecoderLayer and create_decoder_layer. |
| src/mobius/components/_attention.py | Adds shared _apply_qk_norm helper and introduces SinkAttention implementation. |
| src/mobius/components/init.py | Re-exports SinkAttention from the public components API. |
| src/mobius/_testing/torch_reference.py | Extends load_torch_model with revision and attn_implementation pinning + validation. |
| src/mobius/_registry.py | Registers granite_swa model type and adds default test model id mapping. |
| src/mobius/_configs/_base.py | Adds GraniteSwaConfig extraction to match HF defaults and re-assert RoPE for raw config.json path. |
| src/mobius/_configs/init.py | Exports GraniteSwaConfig. |
| scripts/generate_golden.py | Pins HF reference attention backend to eager for granite_swa and forwards Hub revision into reference loads. |
Suppressed comments (1)
tests/integration_test.py:5074
- This test pins the HF weights/build to _GRANITE_SWA_REVISION but loads ArchitectureConfig from the model’s default branch via _get_config(). If the Hub config drifts, the feeds (KV cache layout, sliding_window, layer_types) may stop matching the pinned checkpoint.
config = _get_config(_GRANITE_SWA_MODEL_ID)
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| load_weights=True, | ||
| ) | ||
| torch_model, tokenizer = _load_granite_swa_reference() | ||
| config = _get_config(_GRANITE_SWA_MODEL_ID) |
Comment on lines
+587
to
+596
| if isinstance(attention_bias, GQAContext): | ||
| raise TypeError( | ||
| "SinkAttention cannot emit GroupQueryAttention: the sink logit " | ||
| "must take part in the softmax denominator. Build this model " | ||
| "with a float additive attention bias instead." | ||
| ) | ||
| if static_cache is not None: | ||
| raise NotImplementedError( | ||
| "SinkAttention does not support the opset-24 static KV cache." | ||
| ) |
TestRegistryConsistency::test_config_class_declared_on_model requires that a registry entry declaring a non-default config_class agrees with the model class, so that a model can never silently inherit a different config from its parent. GraniteSwaCausalLMModel was inheriting CausalLMConfig from CausalLMModel while the registry declared GraniteSwaConfig, which would drop layer_rope_theta on any path that resolves the config from the class rather than the registry. The local run that missed this used `-k granite_swa`, which does not select the consistency test; the model type only appears in its assertion message. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchu@microsoft.com>
The shared SinkAttention component unconditionally forced the sink softmax to
float32 on non-float32 builds. That reproduces GraniteSWA, whose eager kernel
does
sink_scale = (lse - sinks).to(torch.float32).sigmoid()
attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32)
but it is wrong for GPT-OSS, which is explicit about staying in the compute
dtype:
probs = F.softmax(combined_logits, dim=-1, dtype=combined_logits.dtype)
So the shared component was changing GPT-OSS float16/bfloat16 numerics and
doubling the size of the largest score tensor, neither of which main does.
The precision contract now lives in the type rather than in a constructor
flag, because it is a property of the architecture and not a tuning knob:
SinkAttention keeps upcast_sink_softmax = False (GPT-OSS), and the new
Float32SinkAttention subclass sets it to True (GraniteSWA). A model cannot
silently acquire or lose the upcast by being constructed differently.
Verified: the tiny gpt_oss graph built at origin/main (059bf5f) and at this
commit has an identical node op_type sequence and identical sorted initializer
names for float32, float16 and bfloat16 alike (Cast counts 8 / 10 / 10 on both
sides), so GPT-OSS is graph-neutral again at every dtype rather than only at
float32. gpt_oss and all granite synthetic parity tests pass.
New regression tests assert both halves of the contract: GPT-OSS has no Cast
on the sink-softmax path at f32/f16/bf16, and GraniteSWA keeps the upcast at
f16/bf16 while staying Cast-free at f32.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchu@microsoft.com>
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.
Adds Mobius support for the HuggingFace architecture
granite_swa(GraniteSWAForCausalLM), validated againstibm-granite/granite-swash-2bpinned at revisionaf1e3227100b61088eead48389ab5409b5d0e39c.What GraniteSWA is
Granite plus three changes, all implemented against the upstream
modular_granite_swa.py(transformers 5.15.0):config.layer_typesselects, per layer, full causal attention or asliding_window-wide local window (128 for this checkpoint; every fourth layer is full attention).config.layer_rope_theta[i], where0means NoPE for that layer.Granite's four scaling multipliers (
embedding_multiplier,attention_multiplier,logits_scaling,residual_multiplier) still apply and are wired through the existingcreate_decoder_layerpath.Implementation
Shared
SinkAttentioncomponent, reused from GPT-OSSUpstream's eager kernel computes an ordinary softmax and rescales the output by
sigmoid(logsumexp(scores) - sink). WritingZ = sum(exp(scores)), that factor isZ / (Z + exp(sink))— which is exactly the mass removed by keeping the sink as one extra softmax column and then dropping it. That column form is already how GPT-OSS expresses sinks in Mobius, so rather than writing a second implementation, the GPT-OSS one is promoted into a reusableSinkAttentioncomponent (subclassingAttention, so projections, Q/K-norm and RoPE setup are shared) and_GptOssAttentionbecomes a thin subclass.DecoderLayer/create_decoder_layergained anattention_classkeyword so a model can swap the attention module without duplicating the layer.Because the sink lives inside the softmax, no fused
AttentionorGroupQueryAttentionop can be used — the score matrix is built explicitly, and the attention bias therefore carries causality, the sliding window and padding (is_causalis not available). This is asserted by a test.The two sink models have different precision contracts
This is a genuine divergence upstream, so it is encoded in the type rather than in a constructor flag:
F.softmax(combined_logits, dim=-1, dtype=combined_logits.dtype). It usesSinkAttention(upcast_sink_softmax = False).(lse - sinks).to(torch.float32).sigmoid()andF.softmax(attn_weights, dim=-1, dtype=torch.float32). It usesFloat32SinkAttention(upcast_sink_softmax = True). In the extra-column formulation thelogsumexp/sigmoidpair and the softmax are the same reduction, so upcasting that one softmax reproduces both upstream upcasts at once.Keeping this in the type means a model cannot silently acquire or lose the upcast by being constructed differently. For float32 builds the two classes emit an identical, Cast-free graph.
GPT-OSS is graph-neutral at every dtype. Verified by building the tiny
gpt_ossgraph at059bf5f1and at this branch: identical nodeop_typesequence and identical sorted initializer names for float32, float16 and bfloat16 (Cast counts 8 / 10 / 10 on both sides). GPT-OSS synthetic parity also still passes. Regression tests now assert both halves of the contract — GPT-OSS has no Cast on the sink-softmax path at any dtype, GraniteSWA has the upcast at f16/bf16 and not at f32.Config
GraniteSwaConfig.from_transformersreproduces HF__post_init__semantics —layer_typesdefaulting to full attention on every fourth layer,layer_rope_thetadefaulting to the globalrope_theta,0preserved as NoPE — and additionally re-asserts RoPE for the rawconfig.jsonpath: the checkpoint carries only a flatrope_theta: 10000with norope_parameters, which the generic extractor otherwise reads as "no RoPE signal at all". Without this, a raw-JSON build would silently produce an all-NoPE model.config_classis declared on both the registry entry and the model class, so the model can never silently fall back toCausalLMConfigand droplayer_rope_theta.Revision and eager-attention pinning
load_torch_modelgainedrevisionandattn_implementationparameters, andscripts/generate_golden.pynow forwards both. The sink is not SDPA-expressible — upstream setsGraniteSWAPreTrainedModel._supports_sdpa = Falsefor exactly this reason — so the reference is pinned toeagerand the resolved implementation is asserted, rather than relying on transformers' default resolution.Evidence
All numbers below are measured, not estimated.
L1
pytest tests/build_graph_test.py -k granite_swa— 4 passedpytest tests/weight_alignment_test.py -k granite_swa— 1 passedpytest src/mobius/models/granite_swa_test.py— 22 passedThe unit tests include an executed check of the sink algebra: a one-layer sink attention graph is run in ORT and compared against upstream's
sigmoid(logsumexp - sink)formula atrtol/atol 1e-4. Others assert the pinnedconfig.jsonextraction, one rotary module per distinct non-zero theta (6RotaryEmbeddingnodes for 3 RoPE layers + 1 NoPE layer), tied embeddings collapsing to a singlemodel.embed_tokens.weightinitializer,model.layers.N.self_attn.sinksmatching HF naming, that no fused attention op is emitted, and the per-dtype precision contract.L2
pytest tests/arch_validation_test.py -m arch_validation -k granite_swa— 3 passed (full-size real HF config, graph build + shape consistency)pytest tests/yaml_schema_test.py— all passedL3
pytest tests/synthetic_parity_test.py -k granite_swa— passed at the defaultatol 1e-3, no tolerance override was needed. The HF reference is forced to eager and asserted. The tiny config deliberately carries 4 layers with mixed spans, two distinct non-zero thetas and one NoPE layer.test_granite_swa_prefill_logits_matchandtest_granite_swa_decode_step_logits_matchboth pass atrtol/atol 1e-3against revision-pinned HF eager. Both use a 169-token prompt against a 128-token window on purpose — below the window a sliding layer and a full-attention layer see the same keys, so a window bug would pass unnoticed.L4 / L5
Goldens generated by an independently invoked HuggingFace pipeline (transformers 5.15.0, float32, eager, revision pinned) and committed.
330, top-22355' "The Last of the Mohicans" by James Fenimore Cooper. It is a poem about the'pytest tests/e2e_golden_test.py -m "golden or generation" -k granite_swash— 2 passed (505 s, CPU fp32)Multi-dtype and multi-EP
CUDA EP creation was verified via
get_providers()and the CUDA-specific memcpy log — not assumed. Reference is HF fp32 CPU, eager, 169-token prompt, full-logit comparison including a native cached decode step:Generation, 20-token greedy, compared against the committed float32 L5 golden token sequence:
bfloat16end to end)So fp32, fp16 and bf16, on both CPU and CUDA, agree semantically; no dtype is degraded or silently downgraded.
CLI build and metadata
mobius build --model ibm-granite/granite-swash-2b --revision af1e3227... --dtype f16 --runtime ort-genaicompletes;genai_config.jsonfaithfully reflects the graph (type: granite_swa,head_size: 128, 20 / 4 heads, 24 layers, correct input/output names and cache templates). The bf16 build completes as well. Docs are auto-generated from the class metadata and pick the model up (docs/models/granite_swa.md,Text Generation).Olive INT4 quantization
OnnxKQuantQuantization(4-bit, block size 32) runs to completion on the f16 export in 1020 s on CPU, and the result loads and generates:MatMulNBitsnodes = 24 layers × 7 projections — every projection was quantized.MatMulremain: the 48 attention score/context matmuls (activation × activation, correctly left alone) plus the LM head.' "The Last of the Mohicans" by James Fenimore Cooper. I hope you enjoy it.'— the first 14 of 20 tokens match the fp32 golden exactly before diverging into an equally coherent continuation, which is the expected behaviour for INT4.Two workarounds were needed on this machine and are worth recording:
olive run --configaborts because Olive 0.13 auto-registers every provider DLL bundled in the ORT GPU wheel andnvinfer_10.dllis absent, so the documented suppression ofolive.systems.local.maybe_register_ep_librariesaroundolive.workflows.runis required; andcupyis not installed here, so the pass ran on CPU.Notes and caveats (no waiver needed, but worth recording)
Sliding-cache representation differs from HuggingFace, by design. HF's
DynamicSlidingWindowLayerphysically drops keys once they leave the window — a 20-token prefill withsliding_window=8returns 7 cached keys for sliding layers and 20 for full-attention layers. The exported graph keeps the full cache and masks the same keys out through the sliding bias. Both are correct for every reachable query (the window only moves forward, so dropped keys can never be attended to again), but the two caches cannot be cross-fed. This is why the decode parity test runs each side against its own cache; the test asserts the trimming difference explicitly so the reason is not lost. Consequence: the exported KV cache grows with sequence length rather than being bounded by the window, which is the existing Mobius behaviour for sliding-window models.414 CUDA Memcpy nodes (performance, not correctness). Attention sinks force a manual score matrix, so sink attention builds per-layer dynamic shape tensors (
Shape/Concat/Expand) that ORT pins to CPU — roughly 17 per layer across 24 layers, on both the f16 and bf16 graphs. This is inherent to any sink model (GPT-OSS has the same shape) and is not introduced by this change.No fp32 GPU golden run.
granite-swash-2bis 2.1B parameters, so an fp32 CUDA session needs ~8.6 GB against ~7.3 GB free on the available A1000 — it OOMs. The f16 and bf16 CUDA runs above target the same committed fp32 golden token sequence and match it exactly, so the GPU evidence is equivalent rather than missing.ORT GenAI 0.15.2 rejects the model type string, but runs the graph. Loading the exported package fails with
Unsupported model_type in config.json: granite_swa. A diagnostic probe (a scratch copy ofgenai_config.jsonwith the type string changed, not a change to the exporter) loads and generates successfully, reproducing the exact golden continuation:So this is purely a gap in that runtime's model-type registry, not a topology or export problem. Per the quality checklist, downstream runtime acceptance is not an export gate, and deliberately no compatibility alias was added to the exporter — claiming
granite_swaisgranitewould be a false alias that hides the sliding window, the sinks and the per-layer theta.A CUDA gotcha on this machine, for anyone reproducing the numbers. A raw
ort.InferenceSession(..., providers=["CUDAExecutionProvider"])silently falls back to CPU here unlesstorchis imported first, becausecublasLt64_12.dllonly reaches the DLL search path via torch's bundled CUDA libraries. All CUDA numbers above were re-verified withget_providers()after hitting exactly this.Pre-existing failures in this environment (not caused by this change)
Verified by
git stashon the base commit059bf5f1. This machine has ORT 1.26.0 whilepyproject.tomlasks for>=1.28.0, which explains the softcap ones:src/mobius/rewrite_rules/_decompose_attention_test.py::TestDecomposeAttentionParity::test_matches_fused_op[softcap]and[decode_softcap_past]src/mobius/models/deepseek_v4_test.py::TestDeepSeekV4QMoEExport::...(2 params)src/mobius/integrations/gguf/_builder_test.py::TestBuildQuantizedGguf::test_gatherblockquantized_zero_point_dequantizes_q4_0scripts/generate_golden_test.py::TestDryRun(2) — Windows cp1252 consoleUnicodeEncodeErroron the→character; passes withPYTHONIOENCODING=utf-8Full non-GPU sweep on this branch: 5084 passed, with only the failures listed above (plus one
qwen_image_testCUDA diffusion test that contended for the GPU with another session).