vLLM DFlash spec decode: 6 fixes needed to get the muse-glimmer image working (with working Dockerfile)

#40
by j4ys0n - opened

The recipe at recipes.vllm.ai gives this for speculative decoding:

--speculative-config '{"method": "dflash", "model": "meta-models/Muse-Glimmer-30B-assistant", "num_speculative_tokens": 15}'

On the current vllm/vllm-openai:muse-glimmer image (0.1.dev19075+gd89ec6d6a, pushed 2026-08-10) this fails at startup, and fixing the first error just reveals the next one. I worked through the whole chain by pulling the image layers and reading the source. Six distinct failures, now all fixed and confirmed serving tokens. Target model works fine without the spec config; everything below only affects DFlash.

Root causes, grouped:

  1. Arch rename past the registry. The assistant checkpoint declares architectures: ["MuseGlimmerAssistantModel"], which IS registered — but EAGLEConfig(method="dflash") prepends DFlash before the registry lookup, producing an unregistered name. First crash.

  2. The muse_glimmer_assistant → Qwen3Config alias (transformers_utils/config.py:106) corrupts two fields, because the muse JSON is Qwen3-shaped but not Qwen3-behaved: (a) no vocab_size in the JSON → Qwen3's 151936 default → broken draft logits sizing against the 202048 target vocab (this is also where the pad_token_id must be within (0, 151935) warnings come from); (b) no use_sliding_window in the JSON → Qwen3Config nulls sliding_window: 2048 → "DFlash sliding attention requires a window size" crash.

  3. Renamed drafter tensors. The checkpoint ships encoder.fc.weight / encoder.output_norm_enc.weight; the loader expects fc / hidden_norm (z-lab DFlash naming) and has no mapping. The registry comment claims "same safetensors as DFlashDraftModel" — not quite; the Onyx→Muse rename touched these two tensors too.

  4. Wrapper-shape assumptions. Muse's get_language_model() returns the decoder itself, but two sites dereference .model on the result expecting a ForCausalLM-style wrapper: interfaces.py (set_aux_hidden_state_layers assert) and spec_decode/dflash/utils.py:57 (embedding tie). Same latent bug exists in the eagle/dspark/gemma4 tie paths.

  5. Config gotcha, not a bug: the recipe's TP=1 command omits --max-num-seqs, and the default (1024) times 14 reserved draft slots per sequence exceeds the 8192 chunked-prefill budget → max_num_scheduled_tokens is set to -6144. Add --max-num-seqs 64 (and optionally --max-num-batched-tokens 16384).

Working fix — derived image, each patch assert-guarded:

FROM vllm/vllm-openai:muse-glimmer
RUN python3 - <<'PY'
from pathlib import Path
base = Path("/usr/local/lib/python3.12/dist-packages/vllm")

def patch(rel, old, new):
    p = base / rel
    src = p.read_text()
    n = src.count(old)
    assert n == 1, f"{rel}: expected 1 occurrence, found {n}"
    p.write_text(src.replace(old, new))
    print(f"patched {rel}")

patch("transformers_utils/configs/eagle.py",
    'arch.startswith("DFlash") or arch.endswith("DFlash")',
    'arch.startswith("DFlash") or arch.endswith("DFlash") or arch == "MuseGlimmerAssistantModel"')

patch("model_executor/models/qwen3_dflash.py",
    'orig_to_new_substr={"midlayer.": "layers.0."},',
    'orig_to_new_substr={"midlayer.": "layers.0.", "encoder.fc": "fc", "encoder.output_norm_enc": "hidden_norm"},')

patch("model_executor/models/qwen3_dflash.py",
    'self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)',
    'self.config.draft_vocab_size = vllm_config.model_config.get_vocab_size()')

patch("model_executor/models/interfaces.py",
    '''        assert hasattr(parent_ref, "model"), (
            "Model instance must have 'model' attribute to set number of layers"
        )''',
    '''        if isinstance(parent_ref, EagleModelMixin):
            parent_ref._set_aux_hidden_state_layers(layers)
            return
        assert hasattr(parent_ref, "model"), (
            "Model instance must have 'model' attribute to set number of layers"
        )''')

patch("model_executor/models/qwen3_dflash.py",
    'self.quant_config = get_draft_quant_config(vllm_config)',
    '''self.quant_config = get_draft_quant_config(vllm_config)
        if getattr(self.config, "sliding_window", None) is None:
            self.config.sliding_window = getattr(
                vllm_config.model_config.hf_text_config, "sliding_window", None
            )''')

patch("v1/worker/gpu/spec_decode/dflash/utils.py",
    'target_inner = target_language_model.model',
    'target_inner = getattr(target_language_model, "model", target_language_model)')
PY

Confirmed working: RTX PRO 6000 Blackwell, BF16 target, TP=1, FA2, full recipe flags plus --max-num-seqs 64 --max-num-batched-tokens 16384. Patches 3 and 5 restore data-faithful values (target vocab / target window) rather than hardcoding, so they shouldn't break other DFlash checkpoints, but I've only tested this pairing.

--model /var/lib/vllm/cache/huggingface/meta-models/Muse-Glimmer-30B --host 10.1.1.80 --port 40006 --served-model-name muse-glimmer-30b --max-model-len=131072 --gpu-memory-utilization=0.92 --enable-auto-tool-choice --tool-call-parser=muse_glimmer --reasoning-parser=muse_glimmer --generation-config=auto --speculative-config={"method": "dflash", "model": "meta-models/Muse-Glimmer-30B-assistant", "num_speculative_tokens": 15} --max-num-seqs=64 --max-num-batched-tokens=16384

Sign up or log in to comment