fix(glm): preserve dense KV-B weights

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 39ff6824-d35f-4d3b-8f5f-043a7119a100

Author: justinchuby

src/mobius/integrations/gguf/_route_census_test.py

@@ -64,8 +64,8 @@ def test_every_route_has_one_actionable_classification() -> None:
     }
     assert Counter(item.category for item in items) == {
         "dependency-or-runtime-abi-blocked": 100,
-        "evidence-only": 107,
-        "immediately-implementable": 53,
+        "evidence-only": 108,
+        "immediately-implementable": 52,
         "intentionally-rejected": 19,
         "artifact-unavailable": 5,
     }
@@ -94,6 +94,7 @@ def test_route_reasons_are_sourced_from_authoritative_records() -> None:
 def test_known_route_boundaries_are_not_collapsed() -> None:
     by_id = {item.route_id: item for item in iter_remaining_route_work()}
     assert by_id["architecture:bitnet"].category == "evidence-only"
+    assert by_id["architecture:glm-dsa"].category == "evidence-only"
     assert by_id["architecture:deepseek4"].category == "dependency-or-runtime-abi-blocked"
     assert by_id["architecture:rwkv6"].category == "dependency-or-runtime-abi-blocked"
     assert by_id["architecture:bailingmoe2"].category == "intentionally-rejected"

src/mobius/models/glm_moe_dsa.py

@@ -673,7 +673,7 @@ def preprocess_weights(
         state_dict = dict(state_dict)
         kv_b_suffix = ".self_attn.kv_b_proj.weight"
         for key in tuple(state_dict):
-            if not key.endswith(kv_b_suffix):
+            if not self.config.use_dsa or not key.endswith(kv_b_suffix):
                 continue
             tensor = state_dict.pop(key)
             expected_rows = self.config.num_attention_heads * (

src/mobius/models/glm_moe_dsa_test.py

@@ -374,17 +374,10 @@ def test_moe_layers_fuse_to_qmoe_when_quantized(self):
 class TestPreprocessWeights:
     def _state_dict(self, config: ArchitectureConfig) -> dict[str, torch.Tensor]:
         model = GlmMoeDsaCausalLMModel(config)
-        state = {
+        return {
             name: torch.zeros(tuple(int(d) for d in p.shape))
             for name, p in model.named_parameters()
         }
-        # Convert Mobius's own attribute-path names into the HF-style
-        # ``model.layers.N.*`` names preprocess_weights expects on input.
-        renamed = {
-            f"model.{name}" if not name.startswith("lm_head") else name: v
-            for name, v in state.items()
-        }
-        return renamed
 
     def test_drops_mtp_layer_weights_with_warning(self, caplog):
         config = _glm_config(num_hidden_layers=4)
@@ -420,6 +413,43 @@ def test_dsa_path_keeps_indexer_weights(self):
         out = model.preprocess_weights(state)
         assert any(".self_attn.indexer." in k for k in out)
 
+    @pytest.mark.parametrize("use_dsa", [True, False], ids=["dsa", "dense"])
+    def test_preprocessed_weights_cover_every_graph_parameter(self, use_dsa: bool):
+        config = _glm_config(
+            num_hidden_layers=1,
+            indexer_types=["full"],
+            use_dsa=use_dsa,
+        )
+        model = GlmMoeDsaCausalLMModel(config)
+        graph = build_from_module(model, config, task="glm-moe-dsa")["model"].graph
+        state = {
+            name: torch.zeros(tuple(int(d) for d in initializer.shape))
+            for name, initializer in graph.initializers.items()
+            if initializer.const_value is None
+        }
+        expected = set(state)
+
+        if use_dsa:
+            prefix = "model.layers.0.self_attn."
+            key_rows = state.pop(f"{prefix}k_b_proj.weight").reshape(
+                config.num_attention_heads,
+                config.qk_nope_head_dim,
+                config.kv_lora_rank,
+            )
+            value_rows = state.pop(f"{prefix}v_b_proj.weight").reshape(
+                config.num_attention_heads,
+                config.v_head_dim,
+                config.kv_lora_rank,
+            )
+            state[f"{prefix}kv_b_proj.weight"] = torch.cat(
+                (key_rows, value_rows), dim=1
+            ).reshape(-1, config.kv_lora_rank)
+
+        out = model.preprocess_weights(state)
+
+        assert expected
+        assert set(out) == expected
+
     def test_splits_hf_fused_kv_b_projection(self):
         config = _glm_config(num_hidden_layers=1, indexer_types=["full"])
         model = GlmMoeDsaCausalLMModel(config)