SWA Granite HF registration (#1017)

## Description

Drops the dependency on the granite-swa-model transformers
[fork](https://github.com/Bharat-Runwal/transformers/tree/granite-swa-model).
Routes `granite_swa` through the generic FMS hf_pretrained/hf_configured
path. Pairs with
[this](foundation-model-stack/foundation-model-stack#538)
FMS registration PR.
<!-- Provide a clear description of your changes. -->

## Related Issues
none
<!-- Link related issues, e.g., `Fixes #` or `Relates to #456` -->

## Test Plan
none
<!-- Describe how you tested your changes. Include commands or steps to
reproduce. -->

## Checklist

- [x] I have read the [contributing
guidelines](https://docs.vllm.ai/projects/spyre/en/latest/contributing)
- [x] My code follows the project's code style (run `bash format.sh`)
- [ ] I have added tests for my changes (if applicable)
- [ ] I have updated the documentation (if applicable)
- [x] My commits include a `Signed-off-by:` line (DCO compliance)

---------

Signed-off-by: Yannick Schnider <Yannick.Schnider1@ibm.com>

Author: yannicks1

sendnn_inference/model_executor/model_loader/spyre.py

@@ -178,7 +178,13 @@ def __init__(
             self.parallel_config
         )
 
-        if self.config.model_type in {"llama", "granite", "granitemoehybrid", "qwen3"}:
+        if self.config.model_type in {
+            "llama",
+            "granite",
+            "granitemoehybrid",
+            "qwen3",
+            "granite_swa",
+        }:
             self.kv_cache_specs["num_layers"] = self.config.num_hidden_layers
             self.kv_cache_specs["head_dim"] = getattr(
                 self.fms_model.config,
@@ -237,6 +243,7 @@ def load_weights(
         # FMS's `hf_configured` path, which fetches only config.json and then
         # random-inits the model via `reset_parameters()`.
         variant: str | None = None
+        source: str | None = None
         if self.load_config.load_format == "dummy":
             logger.info(
                 "Loading model %s with random weights.",
@@ -271,6 +278,7 @@ def load_weights(
                 architecture=architecture,
                 variant=variant,
                 model_path=model_path,
+                source=source,
                 distributed_strategy=distributed_strategy,
                 group=dist.group.WORLD,
                 fused_weights=False,

sendnn_inference/platform.py

@@ -582,6 +582,25 @@ def _get_matching_warmup_shapes(
 
     @classmethod
     def pre_register_and_update(cls, parser: FlexibleArgumentParser | None = None) -> None:
+        # vLLM has no GraniteSWAForCausalLM architecture. FMS loads the real weights
+        # via the hf_pretrained path (model_loader/spyre.py); vLLM only needs
+        # a class for metadata (runner_type, defaults), so alias it to GraniteForCausalLM.
+        from vllm import ModelRegistry
+
+        if "GraniteSWAForCausalLM" not in ModelRegistry.get_supported_archs():
+            ModelRegistry.register_model(
+                "GraniteSWAForCausalLM",
+                "vllm.model_executor.models.granite:GraniteForCausalLM",
+            )
+
+        # Importing fms.models.hf registers the granite_swa model_type with
+        # transformers AutoConfig. This must run before vLLM's frontend parses
+        # config.json (in ModelConfig) so a granite_swa checkpoint loads without a
+        # granite_swa implementation in transformers. Done here (not in
+        # sendnn_inference/__init__.py, which deliberately avoids importing torch
+        # early); pre_register_and_update runs before ModelConfig is built.
+        import fms.models.hf  # noqa: F401
+
         if parser is None:
             return
 

sendnn_inference/v1/worker/spyre_model_runner.py

@@ -144,8 +144,8 @@ def __init__(
                 self.pad_token_id = getattr(self.model_config.hf_config, "pad_token_id", None) or 0
             if self.model_config.get_sliding_window():
                 logger.warning(
-                    "Sliding window is not supported on Spyre. "
-                    "The model will run without sliding window."
+                    "Sliding window is not fully supported on Spyre yet."
+                    "The model will run with experimental sliding window support."
                 )
 
         if vllm_config.device_config is None:

tests/e2e/test_granite_swa_load.py

@@ -0,0 +1,60 @@
+"""Load-path test for the granite_swa model_type.
+
+Constructing the LLM with --load-format dummy on a tiny granite_swa fixture
+flexes every granite_swa code path (each raises on failure):
+  1. SpyrePlatform.pre_register_and_update — GraniteSWAForCausalLM arch alias +
+     `import fms.models.hf` (granite_swa AutoConfig registration).
+  2. vLLM ModelConfig / AutoConfig parse of model_type="granite_swa" + arch resolution.
+  3. SpyreCausalLM.load_weights granite_swa kv-cache branch (no NotImplementedError).
+  4. FMS hf_configured path building a native GraniteSWA via reset_parameters.
+  5. The Spyre warmup that LLM() construction triggers (compile_or_warm_up_model
+     -> _warmup_spyre_dynamic_size), which runs prefill + decode forward passes
+     with dummy requests on eager CPU -- so a forward is exercised without generate().
+
+Temporary: granite_swa is loaded via the FMS AutoConfig shim because transformers
+has no granite_swa implementation yet. Once transformers ships granite_swa natively,
+that code and this test will be removed.
+"""
+
+from pathlib import Path
+
+import pytest
+
+# granite_swa ships only in FMS builds with fms/models/granite_swa.py, not in
+# released ibm-fms (<=1.12.1, which CI pins via uv.lock). Without it,
+# `import fms.models.hf` can't register granite_swa with AutoConfig and the load
+# fails. Skip cleanly where it's absent so CI stays green until FMS ships it.
+pytest.importorskip("fms.models.granite_swa")
+
+from vllm import LLM, ModelRegistry  # noqa: E402
+
+from sendnn_inference import envs as envs_spyre  # noqa: E402
+
+TINY_GRANITE_SWA_DIR = str(
+    Path(__file__).parent.parent / "fixtures" / "model_configs" / "ibm-granite" / "tiny-granite-swa"
+)
+# Reuse the cached micro-g3.3 tokenizer; the fixture's vocab_size (49159) matches it.
+TOKENIZER = "ibm-ai-platform/micro-g3.3-8b-instruct-1b"
+TOKENIZER_REVISION = "6e9c6465a9d7e5e9fa35004a29f0c90befa7d23f"
+
+
+@pytest.mark.cpu
+@pytest.mark.decoder
+def test_granite_swa_dummy_load():
+    """A granite_swa config loads via the dummy (random-init) path. No inference
+    asserted — reaching the assertions means every granite_swa path above succeeded."""
+    envs_spyre.override("SENDNN_INFERENCE_DYNAMO_BACKEND", "eager")
+
+    llm = LLM(
+        model=TINY_GRANITE_SWA_DIR,
+        tokenizer=TOKENIZER,
+        tokenizer_revision=TOKENIZER_REVISION,
+        load_format="dummy",
+        max_model_len=512,
+        max_num_seqs=4,
+        max_num_batched_tokens=128,
+    )
+
+    assert llm is not None
+    # Cheap, version-stable confirmation the granite_swa arch alias was registered.
+    assert "GraniteSWAForCausalLM" in ModelRegistry.get_supported_archs()

tests/fixtures/model_configs/ibm-granite/tiny-granite-swa/config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": ["GraniteSWAForCausalLM"],
+  "model_type": "granite_swa",
+  "vocab_size": 49159,
+  "hidden_size": 128,
+  "intermediate_size": 256,
+  "num_hidden_layers": 4,
+  "num_attention_heads": 4,
+  "num_key_value_heads": 2,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "max_position_embeddings": 2048,
+  "initializer_range": 0.02,
+  "rms_norm_eps": 1e-05,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "mlp_bias": false,
+  "embedding_multiplier": 1.0,
+  "logits_scaling": 1.0,
+  "residual_multiplier": 1.0,
+  "attention_multiplier": 1.0,
+  "sliding_window": 128,
+  "layer_types": ["full_attention", "sliding_attention", "sliding_attention", "full_attention"],
+  "rope_theta": 10000.0,
+  "rope_parameters": {"rope_type": "default", "rope_theta": 10000.0},
+  "tie_word_embeddings": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "pad_token_id": 0,
+  "use_cache": true,
+  "dtype": "float32"
+}