[Spec Decode] Support Muse Glimmer DFlash checkpoints

tests/runner/test_kv_cache_manager.py

@@ -848,6 +848,29 @@ def test_get_kv_cache_spec_with_eagle3_mla(self):
             assert isinstance(spec, MLAAttentionSpec)
             assert spec.num_kv_heads == 1
 
+    def test_get_kv_cache_spec_with_dflash_uses_explicit_head_dim(self):
+        self.runner.vllm_config.compilation_config.static_forward_context = {}
+        mock_speculative_config = MagicMock()
+        mock_speculative_config.method = "dflash"
+        mock_speculative_config.use_gemma4_mtp = MagicMock(return_value=False)
+        mock_hf_config = MagicMock()
+        mock_hf_config.num_key_value_heads = 8
+        mock_hf_config.hidden_size = 6656
+        mock_hf_config.num_attention_heads = 32
+        mock_hf_config.num_hidden_layers = 5
+        mock_hf_config.head_dim = 128
+        mock_speculative_config.draft_model_config.hf_config = mock_hf_config
+        self.runner.speculative_config = mock_speculative_config
+
+        kv_cache_spec = self.runner.get_kv_cache_spec()
+
+        for layer_idx in range(5):
+            draft_spec = kv_cache_spec[f"draft_layer.{layer_idx}"]
+            assert isinstance(draft_spec, FullAttentionSpec)
+            assert draft_spec.num_kv_heads == common_utils.get_padded_num_heads(
+                8, self.runner.mesh.shape["model"])
+            assert draft_spec.head_size == 128
+
     @patch('tpu_inference.models.common.kv_share.compute_mtp_kv_share_map')
     def test_get_kv_cache_spec_with_gemma4_mtp(self, mock_compute_map):
         # tests we create kv cache spec for gemma4 mtp draft model (KV-sharing)

tests/spec_decode/test_dflash.py

@@ -13,14 +13,19 @@
 # limitations under the License.
 """Unit tests for the JAX DFlash speculative decoding proposer."""
 
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import jax
 import jax.numpy as jnp
 import numpy as np
 import pytest
 from flax import nnx
+from transformers import PretrainedConfig
 
+from tpu_inference.models.common.model_loader import _get_model_architecture
+from tpu_inference.models.jax.dflash import (DFlashForCausalLM,
+                                             get_dflash_sliding_window)
 from tpu_inference.spec_decode.jax.dflash import DFlashProposer
 
 
@@ -111,6 +116,22 @@ def __init__(self):
 
 
 # ----- Existing Minimal Tests -----
+def test_dflash_sliding_window_requires_explicit_enablement():
+    config = SimpleNamespace(sliding_window=2048, use_sliding_window=True)
+    assert get_dflash_sliding_window(config) == 2048
+    config.use_sliding_window = False
+    assert get_dflash_sliding_window(config) is None
+
+
+def test_registry_supports_muse_glimmer_assistant_architectures():
+    for architecture in (
+            "MuseGlimmerAssistantModel",
+            "DFlashMuseGlimmerAssistantModel",
+    ):
+        config = PretrainedConfig(architectures=[architecture])
+        assert _get_model_architecture(config) is DFlashForCausalLM
+
+
 def test_propose_uses_target_model_logits():
     proposer = object.__new__(DFlashProposer)
     proposer.mesh = _make_single_device_mesh()

tpu_inference/models/common/model_loader.py

@@ -106,6 +106,8 @@ def _get_model_architecture(config: PretrainedConfig) -> nnx.Module:
     _MODEL_REGISTRY["Gemma4MTPModel"] = Gemma4MTPForCausalLM
     _MODEL_REGISTRY["DFlashForCausalLM"] = DFlashForCausalLM
     _MODEL_REGISTRY["DFlashDraftModel"] = DFlashForCausalLM
+    _MODEL_REGISTRY["MuseGlimmerAssistantModel"] = DFlashForCausalLM
+    _MODEL_REGISTRY["DFlashMuseGlimmerAssistantModel"] = DFlashForCausalLM
 
     architectures = getattr(config, "architectures", [])
     for arch in architectures:

tpu_inference/models/jax/dflash.py

@@ -43,6 +43,22 @@
 _FA_VMEM_LIMIT = 128 * 1024 * 1024
 
 
+def get_dflash_config_value(config: Any,
+                            name: str,
+                            default: Any = None) -> Any:
+    """Read DFlash fields from either nested or top-level draft configs."""
+    dflash_config = getattr(config, "dflash_config", None) or {}
+    if name in dflash_config:
+        return dflash_config[name]
+    return getattr(config, name, default)
+
+
+def get_dflash_sliding_window(config: Any) -> Optional[int]:
+    if not getattr(config, "use_sliding_window", False):
+        return None
+    return getattr(config, "sliding_window", None)
+
+
 class DFlashAttention(nnx.Module):
     """DFlash cross+self attention with on-device KV cache.
 
@@ -68,6 +84,7 @@ def __init__(
         self.rope_theta = getattr(config, "rope_theta", 1000000.0)
         self.rope_scaling = getattr(config, "rope_scaling", None)
         self.rms_norm_eps = getattr(config, "rms_norm_eps", 1e-6)
+        self.sliding_window = get_dflash_sliding_window(config)
 
         self.head_dim_original = getattr(config, "head_dim",
                                          self.hidden_size // self.num_heads)
@@ -180,6 +197,7 @@ def __call__(
             mesh=self.mesh,
             head_dim_original=self.head_dim_original,
             sm_scale=self.head_dim_original**-0.5,
+            attention_chunk_size=self.sliding_window,
             use_causal_mask=True,
             update_kv_cache=True,
         )
@@ -198,6 +216,7 @@ def __call__(
             mesh=self.mesh,
             head_dim_original=self.head_dim_original,
             sm_scale=self.head_dim_original**-0.5,
+            attention_chunk_size=self.sliding_window,
             use_causal_mask=False,  # Noise tokens attend to all KV tokens
             update_kv_cache=True,
         )
@@ -351,8 +370,8 @@ def __init__(
             ) for _ in range(hf_config.num_hidden_layers)
         ])
 
-        dflash_config = getattr(hf_config, "dflash_config", {})
-        target_layer_ids = dflash_config.get("target_layer_ids", None)
+        target_layer_ids = get_dflash_config_value(hf_config,
+                                                   "target_layer_ids")
         num_target_layers = getattr(hf_config, "num_target_layers", None)
         if target_layer_ids is not None:
             num_context_features = len(target_layer_ids)
@@ -450,8 +469,9 @@ def __init__(
         hf_config = spec_config.draft_model_config.hf_config
         self.hf_config = hf_config
         self.block_size = getattr(hf_config, "block_size", 8)
-        dflash_config = getattr(hf_config, "dflash_config", {})
-        self.mask_token_id = dflash_config.get("mask_token_id", 0)
+        dflash_config = getattr(hf_config, "dflash_config", {}) or {}
+        self.mask_token_id = get_dflash_config_value(hf_config,
+                                                     "mask_token_id", 0)
 
         self._position_scheme = dflash_config.get("position_scheme",
                                                   "incremental")
@@ -555,10 +575,14 @@ def load_weights(self, rng_key: jax.Array):
             "model.fc": "model.fc.kernel",
             "fc.weight": "model.fc.kernel",
             "model.fc.weight": "model.fc.kernel",
+            "encoder.fc": "model.fc.kernel",
+            "encoder.fc.weight": "model.fc.kernel",
             "hidden_norm": "model.hidden_norm.scale",
             "model.hidden_norm": "model.hidden_norm.scale",
             "hidden_norm.weight": "model.hidden_norm.scale",
             "model.hidden_norm.weight": "model.hidden_norm.scale",
+            "encoder.output_norm_enc": "model.hidden_norm.scale",
+            "encoder.output_norm_enc.weight": "model.hidden_norm.scale",
             "norm": "model.norm.scale",
             "model.norm": "model.norm.scale",
             "norm.weight": "model.norm.scale",

tpu_inference/runner/kv_cache_manager.py

@@ -575,8 +575,10 @@ def get_kv_cache_spec(self):
                         draft_num_layers = 1
                     num_kv_heads = common_utils.get_padded_num_heads(
                         draft_hf_config.num_key_value_heads, model_cnt)
+                    configured_head_size = vars(draft_hf_config).get(
+                        "head_dim")
                     head_size = common_utils.get_padded_head_dim(
-                        draft_hf_config.hidden_size //
+                        configured_head_size or draft_hf_config.hidden_size //
                         draft_hf_config.num_attention_heads)
                     for i in range(draft_num_layers):
                         if self.use_mla:

tpu_inference/spec_decode/jax/dflash.py

@@ -49,7 +49,8 @@
 ]
 
 TARGET_LM_HEAD_PATHS = [
-    "model.lm_head", "lm_head", "lm_head.weight", "lm_head.kernel",
+    "lm_head.weight", "lm_head.kernel", "model.lm_head.weight",
+    "model.lm_head.kernel", "model.lm_head", "lm_head",
     "model.embed_tokens.weight", "model.embed.embedding",
     "model.embed_tokens.embedding", "embed.embedding", "embed_tokens.embedding"
 ]
@@ -87,8 +88,9 @@ def __init__(
         hf_config = self.draft_model_config.hf_config
         self.block_size = getattr(hf_config, "block_size",
                                   self.num_speculative_tokens + 1)
-        dflash_config = getattr(hf_config, "dflash_config", {})
-        self.mask_token_id = dflash_config.get("mask_token_id", 0)
+        dflash_config = getattr(hf_config, "dflash_config", {}) or {}
+        self.mask_token_id = dflash_config.get(
+            "mask_token_id", getattr(hf_config, "mask_token_id", 0))
         self.hidden_size = hf_config.hidden_size
         self.num_layers = hf_config.num_hidden_layers
 
@@ -149,7 +151,15 @@ def get_target_value(target, paths):
 
                 # 3. Check JAX nnx.State
                 param = _find_param(target, paths)
-                return param.value if param is not None else None
+                if param is None:
+                    return None
+                if hasattr(param, "value"):
+                    return param.value
+                leaves = jax.tree_util.tree_leaves(param)
+                if len(leaves) == 1:
+                    leaf = leaves[0]
+                    return leaf.value if hasattr(leaf, "value") else leaf
+                return None
 
             # Resolve draft and target embeddings
             draft_embed_param = _find_param(self.state, DRAFT_EMBED_PATHS)