Change GraniteSwitch inheritance to MoEShared

.github/workflows/gpu-tests.yaml

@@ -106,17 +106,18 @@ jobs:
     # this is evaluated before any step runs.
     name: GPU Tests (${{ matrix.label }} · ${{ inputs.suite || 'full' }})
     # granite-switch supports two mutually-exclusive vLLM lines and both must be
-    # tested. They cannot share a venv (pyproject declares vllm19/vllm20 as
+    # tested. They cannot share a venv (pyproject declares vllm26/vllm27 as
     # conflicting groups), so each gets its own cluster pod. Only the dev* groups
-    # used because the bare vllm19/vllm20 groups omit pytest.
+    # used because the bare vllm26/vllm27 groups omit pytest. The label encodes
+    # the vLLM minor the leg installs (the GPU harness asserts it matches).
     strategy:
-      fail-fast: false   # a vllm19 failure must not hide the vllm20 result
+      fail-fast: false   # a vllm26 failure must not hide the vllm27 result
       matrix:
         include:
-          - label: vllm19
+          - label: vllm26
             dep_group: dev
-          - label: vllm20
-            dep_group: dev-vllm20
+          - label: vllm27
+            dep_group: dev-vllm27
     runs-on: [self-hosted, gpu]
 
     # MUST be set explicitly. GitHub's default is 360 minutes, which is exactly the

CLAUDE.md

@@ -446,6 +446,21 @@ config keys, an old unfused checkpoint would otherwise load into a fused model a
 mismatch keys, so `GraniteSwitchConfig` rejects `unfused_qkv=True` with a re-compose error.
 See [docs/SR_ARCHITECTURE.md](docs/SR_ARCHITECTURE.md) for details.
 
+### 13. `shared_intermediate_size == 0` Means Skip the Shared MLP, Not Build a Zero-Width One
+
+A pure sparse-MoE base (`granitemoe`) has no dense shared MLP; upstream encodes this as
+`shared_intermediate_size == 0`. `GraniteMoeSharedMLP.__init__` sets `self.hidden_size =
+config.shared_intermediate_size`, so constructing it at 0 would register `[0, H]`/`[H, 0]` weights
+that no checkpoint ships and then add their output into the MoE result. All three decoders (`hf`,
+`vllm/decoder/lora`, `vllm/decoder/shadow_residual`) therefore gate on `shared_intermediate_size >
+0`, set `shared_mlp = None` otherwise, and `config.py` correspondingly drops the
+`shared_input_linear`/`shared_output_linear` LoRA targets so no zero-width LoRA is built either.
+A layer with neither experts nor a shared MLP is invalid: the three decoders raise `ValueError`,
+but **the config itself does not** — the guard is decoder-level. This is why the vLLM-LoRA forward's
+`not has_experts → shared_mlp(x)` branch is safe only because that `ValueError` upstream rules out
+the no-MLP case; protect it if you refactor the gate. Config-level coverage:
+`tests/unit/test_config_edge_cases.py`.
+
 ## Pre-commit
 
 **See [docs/CICD.md](docs/CICD.md) for the full CI/CD setup — pre-commit hook list, setup steps, and what runs on every commit vs. in CI.**

CONTRIBUTING.md

@@ -27,7 +27,7 @@ Or via pip: `pip install uv`
    > ```bash
    > uv sync --frozen --no-default-groups --extra hf --extra compose
    > ```
-   > `--no-default-groups` skips the default `vllm19` group; `--frozen` installs strictly from the
+   > `--no-default-groups` skips the default `vllm26` group; `--frozen` installs strictly from the
    > committed `uv.lock` without modifying it. (What fails on macOS is *installing* the vLLM/CUDA
    > wheels — resolving the lockfile with `uv lock` works fine, so the `uv-lock` pre-commit hook runs
    > locally too; the CPU subset above simply never installs those wheels.) To run the

README.md

@@ -47,15 +47,16 @@ Other install options depending on your use case:
 ```bash
 pip install "granite-switch[compose]"   # Compose modular models
 pip install "granite-switch[hf]"        # HuggingFace inference
-pip install "granite-switch[vllm20]"    # vLLM 0.20+ (requires CUDA 13+)
+pip install "granite-switch[vllm27]"    # newer vLLM line (0.27.x)
 pip install "granite-switch[dev]"       # Everything
 ```
 
-Requires Python 3.10+ and PyTorch 2.0+.
+Requires Python 3.11+ and PyTorch 2.11+.
 
-> **vLLM version note:** This project currently defaults to vLLM 0.19.1 due to vLLM 0.20's
-> dependency on CUDA 13.0+ (via PyTorch 2.11), which is incompatible with many existing
-> environments running CUDA 12.x drivers. Use `.[vllm20]` if your environment supports CUDA 13+.
+> **vLLM version note:** This project requires `transformers>=5.16`, so it pins vLLM to the
+> `0.26.x` line (the earliest vLLM whose model code handles the transformers-5.13+
+> `full_attention` layer-type rename) via the default `[vllm]` extra. `[vllm27]` selects the
+> newer `0.27.x` line. Both require PyTorch 2.11+ (CUDA 13+).
 
 ### Compose a Model
 

docs/AUDIO.md

@@ -19,11 +19,11 @@ non-16 kHz input:
 
 ```bash
 # Serving an audio-enabled checkpoint
-uv sync --extra vllm --extra audio     # or --extra vllm20 --extra audio
+uv sync --extra vllm --extra audio     # or --extra vllm27 --extra audio
 
 # Development / running the test suite (the dev groups include audio already)
-uv sync --group dev                    # vLLM 0.19.x
-uv sync --group dev-vllm20             # vLLM 0.20.x
+uv sync --group dev                    # vLLM 0.26.x
+uv sync --group dev-vllm27             # vLLM 0.27.x
 ```
 
 ## Building an audio-enabled checkpoint

docs/CICD.md

@@ -221,7 +221,7 @@ the full test suite was run to confirm the mechanical changes introduced no
 regressions.
 
 **How:** Submitted to the Vela GPU cluster (namespace `security`, 4 GPUs,
-`vllm19` dependency group) via a job config derived from `tests_on_uv_vllm19.yaml`,
+`vllm26` dependency group) via a job config derived from `tests_on_uv_vllm26.yaml`,
 pointed at the formatted branch. It runs a ruff sanity check followed by all five
 suites with `pytest -n 4`.
 
@@ -259,7 +259,7 @@ very unlikely to stem from PR 1 because:
    integration `test_forward_logit_equivalence`.
 
 **Follow-up:** Confirm the same test also flips a position on unmodified `main`
-(e.g. an integration-only run from `tests_on_uv_vllm19.yaml`). If so, treat it as
+(e.g. an integration-only run from `tests_on_uv_vllm26.yaml`). If so, treat it as
 a flaky near-tie test to be addressed separately (tolerance/top-k handling), and
 consider PR 1 cleared.
 

pyproject.toml

@@ -10,8 +10,8 @@ readme = "README.md"
 license = "Apache-2.0"
 requires-python = ">=3.11,<3.14"
 dependencies = [
-    "torch>=2.10.0",
-    "transformers>=5.5.1,<5.10.0",
+    "torch>=2.11.0",
+    "transformers>=5.16,<=5.17.0",
 ]
 
 [project.urls]
@@ -20,8 +20,8 @@ Documentation = "https://github.com/generative-computing/granite-switch/tree/mai
 
 [project.optional-dependencies]
 hf = ["accelerate>=0.20.0"]
-vllm = ["vllm>=0.19.1,<0.20.0"]
-vllm20 = ["vllm>=0.20.0,<0.21.0"]
+vllm = ["vllm>=0.26.0,<0.27.0"]
+vllm27 = ["vllm>=0.27.0,<0.28.0"]
 compose = ["huggingface_hub", "pyyaml", "tqdm", "safetensors"]
 build = ["huggingface_hub", "pyyaml", "tqdm", "safetensors"]  # Backward compatibility alias for compose
 # Audio (ASR) decode + resample. Reuse vLLM's own audio deps (unversioned, so it
@@ -56,33 +56,33 @@ markers = [
 ]
 
 [dependency-groups]
-vllm19     = ["vllm>=0.19.1,<0.20.0"]
-vllm20     = ["vllm>=0.20.0,<0.21.0"]
+vllm26     = ["vllm>=0.26.0,<0.27.0"]
+vllm27     = ["vllm>=0.27.0,<0.28.0"]
 # `audio` is included so the audio tests can actually run: the ASR path needs
 # vLLM's audio deps (av/soundfile/resampy) at runtime, and no group pulled them
 # in before (integration tests failed with ModuleNotFoundError on a synced pod).
-dev        = ["pytest", "pytest-cov", { include-group = "vllm19" }, "granite-switch[hf,compose,audio]"]
-dev-vllm20 = ["pytest", "pytest-cov", { include-group = "vllm20" }, "granite-switch[hf,compose,audio]"]
+dev        = ["pytest", "pytest-cov", { include-group = "vllm26" }, "granite-switch[hf,compose,audio]"]
+dev-vllm27 = ["pytest", "pytest-cov", { include-group = "vllm27" }, "granite-switch[hf,compose,audio]"]
 test       = ["pytest", "pytest-cov", "bitsandbytes", "optimum-quanto", { include-group = "dev" }]
 
 [tool.uv]
-default-groups = ["vllm19"]
+default-groups = ["vllm26"]
 conflicts = [
     # group-vs-group
-    [{ group = "vllm19" },     { group = "vllm20" }],
-    [{ group = "dev" },        { group = "vllm20" }],
-    [{ group = "dev" },        { group = "dev-vllm20" }],
-    [{ group = "dev-vllm20" }, { group = "vllm19" }],
+    [{ group = "vllm26" },     { group = "vllm27" }],
+    [{ group = "dev" },        { group = "vllm27" }],
+    [{ group = "dev" },        { group = "dev-vllm27" }],
+    [{ group = "dev-vllm27" }, { group = "vllm26" }],
     # group-vs-extra
-    [{ group = "vllm19" },     { extra = "vllm20" }],
-    [{ group = "vllm20" },     { extra = "vllm" }],
-    [{ group = "vllm20" },     { extra = "tutorials" }],
-    [{ group = "dev" },        { extra = "vllm20" }],
-    [{ group = "dev-vllm20" }, { extra = "vllm" }],
-    [{ group = "dev-vllm20" }, { extra = "tutorials" }],
+    [{ group = "vllm26" },     { extra = "vllm27" }],
+    [{ group = "vllm27" },     { extra = "vllm" }],
+    [{ group = "vllm27" },     { extra = "tutorials" }],
+    [{ group = "dev" },        { extra = "vllm27" }],
+    [{ group = "dev-vllm27" }, { extra = "vllm" }],
+    [{ group = "dev-vllm27" }, { extra = "tutorials" }],
     # extra-vs-extra
-    [{ extra = "vllm" },       { extra = "vllm20" }],
-    [{ extra = "tutorials" },  { extra = "vllm20" }],
+    [{ extra = "vllm" },       { extra = "vllm27" }],
+    [{ extra = "tutorials" },  { extra = "vllm27" }],
 ]
 
 [tool.coverage.run]

src/granite_switch/composer/__init__.py

@@ -11,8 +11,8 @@
     granite_dense_arch,
     granite_dense_sr_arch,
     granite_moe_arch,
-    granite_moe_hybrid_arch,
-    granite_moe_hybrid_sr_arch,
+    granite_moe_shared_arch,
+    granite_moe_shared_sr_arch,
     granite_moe_sr_arch,
     resolve_arch,
 )
@@ -28,8 +28,8 @@
     "granite_dense_arch",
     "granite_dense_sr_arch",
     "granite_moe_arch",
-    "granite_moe_hybrid_arch",
-    "granite_moe_hybrid_sr_arch",
+    "granite_moe_shared_arch",
+    "granite_moe_shared_sr_arch",
     "granite_moe_sr_arch",
     "resolve_arch",
 ]

src/granite_switch/composer/arch.py

@@ -220,7 +220,7 @@ def _dense_mlp_to_shared_groups() -> list[ModuleDescriptor]:
     Used for dense Granite models whose base uses ``mlp.gate_proj`` /
     ``mlp.up_proj`` / ``mlp.down_proj`` but whose switch model uses
     ``shared_mlp.input_linear`` / ``shared_mlp.output_linear``
-    (the ``GraniteMoeHybridMLP`` layout).
+    (the ``GraniteMoeSharedMLP`` layout).
     """
     return [
         ModuleDescriptor(
@@ -315,7 +315,7 @@ def _cross_stream_groups() -> list[ModuleDescriptor]:
     "embedding_multiplier": 1.0,
     "logits_scaling": 1.0,
     "attention_multiplier": 1.0,
-    # Granite/GraniteMoeHybrid vLLM classes use separate add-then-norm.
+    # Granite/GraniteMoeShared vLLM classes use separate add-then-norm.
     "fused_add_norm": False,
 }
 
@@ -326,28 +326,24 @@ def _cross_stream_groups() -> list[ModuleDescriptor]:
     "shared_intermediate_size": None,
 }
 
-# Layer type fields (propagated for hybrid models)
-_HYBRID_OPTIONAL_FIELDS: dict[str, Any] = {
-    "layer_types": None,
-    "position_embedding_type": "rope",
-}
-
 
 # ---------------------------------------------------------------------------
 # Architecture factory functions
 # ---------------------------------------------------------------------------
 
 
-def granite_moe_hybrid_arch(base_config=None) -> ArchDescriptor:
-    """GraniteMoeHybrid architecture (model_type ``granitemoehybrid``).
+def granite_moe_shared_arch(base_config=None) -> ArchDescriptor:
+    """GraniteMoeShared architecture (model_type ``granitemoeshared`` /
+    ``granitemoehybrid``).
 
-    GraniteMoeHybrid models use ``shared_mlp`` module naming
+    Granite 4 MoE-with-shared-expert models use ``shared_mlp`` module naming
     (``shared_input_linear``, ``shared_output_linear``), even dense layers
-    with ``num_local_experts=0``.
+    with ``num_local_experts=0``. Real Granite 4.x dense checkpoints are still
+    typed ``granitemoehybrid`` upstream (they carry no mamba layers), so this
+    same descriptor serves both model_type strings.
     """
     optional_fields = dict(_GRANITE_OPTIONAL_FIELDS)
     optional_fields.update(_MOE_OPTIONAL_FIELDS)
-    optional_fields.update(_HYBRID_OPTIONAL_FIELDS)
 
     return ArchDescriptor(
         groups=list(_common_attn_groups()) + list(_moe_shared_mlp_groups()),
@@ -361,18 +357,16 @@ def granite_moe_arch(base_config=None) -> ArchDescriptor:
 
     Pure sparse MoE: every layer has an expert bank and **no** dense
     ``shared_mlp``.  The descriptor is therefore ``_common_attn_groups()`` and
-    nothing else — a strict subset of :func:`granite_moe_hybrid_arch`.
+    nothing else — a strict subset of :func:`granite_moe_shared_arch`.
 
-    The frozen expert tensors (``block_sparse_moe.input_linear`` /
-    ``output_linear`` / ``router.layer``) are named identically in the switch
-    model, so with no shared-MLP group to shadow them they transfer by identity.
+    The frozen expert tensors (``block_sparse_moe.experts.gate_up_proj`` /
+    ``experts.down_proj`` / ``router.weight`` in the transformers-5.16 layout)
+    are named identically in the switch model, so with no shared-MLP group to
+    shadow them they transfer by identity.
 
     ``shared_intermediate_size`` is pinned to ``0``, which is upstream's own
     encoding for "no shared MLP"
     (``granitemoeshared``: ``shared_mlp = None if shared_intermediate_size == 0``).
-    ``position_embedding_type`` is deliberately not propagated:
-    ``GraniteSwitchConfig`` already defaults it to ``"rope"``, which is what
-    granitemoe uses.
     """
     optional_fields = dict(_GRANITE_OPTIONAL_FIELDS)
     optional_fields.update(_MOE_OPTIONAL_FIELDS)
@@ -409,13 +403,13 @@ def granite_dense_arch(base_config=None) -> ArchDescriptor:
 ]
 
 
-def granite_moe_hybrid_sr_arch(base_config=None) -> ArchDescriptor:
-    """GraniteMoeHybrid Shadow Residual architecture.
+def granite_moe_shared_sr_arch(base_config=None) -> ArchDescriptor:
+    """GraniteMoeShared Shadow Residual architecture.
 
-    Identical to :func:`granite_moe_hybrid_arch` (fused projections) plus the
+    Identical to :func:`granite_moe_shared_arch` (fused projections) plus the
     layer-level ``cross_stream`` injection site.
     """
-    arch = granite_moe_hybrid_arch(base_config=base_config)
+    arch = granite_moe_shared_arch(base_config=base_config)
     arch.groups = arch.groups + _cross_stream_groups()
     arch.buffer_keywords = list(_SR_BUFFER_KEYWORDS)
     return arch
@@ -444,7 +438,12 @@ def granite_dense_sr_arch(base_config=None) -> ArchDescriptor:
 _ARCH_REGISTRY = {
     "granite": granite_dense_arch,
     "granitemoe": granite_moe_arch,
-    "granitemoehybrid": granite_moe_hybrid_arch,
+    # granitemoeshared is the de-hybridized base family. The granitemoehybrid
+    # key is retained because real Granite 4.x dense checkpoints are still typed
+    # granitemoehybrid upstream (they carry no mamba layers); both resolve to the
+    # same shared-expert descriptor.
+    "granitemoeshared": granite_moe_shared_arch,
+    "granitemoehybrid": granite_moe_shared_arch,
 }
 
 # Must stay key-for-key in step with _ARCH_REGISTRY: a model_type registered in
@@ -453,7 +452,8 @@ def granite_dense_sr_arch(base_config=None) -> ArchDescriptor:
 _SR_ARCH_REGISTRY = {
     "granite": granite_dense_sr_arch,
     "granitemoe": granite_moe_sr_arch,
-    "granitemoehybrid": granite_moe_hybrid_sr_arch,
+    "granitemoeshared": granite_moe_shared_sr_arch,
+    "granitemoehybrid": granite_moe_shared_sr_arch,
 }
 
 

src/granite_switch/composer/compose_utils.py

@@ -168,7 +168,7 @@ def from_base_and_adapters(
                 lora_rank = built_in_lora_rank
                 adapter_ranks = [built_in_lora_rank] * num_built_in
                 adapter_alphas = {}
-                # Auto-detect lora_target_modules from layer_types
+                # Auto-detect lora_target_modules (attention-only switch model)
                 lora_target_modules = None
                 source_analysis = {}
             else:
@@ -188,34 +188,28 @@ def from_base_and_adapters(
         for field_name, default in arch.optional_config_fields.items():
             config_kwargs[field_name] = getattr(base_config, field_name, default)
 
-        # For Granite 3.x whose arch descriptor doesn't include
-        # shared_intermediate_size, default it to intermediate_size.
-        # GraniteMoeHybridConfig defaults it to 1024 (not None), so
-        # GraniteSwitchConfig's fallback logic doesn't trigger.
+        # For a dense Granite base whose arch descriptor doesn't include
+        # shared_intermediate_size, supply it explicitly from intermediate_size:
+        # a dense Granite layer always has a shared MLP of that width. This makes
+        # the composer the source of truth and does not rely on any parent-class
+        # default (GraniteMoeShared defaults it to 0, which is the "no shared MLP"
+        # sentinel — GraniteSwitchConfig would keep that 0 verbatim if passed).
         if "shared_intermediate_size" not in config_kwargs:
             config_kwargs["shared_intermediate_size"] = config_kwargs[
                 "intermediate_size"
             ]
 
-        # Normalize layer_types: map everything to "attention" (only attention
-        # layers are supported).
-        lt = config_kwargs.get("layer_types")
-        if lt is not None:
-            config_kwargs["layer_types"] = ["attention" for _ in lt]
-
         # When adapters are present, reserve the switch's cache slots at the
         # front: MultiSwitch (coded) owns SWITCH_CACHE_LAYERS == 2 (counting +
         # memory heads). The model subtracts the same count in
         # modeling_granite_switch.py to recover the physical decoder layers.
+        # The switch is attention-only, so no ``layer_types`` is carried:
+        # ``DynamicCache`` derives an all-``full_attention`` layout from the
+        # (inflated) ``num_hidden_layers``.
         if num_total > 0:
             config_kwargs["num_hidden_layers"] = (
                 config_kwargs["num_hidden_layers"] + SWITCH_CACHE_LAYERS
             )
-            if config_kwargs.get("layer_types") is not None:
-                config_kwargs["layer_types"] = [
-                    *(["attention"] * SWITCH_CACHE_LAYERS),
-                    *list(config_kwargs["layer_types"]),
-                ]
 
         # Switch-specific parameters
         config_kwargs.update(