fix(muse_glimmer)!: keep the attention gate out of the qkv fusion

docs/models.md

@@ -84,6 +84,8 @@ unset on PCIe GPUs: wider grids add no bandwidth there.
   its experts as dense weights and no longer loads; reconvert it with `ft checkpoint`.
 - FTW files converted by builds before the quantization refactor may fail to load;
   see [ftw-hotfix.md](ftw-hotfix.md) for the affected checkpoints and the repair tool.
+- Muse-Glimmer keeps the attention gate as its own linear next to `qkv_proj`; an FTW converted
+  when the gate was fused into `qkvg_proj` no longer loads, reconvert it with `ft checkpoint`.
 - An FTW converted before its family served images holds no vision encoder: `ft serve`
   refuses it unless started with `--text-model-only` (or `--mm-disable vision`); reconvert it
   with `ft checkpoint`, or add the encoder in place with [scripts/ftw_hotfix.py](ftw-hotfix.md).

python/freetoken/models/muse_glimmer/attention.py

@@ -18,16 +18,16 @@
 class MuseGlimmerAttention(BaseOP):
     """Gated GQA for one sliding-window or full-attention (NoPE) layer.
 
-        q, k, v, gate = split(qkvg_proj(x))
+        q, k, v = split(qkv_proj(x))
         q = qk_norm(q); k = qk_norm(k)          # weightless per-head RMSNorm
         q, k = rope(q, k)                        # sliding layers only; full layers are NoPE
         attn = paged_attention(q, k, v)          # sm_scale carries qk_scale_factor/sqrt(d)
-        out = o_proj(attn * sigmoid(gate))
+        out = o_proj(attn * sigmoid(gate_proj(x)))
 
     The reference multiplies q by ``qk_scale_factor`` before rope; rope is a rotation, so
     the factor commutes and is folded into ``AttentionSpec.sm_scale`` instead (parse_config
-    sets ``attn_sm_scale = qk_scale_factor / sqrt(head_dim)``). The gate is computed from
-    the same normed layer input as q/k/v, so it rides the fused projection.
+    sets ``attn_sm_scale = qk_scale_factor / sqrt(head_dim)``). The gate stays its own
+    linear: the NVFP4 release quantizes it with a different global scale than q/k/v.
     """
 
     def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
@@ -41,10 +41,14 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = ""):
         self.qo_attn_dim = self.num_q * head_dim
         self.kv_attn_dim = self.num_kv * head_dim
 
-        self._qkvg_split = [self.qo_attn_dim, self.kv_attn_dim, self.kv_attn_dim, self.qo_attn_dim]
-        self.qkvg_proj = LinearColParallelMerged(
-            config.hidden_size, self._qkvg_split, has_bias=False,
-            quant_config=config.quant, prefix=f"{prefix}.qkvg_proj",
+        self._qkv_split = [self.qo_attn_dim, self.kv_attn_dim, self.kv_attn_dim]
+        self.qkv_proj = LinearColParallelMerged(
+            config.hidden_size, self._qkv_split, has_bias=False,
+            quant_config=config.quant, prefix=f"{prefix}.qkv_proj",
+        )
+        self.gate_proj = LinearColParallelMerged(
+            config.hidden_size, [self.qo_attn_dim], has_bias=False,
+            quant_config=config.quant, prefix=f"{prefix}.gate_proj",
         )
         self.o_proj = LinearReplicated(
             self.qo_attn_dim, config.hidden_size, has_bias=False,
@@ -76,12 +80,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         ctx = get_global_ctx()
         T = x.shape[0]
 
-        qkvg = self.qkvg_proj.forward(x)
-        q, k, v, gate = torch.split(qkvg, self._qkvg_split, dim=-1)
-        del qkvg
+        gate = self.gate_proj.forward(x)
+        qkv = self.qkv_proj.forward(x)
+        q, k, v = torch.split(qkv, self._qkv_split, dim=-1)
+        del qkv
         q = self.qk_norm.forward(q.contiguous().view(T, self.num_q, self.head_dim))
         k = self.qk_norm.forward(k.contiguous().view(T, self.num_kv, self.head_dim))
-        v = v.contiguous()  # the split view has the qkvg row stride; the KV store needs contiguous
+        v = v.contiguous()  # the split view has the qkv row stride; the KV store needs contiguous
 
         q = q.reshape(T, self.qo_attn_dim)
         k = k.reshape(T, self.kv_attn_dim)

python/freetoken/models/muse_glimmer/gguf.py

@@ -171,8 +171,9 @@ def t(name: str) -> GgufTensorInfo:
         else:
             qkv = [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
             q, k, v = (matrix(p) for p in qkv)
+        yield from fused(f"{sa}.qkv_proj", qkv, [_unpermute(q, heads), _unpermute(k, kv_heads), v], hidden)
         gate = t(f"{blk}.attn_gate")
-        yield from fused(f"{sa}.qkvg_proj", [*qkv, gate], [_unpermute(q, heads), _unpermute(k, kv_heads), v, matrix(gate)], hidden)
+        yield from fused(f"{sa}.gate_proj", [gate], [matrix(gate)], hidden)
         out = t(f"{blk}.attn_output")
         yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
         mlp = f"{pre}.mlp"

python/freetoken/models/muse_glimmer/weight.py

@@ -26,13 +26,10 @@
 _VISION_QKV_PARTS = (".attn.q_proj", ".attn.k_proj", ".attn.v_proj")
 
 # Fused projections, concatenated on the output dim in this exact order to match the
-# model's merged-linear splits. The attention gate rides the q/k/v fusion (it is computed
-# from the same layer input); ``.self_attn.gate_proj`` and the SwiGLU ``.mlp.gate_proj``
-# are disambiguated by the full suffix.
+# model's merged-linear splits. The full suffix keeps the standalone attention gate
+# ``.self_attn.gate_proj`` out of the SwiGLU ``.mlp.gate_proj`` fusion.
 _FUSIONS: dict[str, tuple[str, ...]] = {
-    ".self_attn.qkvg_proj": (
-        ".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.gate_proj",
-    ),
+    ".self_attn.qkv_proj": (".self_attn.q_proj", ".self_attn.k_proj", ".self_attn.v_proj"),
     ".mlp.gate_up_proj": (".mlp.gate_proj", ".mlp.up_proj"),
 }
 
@@ -134,8 +131,8 @@ def _iter_weights_compressed_tensors(
 
     Every text Linear (q/k/v/o, the attention gate, the MLP) is kept native NVFP4
     (W4A16): ``.weight`` (uint8 packed) + ``.weight_scale`` (fp8 block) + ``.weight_global``
-    (fp16 per-row, the reciprocal of the stored quant-side global). q/k/v/gate fuse into
-    ``qkvg_proj`` and gate/up into ``gate_up_proj`` on the output dim, each part keeping
+    (fp16 per-row, the reciprocal of the stored quant-side global). q/k/v fuse into
+    ``qkv_proj`` and gate/up into ``gate_up_proj`` on the output dim, each part keeping
     its own scales, so the fused FP4 weights are exact. Embeddings, norms, lm_head and the
     vision tower are bf16 (the checkpoint's ignore list). Scale lookups go through the
     shard-map reader: they can land in a different shard than their weight_packed."""

python/freetoken/models/register.py

@@ -80,10 +80,6 @@ class ModelSpec:
 _GEMMA4_UNIFIED_PROCESSOR = "freetoken.mm.processors.gemma4:Gemma4UnifiedMMProcessor"
 _GEMMA4_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),)
 _MUSE_GLIMMER_PROCESSOR = "freetoken.mm.processors.muse_glimmer:MuseGlimmerMMProcessor"
-_MUSE_GLIMMER_PACKED = (
-    ("qkvg_proj", ("q_proj", "k_proj", "v_proj", "gate_proj")),
-    ("gate_up_proj", ("gate_proj", "up_proj")),
-)
 _MUSE_GLIMMER_ENCODERS = (EncoderSpec("vision", "vision_config", ("image",)),)
 _MINIMAX_M3_PACKED = _DENSE_PACKED + (
     ("index_qk_proj", ("index_q_proj", "index_k_proj")),
@@ -316,7 +312,7 @@ class ModelSpec:
         "freetoken.models.muse_glimmer",
         "MuseGlimmerForConditionalGeneration",
         checkpoint_roots=_LANGUAGE_MODEL_ROOT,
-        packed_modules_mapping=_MUSE_GLIMMER_PACKED,
+        packed_modules_mapping=_DENSE_PACKED,
         mm_processor=_MUSE_GLIMMER_PROCESSOR,
         encoders=_MUSE_GLIMMER_ENCODERS,
     ),
@@ -326,7 +322,7 @@ class ModelSpec:
         "MuseGlimmerForCausalLM",
         parse_config="parse_gguf_config",
         iter_weights="iter_gguf_weights",
-        packed_modules_mapping=_MUSE_GLIMMER_PACKED,
+        packed_modules_mapping=_DENSE_PACKED,
     ),
     "MistralForCausalLM": ModelSpec(
         "freetoken.models.mistral",

tests/models/test_muse_glimmer.py

@@ -252,13 +252,12 @@ def test_weight_rename_and_fusion():
     ((key, weight),) = _vision_tensors("vision_tower.layers.0.attn.proj.weight", torch.zeros(2, 2), vbuf)
     assert key == "vision_tower.layers.0.attn.proj.weight" and torch.equal(weight, torch.zeros(2, 2))
 
-    # q/k/v + the attention gate fuse into qkvg_proj in declaration order.
+    # q/k/v fuse into qkv_proj in declaration order; the attention gate stays out of every fusion.
     buf: dict = {}
     parts = {
         "q_proj": torch.full((4, 2), 0.0),
         "k_proj": torch.full((2, 2), 1.0),
         "v_proj": torch.full((2, 2), 2.0),
-        "gate_proj": torch.full((4, 2), 3.0),
     }
     fused = None
     for name, tensor in parts.items():
@@ -268,10 +267,10 @@ def test_weight_rename_and_fusion():
             fused = out[0]
     assert fused is not None and not buf
     key, tensor = fused
-    assert key == "model.layers.0.self_attn.qkvg_proj.weight"
-    assert tensor.shape == (12, 2)
+    assert key == "model.layers.0.self_attn.qkv_proj.weight"
+    assert tensor.shape == (8, 2)
     assert tensor[0, 0] == 0.0 and tensor[4, 0] == 1.0 and tensor[6, 0] == 2.0
-    assert tensor[8, 0] == 3.0
+    assert ct_bf16_fuse("model.layers.0.self_attn.gate_proj", torch.zeros(4, 2), buf, _FUSIONS) is None
 
     # The MLP's own gate_proj is a different fusion (gate|up), not the attention one.
     buf2: dict = {}
@@ -365,7 +364,7 @@ def _bf16_checkpoint_tensors(hf) -> dict:
 @pytest.mark.parametrize("include_vision", [False, True])
 def test_iter_weights_bf16_matches_model_state_dict(tmp_path, monkeypatch, include_vision):
     """The BF16 loader must produce exactly the model's state-dict keys with the
-    right shapes (rename + qkvg / gate_up fusion, norms raw); the image path only
+    right shapes (rename + qkv / gate_up fusion, norms raw); the image path only
     when an encoder is built, then with q/k/v fused into attn.qkv."""
     import torch
 
@@ -399,14 +398,15 @@ def test_iter_weights_bf16_matches_model_state_dict(tmp_path, monkeypatch, inclu
         assert torch.equal(loaded["vision_tower.layers.1.attn.qkv.weight"], torch.cat([tensors[v + f"{n}_proj.weight"] for n in "qkv"]))
         assert torch.equal(loaded["vision_tower.layers.1.attn.qkv.bias"], torch.cat([tensors[v + f"{n}_proj.bias"] for n in "qkv"]))
         assert torch.equal(loaded["vision_tower.projection.weight"], tensors["model.vision_projection.weight"])
-    # fusion order [q, k, v, gate] against the raw parts
-    fused = loaded["model.layers.0.self_attn.qkvg_proj.weight"]
+    # fusion order [q, k, v] against the raw parts; the attention gate loads on its own
+    fused = loaded["model.layers.0.self_attn.qkv_proj.weight"]
     p = "model.language_model.layers.0.self_attn."
     q_dim = tensors[p + "q_proj.weight"].shape[0]
     kv_dim = tensors[p + "k_proj.weight"].shape[0]
     assert torch.equal(fused[:q_dim], tensors[p + "q_proj.weight"])
     assert torch.equal(fused[q_dim : q_dim + kv_dim], tensors[p + "k_proj.weight"])
-    assert torch.equal(fused[-q_dim:], tensors[p + "gate_proj.weight"])
+    assert torch.equal(fused[-kv_dim:], tensors[p + "v_proj.weight"])
+    assert torch.equal(loaded["model.layers.0.self_attn.gate_proj.weight"], tensors[p + "gate_proj.weight"])
 
 
 def test_iter_weights_nvfp4_cross_shard_scales(tmp_path, monkeypatch):
@@ -481,18 +481,19 @@ def nvfp4(base, out_f, in_f, global_scale, shard_for_scales=None):
     loaded = dict(
         iter_weights(str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True)
     )
-    qkvg = "model.layers.0.self_attn.qkvg_proj"
-    assert loaded[qkvg + ".weight"].shape == (2 * q + 2 * kv, H // 2)
-    assert loaded[qkvg + ".weight_scale"].shape == (2 * q + 2 * kv, H // 16)
-    assert loaded[qkvg + ".weight_global"].shape == (2 * q + 2 * kv,)
+    qkv = "model.layers.0.self_attn.qkv_proj"
+    assert loaded[qkv + ".weight"].shape == (q + 2 * kv, H // 2)
+    assert loaded[qkv + ".weight_scale"].shape == (q + 2 * kv, H // 16)
+    assert loaded[qkv + ".weight_global"].shape == (q + 2 * kv,)
     # global is the reciprocal of the stored quant-side scale
-    assert loaded[qkvg + ".weight_global"][0].item() == pytest.approx(0.5)
+    assert loaded[qkv + ".weight_global"][0].item() == pytest.approx(0.5)
+    assert loaded["model.layers.0.self_attn.gate_proj.weight"].shape == (q, H // 2)
     dp = "model.layers.0.mlp.down_proj"
     assert loaded[dp + ".weight"].shape == (H, I // 2)  # cross-shard scales resolved
     assert loaded[dp + ".weight_global"][0].item() == pytest.approx(0.25)
     # the quant-side activation global lands as the dequant-side input_scale the layer declares
     assert loaded[dp + ".input_scale"].item() == pytest.approx(1.0)
-    assert loaded[qkvg + ".input_scale"].shape == ()
+    assert loaded[qkv + ".input_scale"].shape == ()
     assert not any(k.endswith(".input_global_scale") for k in loaded)
     expected_vision = {"vision_tower." + k for k in MuseGlimmerVisionModel(parse_vision_config(hf)).state_dict()}
     assert {k for k in loaded if k.startswith("vision_tower.")} == expected_vision
@@ -503,7 +504,7 @@ def nvfp4(base, out_f, in_f, global_scale, shard_for_scales=None):
     text_only = dict(iter_weights(
         str(tmp_path), torch.device("cpu"), include_moe_experts=False, include_non_moe=True, include_vision=False,
     ))
-    assert not any(k.startswith("vision_tower.") for k in text_only) and qkvg + ".weight" in text_only
+    assert not any(k.startswith("vision_tower.") for k in text_only) and qkv + ".weight" in text_only
 
 
 def test_raw_config_shim_serves_unknown_model_type(tmp_path):
@@ -551,7 +552,8 @@ def test_model_state_dict_matches_loader_keys():
     keys = set(model.state_dict().keys())
     layer0 = {k for k in keys if k.startswith("model.layers.0.")}
     assert layer0 == {
-        "model.layers.0.self_attn.qkvg_proj.weight",
+        "model.layers.0.self_attn.qkv_proj.weight",
+        "model.layers.0.self_attn.gate_proj.weight",
         "model.layers.0.self_attn.o_proj.weight",
         "model.layers.0.mlp.gate_up_proj.weight",
         "model.layers.0.mlp.down_proj.weight",
@@ -588,7 +590,7 @@ def test_model_state_dict_matches_loader_keys():
     qmodel = MuseGlimmerForCausalLM(qcfg)
     attn = qmodel.model.layers.op_list[0].self_attn
     mlp = qmodel.model.layers.op_list[0].mlp
-    for proj in (attn.qkvg_proj, attn.o_proj, mlp.gate_up_proj, mlp.down_proj):
+    for proj in (attn.qkv_proj, attn.gate_proj, attn.o_proj, mlp.gate_up_proj, mlp.down_proj):
         assert type(proj.quant_method) is Nvfp4LinearMethod
     assert type(qmodel.lm_head.quant_method) is UnquantizedLinearMethod  # lm_head stays bf16
     del model, qmodel, torch

tests/models/test_quant_config.py

@@ -257,7 +257,8 @@ def _gemma_tied_head(model, idx):
         "model.layers.{lin}.mlp.down_proj": FP8B, "lm_head": BF16,
     }, strategy="auto") for name in ("Qwen/Qwen3.8-27B-FP8", "Qwen/Qwen3.6-27B-FP8")],
     Case("RedHatAI/Muse-Glimmer-30B-NVFP4", CompressedTensorsConfig, {
-        "model.layers.0.self_attn.qkvg_proj": NVFP4, "model.layers.0.self_attn.o_proj": NVFP4,
+        "model.layers.0.self_attn.qkv_proj": NVFP4, "model.layers.0.self_attn.gate_proj": NVFP4,
+        "model.layers.0.self_attn.o_proj": NVFP4,
         "model.layers.0.mlp.gate_up_proj": NVFP4, "model.layers.0.mlp.down_proj": NVFP4, "lm_head": BF16,
     }, strategy="auto"),
     Case("nvidia/Gemma-4-31B-IT-NVFP4", ModelOptConfig, {