diff --git a/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py b/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py
new file mode 100644
index 00000000000..968f35c2c3c
--- /dev/null
+++ b/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 --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."
+)
diff --git a/examples/speculative_decoding/export/gemma4_dspark_swa_vllm.patch b/examples/speculative_decoding/export/gemma4_dspark_swa_vllm.patch
new file mode 100644
index 00000000000..4f71eb7b5ca
--- /dev/null
+++ b/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 && 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,
diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py
index 255b1d9ab04..7d5f743bd7c 100644
--- a/modelopt/torch/export/plugins/hf_spec_export.py
+++ b/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
diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py
index 6463cb4109d..8be96f2e11c 100644
--- a/modelopt/torch/speculative/plugins/modeling_dflash.py
+++ b/modelopt/torch/speculative/plugins/modeling_dflash.py
@@ -41,6 +41,7 @@
The draft architecture is independent of the target model.
"""
+import copy
from dataclasses import dataclass
import torch
@@ -140,6 +141,18 @@ def __init__(self, config, layer_idx):
self.num_key_value_groups = self.num_heads // self.num_kv_heads
self.scaling = self.head_dim**-0.5
self.attention_dropout = getattr(config, "attention_dropout", 0.0)
+ # DFlash/DSpark drafts attend bidirectionally: a block of draft tokens is
+ # predicted in one shot, so those tokens must see each other. Serving must
+ # agree -- vLLM resolves per-layer causality in
+ # qwen3_dflash._dflash_layer_causal(): an explicit ``dflash_config.causal``
+ # overrides all layers, otherwise a layer is causal only when
+ # ``layer_types[i] == "sliding_attention"``. The exporter emits no
+ # ``causal`` field for a plain full-attention draft, so it stays non-causal
+ # on both sides. With ``dflash_swa_window_size`` set, the exporter instead
+ # emits ``use_swa: True`` + an explicit ``causal: False`` and leaves
+ # ``layer_types`` all-full, which keeps vLLM non-causal too. Only mark
+ # layers ``sliding_attention`` if you also intend them to be CAUSAL at
+ # serving time, and train them that way -- see _build_draft_attention_mask.
self.is_causal = False
attn_bias = getattr(config, "attention_bias", False)
@@ -221,6 +234,117 @@ def forward(self, hidden_states, target_hidden, position_embeddings, attention_m
return self.o_proj(attn_output)
+class DFlashGemma4Attention(DFlashAttention):
+ """DFlash attention for a Gemma4-style draft.
+
+ Two deltas versus the Qwen3-style :class:`DFlashAttention`:
+
+ * ``attention_k_eq_v``: Gemma4 can derive V from the K projection instead of
+ carrying a separate ``v_proj``, halving the KV parameters. vLLM's
+ ``Gemma4DSparkAttention`` does exactly this (``v_src = k`` when
+ ``use_k_eq_v``), and its fused context-KV precompute *asserts* every draft
+ layer is built this way, so a draft trained with a separate ``v_proj``
+ cannot be served by that path at all.
+ * ``v_norm``: applied to V, with **no learnable weight**, mirroring vLLM's
+ ``RMSNorm(..., has_weight=False)``. Plain (Qwen3) DFlash does not norm V.
+
+ ``use_k_eq_v`` follows vLLM: full-attention layers only, and only when the
+ config opts in. Sliding layers keep their own ``v_proj``.
+ """
+
+ def __init__(self, config, layer_idx):
+ """Initialize Gemma4 draft attention with per-layer dims, dropping ``v_proj`` under k_eq_v."""
+ super().__init__(config, layer_idx)
+ layer_types = getattr(config, "layer_types", None)
+ is_full = layer_types is None or layer_types[layer_idx] == "full_attention"
+ self.use_k_eq_v = is_full and getattr(config, "attention_k_eq_v", False)
+
+ # Gemma4's attention dims are PER LAYER: full-attention layers use a larger
+ # ``global_head_dim`` and, under k_eq_v, a smaller ``num_global_key_value_heads``.
+ # This mirrors vLLM's ``gemma4_layer_config`` (transformers_utils/configs/gemma4.py),
+ # which Gemma4DSparkAttention calls to size q/k/o. Getting this wrong is silent:
+ # the DSpark weight loader only fills names it finds, so a mis-shaped k_proj is
+ # simply left randomly initialized.
+ if is_full:
+ self.head_dim = getattr(config, "global_head_dim", None) or self.head_dim
+ if self.use_k_eq_v:
+ self.num_kv_heads = (
+ getattr(config, "num_global_key_value_heads", None) or self.num_kv_heads
+ )
+ self.num_key_value_groups = self.num_heads // self.num_kv_heads
+ self.scaling = self.head_dim**-0.5
+ attn_bias = getattr(config, "attention_bias", False)
+ self.q_proj = nn.Linear(
+ config.hidden_size, self.num_heads * self.head_dim, bias=attn_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, self.num_kv_heads * self.head_dim, bias=attn_bias
+ )
+ self.o_proj = nn.Linear(
+ self.num_heads * self.head_dim, config.hidden_size, bias=attn_bias
+ )
+ self.q_norm = _NORM_CLS(self.head_dim, eps=config.rms_norm_eps)
+ self.k_norm = _NORM_CLS(self.head_dim, eps=config.rms_norm_eps)
+
+ if self.use_k_eq_v:
+ # Registered by the parent; drop it so it is neither trained nor exported.
+ del self.v_proj
+ self.v_proj = None
+ elif is_full:
+ self.v_proj = nn.Linear(
+ config.hidden_size,
+ self.num_kv_heads * self.head_dim,
+ bias=getattr(config, "attention_bias", False),
+ )
+ # vLLM builds this as ``RMSNorm(..., has_weight=False)`` and the reference
+ # checkpoint ships NO v_norm tensor, so keep the scale fixed at ones and
+ # non-persistent: it must not appear in the exported state_dict.
+ self.v_norm = _NORM_CLS(self.head_dim, eps=config.rms_norm_eps)
+ del self.v_norm.weight
+ self.v_norm.register_buffer("weight", torch.ones(self.head_dim), persistent=False)
+
+ def _project_v(self, target_hidden, hidden_states, k_ctx, k_noise):
+ """Return the V sequence, from K under k_eq_v or from ``v_proj`` otherwise."""
+ if self.use_k_eq_v:
+ return k_ctx, k_noise
+ return self.v_proj(target_hidden), self.v_proj(hidden_states)
+
+ def forward(self, hidden_states, target_hidden, position_embeddings, attention_mask=None):
+ """Forward with KV injection; V is normed and, under k_eq_v, shares K's projection."""
+ bsz, q_len, _ = hidden_states.shape
+ ctx_len = target_hidden.shape[1]
+
+ q = self.q_proj(hidden_states).view(bsz, q_len, -1, self.head_dim)
+ q = self.q_norm(q).transpose(1, 2)
+
+ k_ctx = self.k_proj(target_hidden)
+ k_noise = self.k_proj(hidden_states)
+ k = torch.cat([k_ctx, k_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim)
+ k = self.k_norm(k).transpose(1, 2)
+
+ v_ctx, v_noise = self._project_v(target_hidden, hidden_states, k_ctx, k_noise)
+ v = torch.cat([v_ctx, v_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim)
+ # vLLM norms V (no RoPE on V), unlike the Qwen3-style path.
+ v = self.v_norm(v).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+ attn_fn = self._get_attn_fn()
+ attn_output, _ = attn_fn(
+ self,
+ q,
+ k,
+ v,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ sliding_window=self.sliding_window,
+ )
+ attn_output = attn_output.reshape(bsz, q_len, -1)
+ return self.o_proj(attn_output)
+
+
class DFlashDecoderLayer(nn.Module):
"""Draft decoder layer with KV injection."""
@@ -248,6 +372,56 @@ def forward(self, hidden_states, target_hidden, position_embeddings, attention_m
return hidden_states
+class DFlashGemma4DecoderLayer(nn.Module):
+ """Draft decoder layer matching Gemma4's block, with KV injection.
+
+ Gemma4 wraps each sub-block in a *pair* of norms ("sandwich norm") and scales
+ the layer output by a learned ``layer_scalar``, where Qwen3 uses a single
+ pre-norm per sub-block. vLLM's ``Gemma4MTPDecoderLayer`` -- which
+ ``Gemma4DSparkDecoderLayer`` inherits -- looks up
+ ``pre_feedforward_layernorm`` / ``post_feedforward_layernorm`` /
+ ``layer_scalar`` by name, and its DSpark weight loader silently leaves any
+ parameter it cannot find randomly initialized. A Qwen3-shaped draft
+ therefore *loads without error* and produces garbage, so the shapes must
+ match exactly.
+
+ The residual/norm order below mirrors ``Gemma4MTPDecoderLayer.forward``:
+ norm -> attn -> norm -> +residual -> norm -> mlp -> norm -> +residual, then
+ scale by ``layer_scalar``.
+ """
+
+ def __init__(self, config, layer_idx):
+ """Initialize a Gemma4-style draft layer (sandwich norms + layer scalar)."""
+ super().__init__()
+ self.self_attn = DFlashGemma4Attention(config, layer_idx)
+ self.mlp = _MLP_CLS(config)
+ self.input_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps)
+ self.pre_feedforward_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_feedforward_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps)
+ # A buffer (not a parameter) to match vLLM's `register_buffer`, so the
+ # exported tensor name and shape line up with the reference checkpoint.
+ self.register_buffer("layer_scalar", torch.ones(1))
+
+ def forward(self, hidden_states, target_hidden, position_embeddings, attention_mask=None):
+ """Forward with sandwich norms, KV injection, and the layer scalar."""
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ hidden_states = self.self_attn(
+ hidden_states, target_hidden, position_embeddings, attention_mask
+ )
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = hidden_states + residual
+
+ residual = hidden_states
+ hidden_states = self.pre_feedforward_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = self.post_feedforward_layernorm(hidden_states)
+ hidden_states = hidden_states + residual
+
+ return hidden_states * self.layer_scalar
+
+
class DFlashModule(nn.Module):
"""DFlash draft module using Qwen3 components (MLP, RMSNorm, RotaryEmbedding)."""
@@ -263,11 +437,20 @@ def __init__(self, config):
self.hidden_norm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps)
# Decoder layers
+ # Gemma4 drafts need Gemma4's block shape (sandwich norms + layer_scalar,
+ # optional k_eq_v); everything else keeps the Qwen3-style block.
+ layer_cls = (
+ DFlashGemma4DecoderLayer
+ if str(getattr(config, "model_type", "")).startswith("gemma4")
+ else DFlashDecoderLayer
+ )
self.layers = nn.ModuleList(
- [DFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ [layer_cls(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
)
self.norm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps)
self._rotary_config = config # Used by _maybe_init_rotary_emb
+ self._gemma4_rope_kinds = self._build_gemma4_rope_kinds(config)
+ self._layer_types = list(getattr(config, "layer_types", []) or [])
# Explicit weight init is needed because DFlashModule is instantiated via
# mtsp.convert() AFTER the base model's post_init() has already run, so HF's
@@ -280,9 +463,46 @@ def _maybe_init_rotary_emb(self, device=None):
Same pattern as EAGLE3's _maybe_init_rope. Avoids creating rotary_emb
during __init__ (which runs on meta device during from_pretrained),
preventing the meta-tensor inv_freq issue on checkpoint resume.
+
+ Gemma4 needs one module PER attention kind, not one for the whole draft:
+ its full-attention layers use ``global_head_dim`` while sliding layers use
+ ``head_dim``, and the two kinds carry different ``rope_parameters`` (theta
+ 1e6 vs 1e4). vLLM builds RoPE per layer for exactly this reason; a single
+ shared module silently mismatches the head dim on one of the two kinds.
"""
if not hasattr(self, "rotary_emb"):
self.rotary_emb = _ROTARY_CLS(config=self._rotary_config, device=device)
+ if self._gemma4_rope_kinds and not hasattr(self, "rotary_emb_by_kind"):
+ self.rotary_emb_by_kind = nn.ModuleDict(
+ {
+ kind: _ROTARY_CLS(config=cfg, device=device)
+ for kind, cfg in self._gemma4_rope_kinds.items()
+ }
+ )
+
+ @staticmethod
+ def _build_gemma4_rope_kinds(config):
+ """Per-attention-kind rotary configs for a Gemma4 draft, or ``{}`` otherwise.
+
+ Returns a shallow copy of ``config`` per distinct ``layer_types`` entry with
+ ``head_dim`` and ``rope_parameters`` resolved for that kind.
+ """
+ if not str(getattr(config, "model_type", "")).startswith("gemma4"):
+ return {}
+ layer_types = getattr(config, "layer_types", None)
+ if not layer_types:
+ return {}
+ rope_params = getattr(config, "rope_parameters", None)
+ kinds = {}
+ for kind in dict.fromkeys(layer_types):
+ cfg = copy.copy(config)
+ if kind == "full_attention":
+ cfg.head_dim = getattr(config, "global_head_dim", None) or config.head_dim
+ if isinstance(rope_params, dict) and isinstance(rope_params.get(kind), dict):
+ cfg.rope_parameters = dict(rope_params[kind])
+ cfg.rope_theta = cfg.rope_parameters.get("rope_theta", config.rope_theta)
+ kinds[kind] = cfg
+ return kinds
def _init_weights(self, config):
"""Initialize weights matching HF PreTrainedModel._init_weights."""
@@ -299,8 +519,15 @@ def forward(self, noise_embedding, target_hidden, position_ids, attention_mask=N
target_hidden = self.hidden_norm(self.fc(target_hidden))
self._maybe_init_rotary_emb(device=hidden_states.device)
position_embeddings = self.rotary_emb(hidden_states, position_ids)
-
- for layer in self.layers:
- hidden_states = layer(hidden_states, target_hidden, position_embeddings, attention_mask)
+ per_kind = {
+ kind: emb(hidden_states, position_ids)
+ for kind, emb in getattr(self, "rotary_emb_by_kind", {}).items()
+ }
+
+ for layer_idx, layer in enumerate(self.layers):
+ layer_pos = position_embeddings
+ if per_kind and layer_idx < len(self._layer_types):
+ layer_pos = per_kind.get(self._layer_types[layer_idx], position_embeddings)
+ hidden_states = layer(hidden_states, target_hidden, layer_pos, attention_mask)
return self.norm(hidden_states)
diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py
index 2b5fe989c03..ac787c789f1 100644
--- a/modelopt/torch/speculative/plugins/modeling_fakebase.py
+++ b/modelopt/torch/speculative/plugins/modeling_fakebase.py
@@ -72,6 +72,40 @@
_SAFETENSORS_SINGLE_FILENAMES = ["model.safetensors", "consolidated.safetensors"]
+def _resolve_rope_theta(base_cfg, attn_kind: str = "sliding_attention") -> float | None:
+ """Return the base model's RoPE theta, handling nested ``rope_parameters``.
+
+ Most models expose a flat ``rope_theta``. Gemma 4 instead nests per-attention-kind RoPE
+ settings under ``rope_parameters``, e.g.::
+
+ {"full_attention": {"rope_theta": 1e6, "rope_type": "proportional",
+ "partial_rotary_factor": 0.25},
+ "sliding_attention": {"rope_theta": 1e4, "rope_type": "default"}}
+
+ A flat ``getattr(base_cfg, "rope_theta", None)`` returns ``None`` there, and the draft then
+ silently trains on the draft class's default theta instead of the base's — training loss and
+ accuracy still improve while MT-Bench AAL is capped, because RoPE frequencies get baked into
+ the trained weights.
+
+ ``attn_kind`` selects which entry to read; it must match the attention the DRAFT uses. The
+ default is ``sliding_attention`` because SWA drafts are the common case for Gemma 4, and its
+ ``rope_type`` is plain ``default`` (the ``full_attention`` entry uses ``proportional`` rope
+ with ``partial_rotary_factor``, which the draft classes do not implement).
+ """
+ theta = getattr(base_cfg, "rope_theta", None)
+ if theta is not None:
+ return theta
+ params = getattr(base_cfg, "rope_parameters", None)
+ if not isinstance(params, dict):
+ return None
+ entry = params.get(attn_kind)
+ if entry is None:
+ # Single-kind nested form, or an unknown kind name: fall back to the sole entry.
+ values = [v for v in params.values() if isinstance(v, dict) and "rope_theta" in v]
+ entry = values[0] if len(values) == 1 else None
+ return entry.get("rope_theta") if isinstance(entry, dict) else None
+
+
class FakeBaseConfig(PretrainedConfig):
"""Minimal config for FakeBaseModel that supports offline speculative decoding training."""
@@ -203,7 +237,7 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM
num_key_value_heads=getattr(base_cfg, "num_key_value_heads", None),
intermediate_size=getattr(base_cfg, "intermediate_size", None),
rms_norm_eps=getattr(base_cfg, "rms_norm_eps", 1e-6),
- rope_theta=getattr(base_cfg, "rope_theta", None),
+ rope_theta=_resolve_rope_theta(base_cfg),
final_norm_type=_select_final_norm_type(
getattr(base_cfg, "model_type", None), base_cfg
),
diff --git a/modelopt/torch/speculative/plugins/modeling_final_norm.py b/modelopt/torch/speculative/plugins/modeling_final_norm.py
index 718d591b662..9e4a7ad6dd0 100644
--- a/modelopt/torch/speculative/plugins/modeling_final_norm.py
+++ b/modelopt/torch/speculative/plugins/modeling_final_norm.py
@@ -93,6 +93,14 @@ def extra_repr(self):
# M3's final norm is always gemma-style; map it here too so a config that lost its
# use_gemma_norm flag still gets the correct flavor instead of silently dropping the +1.
"minimax_m3_vl_text": "gemma_rmsnorm",
+ # Gemma 4 VLM nests the LLM as text_config with model_type "gemma4_text"; from_source
+ # reads the NESTED config, so a "gemma4" key alone would never match. Verified numerically
+ # on gemma-4-E4B-it that Gemma4RMSNorm is plain ``normed * weight`` — NOT the ``(1 + weight)``
+ # form used by Gemma 2/3 — reproducing HF ``hidden_states[-1]`` at cos=0.999999 (vs 0.9719
+ # and maxabs_err 47.6 for the ``(1 + weight)`` form), so plain ``rmsnorm`` is correct here
+ # and ``gemma_rmsnorm`` would be wrong. Both keys listed so a text-only checkpoint works too.
+ "gemma4_text": "rmsnorm",
+ "gemma4": "rmsnorm",
# gpt_oss intentionally DISABLED: GptOssRMSNorm uses an fp32 weight + multiply-then-cast,
# unlike _FinalRMSNorm's bf16 weight, so reusing it would silently bias reconstructed logits.
# Re-enable once a gpt_oss-style class (fp32 weight, multiply-then-cast) is in _FINAL_NORM_CLASSES.
diff --git a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml
new file mode 100644
index 00000000000..6bbe32a9935
--- /dev/null
+++ b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml
@@ -0,0 +1,161 @@
+# DSpark (full-train, from scratch) recipe for Gemma-4-E4B-it.
+#
+# Streaming: the real Gemma-4-E4B-it base is served by vLLM; the trainer uses a
+# fake base (FakeBaseModel carries embed_tokens + the final norm). DSpark
+# reconstructs the base teacher distribution from the captured PRE-norm hidden
+# and re-applies the base final norm before lm_head.
+#
+# Gemma-4-E4B-specific notes (all verified on PDX 2026-08-12):
+#
+# * FINAL NORM: Gemma 4 nests the LLM under text_config with model_type
+# "gemma4_text", so modeling_final_norm.py needed BOTH "gemma4_text" and
+# "gemma4" added to _FINAL_NORM_TYPE_BY_MODEL_TYPE. Verified numerically that
+# Gemma4RMSNorm is plain `normed * weight` (NOT Gemma 2/3's `(1 + weight)`),
+# reproducing HF hidden_states[-1] at cos=0.999999, so the existing
+# _FinalRMSNorm is correct as-is.
+#
+# * ROPE: Gemma 4 has NO flat `rope_theta`; it nests per-attention-kind settings
+# under `rope_parameters` (full_attention: theta 1e6 + rope_type
+# "proportional" + partial_rotary_factor 0.25; sliding_attention: theta 1e4 +
+# rope_type "default"). modeling_fakebase.py needed _resolve_rope_theta() to
+# read the nested form. hf_dflash.py ENFORCES rope_theta from the base config
+# and overwrites any value set here, so this could NOT be fixed from the yaml.
+# We take the sliding_attention entry -> rope_theta 10000.0. Note this is the
+# LOCAL theta even though the draft below is all full_attention: the
+# full_attention entry's "proportional" rope + partial_rotary_factor 0.25 is
+# not implemented in the draft classes, and the reference drafter
+# (deepseek-ai/dspark_gemma4_12b_block7) also carries the default/1e4 form.
+# The draft is trained and served with the same value, so it is self-consistent.
+#
+# * CAPTURE IDS: vLLM's EagleModelMixin captures POST-layer with residual added
+# and indexes as layer_idx+1, so valid ids are 1..42 and 42 is the TRUE final
+# layer (verified: cos(id42, final_norm_INPUT) = 1.0000, and ids 9/18/27/36
+# match HF hidden_states[id] at cos=1.0000 with off-by-one dropping to
+# 0.55-0.98). Gemma 4 repeats 5x sliding + 1x full attention, so the
+# full_attention layers (0-based [5,11,17,23,29,35,41]) correspond to capture
+# ids [6,12,18,24,30,36,42]; we sample those to land on residual-stream
+# boundaries rather than spacing uniformly.
+#
+# * MASK TOKEN: Gemma 4's vocab is fully packed (no free/unused ids), but it
+# ships a native `` token at id 4 — used directly.
+#
+# * DRAFT ATTENTION IS NON-CAUSAL (bidirectional), like every other DSpark head.
+# vLLM resolves this per layer in qwen3_dflash._dflash_layer_causal: an
+# explicit `dflash_config.causal` overrides everything, otherwise a layer is
+# causal only when layer_types[i] == "sliding_attention". This draft sets
+# neither, and every layer is full_attention, so all 5 layers come out
+# causal=False. That is what a block-parallel (semi-autoregressive) drafter
+# wants: the block's tokens are predicted at once and must see each other.
+# ModelOpt matches this -- DFlashAttention is non-causal, and the exporter
+# does not emit a `causal` field. See "SERVING" below for the consequence.
+#
+# * SERVING requires speculative_config.attention_backend = "FLASHINFER", for
+# TWO independent reasons; either alone is fatal:
+# 1. head dim -- full_attention layers use global_head_dim=512 and the
+# auto-selected FLASH_ATTN (FA2) caps head dimension at 256.
+# 2. causality -- load_dspark_model sets attention_config.use_non_causal
+# from dflash_has_any_non_causal(), which is True here, and the backend
+# must be able to serve a non-causal mask.
+# VLLM_ATTENTION_BACKEND does NOT reach the draft; it is read from
+# speculative_config.attention_backend. Confirmed at runtime by the log line
+# "Using FlashInfer for draft model non-causal attention".
+
+metadata:
+ recipe_type: speculative_dflash
+ description: DSpark (DFlash backbone + Markov + confidence head) for Gemma-4-E4B-it,
+ 5-layer full-attention draft aligned with the official reference checkpoint.
+
+model:
+ model_name_or_path:
+ trust_remote_code: true
+ use_fake_base_for_offline: true
+
+data:
+ mode: streaming
+ data_path:
+ offline_data_path:
+ chat_template:
+
+training:
+ output_dir:
+ num_train_epochs: 1
+ per_device_train_batch_size: 4
+ gradient_accumulation_steps: 1
+ learning_rate: 1.0e-4
+ warmup_steps: 500
+ training_seq_len: 4096
+ logging_steps: 20
+ save_steps: 1000
+ cp_size: 1
+ dp_shard_size: 1
+ disable_tqdm: true
+ # Eval runs the DFlash backbone only (Markov head not applied in eval forward),
+ # so AR would misreport. Compare via export + offline AL harness instead.
+ estimate_ar: false
+ ar_validate_steps: 0
+ answer_only_loss: true
+ do_eval: false
+ lr_scheduler_type: linear
+ save_strategy: steps
+ weight_decay: 0.0
+ max_grad_norm: 1.0
+ dataloader_drop_last: true
+ bf16: true
+ tf32: true
+ remove_unused_columns: false
+ ddp_find_unused_parameters: true
+ ddp_timeout: 1800
+ report_to: none
+
+dflash:
+ dflash_block_size: 8
+ dflash_num_anchors: 512
+ dflash_use_torch_compile: false
+ dflash_self_logit_distillation: false
+ # block_size=8 -> decay gamma 4 (matches the K2.6 DSpark regime).
+ dflash_loss_decay_factor: 4.0
+ # Gemma 4 ships a native token at id 4 (vocab is fully packed, so there
+ # is no spare/unused id to borrow the way Kimi's 163838 was).
+ dflash_mask_token_id: 4
+ # --- DSpark three-term loss (DeepSpec L1/TVD-dominant defaults) ---
+ dflash_ce_loss_alpha: 0.1
+ dflash_l1_loss_alpha: 0.9
+ dflash_confidence_head_alpha: 1.0
+ dflash_architecture_config:
+ # Aligned with the official reference drafter deepseek-ai/dspark_gemma4_12b_block7
+ # (the checkpoint vLLM's gemma4_dspark.py was written for). Its backbone is
+ # 5 full-attention layers with attention_k_eq_v.
+ #
+ # Why not SWA here: Gemma4DSparkModel._build_fused_kv_buffers asserts
+ # `all(a.use_k_eq_v for a in layers_attn)`, and Gemma4DSparkAttention only
+ # sets use_k_eq_v when layer_type == "full_attention". So the assert is
+ # really about k_eq_v uniformity, but because the two are coupled, ANY
+ # sliding layer trips it. Mixed sliding/full would additionally need the V2
+ # model runner (see _resolve_layer_attention in qwen3_dflash.py).
+ num_hidden_layers: 5
+ num_attention_heads: 16
+ num_key_value_heads: 8
+ # Full-attention layers use global_head_dim (512) and, under k_eq_v,
+ # num_global_key_value_heads (1) -- see gemma4_layer_config in vLLM.
+ head_dim: 256
+ global_head_dim: 512
+ num_global_key_value_heads: 1
+ intermediate_size: 10240
+ projector_type: dspark
+ markov_rank: 256
+ markov_head_type: vanilla
+ use_confidence_head: true
+ # --- Gemma4 block shape (sandwich norms + layer_scalar) + k_eq_v ---
+ # model_type gemma4* selects DFlashGemma4DecoderLayer, which adds
+ # pre/post_feedforward_layernorm and layer_scalar. vLLM's
+ # Gemma4MTPDecoderLayer looks those up BY NAME and its DSpark loader leaves
+ # anything it cannot find randomly initialized -- silently. A Qwen3-shaped
+ # draft therefore loads without error and produces garbage.
+ model_type: gemma4_text
+ attention_k_eq_v: true
+ layer_types:
+ - full_attention
+ - full_attention
+ - full_attention
+ - full_attention
+ - full_attention
diff --git a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml
new file mode 100644
index 00000000000..d87badefad1
--- /dev/null
+++ b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml
@@ -0,0 +1,202 @@
+# DSpark recipe for Gemma-4-E4B-it -- SLIDING-WINDOW (SWA) VARIANT.
+#
+# Derived from dspark_gemma4_e4b.yaml (the validated full-attention baseline);
+# the ONLY delta is `dflash_swa_window_size` below. Keep the two in sync.
+#
+# !! NOT YET SERVABLE -- requires a vLLM fix. Read this before running. !!
+#
+# Training side is ready: _build_draft_attention_mask() already windows the
+# context (`kv > q_real_pos - window`) while leaving block-internal attention
+# bidirectional, and hf_dspark.py passes dflash_swa_window_size through. The
+# exporter emits dflash_config.use_swa/swa_window_size + a top-level
+# sliding_window, and pins causal: False.
+#
+# The serving gap is Gemma4-specific. vLLM's Qwen3 DFlash layer resolves its
+# window via _resolve_layer_attention(), which honours dflash_config.use_swa.
+# Gemma4DSparkAttention does NOT: it inherits Gemma4MTPAttention.__init__,
+# which derives the window purely from the base-model layer pattern --
+#
+# self.is_sliding = layer_type == "sliding_attention"
+# sliding_window = config.sliding_window if self.is_sliding else None
+#
+# Our draft is deliberately all "full_attention" (that is what keeps
+# attention_k_eq_v uniform and the fused-KV precompute assert satisfied), so
+# is_sliding is False on every layer and per_layer_sliding_window stays None.
+# The draft would therefore TRAIN with a 512-token window and SERVE with full
+# attention -- a silent train/inference mismatch that only shows up as a
+# depressed acceptance length, with no error raised.
+#
+# Fix (vLLM side, small): have Gemma4DSparkAttention resolve its window through
+# _resolve_layer_attention(config, layer_idx) like the Qwen3 DFlash layer does,
+# and pass the result as per_layer_sliding_window. Verified by config
+# simulation that this keeps all 5 layers use_k_eq_v=True (fused-KV assert
+# passes), causal=False, and does not require the V2 model runner.
+#
+# Until that lands, train with dspark_gemma4_e4b.yaml instead.
+#
+#
+# Streaming: the real Gemma-4-E4B-it base is served by vLLM; the trainer uses a
+# fake base (FakeBaseModel carries embed_tokens + the final norm). DSpark
+# reconstructs the base teacher distribution from the captured PRE-norm hidden
+# and re-applies the base final norm before lm_head.
+#
+# Gemma-4-E4B-specific notes (all verified on PDX 2026-08-12):
+#
+# * FINAL NORM: Gemma 4 nests the LLM under text_config with model_type
+# "gemma4_text", so modeling_final_norm.py needed BOTH "gemma4_text" and
+# "gemma4" added to _FINAL_NORM_TYPE_BY_MODEL_TYPE. Verified numerically that
+# Gemma4RMSNorm is plain `normed * weight` (NOT Gemma 2/3's `(1 + weight)`),
+# reproducing HF hidden_states[-1] at cos=0.999999, so the existing
+# _FinalRMSNorm is correct as-is.
+#
+# * ROPE: Gemma 4 has NO flat `rope_theta`; it nests per-attention-kind settings
+# under `rope_parameters` (full_attention: theta 1e6 + rope_type
+# "proportional" + partial_rotary_factor 0.25; sliding_attention: theta 1e4 +
+# rope_type "default"). modeling_fakebase.py needed _resolve_rope_theta() to
+# read the nested form. hf_dflash.py ENFORCES rope_theta from the base config
+# and overwrites any value set here, so this could NOT be fixed from the yaml.
+# We take the sliding_attention entry -> rope_theta 10000.0. Note this is the
+# LOCAL theta even though the draft below is all full_attention: the
+# full_attention entry's "proportional" rope + partial_rotary_factor 0.25 is
+# not implemented in the draft classes, and the reference drafter
+# (deepseek-ai/dspark_gemma4_12b_block7) also carries the default/1e4 form.
+# The draft is trained and served with the same value, so it is self-consistent.
+#
+# * CAPTURE IDS: vLLM's EagleModelMixin captures POST-layer with residual added
+# and indexes as layer_idx+1, so valid ids are 1..42 and 42 is the TRUE final
+# layer (verified: cos(id42, final_norm_INPUT) = 1.0000, and ids 9/18/27/36
+# match HF hidden_states[id] at cos=1.0000 with off-by-one dropping to
+# 0.55-0.98). Gemma 4 repeats 5x sliding + 1x full attention, so the
+# full_attention layers (0-based [5,11,17,23,29,35,41]) correspond to capture
+# ids [6,12,18,24,30,36,42]; we sample those to land on residual-stream
+# boundaries rather than spacing uniformly.
+#
+# * MASK TOKEN: Gemma 4's vocab is fully packed (no free/unused ids), but it
+# ships a native `` token at id 4 — used directly.
+#
+# * DRAFT ATTENTION IS NON-CAUSAL (bidirectional), like every other DSpark head.
+# vLLM resolves this per layer in qwen3_dflash._dflash_layer_causal: an
+# explicit `dflash_config.causal` overrides everything, otherwise a layer is
+# causal only when layer_types[i] == "sliding_attention". This draft sets
+# neither, and every layer is full_attention, so all 5 layers come out
+# causal=False. That is what a block-parallel (semi-autoregressive) drafter
+# wants: the block's tokens are predicted at once and must see each other.
+# ModelOpt matches this -- DFlashAttention is non-causal, and the exporter
+# does not emit a `causal` field. See "SERVING" below for the consequence.
+#
+# * SERVING requires speculative_config.attention_backend = "FLASHINFER", for
+# TWO independent reasons; either alone is fatal:
+# 1. head dim -- full_attention layers use global_head_dim=512 and the
+# auto-selected FLASH_ATTN (FA2) caps head dimension at 256.
+# 2. causality -- load_dspark_model sets attention_config.use_non_causal
+# from dflash_has_any_non_causal(), which is True here, and the backend
+# must be able to serve a non-causal mask.
+# VLLM_ATTENTION_BACKEND does NOT reach the draft; it is read from
+# speculative_config.attention_backend. Confirmed at runtime by the log line
+# "Using FlashInfer for draft model non-causal attention".
+
+metadata:
+ recipe_type: speculative_dflash
+ description: DSpark for Gemma-4-E4B-it, 5-layer draft with non-causal sliding-window
+ attention (window 512). Trains today; needs a vLLM fix before it can be served.
+
+model:
+ model_name_or_path:
+ trust_remote_code: true
+ use_fake_base_for_offline: true
+
+data:
+ mode: streaming
+ data_path:
+ offline_data_path:
+ chat_template:
+
+training:
+ output_dir:
+ num_train_epochs: 1
+ per_device_train_batch_size: 4
+ gradient_accumulation_steps: 1
+ learning_rate: 1.0e-4
+ warmup_steps: 500
+ training_seq_len: 4096
+ logging_steps: 20
+ save_steps: 1000
+ cp_size: 1
+ dp_shard_size: 1
+ disable_tqdm: true
+ # Eval runs the DFlash backbone only (Markov head not applied in eval forward),
+ # so AR would misreport. Compare via export + offline AL harness instead.
+ estimate_ar: false
+ ar_validate_steps: 0
+ answer_only_loss: true
+ do_eval: false
+ lr_scheduler_type: linear
+ save_strategy: steps
+ weight_decay: 0.0
+ max_grad_norm: 1.0
+ dataloader_drop_last: true
+ bf16: true
+ tf32: true
+ remove_unused_columns: false
+ ddp_find_unused_parameters: true
+ ddp_timeout: 1800
+ report_to: none
+
+dflash:
+ dflash_block_size: 8
+ dflash_num_anchors: 512
+ # THE ONLY FUNCTIONAL DELTA vs dspark_gemma4_e4b.yaml.
+ # Non-causal sliding window over the CONTEXT only; block-internal attention
+ # stays bidirectional and un-windowed (config.py enforces window >=
+ # dflash_block_size, so a full block always fits). 512 matches the base
+ # model's own sliding_window.
+ dflash_swa_window_size: 512
+ dflash_use_torch_compile: false
+ dflash_self_logit_distillation: false
+ # block_size=8 -> decay gamma 4 (matches the K2.6 DSpark regime).
+ dflash_loss_decay_factor: 4.0
+ # Gemma 4 ships a native token at id 4 (vocab is fully packed, so there
+ # is no spare/unused id to borrow the way Kimi's 163838 was).
+ dflash_mask_token_id: 4
+ # --- DSpark three-term loss (DeepSpec L1/TVD-dominant defaults) ---
+ dflash_ce_loss_alpha: 0.1
+ dflash_l1_loss_alpha: 0.9
+ dflash_confidence_head_alpha: 1.0
+ dflash_architecture_config:
+ # Aligned with the official reference drafter deepseek-ai/dspark_gemma4_12b_block7
+ # (the checkpoint vLLM's gemma4_dspark.py was written for). Its backbone is
+ # 5 full-attention layers with attention_k_eq_v.
+ #
+ # Why not SWA here: Gemma4DSparkModel._build_fused_kv_buffers asserts
+ # `all(a.use_k_eq_v for a in layers_attn)`, and Gemma4DSparkAttention only
+ # sets use_k_eq_v when layer_type == "full_attention". So the assert is
+ # really about k_eq_v uniformity, but because the two are coupled, ANY
+ # sliding layer trips it. Mixed sliding/full would additionally need the V2
+ # model runner (see _resolve_layer_attention in qwen3_dflash.py).
+ num_hidden_layers: 5
+ num_attention_heads: 16
+ num_key_value_heads: 8
+ # Full-attention layers use global_head_dim (512) and, under k_eq_v,
+ # num_global_key_value_heads (1) -- see gemma4_layer_config in vLLM.
+ head_dim: 256
+ global_head_dim: 512
+ num_global_key_value_heads: 1
+ intermediate_size: 10240
+ projector_type: dspark
+ markov_rank: 256
+ markov_head_type: vanilla
+ use_confidence_head: true
+ # --- Gemma4 block shape (sandwich norms + layer_scalar) + k_eq_v ---
+ # model_type gemma4* selects DFlashGemma4DecoderLayer, which adds
+ # pre/post_feedforward_layernorm and layer_scalar. vLLM's
+ # Gemma4MTPDecoderLayer looks those up BY NAME and its DSpark loader leaves
+ # anything it cannot find randomly initialized -- silently. A Qwen3-shaped
+ # draft therefore loads without error and produces garbage.
+ model_type: gemma4_text
+ attention_k_eq_v: true
+ layer_types:
+ - full_attention
+ - full_attention
+ - full_attention
+ - full_attention
+ - full_attention
diff --git a/tools/launcher/examples/google/gemma-4-E4B-it/chat_template_train.jinja b/tools/launcher/examples/google/gemma-4-E4B-it/chat_template_train.jinja
new file mode 100644
index 00000000000..f4b962ede6c
--- /dev/null
+++ b/tools/launcher/examples/google/gemma-4-E4B-it/chat_template_train.jinja
@@ -0,0 +1,390 @@
+{#
+ Template: Google Gemma 4 Canonical Chat Template
+ Author: Google Gemma Engineering Team
+ Published: 2026-07-09
+ Context: Fixed tool-calling loops, turn closures, and thinking content-ordering.
+#}
+{%- macro format_parameters(properties, required, filter_keys=false) -%}
+ {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+ {%- set ns = namespace(found_first=false) -%}
+ {%- for key, value in properties | dictsort -%}
+ {%- set add_comma = false -%}
+ {%- if not filter_keys or key not in standard_keys -%}
+ {%- if ns.found_first %},{% endif -%}
+ {%- set ns.found_first = true -%}
+ {{ key }}:{
+ {%- if value['description'] -%}
+ description:<|"|>{{ value['description'] }}<|"|>
+ {%- set add_comma = true -%}
+ {%- endif -%}
+ {%- if value['type'] | upper == 'STRING' -%}
+ {%- if value['enum'] -%}
+ {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+ enum:{{ format_argument(value['enum']) }}
+ {%- endif -%}
+ {%- elif value['type'] | upper == 'ARRAY' -%}
+ {%- if value['items'] is mapping and value['items'] -%}
+ {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+ items:{
+ {%- set ns_items = namespace(found_first=false) -%}
+ {%- for item_key, item_value in value['items'] | dictsort -%}
+ {%- if item_value is not none -%}
+ {%- if ns_items.found_first %},{% endif -%}
+ {%- set ns_items.found_first = true -%}
+ {%- if item_key == 'properties' -%}
+ properties:{
+ {%- if item_value is mapping -%}
+ {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+ {%- endif -%}
+ }
+ {%- elif item_key == 'required' -%}
+ required:[
+ {%- for req_item in item_value -%}
+ <|"|>{{- req_item -}}<|"|>
+ {%- if not loop.last %},{% endif -%}
+ {%- endfor -%}
+ ]
+ {%- elif item_key == 'type' -%}
+ {%- if item_value is string -%}
+ type:{{ format_argument(item_value | upper) }}
+ {%- else -%}
+ type:{{ format_argument(item_value | map('upper') | list) }}
+ {%- endif -%}
+ {%- else -%}
+ {{ item_key }}:{{ format_argument(item_value) }}
+ {%- endif -%}
+ {%- endif -%}
+ {%- endfor -%}
+ }
+ {%- endif -%}
+ {%- endif -%}
+ {%- if value['nullable'] %}
+ {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+ nullable:true
+ {%- endif -%}
+ {%- if value['type'] | upper == 'OBJECT' -%}
+ {%- if value['properties'] is defined and value['properties'] is mapping -%}
+ {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+ properties:{
+ {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+ }
+ {%- elif value is mapping -%}
+ {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+ properties:{
+ {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
+ }
+ {%- endif -%}
+ {%- if value['required'] -%}
+ {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+ required:[
+ {%- for item in value['required'] | default([]) -%}
+ <|"|>{{- item -}}<|"|>
+ {%- if not loop.last %},{% endif -%}
+ {%- endfor -%}
+ ]
+ {%- endif -%}
+ {%- endif -%}
+ {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+ type:<|"|>{{ value['type'] | upper }}<|"|>}
+ {%- endif -%}
+ {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+ declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+ {%- set params = tool_data['function']['parameters'] -%}
+ {%- if params -%}
+ ,parameters:{
+ {%- if params['properties'] -%}
+ properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+ {%- endif -%}
+ {%- if params['required'] -%}
+ required:[
+ {%- for item in params['required'] -%}
+ <|"|>{{- item -}}<|"|>
+ {{- ',' if not loop.last -}}
+ {%- endfor -%}
+ ],
+ {%- endif -%}
+ {%- if params['type'] -%}
+ type:<|"|>{{- params['type'] | upper -}}<|"|>}
+ {%- endif -%}
+ {%- endif -%}
+ {%- if 'response' in tool_data['function'] -%}
+ {%- set response_declaration = tool_data['function']['response'] -%}
+ ,response:{
+ {%- if response_declaration['description'] -%}
+ description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+ {%- endif -%}
+ {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+ type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+ {%- endif -%}
+ {%- endif -%}
+ }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+ {%- if argument is none -%}
+ {{- 'null' -}}
+ {%- elif argument is string -%}
+ {{- '<|"|>' + argument + '<|"|>' -}}
+ {%- elif argument is boolean -%}
+ {{- 'true' if argument else 'false' -}}
+ {%- elif argument is mapping -%}
+ {{- '{' -}}
+ {%- set ns = namespace(found_first=false) -%}
+ {%- for key, value in argument | dictsort -%}
+ {%- if ns.found_first %},{% endif -%}
+ {%- set ns.found_first = true -%}
+ {%- if escape_keys -%}
+ {{- '<|"|>' + key + '<|"|>' -}}
+ {%- else -%}
+ {{- key -}}
+ {%- endif -%}
+ :{{- format_argument(value, escape_keys=escape_keys) -}}
+ {%- endfor -%}
+ {{- '}' -}}
+ {%- elif argument is sequence -%}
+ {{- '[' -}}
+ {%- for item in argument -%}
+ {{- format_argument(item, escape_keys=escape_keys) -}}
+ {%- if not loop.last %},{% endif -%}
+ {%- endfor -%}
+ {{- ']' -}}
+ {%- else -%}
+ {{- argument -}}
+ {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+ {%- set ns = namespace(result='') -%}
+ {%- for part in text.split('') -%}
+ {%- if '<|channel>' in part -%}
+ {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+ {%- else -%}
+ {%- set ns.result = ns.result + part -%}
+ {%- endif -%}
+ {%- endfor -%}
+ {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+ {{- '<|tool_response>' -}}
+ {%- if response is mapping -%}
+ {{- 'response:' + tool_name + '{' -}}
+ {%- for key, value in response | dictsort -%}
+ {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+ {%- if not loop.last %},{% endif -%}
+ {%- endfor -%}
+ {{- '}' -}}
+ {%- else -%}
+ {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+ {%- endif -%}
+ {{- '' -}}
+{%- endmacro -%}
+
+{#- ===== SETUP ===== -#}
+{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%}
+{%- set loop_messages = messages -%}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set preserve_thinking = preserve_thinking | default(false) -%}
+{{- bos_token -}}
+{#- Handle System/Tool Definitions Block -#}
+{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%}
+ {{- '<|turn>system\n' -}}
+ {#- Inject Thinking token at the very top of the FIRST system turn -#}
+ {%- if enable_thinking -%}
+ {{- '<|think|>\n' -}}
+ {%- set ns.prev_message_type = 'think' -%}
+ {%- endif -%}
+ {%- if messages and messages[0]['role'] in ['system', 'developer'] -%}
+ {%- if messages[0]['content'] is string -%}
+ {{- messages[0]['content'] | trim -}}
+ {%- elif messages[0]['content'] is sequence -%}
+ {%- for item in messages[0]['content'] -%}
+ {{- item['text'] | trim + ' '-}}
+ {%- endfor -%}
+ {%- endif -%}
+ {%- set loop_messages = messages[1:] -%}
+ {%- endif -%}
+ {%- if tools -%}
+ {%- for tool in tools %}
+ {{- '<|tool>' -}}
+ {{- format_function_declaration(tool) | trim -}}
+ {{- '' -}}
+ {%- endfor %}
+ {%- set ns.prev_message_type = 'tool' -%}
+ {%- endif -%}
+ {{- '\n' -}}
+{%- endif %}
+
+{#- Pre-scan: find last user message index for reasoning guard -#}
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+ {%- if loop_messages[i]['role'] == 'user' -%}
+ {%- set ns_turn.last_user_idx = i -%}
+ {%- endif -%}
+{%- endfor -%}
+
+{#- Loop through messages -#}
+{%- for message in loop_messages -%}
+ {%- if message['role'] != 'tool' -%}
+ {%- set ns.prev_message_type = None -%}
+ {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+ {#- Detect continuation using tracked state — O(1) instead of O(n) backward scan -#}
+ {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%}
+ {%- if not continue_same_model_turn -%}
+ {{- '<|turn>' + role + '\n' }}
+ {%- endif -%}
+
+ {#- Render reasoning/reasoning_content as thinking channel -#}
+ {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
+ {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or (preserve_thinking and message.get('tool_calls')) -%}
+ {%- if thinking_text and thinking_gate -%}
+ {{- '<|channel>thought\n' + thinking_text + '\n' -}}
+ {%- endif -%}
+
+ {%- if message.get('tool_calls') -%}
+ {%- for tool_call in message.get('tool_calls') -%}
+ {%- set function = tool_call['function'] -%}
+ {{- '<|tool_call>call:' + function['name'] + '{' -}}
+ {%- if function['arguments'] is mapping -%}
+ {%- set ns_args = namespace(found_first=false) -%}
+ {%- for key, value in function['arguments'] | dictsort -%}
+ {%- if ns_args.found_first %},{% endif -%}
+ {%- set ns_args.found_first = true -%}
+ {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+ {%- endfor -%}
+ {%- elif function['arguments'] is none -%}
+ {%- else -%}
+ {{- raise_exception(
+ "chat_template: tool_calls[].function.arguments must be a "
+ "JSON object (mapping), not a string. Deserialize arguments "
+ "before passing to the template."
+ ) -}}
+ {%- endif -%}
+ {{- '}' -}}
+ {%- endfor -%}
+ {%- set ns.prev_message_type = 'tool_call' -%}
+ {%- endif -%}
+
+ {%- set ns_tr_out = namespace(flag=false) -%}
+ {%- if message.get('tool_responses') -%}
+ {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
+ {%- for tool_response in message.get('tool_responses') -%}
+ {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}}
+ {%- set ns_tr_out.flag = true -%}
+ {%- set ns.prev_message_type = 'tool_response' -%}
+ {%- endfor -%}
+ {%- elif message.get('tool_calls') -%}
+ {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
+ {%- set ns_tool_scan = namespace(stopped=false) -%}
+ {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+ {%- if ns_tool_scan.stopped -%}
+ {%- elif loop_messages[k]['role'] != 'tool' -%}
+ {%- set ns_tool_scan.stopped = true -%}
+ {%- else -%}
+ {%- set follow = loop_messages[k] -%}
+ {#- Resolve tool_call_id to function name -#}
+ {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%}
+ {%- for tc in message.get('tool_calls') -%}
+ {%- if tc.get('id') == follow.get('tool_call_id') -%}
+ {%- set ns_tname.name = tc['function']['name'] -%}
+ {%- endif -%}
+ {%- endfor -%}
+ {#- Handle content as string or content-parts array -#}
+ {%- set tool_body = follow.get('content') -%}
+ {%- if tool_body is string -%}
+ {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+ {%- elif tool_body is sequence and tool_body is not string -%}
+ {%- set ns_txt = namespace(s='') -%}
+ {%- for part in tool_body -%}
+ {%- if part.get('type') == 'text' -%}
+ {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+ {%- endif -%}
+ {%- endfor -%}
+ {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+ {%- for part in tool_body -%}
+ {%- if part.get('type') in ['image', 'image_url'] -%}
+ {{- '<|image|>' -}}
+ {%- elif part.get('type') in ['audio', 'input_audio'] -%}
+ {{- '<|audio|>' -}}
+ {%- elif part.get('type') == 'video' -%}
+ {{- '<|video|>' -}}
+ {%- endif -%}
+ {%- endfor -%}
+ {%- else -%}
+ {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+ {%- endif -%}
+ {%- set ns_tr_out.flag = true -%}
+ {%- set ns.prev_message_type = 'tool_response' -%}
+ {%- endif -%}
+ {%- endfor -%}
+ {%- endif -%}
+
+ {%- set captured_content -%}
+ {%- if message.get('content') is string -%}
+ {%- if role == 'model' -%}
+ {{- strip_thinking(message['content']) -}}
+ {%- else -%}
+ {{- message['content'] | trim -}}
+ {%- endif -%}
+ {%- elif message.get('content') is sequence -%}
+ {%- for item in message['content'] -%}
+ {%- if item.get('type') == 'text' -%}
+ {%- if role == 'model' -%}
+ {{- strip_thinking(item['text']) -}}
+ {%- else -%}
+ {{- item['text'] | trim -}}
+ {%- endif -%}
+ {%- elif item.get('type') in ['image', 'image_url'] -%}
+ {{- '<|image|>' -}}
+ {%- elif item.get('type') in ['audio', 'input_audio'] -%}
+ {{- '<|audio|>' -}}
+ {%- elif item.get('type') == 'video' -%}
+ {{- '<|video|>' -}}
+ {%- endif -%}
+ {%- endfor -%}
+ {%- endif -%}
+ {%- endset -%}
+
+ {%- if role == 'model' -%}
+ {%- generation -%}{{- captured_content -}}{%- endgeneration -%}
+ {%- else -%}
+ {{- captured_content -}}
+ {%- endif -%}
+ {%- set has_content = captured_content | trim | length > 0 -%}
+
+ {#- Forward-scan: find next non-tool message role for continuation detection -#}
+ {%- set next_nt = namespace(role=None, found=false) -%}
+ {%- for j in range(loop.index0 + 1, loop_messages | length) -%}
+ {%- if not next_nt.found -%}
+ {%- if loop_messages[j]['role'] != 'tool' -%}
+ {%- set next_nt.role = loop_messages[j]['role'] -%}
+ {%- set next_nt.found = true -%}
+ {%- endif -%}
+ {%- endif -%}
+ {%- endfor -%}
+
+ {%- set continues_into_next = (
+ role == 'model'
+ and next_nt.role == 'assistant'
+ and (not message.get('tool_calls') or ns_tr_out.flag)
+ ) -%}
+
+ {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
+ {{- '<|tool_response>' -}}
+ {%- elif continues_into_next -%}
+ {%- elif not (ns_tr_out.flag and not has_content and not next_nt.found) -%}
+ {{- '\n' -}}
+ {%- endif -%}
+
+ {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#}
+ {%- set ns.prev_non_tool_role = message['role'] -%}
+ {%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+ {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+ {{- '<|turn>model\n' -}}
+ {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%}
+ {{- '<|channel>thought\n' -}}
+ {%- endif -%}
+{%- endif -%}
diff --git a/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_full.yaml b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_full.yaml
new file mode 100644
index 00000000000..12b8ea36ff0
--- /dev/null
+++ b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_full.yaml
@@ -0,0 +1,96 @@
+# DSpark FULL training run for Gemma-4-E4B-it on AWS-PDX.
+#
+# First real training run (the earlier 20-step smoke was a pipeline test only).
+# Corpus: 1,130,852 synthesized rows from ~/lustre/g4_corpus_train.
+#
+# TOPOLOGY: 4 nodes on the interactive QOS (its hard cap), split 2 serve + 2
+# trainer. The launcher does NOT support "own serve on every node" -- for
+# nodes >= 2 it splits WHOLE nodes (see train_eagle_streaming.sh header), so
+# 2:2 is what yields a clean global batch of 64:
+#
+# 2 trainer nodes x 8 GPU = 16 DP ranks
+# per_device_train_batch_size 4 x grad_accum 1
+# -> global batch = 16 * 4 * 1 = 64
+#
+# 2 epochs: 1,130,852 * 2 / 64 ~= 35,340 steps
+#
+# CORPUS: must be the CLEANED copy (g4_corpus_train). The raw synthesis output
+# carries a user-only `messages` column, and hf_streaming_dataset prefers
+# `messages` over `conversations` -- that yields an all-zero loss mask under
+# answer_only_loss, every row is rejected, and streaming SILENTLY HANGS with no
+# error. Verified 200/200 raw rows were user-only; the cleaned copy has 0.
+#
+# Gemma-4-E4B specifics (verified on PDX 2026-08-12, unchanged):
+# * EAGLE_CAPTURE_IDS: vLLM captures POST-layer with the residual added and
+# indexes as layer_idx+1, so valid ids are 1..42 and 42 is the TRUE final
+# layer. Gemma 4 repeats 5x sliding + 1x full attention, so full-attention
+# layers sit at ids [6,12,18,24,30,36,42]; we sample 6 of those (must be
+# num_draft_layers + 1 = 6) to land on residual-stream boundaries.
+# * No --trust-remote-code needed and no vLLM patch; gemma4 is native in the
+# 2026-08-11 nightly.
+
+job_name: Gemma-4-E4B_DSpark_full_2ep
+pipeline:
+ allow_to_fail: false
+ skip: false
+ note:
+
+ global_vars:
+ hf_model: /hf-local/gemma-4-E4B-it
+
+ task_0:
+ script: common/eagle3/train_eagle_streaming.sh
+ args:
+ - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml
+ - model.model_name_or_path=<>
+ - model.use_fake_base_for_offline=true
+ - data.mode=streaming
+ # Gemma 4's stock chat template has NO generation markers, so
+ # answer_only_loss silently yields an all-zero loss_mask and every row is
+ # rejected. This copy adds them; see the template file header.
+ - data.chat_template=examples/google/gemma-4-E4B-it/chat_template_train.jinja
+ - data.data_path=/traindata
+ - training.output_dir=/scratchspace/dspark_full_2ep
+ - training.training_seq_len=4096
+ - training.disable_tqdm=true
+ - training.ar_validate_steps=0
+ - training.num_train_epochs=2
+ - training.per_device_train_batch_size=4
+ - training.gradient_accumulation_steps=1
+ - training.learning_rate=1.0e-4
+ - training.warmup_steps=500
+ - training.logging_steps=20
+ - training.save_steps=1000
+ # All checkpoints are kept (user's call): HF save_total_limit defaults to
+ # None, which retains every one, so no override is needed.
+ - training.answer_only_loss=true
+ - training.report_to=none
+ environment:
+ - HF_MODEL_CKPT: <>
+ # MUST be exactly num_draft_layers + 1 (5 draft layers -> 6 ids).
+ - EAGLE_CAPTURE_IDS: "[6,12,18,24,36,42]"
+ # 2 serve nodes + 2 trainer nodes (of the 4 allocated).
+ - SERVE_NODES: "2"
+ - SERVE_GPU_MEM_UTIL: "0.9"
+ - SERVE_MAX_MODEL_LEN: "4352"
+ - SERVE_MAX_NUM_SEQS: "64"
+ - SERVE_READY_TIMEOUT: "2400"
+ - STREAMING_NUM_WORKERS: "4"
+ - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200"
+ - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200"
+ # NIXL hidden-state transport. Without these NIXL falls back to UCX, which
+ # reports "UCX CUDA support was not found" and dies with
+ # NIXL_ERR_REMOTE_DISCONNECT the moment the trainer pulls hidden states.
+ - NIXL_BACKENDS: LIBFABRIC
+ - FI_PROVIDER: efa
+ slurm_config:
+ _factory_: "slurm_factory"
+ nodes: 4
+ ntasks_per_node: 1
+ gpus_per_node: 8
+ # The auxfix image carries the AWS libfabric stack that NIXL needs; the
+ # nem35 image has no libfabric and no NIXL plugins.
+ container: /home/haoguo/lustre/containers/vllm-nightly-efa-x86_64-auxfix.sqsh
+ container_mounts:
+ - /home/haoguo/lustre/hf-local:/hf-local
+ - /home/haoguo/lustre/g4_corpus_train:/traindata
diff --git a/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke.yaml b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke.yaml
new file mode 100644
index 00000000000..5ff232145ea
--- /dev/null
+++ b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke.yaml
@@ -0,0 +1,98 @@
+# DSpark streaming SMOKE run (co-located single node) for Gemma-4-E4B-it on AWS-PDX.
+#
+# Purpose: prove the streaming chain end-to-end (vLLM serve -> aux hidden-state
+# capture -> NIXL transfer -> fake-base trainer) for a NEW base model. Only ~20
+# steps on a 1024-row corpus; the numbers are meaningless, the point is that the
+# pipeline runs and the loss is finite and decreasing.
+#
+# Topology: serve TP=1 on GPU 0 (E4B is 16 GB, fits one B300), DSpark trainer on
+# GPUs 4-7. Intra-node NIXL, no cross-node EFA.
+#
+# Gemma-4-E4B specifics verified on PDX 2026-08-12:
+# * EAGLE_CAPTURE_IDS: vLLM's EagleModelMixin captures POST-layer with the
+# residual added and indexes as layer_idx+1, so valid ids are 1..42 and 42 is
+# the TRUE final layer (verified cos=1.0000 against HF hidden_states, with
+# off-by-one dropping to 0.55-0.98). Gemma 4 repeats 5x sliding + 1x full
+# attention, so full-attention layers sit at capture ids
+# [6,12,18,24,30,36,42]; we sample those to land on residual-stream
+# boundaries instead of spacing uniformly (deep layers are near-redundant:
+# adjacent cosine ~0.98-0.99 around id 36 vs 0.55-0.70 around id 9).
+# * NO --trust-remote-code needed (the repo ships no .py) and no vLLM patch:
+# gemma4 is natively supported in the 2026-08-11 nightly container.
+# * Corpus has its user-only `messages` column DROPPED — hf_streaming_dataset
+# prefers `messages` over `conversations` and a user-only one makes streaming
+# SILENTLY HANG.
+
+job_name: Gemma-4-E4B_DSpark_streaming_smoke
+pipeline:
+ allow_to_fail: false
+ skip: false
+ note:
+
+ global_vars:
+ hf_model: /hf-local/gemma-4-E4B-it
+
+ task_0:
+ script: common/eagle3/train_eagle_streaming.sh
+ args:
+ - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml
+ - model.model_name_or_path=<>
+ - model.use_fake_base_for_offline=true
+ - data.mode=streaming
+ # Gemma 4's stock chat template has NO generation markers, so
+ # answer_only_loss silently yields an all-zero loss_mask and EVERY row is
+ # rejected ("no fetchable sample found in the entire corpus"). This copy
+ # adds them; see the template file header for details.
+ - data.chat_template=examples/google/gemma-4-E4B-it/chat_template_train.jinja
+ - data.data_path=/smokedata
+ - training.output_dir=/scratchspace/dspark_smoke
+ - training.training_seq_len=2048
+ - training.disable_tqdm=true
+ - training.ar_validate_steps=500000
+ - training.num_train_epochs=1
+ - training.max_steps=20
+ - training.logging_steps=1
+ - training.save_steps=20
+ - training.per_device_train_batch_size=1
+ - training.answer_only_loss=true
+ - training.report_to=none
+ environment:
+ - HF_MODEL_CKPT: <>
+ # MUST be exactly num_draft_layers + 1 entries (5 draft layers -> 6 ids):
+ # the projector is sized from the DRAFT's num_hidden_layers, so 7 ids gave
+ # "mat1 and mat2 shapes cannot be multiplied (2048x15360 and 12800x2560)".
+ # Confirmed against the working runs: K2.6 6 layers -> 7 ids, gpt-oss 5 -> 6.
+ # Chosen from the full-attention layers of the 5:1 sliding/full cycle
+ # ([6,12,18,24,30,36,42]), dropping 30 to keep both ends and the true final 42.
+ - EAGLE_CAPTURE_IDS: "[6,12,18,24,36,42]"
+ - SERVE_GPU: "0"
+ - SERVE_TP: "1"
+ - SERVE_GPU_MEM_UTIL: "0.85"
+ - STREAMING_NUM_WORKERS: "1"
+ - SERVE_MAX_MODEL_LEN: "2176"
+ - SERVE_MAX_NUM_SEQS: "4"
+ - SERVE_READY_TIMEOUT: "2400"
+ - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200"
+ - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200"
+ # NIXL hidden-state transport. Without these, NIXL falls back to the UCX
+ # backend, which reports "8 NVIDIA GPU(s) were detected, but UCX CUDA
+ # support was not found" and then dies with NIXL_ERR_REMOTE_DISCONNECT the
+ # moment the trainer tries to pull hidden states. LIBFABRIC+efa is what the
+ # working Kimi-K2.6 streaming runs use.
+ - NIXL_BACKENDS: LIBFABRIC
+ - FI_PROVIDER: efa
+ slurm_config:
+ _factory_: "slurm_factory"
+ nodes: 1
+ ntasks_per_node: 1
+ gpus_per_node: 8
+ # The 2026-08-11 nem35 image has NO libfabric and no NIXL plugins, so
+ # NIXL_BACKENDS=LIBFABRIC dies with NIXL_ERR_NOT_FOUND and the UCX fallback
+ # dies with NIXL_ERR_REMOTE_DISCONNECT ("UCX CUDA support was not found").
+ # The auxfix image carries the AWS libfabric stack and is what the working
+ # Kimi-K2.6 streaming runs use. Its vLLM is older (transformers 5.12.1) --
+ # gemma4 support must be re-verified in THIS image.
+ container: /home/haoguo/lustre/containers/vllm-nightly-efa-x86_64-auxfix.sqsh
+ container_mounts:
+ - /home/haoguo/lustre/hf-local:/hf-local
+ - /home/haoguo/lustre/g4_smoke_corpus:/smokedata
diff --git a/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke_swa.yaml b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke_swa.yaml
new file mode 100644
index 00000000000..f4cd8f0198c
--- /dev/null
+++ b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke_swa.yaml
@@ -0,0 +1,111 @@
+# DSpark streaming SMOKE run (co-located single node) for Gemma-4-E4B-it on
+# AWS-PDX -- SLIDING-WINDOW (SWA) VARIANT.
+#
+# Derived from hf_streaming_dspark_smoke.yaml; the only deltas are the recipe
+# (dspark_gemma4_e4b_swa.yaml, which sets dflash_swa_window_size: 512) and the
+# output dir. Keep the two in sync.
+#
+# TRAINING ONLY. The resulting checkpoint is NOT servable as-is: vLLM's
+# Gemma4DSparkAttention derives its sliding window from the BASE layer pattern
+# instead of dflash_config.use_swa, so an SWA-trained draft would be served with
+# FULL attention -- silently, showing up only as depressed acceptance. See
+# examples/speculative_decoding/export/gemma4_dspark_swa_vllm.patch. The point of
+# this run is to prove SWA training converges and to produce a checkpoint.
+#
+#
+# Purpose: prove the streaming chain end-to-end (vLLM serve -> aux hidden-state
+# capture -> NIXL transfer -> fake-base trainer) for a NEW base model. Only ~20
+# steps on a 1024-row corpus; the numbers are meaningless, the point is that the
+# pipeline runs and the loss is finite and decreasing.
+#
+# Topology: serve TP=1 on GPU 0 (E4B is 16 GB, fits one B300), DSpark trainer on
+# GPUs 4-7. Intra-node NIXL, no cross-node EFA.
+#
+# Gemma-4-E4B specifics verified on PDX 2026-08-12:
+# * EAGLE_CAPTURE_IDS: vLLM's EagleModelMixin captures POST-layer with the
+# residual added and indexes as layer_idx+1, so valid ids are 1..42 and 42 is
+# the TRUE final layer (verified cos=1.0000 against HF hidden_states, with
+# off-by-one dropping to 0.55-0.98). Gemma 4 repeats 5x sliding + 1x full
+# attention, so full-attention layers sit at capture ids
+# [6,12,18,24,30,36,42]; we sample those to land on residual-stream
+# boundaries instead of spacing uniformly (deep layers are near-redundant:
+# adjacent cosine ~0.98-0.99 around id 36 vs 0.55-0.70 around id 9).
+# * NO --trust-remote-code needed (the repo ships no .py) and no vLLM patch:
+# gemma4 is natively supported in the 2026-08-11 nightly container.
+# * Corpus has its user-only `messages` column DROPPED — hf_streaming_dataset
+# prefers `messages` over `conversations` and a user-only one makes streaming
+# SILENTLY HANG.
+
+job_name: Gemma-4-E4B_DSpark_streaming_smoke_swa
+pipeline:
+ allow_to_fail: false
+ skip: false
+ note:
+
+ global_vars:
+ hf_model: /hf-local/gemma-4-E4B-it
+
+ task_0:
+ script: common/eagle3/train_eagle_streaming.sh
+ args:
+ - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml
+ - model.model_name_or_path=<>
+ - model.use_fake_base_for_offline=true
+ - data.mode=streaming
+ # Gemma 4's stock chat template has NO generation markers, so
+ # answer_only_loss silently yields an all-zero loss_mask and EVERY row is
+ # rejected ("no fetchable sample found in the entire corpus"). This copy
+ # adds them; see the template file header for details.
+ - data.chat_template=examples/google/gemma-4-E4B-it/chat_template_train.jinja
+ - data.data_path=/smokedata
+ - training.output_dir=/scratchspace/dspark_smoke_swa
+ - training.training_seq_len=2048
+ - training.disable_tqdm=true
+ - training.ar_validate_steps=500000
+ - training.num_train_epochs=1
+ - training.max_steps=20
+ - training.logging_steps=1
+ - training.save_steps=20
+ - training.per_device_train_batch_size=1
+ - training.answer_only_loss=true
+ - training.report_to=none
+ environment:
+ - HF_MODEL_CKPT: <>
+ # MUST be exactly num_draft_layers + 1 entries (5 draft layers -> 6 ids):
+ # the projector is sized from the DRAFT's num_hidden_layers, so 7 ids gave
+ # "mat1 and mat2 shapes cannot be multiplied (2048x15360 and 12800x2560)".
+ # Confirmed against the working runs: K2.6 6 layers -> 7 ids, gpt-oss 5 -> 6.
+ # Chosen from the full-attention layers of the 5:1 sliding/full cycle
+ # ([6,12,18,24,30,36,42]), dropping 30 to keep both ends and the true final 42.
+ - EAGLE_CAPTURE_IDS: "[6,12,18,24,36,42]"
+ - SERVE_GPU: "0"
+ - SERVE_TP: "1"
+ - SERVE_GPU_MEM_UTIL: "0.85"
+ - STREAMING_NUM_WORKERS: "1"
+ - SERVE_MAX_MODEL_LEN: "2176"
+ - SERVE_MAX_NUM_SEQS: "4"
+ - SERVE_READY_TIMEOUT: "2400"
+ - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200"
+ - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200"
+ # NIXL hidden-state transport. Without these, NIXL falls back to the UCX
+ # backend, which reports "8 NVIDIA GPU(s) were detected, but UCX CUDA
+ # support was not found" and then dies with NIXL_ERR_REMOTE_DISCONNECT the
+ # moment the trainer tries to pull hidden states. LIBFABRIC+efa is what the
+ # working Kimi-K2.6 streaming runs use.
+ - NIXL_BACKENDS: LIBFABRIC
+ - FI_PROVIDER: efa
+ slurm_config:
+ _factory_: "slurm_factory"
+ nodes: 1
+ ntasks_per_node: 1
+ gpus_per_node: 8
+ # The 2026-08-11 nem35 image has NO libfabric and no NIXL plugins, so
+ # NIXL_BACKENDS=LIBFABRIC dies with NIXL_ERR_NOT_FOUND and the UCX fallback
+ # dies with NIXL_ERR_REMOTE_DISCONNECT ("UCX CUDA support was not found").
+ # The auxfix image carries the AWS libfabric stack and is what the working
+ # Kimi-K2.6 streaming runs use. Its vLLM is older (transformers 5.12.1) --
+ # gemma4 support must be re-verified in THIS image.
+ container: /home/haoguo/lustre/containers/vllm-nightly-efa-x86_64-auxfix.sqsh
+ container_mounts:
+ - /home/haoguo/lustre/hf-local:/hf-local
+ - /home/haoguo/lustre/g4_smoke_corpus:/smokedata