Add multi-component export support for CLM and KEV

docs/api/model_package.md

@@ -114,6 +114,35 @@ pkg.save("output/llama/", external_data="safetensors")
 pkg.save("output/llama-serial/", max_workers=1)
 ```
 
+### Generic backbone-and-head packages
+
+Use `MultiComponentModelTask` when one non-generative package contains exactly
+one backbone or encoder and one or more named head graphs:
+
+```python
+from mobius import (
+    ComponentConfig,
+    ComponentRole,
+    ComponentSpec,
+    MultiComponentModelTask,
+)
+
+
+class EncoderWithHeadsTask(MultiComponentModelTask):
+    components = ComponentSpec(
+        encoder=ComponentConfig("encoder", ComponentRole.ENCODER),
+        classifier=ComponentConfig("heads.classifier", ComponentRole.HEAD),
+    )
+
+    def build_component(self, name, component, module, config):
+        ...
+```
+
+The component names become `ModelPackage` keys. Roles are also exposed through
+`model_roles`, so existing inspection and optimization code treats head graphs
+as heads rather than decoder models. Dotted module paths allow components such
+as `heads.classifier` to be resolved from the root module.
+
 ## Output Layout
 
 - **Single model**: `directory/model.onnx` + `directory/model.onnx.data`

docs/model-catalog.md

@@ -1,5 +1,92 @@
 # Model Catalog
 
+## Non-generative decision models
+
+Mobius exports `Contrastive-LM/CLM-v0.1-8B` with one `ModelPackage`
+containing its headless Qwen3 encoder, state/action projection heads, and
+scaled-cosine scorer (`CLMModel` + `CLMTask`). The CLM release did **not** pin a
+Qwen3-8B base revision: generic construction records the base as `unpinned`;
+`build_clm_package(..., base_revision=<sha>)` requires the caller's selected
+revision and records it as reproducible provenance.
+
+The v0.1 heads first L2-normalize each selected Qwen embedding, then use depth
+3: input `Linear + GELU`, one
+`Linear + LayerNorm + GELU` hidden block (optionally residual), output
+`Linear`, then L2 normalization. The headless encoder emits the complete
+`token_hidden_states` tensor, contains no vocabulary projection, and has no
+generation-cache inputs or outputs. CLM uses
+last-token pooling, but neither the upstream padding side nor its EOS behavior
+is pinned. Callers must select the last attended position using their
+attention mask (the deterministic `clm_last_token_indices` helper is provided)
+before passing `[items, 4096]` embeddings to either projection head. Mobius
+does not assume that sequence position `-1` is a real token.
+
+`jaredpalmer/kev-4b` similarly exports its headless Qwen3.5 backbone and grouped
+pointer head (`KevModel` + `KevTask`). Its published Qwen3.5 base revision is
+pinned. The backbone package exposes only input IDs, attention/position inputs,
+and token hidden states; KV, convolution, and recurrent generation caches are
+not part of this scoring export. `build_kev_package` accepts a dense state dict
+after the published PEFT
+adapter has been merged. Mobius does not currently merge arbitrary PEFT LoRA
+artifacts into Qwen3.5: passing separate base/adapter mappings fails with an
+actionable instruction to use PEFT `merge_and_unload()` first, rather than
+silently exporting the unadapted base.
+
+The deterministic request adapters in `mobius.models.decision` implement each
+model's verified structured rendering, candidate mapping, stable grouped
+softmax, confidence, and ranking contracts. They require a caller-provided
+tokenizer and local tensors; export and preprocessing do not use an HTTP or
+session API. Checkpoint mapping helpers accept already-deserialized mappings,
+so applications retain control over pickle deserialization. Production build
+helpers validate the complete published checkpoint metadata, exact head key
+sets, tensor shapes, and scalar calibration before constructing any graph.
+Flexible `synthetic_*_checkpoint` helpers are test fixtures and intentionally
+do not pass these production contracts.
+
+The Kev adapter enforces the published 8,192-token causal-row limit, 1–255
+option bounds, and a non-empty question set. Non-strict mode truncates state
+user tokens to 8,191 before adding the state control token; strict mode rejects
+that overflow. Every question remains an independent state-plus-question row;
+`batch_kev_rows` produces padded backbone inputs and matching local
+decide/option indices for the pointer graph. Score confidence is
+`max(0, 1 - Σ pᵢ|i-mode|/(L-1))`, matching the public adapter.
+
+### Decision-model parity tests
+
+`tests/kev_parity_test.py` and `tests/clm_parity_test.py` contain fast synthetic
+ONNX-versus-PyTorch tests that run in the normal CPU test lane:
+
+```bash
+pytest tests/kev_parity_test.py tests/clm_parity_test.py -m "not integration"
+```
+
+The full tests are opt-in because they load the exported FP32 package and its
+reference model on CUDA. Keep user-site packages disabled when the selected
+Conda environment contains the intended ONNX Runtime GPU build.
+
+```bash
+# Pinned KEV base, adapter, preprocessing, hidden-state readouts, pointer head,
+# probabilities, and typed answers.
+PYTHONNOUSERSITE=1 \
+MOBIUS_DECISION_EXPORT_ROOT=/path/to/kev-4b-fp32 \
+MOBIUS_KEV_ALLOW_DOWNLOAD=1 \
+pytest tests/kev_parity_test.py -m integration
+
+# CLM with an explicitly selected Qwen3 revision.
+PYTHONNOUSERSITE=1 \
+CLM_RUN_FULL_PARITY=1 \
+CLM_BASE_REVISION=<immutable-qwen3-revision> \
+MOBIUS_DECISION_EXPORT_ROOT=/path/to/clm-v0.1-8b-fp32 \
+pytest tests/clm_parity_test.py -m integration -k exported_clm
+```
+
+CLM's official serving path remains a separate qualification because upstream
+did not pin Qwen3, vLLM, or tokenizer/pooling versions. Start and record an
+external `vllm serve ... --runner pooling` instance, then set
+`CLM_RUN_VLLM_PARITY=1`, `CLM_EMBED_URL`, `CLM_EMBED_MODEL`, and
+`CLM_VLLM_METADATA` to compare its normalized embeddings and probabilities
+against the same exported package.
+
 This user-facing catalog groups registered HuggingFace model types by task and lists
 their module classes and example model IDs. Diffusers pipeline components are described
 separately because they are not entries in the model registry. Neither list implies GGUF

src/mobius/__init__.py

@@ -21,6 +21,12 @@
     "BaseModelConfig",
     "CausalLMConfig",
     "CausalLMTask",
+    "CLMModel",
+    "CLMProjectionHead",
+    "CLMTask",
+    "ComponentConfig",
+    "ComponentRole",
+    "ComponentSpec",
     "ComponentInfo",
     "SharedWeightEndpoint",
     "SharedWeightInfo",
@@ -42,6 +48,10 @@
     "ModelRegistration",
     "ModelRegistry",
     "ModelTask",
+    "MultiComponentModelTask",
+    "KevModel",
+    "KevPointerHead",
+    "KevTask",
     "MLPWorldModel",
     "MMSConfig",
     "OPSET_VERSION",
@@ -58,11 +68,13 @@
     "adapter_source_from_onnx_adapter",
     "attach_peft_adapter",
     "build",
+    "build_clm_package",
     "build_context",
     "build_diffusers_pipeline",
     "build_from_gguf",
     "build_from_module",
     "build_from_nemo",
+    "build_kev_package",
     "compose_adapter_deltas",
     "components",
     "ep_capabilities",
@@ -154,5 +166,23 @@
 from mobius.integrations.gguf import build_from_gguf
 from mobius.integrations.nemo import build_from_nemo
 from mobius.integrations.transformers import build
-from mobius.models import MLPWorldModel
-from mobius.tasks import CausalLMTask, ModelTask, WorldModelTask
+from mobius.models import (
+    CLMModel,
+    CLMProjectionHead,
+    KevModel,
+    KevPointerHead,
+    MLPWorldModel,
+    build_clm_package,
+    build_kev_package,
+)
+from mobius.tasks import (
+    CausalLMTask,
+    CLMTask,
+    ComponentConfig,
+    ComponentRole,
+    ComponentSpec,
+    KevTask,
+    ModelTask,
+    MultiComponentModelTask,
+    WorldModelTask,
+)

src/mobius/_model_package.py

@@ -1415,7 +1415,7 @@ def apply_weights(
         prefix_map: dict[str, str] | None = None,
         *,
         fold_constants: bool = True,
-    ) -> None:
+    ) -> set[str]:
         """Apply weights from a state dict across component models.
 
         For single-component packages, all weights are applied to the sole
@@ -1432,6 +1432,9 @@ def apply_weights(
                 Weights whose name starts with a prefix are applied to the
                 named component (with the prefix stripped). Unmatched weights
                 are applied to all components.
+
+        Returns:
+            The original state-dict names that matched package initializers.
         """
         applied: set[str] = set()
 
@@ -1448,15 +1451,15 @@ def apply_weights(
             routed: dict[str, dict[str, torch.Tensor]] = {name: {} for name in self.data}
             unmatched: dict[str, torch.Tensor] = {}
             # Track original HF names for weights that get stripped
-            stripped_to_original: dict[str, str] = {}
+            stripped_to_original: dict[str, dict[str, str]] = {name: {} for name in self.data}
 
             for weight_name, tensor in state_dict.items():
                 matched = False
                 for prefix, component in prefix_map.items():
                     if weight_name.startswith(prefix):
                         stripped = weight_name[len(prefix) :].lstrip(".")
                         routed[component][stripped] = tensor
-                        stripped_to_original[stripped] = weight_name
+                        stripped_to_original[component][stripped] = weight_name
                         matched = True
                         break
                 if not matched:
@@ -1467,7 +1470,7 @@ def apply_weights(
                     self.data[component_name], component_weights
                 )
                 for s in applied_stripped:
-                    applied.add(stripped_to_original.get(s, s))
+                    applied.add(stripped_to_original[component_name].get(s, s))
 
             # Try unmatched weights against all models
             if unmatched:
@@ -1477,13 +1480,14 @@ def apply_weights(
         _log_weight_mapping(state_dict, applied)
 
         if not fold_constants:
-            return
+            return applied
 
         # Fold constants now that weights have been loaded.
         # PackQKV emits Concat(w_q, w_k, w_v) in the graph; those nodes can only
         # be constant-folded once the weight tensors carry their const_value.
         for model in self.data.values():
             fold_initializers_after_weights(model)
+        return applied
 
 
 @contextmanager

src/mobius/_model_package_test.py

@@ -1477,6 +1477,26 @@ def test_multi_component_with_prefix_map(self):
         )
         assert model1.graph.initializers[init_name].const_value is not None
 
+    def test_prefix_routes_report_colliding_stripped_names(self):
+        config = make_config()
+        pkg1 = build_from_module(CausalLMModel(config), config)
+        pkg2 = build_from_module(CausalLMModel(config), config)
+        model1 = pkg1["model"]
+        model2 = pkg2["model"]
+        pkg = ModelPackage({"text": model1, "vision": model2})
+
+        init_name = next(iter(model1.graph.initializers.keys()))
+        shape = list(model1.graph.initializers[init_name].shape)
+        applied = pkg.apply_weights(
+            {
+                f"text.{init_name}": torch.ones(shape),
+                f"vision.{init_name}": torch.zeros(shape),
+            },
+            prefix_map={"text.": "text", "vision.": "vision"},
+        )
+
+        assert applied == {f"text.{init_name}", f"vision.{init_name}"}
+
 
 class TestBuildPackageFromModule:
     def test_returns_model_package(self):

src/mobius/components/_gated_deltanet.py

@@ -164,8 +164,8 @@ def forward(
         self,
         op: OpBuilder,
         hidden_states: ir.Value,
-        conv_state: ir.Value,
-        recurrent_state: ir.Value,
+        conv_state: ir.Value | None,
+        recurrent_state: ir.Value | None,
     ):
         """Forward pass for the Gated DeltaNet layer.
 
@@ -184,6 +184,38 @@ def forward(
             new_recurrent_state: (batch, num_v_heads, k_dim, v_dim)
         """
         batch_dim = op.Shape(hidden_states, start=0, end=1)
+        if conv_state is None:
+            conv_state = op.Expand(
+                op.CastLike(op.Constant(value_float=0.0), hidden_states),
+                op.Concat(
+                    batch_dim,
+                    op.Constant(
+                        value_ints=[
+                            self.conv_dim,
+                            self.conv_kernel_size - 1,
+                        ]
+                    ),
+                    axis=0,
+                ),
+            )
+        if recurrent_state is None:
+            recurrent_state = op.Expand(
+                op.Cast(
+                    op.Constant(value_float=0.0),
+                    to=self._stash_type,
+                ),
+                op.Concat(
+                    batch_dim,
+                    op.Constant(
+                        value_ints=[
+                            self.num_v_heads,
+                            self.head_k_dim,
+                            self.head_v_dim,
+                        ]
+                    ),
+                    axis=0,
+                ),
+            )
 
         # === Projections ===
         mixed_qkv = self.in_proj_qkv(op, hidden_states)

src/mobius/components/_gated_deltanet_test.py

@@ -87,6 +87,19 @@ def test_forward_builds_graph(self):
         assert count_op_type(graph, "Scan") == 0
         assert count_op_type(graph, "Conv") == 0
 
+    def test_forward_initializes_missing_states(self):
+        config = self._make_deltanet_config()
+        dn = GatedDeltaNet(config)
+        builder, op, graph = create_test_builder()
+        hidden = create_test_input(builder, "hidden", [1, 4, 64])
+
+        output, new_conv, new_rec = dn(op, hidden, None, None)
+        builder._adapt_outputs([output, new_conv, new_rec], "")
+
+        assert count_op_type(graph, "Expand") >= 2
+        assert count_op_type(graph, "CausalConvWithState") == 1
+        assert count_op_type(graph, "LinearAttention") == 1
+
     def test_forward_has_sigmoid_for_beta(self):
         config = self._make_deltanet_config()
         dn = GatedDeltaNet(config)