feat(speculative): support Gemma-4-E4B as a streaming DFlash/DSpark target

examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+"""Convert a ModelOpt DSpark drafter export into the layout vLLM's
+Gemma4DSparkForCausalLM expects.
+
+Three classes of mismatch, all on the ModelOpt side (vLLM needs no patch):
+
+1. config.json
+   - architectures: DFlashDraftModel -> Gemma4DSparkModel   (registry key)
+   - model_type:    qwen3            -> gemma4_text          (Gemma4DSparkAttention
+     reads layer_types/head_dim/global_head_dim off a Gemma4 text config)
+   - target_layer_ids / markov_rank are read as TOP-LEVEL attrs by
+     Gemma4DSparkModel.__init__, but ModelOpt nests them under dflash_config.
+
+2. weight names
+   - markov_w1/markov_w2 -> markov_head.markov_w1/markov_head.markov_w2
+     (DSparkMarkovHead registers them under a `markov_head` submodule)
+
+3. missing tensors  <-- the one that silently destroys AL
+   - Gemma4DSparkForCausalLM builds its OWN ParallelLMHead and VocabParallelEmbedding
+     and its load_weights() only fills names it finds; anything absent stays
+     RANDOMLY INITIALIZED with no error. The ModelOpt export ships neither
+     lm_head nor embed_tokens, so both must be baked in from the base model.
+     Gemma 4 is tie_word_embeddings=true, so both come from the same tensor:
+     model.language_model.embed_tokens.weight
+"""
+
+import argparse
+import json
+import os
+import shutil
+
+from safetensors.torch import load_file, save_file
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--drafter", required=True, help="ModelOpt export dir")
+ap.add_argument("--base", required=True, help="base Gemma-4-E4B-it dir (for lm_head/embed)")
+ap.add_argument("--out", required=True)
+args = ap.parse_args()
+
+os.makedirs(args.out, exist_ok=True)
+cfg = json.load(open(os.path.join(args.drafter, "config.json")))
+df = cfg.get("dflash_config", {}) or {}
+
+new = dict(cfg)
+new["architectures"] = ["Gemma4DSparkModel"]
+new["model_type"] = "gemma4_text"
+# promote what Gemma4DSparkModel.__init__ reads off the top level
+new["target_layer_ids"] = df.get("target_layer_ids", cfg.get("target_layer_ids"))
+new["markov_rank"] = df.get("markov_rank", 256)
+new["mask_token_id"] = df.get("mask_token_id")
+new["block_size"] = cfg.get("block_size")
+new.setdefault("draft_vocab_size", cfg["vocab_size"])
+# Gemma4 attention knobs consulted by gemma4_layer_config()/Gemma4DSparkAttention
+# Do NOT default these: Gemma4 sizes attention per layer, and quietly falling back to
+# the sliding-layer values rebuilds the draft with the wrong q/k/o and q/k-norm shapes.
+# They must come from the training config (see hf_spec_export.py).
+for _req in ("global_head_dim", "attention_k_eq_v"):
+    if _req not in cfg:
+        raise SystemExit(
+            f"drafter config.json is missing {_req!r}. Re-export with a ModelOpt that "
+            "propagates the Gemma4 per-layer attention fields, or add it by hand."
+        )
+if cfg.get("attention_k_eq_v") and "num_global_key_value_heads" not in cfg:
+    raise SystemExit(
+        "attention_k_eq_v is set but num_global_key_value_heads is missing; "
+        "vLLM would size k_proj with the sliding-layer KV head count."
+    )
+new.setdefault("sliding_window", 512)
+new.setdefault("final_logit_softcapping", None)
+json.dump(new, open(os.path.join(args.out, "config.json"), "w"), indent=2)
+print(
+    "config: architectures={} model_type={} target_layer_ids={} markov_rank={}".format(
+        new["architectures"], new["model_type"], new["target_layer_ids"], new["markov_rank"]
+    )
+)
+
+sd = load_file(os.path.join(args.drafter, "model.safetensors"))
+print(f"loaded {len(sd)} drafter tensors")
+
+out = {}
+renamed = 0
+for k, v in sd.items():
+    nk = k
+    if k.startswith(("markov_w1", "markov_w2")):
+        nk = "markov_head." + k
+        renamed += 1
+    out[nk] = v
+print(f"renamed {renamed} markov tensors -> markov_head.*")
+
+# --- bake in lm_head + embed_tokens from the base (tied on Gemma 4) ---
+idx_path = os.path.join(args.base, "model.safetensors.index.json")
+if os.path.exists(idx_path):
+    wm = json.load(open(idx_path))["weight_map"]
+    key = next(k for k in wm if k.endswith("language_model.embed_tokens.weight"))
+    base_sd = load_file(os.path.join(args.base, wm[key]))
+else:
+    base_sd = load_file(os.path.join(args.base, "model.safetensors"))
+    key = next(k for k in base_sd if k.endswith("language_model.embed_tokens.weight"))
+emb = base_sd[key]
+print(f"base embed_tokens {tuple(emb.shape)} from {key!r}")
+assert emb.shape[0] == cfg["vocab_size"] and emb.shape[1] == cfg["hidden_size"], (
+    f"base embed {tuple(emb.shape)} does not match draft vocab/hidden "
+    f"({cfg['vocab_size']},{cfg['hidden_size']})"
+)
+
+out["lm_head.weight"] = emb.clone()
+out["embed_tokens.weight"] = emb.clone()
+print("baked lm_head.weight + embed_tokens.weight (tie_word_embeddings=true on Gemma 4)")
+
+save_file(out, os.path.join(args.out, "model.safetensors"), metadata={"format": "pt"})
+for f in ("tokenizer.json", "tokenizer_config.json"):
+    src = os.path.join(args.base, f)
+    if os.path.exists(src):
+        shutil.copy(src, os.path.join(args.out, f))
+print(f"wrote {len(out)} tensors -> {args.out}")
+print()
+print("Serve with (note the draft attention backend):")
+_spec = (
+    f'{{"model": "{args.out}", "num_speculative_tokens": 3, '
+    '"method": "dspark", "attention_backend": "FLASHINFER"}'
+)
+print(f"  vllm serve <base> --speculative-config '{_spec}'")
+print(
+    "  FLASHINFER is REQUIRED for TWO independent reasons; either alone is fatal:\n"
+    "    1. head dim  -- the draft re-runs backend auto-selection and lands on FLASH_ATTN,\n"
+    "       whose FA2 kernel caps head dimension at 256, but Gemma4 full-attention layers\n"
+    "       use global_head_dim=512.\n"
+    "    2. causality -- the draft is NON-CAUSAL (bidirectional) on every layer, as DSpark\n"
+    "       heads are: _dflash_layer_causal() only marks sliding_attention layers causal,\n"
+    "       and this draft is all full_attention with no dflash_config.causal override.\n"
+    "       load_dspark_model therefore sets use_non_causal=True, which the backend must\n"
+    "       support. Runtime confirms with: 'Using FlashInfer for draft model non-causal\n"
+    "       attention'.\n"
+    "  The VLLM_ATTENTION_BACKEND env var does NOT reach the draft -- it is read from\n"
+    "  speculative_config.attention_backend."
+)

examples/speculative_decoding/export/gemma4_dspark_swa_vllm.patch

@@ -0,0 +1,108 @@
+Make Gemma4 DSpark drafts honour dflash_config.use_swa / swa_window_size.
+
+Applies to (vLLM): vllm/model_executor/models/gemma4_dspark.py
+
+WHY
+    A DFlash/DSpark draft declares its sliding window through
+    dflash_config.use_swa + swa_window_size, and keeps layer_types all
+    "full_attention" on purpose -- that is what keeps attention_k_eq_v uniform,
+    which Gemma4DSparkModel._build_fused_kv_buffers() asserts.
+
+    vLLM's Qwen3 DFlash layer resolves its window through
+    _resolve_layer_attention(), which reads exactly those fields.
+    Gemma4DSparkAttention does not: it inherits Gemma4MTPAttention.__init__,
+    which derives the window from the BASE model's layer pattern --
+
+        self.is_sliding = layer_type == "sliding_attention"
+        sliding_window  = config.sliding_window if self.is_sliding else None
+
+    so for an all-full_attention draft, per_layer_sliding_window is None on
+    every layer. A draft TRAINED with a window is then SERVED with full
+    attention. Nothing errors; acceptance length just quietly drops.
+
+WHAT
+    Resolve the window via _resolve_layer_attention(config, layer_idx) -- the
+    same call the Qwen3 DFlash layer makes -- and rebuild self.attn with
+    per_layer_sliding_window when a window applies. self.attn is already
+    replaced further down for the k_proj/v_proj rework, so rebuilding it here
+    follows the class's existing pattern; the stale static_forward_context entry
+    is popped first, exactly as Gemma4DSparkDecoderLayer does for
+    "{prefix}.self_attn.attn".
+
+    Layers with no window are untouched, so the existing full-attention path
+    (deepseek-ai/dspark_gemma4_12b_block7 and the recipe it matches) is
+    bit-identical.
+
+VERIFIED (by config simulation, against the config ModelOpt's exporter emits
+for modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml)
+    - all 5 layers get sliding_window=512
+    - all 5 layers stay causal=False  (dflash_config.causal is pinned False)
+    - all 5 layers keep use_k_eq_v=True -> the fused-KV assert still passes
+    - layer_types stays uniform -> the V2 model runner is NOT required
+    NOT yet validated end-to-end on GPU; no acceptance-length number exists.
+
+APPLY
+    cd <vllm> && git apply /path/to/gemma4_dspark_swa_vllm.patch
+
+--- a/vllm/model_executor/models/gemma4_dspark.py
++++ b/vllm/model_executor/models/gemma4_dspark.py
+@@ -8,8 +8,9 @@
+ import torch.nn as nn
+ import torch.nn.functional as F
+
+ from vllm import _custom_ops as ops
++from vllm.attention import Attention
+ from vllm.compilation.decorators import support_torch_compile
+ from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+ from vllm.model_executor.layers.layernorm import RMSNorm
+ from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear
+@@ -22,9 +23,13 @@
+ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+ from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config
+
+ from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer
+-from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal
++from .qwen3_dflash import (
++    DFlashQwen3Model,
++    _dflash_layer_causal,
++    _resolve_layer_attention,
++)
+ from .qwen3_dspark import DSparkMarkovHead, Qwen3DSparkForCausalLM
+ from .utils import extract_layer_index, maybe_prefix
+
+
+@@ -61,8 +66,34 @@
+         )
+         self.is_kv_shared_layer = False
+         self.causal = _dflash_layer_causal(config, layer_idx)
+         self.kv_size = self.num_kv_heads * self.head_dim
++
++        # Gemma4MTPAttention derives the sliding window from the BASE model's
++        # layer pattern (layer_type == "sliding_attention"). That is wrong for a
++        # DFlash/DSpark draft: the draft declares its window via
++        # dflash_config.use_swa / swa_window_size, and deliberately keeps
++        # layer_types all "full_attention" so that attention_k_eq_v stays uniform
++        # for the fused context-KV precompute. Without this, a draft TRAINED with
++        # a sliding window is SERVED with full attention -- no error is raised,
++        # acceptance length just quietly drops. Resolve the window the same way
++        # the Qwen3 DFlash layer does, and rebuild self.attn when it applies.
++        sliding_window, _ = _resolve_layer_attention(config, layer_idx)
++        if sliding_window is not None:
++            get_current_vllm_config().compilation_config.static_forward_context.pop(
++                f"{prefix}.attn", None
++            )
++            self.attn = Attention(
++                self.num_heads,
++                self.head_dim,
++                self.scaling,
++                num_kv_heads=self.num_kv_heads,
++                cache_config=cache_config,
++                quant_config=quant_config,
++                logits_soft_cap=getattr(config, "attn_logit_softcapping", None),
++                per_layer_sliding_window=sliding_window,
++                prefix=f"{prefix}.attn",
++            )
+         attn_bias = getattr(config, "attention_bias", False)
+         self.k_proj = ColumnParallelLinear(
+             config.hidden_size,
+             self.total_num_kv_heads * self.head_dim,

modelopt/torch/export/plugins/hf_spec_export.py

@@ -412,6 +412,16 @@ def _export_config(self):
         else:
             config["layer_types"] = ["full_attention"] * draft_config.num_hidden_layers
 
+        # Gemma4 sizes attention PER LAYER, so the serving side cannot reconstruct the
+        # draft's shapes from head_dim / num_key_value_heads alone: full-attention layers
+        # use ``global_head_dim`` and, under ``attention_k_eq_v``, ``num_global_key_value_heads``
+        # (see gemma4_layer_config in vLLM). Omitting these makes vLLM rebuild the draft with
+        # the sliding-layer dims and fail with a shape mismatch on q/k/o and the q/k norms.
+        for _attr in ("global_head_dim", "num_global_key_value_heads", "attention_k_eq_v"):
+            _val = getattr(draft_config, _attr, None)
+            if _val is not None:
+                config[_attr] = _val
+
         # Sliding-window attention: all draft layers use non-causal SWA (MiMo-style). vLLM's
         # _resolve_layer_attention reads dflash_config.use_swa + swa_window_size; with
         # layer_types left all "full_attention" it applies a non-causal sliding window to