[WIP] experimental GraniteSWAForCausalLM support

sendnn_inference/model_executor/model_loader/spyre.py

@@ -105,6 +105,23 @@
     return mm_device
 
 
+def _granite_swa_variant_from_hf(cfg: PretrainedConfig) -> str:
+    """Pick the FMS granite_swa variant string from the HF config dimensions.
+
+    Mirrors the registrations in fms/models/granite_swa.py:
+      3b:  hidden_size=2560, num_hidden_layers=24
+      20b: hidden_size=4096, num_hidden_layers=44
+    """
+    if cfg.hidden_size == 4096 and cfg.num_hidden_layers == 44:
+        return "20b"
+    if cfg.hidden_size == 2560 and cfg.num_hidden_layers == 24:
+        return "3b"
+    raise ValueError(
+        f"No registered granite_swa variant for "
+        f"hidden_size={cfg.hidden_size}, num_hidden_layers={cfg.num_hidden_layers}"
+    )
+
+
 class SpyreCausalLM(nn.Module):
     def __init__(
         self,
@@ -178,7 +195,7 @@
             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 +254,7 @@
         # 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.",
@@ -258,6 +276,16 @@
                     revision=model_config.revision,
                 )
 
+        # granite_swa is not auto-detected by FMS's hf_pretrained/hf_configured
+        # inference, so route it through its explicitly-registered architecture.
+        # source="hf" converts the HF checkpoint; for dummy it stays None (with
+        # model_path=None) so FMS random-inits via `reset_parameters()`.
+        if self.config.model_type == "granite_swa":
+            architecture = "granite_swa"
+            variant = _granite_swa_variant_from_hf(self.config)
+            if self.load_config.load_format != "dummy":
+                source = "hf"
+
         # Get any fixes needed that must be patched into the kwargs;
         # currently this is only use for multimodal models / llava next
         model_kwargs = spyre_mm.get_mm_specific_load_overrides(self.config)
@@ -271,6 +299,7 @@
                 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,17 @@ 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 explicit "granite_swa" 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",
+            )
+
         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: