Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+learning_curve.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,180 @@
+---
+license: mit
+base_model: z-lab/Qwen3.6-27B-DFlash
+tags:
+- speculative-decoding
+- draft-model
+- dflash
+- dspark
+- qwen
+---
+
+# DSpark-DFlash Draft Head for Qwen3.6-27B-AEON
+
+A speculative-decoding **draft head** for the **Qwen3.6-27B-AEON** family, built by
+reproducing and adapting DeepSeek's **DSpark** recipe on top of z-lab's public
+**DFlash** block-diffusion drafter. It combines two drafting paths:
+
+1. **DFlash block-diffusion backbone** — the public [`z-lab/Qwen3.6-27B-DFlash`](https://huggingface.co/z-lab/Qwen3.6-27B-DFlash)
+   head (MIT), fine-tuned on on-policy AEON traces.
+2. **VanillaMarkov sequential head** — a rank-256 head with a bigram bias term, added
+   on the DSpark-style semi-autoregressive drafting path.
+
+The head is distilled against the **self-generated (on-policy) logits** of a
+Qwen3.6-27B-AEON target, so it is specialized to that target family rather than to
+stock `Qwen/Qwen3.6-27B`.
+
+![Learning curve](learning_curve.png)
+
+---
+
+## Results
+
+All numbers below are reported **with their measurement conditions**. `measured`
+means observed on the described harness; nothing here is extrapolated.
+
+### Offline paired acceptance (draft quality)
+
+Paired evaluation on **on-policy AEON-generated text** (the draft head scores the
+same target continuations under both drafters), **n = 4052 anchors / 176 sequences**,
+**cluster-bootstrap 95% CI**, **eval K = 8**.
+
+Versus the stock z-lab DFlash head:
+
+| Metric | This head vs z-lab DFlash |
+|---|---|
+| Accept, **T = 1.0 sampling** | **+32.4%** relative [CI +27%, +38%] |
+| Accept, **greedy** | **+17.7%** relative |
+
+Per-domain absolute accepted-length gain (all CIs exclude 0):
+
+| Domain | Δ accepted length |
+|---|---|
+| toolcall | +2.83 |
+| chat | +0.80 |
+| code | +0.57 |
+| math | +0.50 |
+
+### End-to-end serving throughput (measured)
+
+**vLLM 0.23.0**, **ABBA × 3 rounds**, **K = 8**, **T = 1.0**, **NVFP4 target**,
+**single RTX PRO 6000**.
+
+| | This head | z-lab DFlash | Δ |
+|---|---|---|---|
+| Aggregate throughput | **194.8 tok/s** | 175.5 tok/s | **+11.0%** [CI +13.6, +26.5 tok/s] |
+| Accept rate | 0.420 | 0.342 | — |
+
+Per-domain throughput gain (all CIs exclude 0):
+
+| Domain | Δ tok/s |
+|---|---|
+| code | +15.0% |
+| toolcall | +14.1% |
+| chat | +8.1% |
+| math | +7.1% |
+
+---
+
+## Training
+
+| Setting | Value |
+|---|---|
+| Loss | L1 distribution loss (0.9) + teacher-argmax CE (0.1) |
+| `loss_decay_gamma` | 6.0 |
+| `block_size` | 11 |
+| `max_context` | 1024 |
+| Anchors / sequence | 32 |
+| Learning rate | 6e-4, cosine schedule |
+| Steps | 6000 (converged at **~4500** in practice) |
+| Head dtype | bf16 |
+| Target data | **target self-generated, on-policy** |
+
+**Corpus** — coding/toolcall-heavy mix, **15,936 sequences** total:
+
+- **40%** toolcall (AEON self-play)
+- **25%** real agent sessions (tool-use traces)
+- **35%** general (of which **57%** is code)
+
+---
+
+## Usage
+
+### Target model
+
+Designed for the **Qwen3.6-27B-AEON** family (**vocab 248320**). Verified against:
+
+- [`AEON-7/Qwen3.6-27B-AEON-Ultimate-Uncensored`](https://huggingface.co/AEON-7/Qwen3.6-27B-AEON-Ultimate-Uncensored) (BF16)
+- [`AEON-7/Qwen3.6-27B-AEON-Ultimate-Uncensored-NVFP4`](https://huggingface.co/AEON-7/Qwen3.6-27B-AEON-Ultimate-Uncensored-NVFP4) (NVFP4, used for the serving numbers above)
+
+More broadly, this head is compatible with **Qwen3.6-27B-AEON merges (vocab 248320)**.
+It is distilled to AEON logits and is **not** intended as a drop-in drafter for stock
+`Qwen/Qwen3.6-27B`.
+
+### Install the vLLM patches (required)
+
+This head uses a **Markov semi-autoregressive** drafting path that **stock vLLM does
+not implement**, so the two bundled patch files are **required**. They are written
+against **vLLM 0.23.0** — do not apply them to other versions.
+
+Overwrite-copy the two files from `vllm_patches/` in this repo into your vLLM install:
+
+```bash
+# from the root of this repo, into your vLLM 0.23.0 site-packages
+cp vllm_patches/qwen3_dflash.py       "$VLLM/vllm/model_executor/models/"
+cp vllm_patches/llm_base_proposer.py  "$VLLM/vllm/v1/spec_decode/"
+```
+
+(`$VLLM` = the directory containing your installed `vllm` package.)
+
+### Serve
+
+```bash
+vllm serve <target> \
+  --speculative-config '{"method":"dflash","model":"<this repo>","num_speculative_tokens":8,"draft_sample_method":"probabilistic"}' \
+  --mamba-cache-dtype float32 \
+  --attention-backend flash_attn
+```
+
+Replace `<target>` with an AEON target (e.g.
+`AEON-7/Qwen3.6-27B-AEON-Ultimate-Uncensored-NVFP4`) and `<this repo>` with this
+draft-head repo id.
+
+### Optional environment variables
+
+- **`DSPARK_MARKOV_TOPN`** — experimental. Truncates the Markov head to its top-N
+  candidates to cut per-step overhead in single-stream serving. It measurably
+  **lowers acceptance on toolcall**, so it is **off by default** and recommended to
+  leave off.
+
+---
+
+## Honest limitations
+
+- **Target-specific.** The head is distilled to the logits of a Qwen3.6-27B-AEON
+  merge. Gains do not transfer to unrelated targets, including stock
+  `Qwen/Qwen3.6-27B`.
+- **Evaluated at K = 8.** The block size is 11 (supports K ≤ 10); all reported
+  acceptance and throughput numbers are at K = 8 and are not guaranteed to hold at
+  other draft lengths.
+- **Single-GPU numbers.** Serving throughput was measured on one RTX PRO 6000 with
+  the NVFP4 target; other hardware, batch regimes, or target quantizations will differ.
+- **Sampling is the strong regime.** The largest wins are under **T = 1.0 sampling**
+  (accept +32.4% offline) rather than greedy (+17.7% offline). Greedy-heavy workloads
+  will see smaller gains.
+
+---
+
+## Acknowledgements
+
+- **z-lab** — the [DFlash](https://github.com/z-lab/dflash) block-diffusion drafter
+  ([`z-lab/Qwen3.6-27B-DFlash`](https://huggingface.co/z-lab/Qwen3.6-27B-DFlash), MIT),
+  which this head fine-tunes.
+  DFlash: Block Diffusion for Flash Speculative Decoding
+  ([arXiv:2602.06036](https://arxiv.org/abs/2602.06036)).
+- **DeepSeek** — the **DSpark** paper and the **DeepSpec** reference implementation,
+  whose recipe this work reproduces and adapts for the Qwen3.6-27B-AEON target.
+
+## License
+
+MIT (inherited from the z-lab DFlash head).

config.json

@@ -0,0 +1,53 @@
+{
+  "architectures": [
+    "DFlashDraftModel"
+  ],
+  "auto_map": {
+    "AutoModel": "dflash.DFlashDraftModel"
+  },
+  "model_type": "qwen3",
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "block_size": 11,
+  "dflash_config": {
+    "causal_head": false,
+    "causal": false,
+    "mask_token_id": 248070,
+    "target_layer_ids": [
+      1,
+      16,
+      31,
+      46,
+      61
+    ]
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 248046,
+  "pad_token_id": 248044,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 17408,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 8192,
+  "num_attention_heads": 32,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 64,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 10000000.0,
+  "rope_scaling": null,
+  "tie_word_embeddings": false,
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 248320,
+  "markov_rank": 256,
+  "markov_head_type": "vanilla"
+}

dflash.py

@@ -0,0 +1,494 @@
+from typing import Callable, Optional
+
+import torch
+from torch import nn
+from transformers import DynamicCache
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.qwen3.modeling_qwen3 import (
+    ALL_ATTENTION_FUNCTIONS,
+    FlashAttentionKwargs,
+    GradientCheckpointingLayer,
+    Qwen3Config,
+    Qwen3MLP,
+    Qwen3PreTrainedModel,
+    Qwen3RMSNorm,
+    Qwen3RotaryEmbedding,
+    eager_attention_forward,
+    rotate_half,
+)
+from typing_extensions import Tuple, Unpack
+
+# Optional DSpark-style Markov sequential head (vendored next to this file).
+# Guarded so the rank-0 path -- and any environment lacking markov_head.py --
+# still imports and behaves byte-identically to the original DFlash head.
+try:
+    from markov_head import build_markov_head
+except Exception:  # pragma: no cover - import guard
+    try:
+        from .markov_head import build_markov_head  # type: ignore
+    except Exception:
+        build_markov_head = None
+
+
+def sample(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor:
+    if temperature < 1e-5:
+        return torch.argmax(logits, dim=-1)
+    bsz, seq_len, vocab_size = logits.shape
+    logits = logits.view(-1, vocab_size)
+    logits = logits / temperature
+    probs = torch.softmax(logits, dim=-1)
+    return torch.multinomial(probs, num_samples=1).view(bsz, seq_len)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_len = q.size(-2)
+    q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :])
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def _to_additive_attention_mask(
+    attention_mask: torch.Tensor,
+    *,
+    query_dtype: torch.dtype,
+    device: torch.device,
+    key_len: int,
+) -> torch.Tensor:
+    if attention_mask.ndim == 4:
+        attention_mask = attention_mask[:, :, :, :key_len]
+    if attention_mask.dtype == torch.bool:
+        additive_mask = torch.zeros_like(attention_mask, dtype=query_dtype, device=device)
+        return additive_mask.masked_fill(
+            attention_mask.logical_not().to(device=device),
+            torch.finfo(query_dtype).min,
+        )
+    return attention_mask.to(device=device, dtype=query_dtype)
+
+
+def _build_dflash_causal_attention_mask(
+    *,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cached_kv_len: int,
+    ctx_len: int,
+) -> torch.Tensor:
+    q_len = query.shape[-2]
+    kv_len = key.shape[-2]
+    key_positions = torch.arange(kv_len, device=query.device)
+    query_positions = cached_kv_len + ctx_len + torch.arange(q_len, device=query.device)
+    can_attend = key_positions.unsqueeze(0) <= query_positions.unsqueeze(1)
+    mask = torch.zeros((1, 1, q_len, kv_len), dtype=query.dtype, device=query.device)
+    return mask.masked_fill(can_attend.logical_not().unsqueeze(0).unsqueeze(0), torch.finfo(query.dtype).min)
+
+
+class Qwen3DFlashAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        dflash_config = getattr(config, "dflash_config", {}) or {}
+        self.is_causal = bool(dflash_config.get("causal_head", False))
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=config.attention_bias,
+        )
+        self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.sliding_window = (
+            config.sliding_window
+            if config.layer_types[layer_idx] == "sliding_attention"
+            else None
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        target_hidden: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        bsz, q_len = hidden_states.shape[:-1]
+        ctx_len = target_hidden.shape[1]
+        is_causal = kwargs.pop("is_causal", None)
+        if is_causal is None:
+            is_causal = self.is_causal
+        q = self.q_proj(hidden_states)
+        q = q.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)
+        v_ctx = self.v_proj(target_hidden)
+        v_noise = self.v_proj(hidden_states)
+        k = torch.cat([k_ctx, k_noise], dim=1).view(
+            bsz, ctx_len + q_len, -1, self.head_dim
+        )
+        v = torch.cat([v_ctx, v_noise], dim=1).view(
+            bsz, ctx_len + q_len, -1, self.head_dim
+        )
+        k = self.k_norm(k).transpose(1, 2)
+        v = v.transpose(1, 2)
+        cos, sin = position_embeddings
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        cached_kv_len = past_key_values.get_seq_length() if past_key_values is not None else 0
+        if past_key_values is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs)
+
+        attn_backend = self.config._attn_implementation
+        use_explicit_dflash_causal_mask = bool(is_causal) and attn_backend in {"eager", "sdpa"}
+        if use_explicit_dflash_causal_mask:
+            dflash_causal_mask = _build_dflash_causal_attention_mask(
+                query=q, key=k, cached_kv_len=cached_kv_len, ctx_len=ctx_len,
+            )
+            if attention_mask is not None:
+                dflash_causal_mask = dflash_causal_mask + _to_additive_attention_mask(
+                    attention_mask, query_dtype=q.dtype, device=q.device, key_len=k.shape[-2],
+                )
+            attention_mask = dflash_causal_mask
+            is_causal = False
+
+        kwargs["is_causal"] = is_causal
+
+        attn_fn: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        attn_output, attn_weights = 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,
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(bsz, q_len, -1)
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Qwen3Config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = Qwen3DFlashAttention(config=config, layer_idx=layer_idx)
+        self.mlp = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        target_hidden: Optional[torch.Tensor] = None,
+        hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[
+            Tuple[torch.Tensor, torch.Tensor]
+        ] = None,  # necessary, but kept here for BC
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states=hidden_states,
+            target_hidden=target_hidden,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )[0]
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+def build_target_layer_ids(num_target_layers: int, num_draft_layers: int):
+    if num_draft_layers == 1:
+        return [(num_target_layers // 2)]
+    start = 1
+    end = num_target_layers - 3
+    span = end - start
+    target_layer_ids = [
+        int(round(start + (i * span) / (num_draft_layers - 1)))
+        for i in range(num_draft_layers)
+    ]
+    return target_layer_ids
+
+
+def extract_context_feature(
+    hidden_states: list[torch.Tensor],
+    layer_ids: Optional[list[int]],
+) -> torch.Tensor:
+    offset = 1
+    selected_states = []
+    for layer_id in layer_ids:
+        selected_states.append(hidden_states[layer_id + offset])
+    target_hidden = torch.cat(selected_states, dim=-1)
+    return target_hidden
+
+
+class DFlashDraftModel(Qwen3PreTrainedModel):
+    config_class = Qwen3Config
+    _no_split_modules = ["Qwen3DFlashDecoderLayer"]
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.config = config
+        if not hasattr(self.config, "dflash_config") or self.config.dflash_config is None:
+            self.config.dflash_config = {}
+        self.causal_head = bool(self.config.dflash_config.get("causal_head", False))
+        self.layers = nn.ModuleList(
+            [
+                Qwen3DFlashDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        dflash_config = getattr(config, "dflash_config", {}) or {}
+        self.target_layer_ids = dflash_config.get(
+            "target_layer_ids",
+            build_target_layer_ids(config.num_target_layers, config.num_hidden_layers),
+        )
+        self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen3RotaryEmbedding(config)
+        self.fc = nn.Linear(
+            len(self.target_layer_ids) * config.hidden_size,
+            config.hidden_size,
+            bias=False,
+        )
+        self.hidden_norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.block_size = config.block_size
+        self.mask_token_id = dflash_config.get("mask_token_id", None)
+        # Optional DSpark Markov sequential head. markov_rank == 0 (default)
+        # -> markov_head is None and the model is byte-identical to plain
+        # DFlash (same params, same forward, same checkpoint).
+        self.markov_rank = int(getattr(self.config, "markov_rank", 0) or 0)
+        self.markov_head_type = getattr(self.config, "markov_head_type", "vanilla")
+        self.markov_head = None
+        if self.markov_rank > 0:
+            if build_markov_head is None:
+                raise RuntimeError(
+                    "config.markov_rank > 0 but markov_head.py could not be imported; "
+                    "ensure markov_head.py sits next to dflash.py."
+                )
+            self.markov_head = build_markov_head(
+                markov_rank=self.markov_rank,
+                vocab_size=config.vocab_size,
+                hidden_size=config.hidden_size,
+                head_type=self.markov_head_type,
+            )
+        # Optional DSpark confidence (accept-rate) head. Disabled by default
+        # (config.confidence_head missing/False) -> confidence_head is None and
+        # the model is byte-identical to the markov/plain DFlash head (same
+        # params, same forward, same checkpoint, same config.json). When enabled
+        # it is a single Linear(H [+ markov_rank], 1), matching the official
+        # DeepSpec AcceptRatePredictor (common.py:43-49). The forward contract
+        # is UNCHANGED: the confidence logit is computed by the caller
+        # (train_head / feature-cache / vLLM), never inside forward().
+        self.confidence_head = None
+        if getattr(self.config, "confidence_head", False):
+            conf_in = config.hidden_size + (
+                self.markov_rank
+                if getattr(self.config, "confidence_head_with_markov", False)
+                else 0
+            )
+            self.confidence_head = nn.Linear(conf_in, 1, bias=True)
+        self.post_init()
+
+    def resolve_causal_head(self, head_type: str = "auto") -> bool:
+        if head_type == "auto":
+            return bool(self.causal_head)
+        if head_type == "bidirectional":
+            return False
+        if head_type == "causal":
+            return True
+        raise ValueError(
+            f"Unsupported head_type={head_type!r}. Expected one of: auto, bidirectional, causal."
+        )
+
+    def forward(
+        self,
+        position_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        noise_embedding: Optional[torch.Tensor] = None,
+        target_hidden: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden_states = noise_embedding
+        target_hidden = self.hidden_norm(self.fc(target_hidden))
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states=hidden_states,
+                target_hidden=target_hidden,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                use_cache=use_cache,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+        return self.norm(hidden_states)
+
+    @torch.inference_mode()
+    def spec_generate(
+        self,
+        target: nn.Module,
+        input_ids: torch.LongTensor,
+        max_new_tokens: int,
+        stop_token_ids: list[int],
+        temperature: float,
+    ):
+        self.eval()
+        num_input_tokens = input_ids.shape[1]
+        max_length = num_input_tokens + max_new_tokens
+
+        block_size = self.block_size
+        output_ids = torch.full(
+            (1, max_length + block_size),
+            self.mask_token_id,
+            dtype=torch.long,
+            device=target.device,
+        )
+        position_ids = torch.arange(
+            output_ids.shape[1], device=target.device
+        ).unsqueeze(0)
+
+        past_key_values_target = DynamicCache()
+        past_key_values_draft = DynamicCache()
+
+        # Prefill stage
+        output = target(
+            input_ids,
+            position_ids=position_ids[:, :num_input_tokens],
+            past_key_values=past_key_values_target,
+            use_cache=True,
+            logits_to_keep=1,
+            output_hidden_states=True,
+        )
+
+        output_ids[:, :num_input_tokens] = input_ids
+        output_ids[:, num_input_tokens : num_input_tokens + 1] = sample(
+            output.logits, temperature
+        )
+        target_hidden = extract_context_feature(
+            output.hidden_states, self.target_layer_ids
+        )
+
+        # Decode stage
+        acceptance_lengths = []
+        start = input_ids.shape[1]
+        while start < max_length:
+            block_output_ids = output_ids[:, start : start + block_size].clone()
+            block_position_ids = position_ids[:, start : start + block_size]
+            noise_embedding = target.model.embed_tokens(block_output_ids)
+            draft_logits = target.lm_head(
+                self(
+                    target_hidden=target_hidden,
+                    noise_embedding=noise_embedding,
+                    position_ids=position_ids[
+                        :, past_key_values_draft.get_seq_length() : start + block_size
+                    ],
+                    past_key_values=past_key_values_draft,
+                    use_cache=True,
+                    is_causal=False,
+                )[:, -block_size + 1 :, :]
+            )
+            past_key_values_draft.crop(start)
+            block_output_ids[:, 1:] = sample(draft_logits)
+            output_ids[:, start + 1 : start + block_size] = block_output_ids[:, 1:]
+
+            output = target(
+                block_output_ids,
+                position_ids=block_position_ids,
+                past_key_values=past_key_values_target,
+                use_cache=True,
+                output_hidden_states=True,
+            )
+
+            posterior = sample(output.logits, temperature)
+            acceptance_length = (
+                (block_output_ids[:, 1:] == posterior[:, :-1])
+                .cumprod(dim=1)
+                .sum(dim=1)[0]
+                .item()
+            )
+            output_ids[:, start : start + acceptance_length + 1] = block_output_ids[
+                :, : acceptance_length + 1
+            ]
+            output_ids[:, start + acceptance_length + 1] = posterior[
+                :, acceptance_length
+            ]
+            start += acceptance_length + 1
+            past_key_values_target.crop(start)
+            target_hidden = extract_context_feature(
+                output.hidden_states, self.target_layer_ids
+            )[:, : acceptance_length + 1, :]
+            acceptance_lengths.append(acceptance_length + 1)
+            if stop_token_ids is not None and any(
+                stop_token_id in output_ids[:, num_input_tokens:]
+                for stop_token_id in stop_token_ids
+            ):
+                break
+        output_ids = output_ids[:, :max_length]
+        output_ids = output_ids[:, output_ids[0] != self.mask_token_id]
+        if stop_token_ids is not None:
+            stop_token_ids = torch.tensor(stop_token_ids, device=output_ids.device)
+            stop_token_indices = torch.isin(
+                output_ids[0][num_input_tokens:], stop_token_ids
+            ).nonzero(as_tuple=True)[0]
+            if stop_token_indices.numel() > 0:
+                output_ids = output_ids[
+                    :, : num_input_tokens + stop_token_indices[0] + 1
+                ]
+
+        return output_ids

markov_head.py

@@ -0,0 +1,229 @@
+"""DSpark-style low-rank Markov sequential head (vendored).
+
+=============================================================================
+ATTRIBUTION
+=============================================================================
+Adapted from DeepSpec (https://github.com/deepseek-ai/DeepSpec),
+file ``deepspec/modeling/dspark/markov_head.py`` (DSpark / DeepSeek-V4 draft
+head).  DeepSpec is released under the MIT License, Copyright (c) 2026 The
+DeepSpec Authors.  Only the *VanillaMarkov* head is vendored here (the
++16-18% accepted-length default that ships in DeepSeek-V4); the gated / RNN
+variants are intentionally omitted to keep the surface minimal.
+
+=============================================================================
+WHAT THIS IS
+=============================================================================
+A parallel block-drafter (our DFlashDraftModel) predicts every block position
+in ONE forward from mask-token inputs, so position k cannot see what was
+actually sampled at position k-1 -- this is the "suffix decay" we measured
+([72,57,45,35,27,22,19,16] top-1 by position).
+
+The Markov head fixes that *cheaply* by adding a per-position logit bias that
+conditions on the previous token only:
+
+    B(x_{k-1}, :) = W2( W1[x_{k-1}] )        W1 in R^{V x r},  W2 in R^{r x V}
+
+The corrected logit for position k is  U_k + B(x_{k-1}, :)  where U_k is the
+backbone's base logit (lm_head(hidden_k)).  At TRAIN time x_{k-1} is the
+teacher-forced ground-truth predecessor (apply_block_logits); at INFERENCE
+time x_{k-1} is the actually-sampled draft token, so the block is sampled
+LEFT-TO-RIGHT (sample_block_tokens).  This is CHAIN mode (single-block verify)
+-> hybrid-safe (no per-branch SSM-state-fork tax).
+
+The head is fully self-contained: it carries its OWN W1/W2 and never touches
+the backbone's (borrowed) embed_tokens / lm_head.
+=============================================================================
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import nn
+
+
+def _sample_tokens(logits: torch.Tensor, temperature: float = 0.0) -> torch.Tensor:
+    """Greedy (temperature < 1e-5) or temperature multinomial sample.
+
+    logits: (..., vocab) -> returns (...) long token ids. Mirrors dflash.sample
+    semantics so the Markov resample matches the rest of the pipeline."""
+    if temperature is None or temperature < 1e-5:
+        return torch.argmax(logits, dim=-1)
+    *lead, vocab = logits.shape
+    flat = (logits / temperature).reshape(-1, vocab)
+    probs = torch.softmax(flat, dim=-1)
+    return torch.multinomial(probs, num_samples=1).reshape(*lead)
+
+
+class VanillaMarkov(nn.Module):
+    """Memoryless low-rank transition bias B(x_{k-1}) = W2(W1[x_{k-1}])."""
+
+    def __init__(self, *, vocab_size: int, markov_rank: int):
+        super().__init__()
+        self.vocab_size = int(vocab_size)
+        self.markov_rank = int(markov_rank)
+        self.markov_head_type = "vanilla"
+        assert self.markov_rank > 0, (
+            f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}."
+        )
+        self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank)
+        self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False)
+
+    def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor:
+        return self.markov_w1(token_ids.long())
+
+    def project_bias(self, latent_states: torch.Tensor) -> torch.Tensor:
+        return self.markov_w2(latent_states)
+
+    def compute_step_bias(
+        self,
+        token_ids: torch.Tensor,
+        hidden_states: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # vanilla head ignores hidden_states (pure function of the prev token)
+        del hidden_states
+        return self.project_bias(self.get_prev_embeddings(token_ids))
+
+    def apply_step_logits(
+        self,
+        logits: torch.Tensor,  # (B, V)
+        *,
+        token_ids: torch.Tensor,  # (B,)
+        hidden_states: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        bias = self.compute_step_bias(token_ids, hidden_states)
+        return logits + bias.to(logits.dtype)
+
+    def apply_block_logits(
+        self,
+        base_logits: torch.Tensor,  # (..., M, V)
+        *,
+        token_ids: torch.Tensor,  # (..., M)  teacher-forced prev tokens
+        hidden_states: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Train-time teacher-forced bias. Shape-agnostic over leading dims:
+        works for our (R, M, V) layout AND the DSpark (B, num_blocks, bs, V)
+        layout, since W1/W2 act only on the last dim."""
+        if base_logits.numel() == 0 or base_logits.shape[-2] == 0:
+            return base_logits
+        bias = self.compute_step_bias(token_ids, hidden_states)
+        return base_logits + bias.to(base_logits.dtype)
+
+    @torch.no_grad()
+    def sample_block_tokens(
+        self,
+        base_logits: torch.Tensor,  # (B, M, V) backbone base logits
+        *,
+        first_prev_token_ids: torch.Tensor,  # (B,) verified token before pos 0
+        hidden_states: Optional[torch.Tensor] = None,
+        temperature: float = 0.0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Inference-time LEFT-TO-RIGHT block sampling. Each position's logit is
+        biased by the token actually sampled at the previous position.
+
+        Returns (sampled_tokens (B, M), corrected_logits (B, M, V))."""
+        batch_size, proposal_len = base_logits.shape[:2]
+        if proposal_len == 0:
+            empty = torch.empty(batch_size, 0, dtype=torch.long, device=base_logits.device)
+            return empty, base_logits
+
+        sampled_tokens = []
+        corrected_logits = []
+        prev_token_ids = first_prev_token_ids.long()
+        for step_idx in range(proposal_len):
+            step_logits = self.apply_step_logits(
+                base_logits[:, step_idx, :],
+                token_ids=prev_token_ids,
+                # Forward the per-position hidden so a GatedMarkovHead actually
+                # gates during LEFT-TO-RIGHT sampling. Was hard-coded None, which
+                # silently dropped the gate (vanilla fallback) even when the
+                # caller had the hidden -> offline accept sims measured gated as
+                # vanilla. None-passing callers keep the vanilla path unchanged.
+                hidden_states=(
+                    hidden_states[:, step_idx, :]
+                    if hidden_states is not None
+                    else None
+                ),
+            )
+            corrected_logits.append(step_logits.unsqueeze(1))
+            next_token_ids = _sample_tokens(step_logits, temperature=temperature)
+            sampled_tokens.append(next_token_ids)
+            prev_token_ids = next_token_ids
+        return torch.stack(sampled_tokens, dim=1), torch.cat(corrected_logits, dim=1)
+
+
+class GatedMarkovHead(VanillaMarkov):
+    """Gated DSpark Markov head (official DeepSpec GatedMarkovHead).
+
+    Uses a sigmoid gate conditioned on [hidden_state; prev_embedding] to
+    modulate the markov bias. Unlike VanillaMarkov which ignores hidden_states,
+    GatedMarkovHead uses the backbone hidden state to adaptively gate the
+    bigram bias -- stronger when the backbone is uncertain, weaker when it's
+    confident. This should help with the serve pos0 gap we observed.
+    """
+
+    def __init__(self, *, vocab_size: int, markov_rank: int, hidden_size: int):
+        super().__init__(vocab_size=vocab_size, markov_rank=markov_rank)
+        self.markov_head_type = "gated"
+        self.gate_proj = nn.Linear(hidden_size + markov_rank, markov_rank)
+
+    def compute_gate(
+        self,
+        token_ids: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        prev_embeddings = self.get_prev_embeddings(token_ids)
+        # Defensive dtype align: nn.Linear (gate_proj) requires its input in the
+        # weight dtype, and torch.cat requires both operands to share a dtype.
+        # The caller may hand us hidden_states in a dtype that differs from the
+        # head params (e.g. a float32 --head-dtype head fed a bf16 forward hidden,
+        # or the serve-side bf16 gate fed an fp32 draft hidden). Cast BOTH cat
+        # operands to gate_proj.weight.dtype so neither the concat nor the matmul
+        # can raise a dtype mismatch. This is a no-op in the intended paths (whole
+        # head is a single dtype), so it changes no numerics.
+        w_dtype = self.gate_proj.weight.dtype
+        gate_inputs = torch.cat(
+            [hidden_states.to(w_dtype), prev_embeddings.to(w_dtype)], dim=-1
+        )
+        return torch.sigmoid(self.gate_proj(gate_inputs))
+
+    def compute_step_bias(
+        self,
+        token_ids: torch.Tensor,
+        hidden_states: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if hidden_states is None:
+            # Fallback to vanilla when hidden_states not available (e.g., offline eval)
+            return self.project_bias(self.get_prev_embeddings(token_ids))
+        prev_embeddings = self.get_prev_embeddings(token_ids)
+        gate = self.compute_gate(token_ids, hidden_states).to(dtype=prev_embeddings.dtype)
+        return self.project_bias(gate * prev_embeddings)
+
+
+def build_markov_head(
+    *,
+    markov_rank: int,
+    vocab_size: int,
+    hidden_size: Optional[int] = None,
+    head_type: str = "vanilla",
+) -> Optional[nn.Module]:
+    """Return a Markov head, or None when markov_rank == 0 (head disabled)."""
+    markov_rank = int(markov_rank)
+    assert markov_rank >= 0, f"markov_rank must be >= 0, got {markov_rank}"
+    if markov_rank == 0:
+        return None
+    head_type = str(head_type).lower()
+    if head_type == "vanilla":
+        return VanillaMarkov(vocab_size=vocab_size, markov_rank=markov_rank)
+    if head_type == "gated":
+        assert hidden_size is not None, "GatedMarkovHead requires hidden_size"
+        return GatedMarkovHead(
+            vocab_size=vocab_size, markov_rank=markov_rank, hidden_size=hidden_size
+        )
+    raise ValueError(
+        f"Unsupported markov_head_type={head_type!r}; only 'vanilla' and 'gated' are vendored."
+    )
+
+
+__all__ = ["VanillaMarkov", "build_markov_head"]

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e201370eb0b0fe9a5256a5a2c6ed600503e81e4f1bdad3c6a0ac5dea84fc544
+size 3714712368

vllm_patches/llm_base_proposer.py

@@ -0,0 +1,1907 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
+from importlib.util import find_spec
+from typing import Any, cast
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from vllm.config import (
+    CUDAGraphMode,
+    VllmConfig,
+    get_layers_from_vllm_config,
+    replace,
+)
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.forward_context import set_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.model_loader import get_model
+from vllm.model_executor.models import supports_multimodal
+from vllm.model_executor.models.deepseek_eagle3 import Eagle3DeepseekV2ForCausalLM
+from vllm.model_executor.models.interfaces import SupportsMultiModal
+from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
+from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.platforms import current_platform
+from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
+from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
+from vllm.v1.kv_cache_interface import KVCacheConfig, UniformTypeKVCacheSpecs
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.ops.topk_topp_sampler import (
+    empty_exponential_noise_like,
+    sample_with_exponential_noise,
+)
+from vllm.v1.sample.sampler import _SAMPLING_EPS
+from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
+from vllm.v1.spec_decode.utils import (
+    PADDING_SLOT_ID,
+    compute_new_slot_mapping,
+    copy_and_expand_eagle_inputs_kernel,
+    eagle_prepare_inputs_padded_kernel,
+    eagle_prepare_next_token_padded_kernel,
+    eagle_step_update_slot_mapping_and_metadata,
+    extend_all_queries_by_N,
+    next_power_of_2,
+)
+from vllm.v1.utils import CpuGpuBuffer
+from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
+from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
+from vllm.v1.worker.utils import AttentionGroup
+
+logger = init_logger(__name__)
+
+
+class SpecDecodeBaseProposer:
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        device: torch.device,
+        pass_hidden_states_to_model: bool,
+        runner=None,
+    ):
+        self.vllm_config = vllm_config
+        assert vllm_config.speculative_config is not None
+        self.speculative_config = vllm_config.speculative_config
+        self.draft_model_config = self.speculative_config.draft_model_config
+        self.method = self.speculative_config.method
+        self.pass_hidden_states_to_model = pass_hidden_states_to_model
+        self._share_mtp_indices = False
+
+        self.device = device
+        self.dtype = vllm_config.model_config.dtype
+        self.max_model_len = vllm_config.model_config.max_model_len
+        self.dp_rank = vllm_config.parallel_config.data_parallel_rank
+        self.num_speculative_tokens = self.speculative_config.num_speculative_tokens
+
+        # We need to get the hidden size from the draft model config because
+        # the draft model's hidden size can be different from the target model's
+        # hidden size (e.g., Llama 3.3 70B).
+        self.hidden_size = self.draft_model_config.get_hidden_size()
+        self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size()
+
+        # DeepSeek V4 MTP consumes the target's pre-hc_head residual stream,
+        # shape (T, hc_mult * hidden_size). Expand the hidden_states buffer
+        # so target_hidden_states fits; detect DeepseekV4 via draft hf_config.
+        draft_hf_config = self.draft_model_config.hf_config
+        if hasattr(draft_hf_config, "compress_ratios") and hasattr(
+            draft_hf_config, "hc_mult"
+        ):
+            self.hidden_size = self.hidden_size * draft_hf_config.hc_mult
+
+        # Unifying eagle, draft model, and parallel drafting support.
+        # DFlash always uses parallel drafting (all tokens in one pass),
+        # but has an additional slot for the next_token_id (does not shift like EAGLE)
+        self.parallel_drafting: bool = self.speculative_config.parallel_drafting
+        self.extra_slots_per_request = (
+            1 if not self.parallel_drafting else self.num_speculative_tokens
+        )
+        self.net_num_new_slots_per_request = self.extra_slots_per_request - (
+            1 if (self.pass_hidden_states_to_model and self.method != "dflash") else 0
+        )
+        self.needs_extra_input_slots = self.net_num_new_slots_per_request > 0
+
+        # When True, all draft steps reuse the same position as the
+        # first step instead of advancing by one each iteration.
+        # Used by draft models with Q-only attention that share KV
+        # with the target and always predict from the same position.
+        self.constant_draft_positions: bool = False
+
+        self.parallel_drafting_token_id: int = 0
+        self.parallel_drafting_hidden_state_tensor: torch.Tensor | None = None
+        if self.parallel_drafting:
+            self._init_parallel_drafting_params()
+        self.use_local_argmax_reduction: bool = (
+            self.speculative_config.use_local_argmax_reduction
+        )
+        self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel
+
+        self.max_batch_size = vllm_config.scheduler_config.max_num_seqs
+        self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        self.token_arange_np = np.arange(self.max_num_tokens, dtype=np.int32)
+
+        # Can be specialized by methods like DFlash to reduce the limit
+        self.max_query_tokens = self.max_num_tokens
+        self.max_positions = self.max_num_tokens
+
+        # Multi-modal data support
+        self.mm_registry = MULTIMODAL_REGISTRY
+        self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
+            vllm_config.model_config
+        )
+
+        self.draft_attn_groups: list[AttentionGroup] = []
+        self.kv_cache_gid: int = -1
+        self.eagle3_use_aux_hidden_state: bool = (
+            self._get_eagle3_use_aux_hidden_state_from_config()
+        )
+
+        self.compilation_config = self.vllm_config.compilation_config
+
+        # Cudagraph dispatcher for PIECEWISE-only dispatching in eagle.
+        # Keys are initialized later via initialize_cudagraph_keys() called from
+        # gpu_model_runner._check_and_update_cudagraph_mode after
+        # adjust_cudagraph_sizes_for_spec_decode is called.
+        self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)
+
+        # persistent buffers for cuda graph
+        self.input_ids = torch.zeros(
+            self.max_num_tokens, dtype=torch.int32, device=device
+        )
+        # Use draft model's M-RoPE setting, not target model's
+        # Draft models may be text-only even if target is multimodal
+        self.uses_mrope = self.draft_model_config.uses_mrope
+        self.uses_xdrope_dim = self.vllm_config.model_config.uses_xdrope_dim
+        self.draft_uses_xdrope_dim = self.draft_model_config.uses_xdrope_dim
+        if self.uses_mrope:
+            # NOTE: `mrope_positions` is implemented with one additional dummy
+            # position on purpose to make it non-contiguous so that it can work
+            # with torch compile.
+            # See detailed explanation in https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923
+
+            # NOTE: When M-RoPE is enabled, position ids are 3D regardless of
+            # the modality of inputs. For text-only inputs, each dimension has
+            # identical position IDs, making M-RoPE functionally equivalent to
+            # 1D-RoPE.
+            # See page 5 of https://arxiv.org/abs/2409.12191
+            self.mrope_positions = torch.zeros(
+                (3, self.max_positions + 1), dtype=torch.int64, device=device
+            )
+        elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
+            self.xdrope_positions = torch.zeros(
+                (self.uses_xdrope_dim, self.max_positions + 1),
+                dtype=torch.int64,
+                device=device,
+            )
+        else:
+            # RoPE need (max_num_tokens,)
+            self.positions = torch.zeros(
+                self.max_positions,
+                dtype=torch.int64,
+                device=device,
+            )
+        self.hidden_states = torch.zeros(
+            (self.max_num_tokens, self.hidden_size), dtype=self.dtype, device=device
+        )
+
+        # Will be set when we initialize the attention backend
+        self.block_size: int = -1
+
+        # We need +1 here because the arange is used to set query_start_loc,
+        # which has one more element than batch_size.
+        max_num_slots_for_arange = max(self.max_batch_size + 1, self.max_num_tokens)
+        self.arange = torch.arange(
+            max_num_slots_for_arange, device=device, dtype=torch.int32
+        )
+
+        if self.needs_extra_input_slots:
+            self._raise_if_padded_drafter_batch_disabled()
+            self._warn_if_multimodal()
+            self._raise_if_mrope()
+
+        self.is_rejected_token_mask: torch.Tensor | None = None
+        self.is_masked_token_mask: torch.Tensor | None = None
+        if self.needs_extra_input_slots:
+            # For draft models and parallel drafting, we need to keep track of
+            # which tokens are rejected to update the slot mapping with padding slots.
+            self.is_rejected_token_mask = torch.zeros(
+                (self.max_num_tokens,), dtype=torch.bool, device=device
+            )
+            # For parallel drafting, we also need to keep track of which tokens
+            # are parallel-padding tokens used to sample at later positions.
+            # We populate this tensor even when using draft models for simplicity.
+            self.is_masked_token_mask = torch.zeros(
+                (self.max_num_tokens,), dtype=torch.bool, device=device
+            )
+
+        self.inputs_embeds = torch.zeros(
+            (self.max_num_tokens, self.inputs_embeds_size),
+            dtype=self.dtype,
+            device=device,
+        )
+
+        self.backup_next_token_ids = CpuGpuBuffer(
+            self.max_batch_size,
+            dtype=torch.int32,
+            pin_memory=is_pin_memory_available(),
+            device=device,
+            with_numpy=True,
+        )
+        self._enable_probabilistic_draft_probs = (
+            self.speculative_config.rejection_sample_method == "standard"
+            and self.speculative_config.draft_sample_method == "probabilistic"
+        )
+        self._last_draft_probs: torch.Tensor | None = None
+        # DSpark confidence-head dynamic-K threshold (serve-time knob, env var).
+        # 0.0 (default, unset) => no truncation: full num_spec block every step,
+        # i.e. byte-identical to the no-confidence behavior (official
+        # _confident_prefix_length returns block_size when threshold<=0). A
+        # positive threshold trims each draft row to its confident PREFIX.
+        try:
+            self._conf_threshold: float = float(
+                os.environ.get("DSPARK_CONF_THRESHOLD", "0.0") or 0.0
+            )
+        except (TypeError, ValueError):
+            self._conf_threshold = 0.0
+        # Per-row confident prefix lengths for the last proposed block, in
+        # input_batch row order ([B] list of ints), or None when no truncation.
+        self._last_draft_prefix_lengths: list[int] | None = None
+
+        self._slot_mapping_buffer = torch.zeros(
+            self.max_positions,
+            dtype=torch.int64,
+            device=device,
+        )
+
+        # Determine allowed attention backends once during initialization.
+        self.allowed_attn_types: tuple | None = None
+        if current_platform.is_rocm():
+            from vllm.models.deepseek_v4.amd.rocm import (
+                DeepseekV4ROCMAiterMLASparseMetadata,
+                DeepseekV4ROCMAiterSparseSWAMetadata,
+            )
+            from vllm.v1.attention.backends.mla.indexer import (
+                DeepseekV32IndexerMetadata,
+            )
+            from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
+                ROCMAiterMLASparseMetadata,
+            )
+            from vllm.v1.attention.backends.rocm_attn import RocmAttentionMetadata
+
+            rocm_types = [
+                TritonAttentionMetadata,
+                RocmAttentionMetadata,
+                ROCMAiterMLASparseMetadata,
+                DeepseekV4ROCMAiterMLASparseMetadata,
+                DeepseekV4ROCMAiterSparseSWAMetadata,
+                DeepseekV32IndexerMetadata,
+            ]
+            # ROCM_AITER_FA is an optional backend
+            # We check is_enabled() here to avoid importing the backend module during
+            # auto-discovery when VLLM_ROCM_USE_AITER=0, which would trigger aiter
+            # import and JIT compilation warnings. Explicit backend selection via
+            # attention_config still works because the backend module is loaded
+            # directly when selected, not through this auto-discovery path.
+            # Check if backend module exists to allow explicit selection
+            if find_spec(
+                AttentionBackendEnum.ROCM_AITER_FA.get_path(include_classname=False)
+            ):
+                from vllm.v1.attention.backends.rocm_aiter_fa import (
+                    AiterFlashAttentionMetadata,
+                )
+
+                rocm_types.append(AiterFlashAttentionMetadata)
+
+            # TRITON_MLA backend support for MLA models (e.g., DeepSeek)
+            from vllm.model_executor.layers.attention.mla_attention import (
+                MLACommonMetadata,
+            )
+
+            rocm_types.append(MLACommonMetadata)
+
+            # FlexAttention backend support
+            from vllm.v1.attention.backends.flex_attention import FlexAttentionMetadata
+
+            rocm_types.append(FlexAttentionMetadata)
+
+            self.allowed_attn_types = tuple(rocm_types)
+
+    def _raise_if_padded_drafter_batch_disabled(self):
+        if self.speculative_config.disable_padded_drafter_batch:
+            raise NotImplementedError(
+                "Speculative Decoding with draft models or parallel drafting only "
+                "supports padded drafter batch. Please unset "
+                "disable_padded_drafter_batch in the speculative_config."
+            )
+
+    def _warn_if_multimodal(self):
+        if self.supports_mm_inputs:
+            logger.warning(
+                "Speculative Decoding with draft models or parallel drafting "
+                "does not fully support multimodal models yet. "
+                "Proceeding with text-only speculative decoding."
+            )
+
+    def _raise_if_mrope(self):
+        if self.draft_model_config.uses_mrope:
+            raise NotImplementedError(
+                "Speculative Decoding with draft models or parallel drafting "
+                "does not support M-RoPE yet"
+            )
+
+    def _init_parallel_drafting_params(self):
+        # For parallel drafting, we need the token ID to use for masked slots
+        # And for EAGLE + parallel drafting, we need the hidden state tensor to use
+        # for those masked slots.
+
+        model_hf_config = self.draft_model_config.hf_config
+        # DFlash stores mask_token_id in dflash_config
+        dflash_config = getattr(model_hf_config, "dflash_config", None)
+        if dflash_config and "mask_token_id" in dflash_config:
+            self.parallel_drafting_token_id = dflash_config["mask_token_id"]
+        elif hasattr(model_hf_config, "pard_token"):
+            self.parallel_drafting_token_id = model_hf_config.pard_token
+        elif hasattr(model_hf_config, "ptd_token_id"):
+            self.parallel_drafting_token_id = model_hf_config.ptd_token_id
+        else:
+            raise ValueError(
+                "For parallel drafting, the draft model config must have "
+                "`pard_token`, `ptd_token_id`, or "
+                "`dflash_config.mask_token_id` specified in its config.json."
+            )
+
+        if self.pass_hidden_states_to_model:
+            self.parallel_drafting_hidden_state_tensor = torch.empty(
+                self.hidden_size, dtype=self.dtype, device=self.device
+            )
+
+    def _get_positions(self, num_tokens: int):
+        if self.uses_mrope:
+            return self.mrope_positions[:, :num_tokens]
+        if self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
+            return self.xdrope_positions[:, :num_tokens]
+        return self.positions[:num_tokens]
+
+    def _set_positions(self, num_tokens: int, positions: torch.Tensor):
+        if self.uses_mrope:
+            self.mrope_positions[:, :num_tokens] = positions
+        elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
+            self.xdrope_positions[:, :num_tokens] = positions
+        else:
+            # Convert M-RoPE positions if target model uses M-RoPE
+            # but draft doesn't, For text inputs, all M-RoPE
+            # dimensions are identical
+            if self.vllm_config.model_config.uses_mrope:
+                positions = positions[0]
+            self.positions[:num_tokens] = positions
+
+    def _get_slot_mapping(
+        self,
+        num_tokens: int,
+        slot_mapping: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Return slot_mapping dict for EAGLE layers.
+
+        If slot_mapping is provided, copies it into the buffer first.
+        """
+        if slot_mapping is not None:
+            num_actual = slot_mapping.shape[0]
+            self._slot_mapping_buffer[:num_actual].copy_(slot_mapping)
+            if num_tokens > num_actual:
+                self._slot_mapping_buffer[num_actual:num_tokens].fill_(PADDING_SLOT_ID)
+
+        view = self._slot_mapping_buffer[:num_tokens]
+        return {name: view for name in self._draft_attn_layer_names}
+
+    def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None:
+        """Initialize cudagraph dispatcher keys for the drafter.
+
+        Only supports PIECEWISE cudagraphs (via mixed_mode).
+        This should be called after adjust_cudagraph_sizes_for_spec_decode.
+        """
+        if (
+            not self.speculative_config.enforce_eager
+            and cudagraph_mode.mixed_mode()
+            in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]
+        ):
+            eagle_cudagraph_mode = CUDAGraphMode.PIECEWISE
+        else:
+            eagle_cudagraph_mode = CUDAGraphMode.NONE
+
+        self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode)
+
+    def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Greedy-sample draft tokens from hidden states."""
+        if self.use_local_argmax_reduction:
+            return self.model.get_top_tokens(hidden_states)
+        return self.model.compute_logits(hidden_states).argmax(dim=-1)
+
+    def _sample_from_logits(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if not self._enable_probabilistic_draft_probs:
+            return logits.argmax(dim=-1), None
+        if sampling_metadata.all_greedy:
+            return logits.argmax(dim=-1), None
+        return compute_probs_and_sample_next_token(
+            logits, sampling_metadata, self.use_fp64_gumbel
+        )
+
+    def _sample_draft_tokens(
+        self,
+        hidden_states: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if not self._enable_probabilistic_draft_probs or sampling_metadata.all_greedy:
+            return self._greedy_sample(hidden_states), None
+        logits = self.model.compute_logits(hidden_states)
+        return self._sample_from_logits(logits, sampling_metadata)
+
+    def take_last_draft_probs(self) -> torch.Tensor | None:
+        return self._last_draft_probs
+
+    def take_last_draft_prefix_lengths(self) -> list[int] | None:
+        """Per-row confident PREFIX lengths for the just-proposed block ([B] ints
+        in input_batch row order), or None when dynamic-K is inactive. Read by
+        the model runner right after propose(), mirroring take_last_draft_probs."""
+        return self._last_draft_prefix_lengths
+
+    def _compute_conf_prefix_lengths(
+        self,
+        sample_hidden_states: torch.Tensor,
+        draft_block: torch.Tensor,
+        first_prev_token_ids: torch.Tensor,
+    ) -> None:
+        """DSpark confidence-head dynamic-K: per-row confident PREFIX length.
+
+        Sets self._last_draft_prefix_lengths to a [B] list giving, for each row,
+        the number of LEADING draft positions whose accept-rate confidence is
+        >= threshold (== first position with sigmoid(conf) < threshold; official
+        draft_ops._confident_prefix_length). No-op (None) unless a positive
+        threshold is configured AND the draft model carries a confidence head.
+
+        The lengths are applied downstream as a PREFIX slice (drop the tail only)
+        in GPUModelRunner.take_draft_token_ids, which keeps the draft_probs
+        [row, :num_draft] slice and the ragged rejection sampler structurally
+        aligned -> lossless regardless of confidence quality.
+        """
+        self._last_draft_prefix_lengths = None
+        _dbg = bool(os.environ.get("DSPARK_DEBUG"))
+        # MEASURE mode: compute + log the head's per-position predicted accept
+        # (sigmoid) even when threshold<=0 (no truncation), to compare serve-path
+        # head calibration against the realized per-position acceptance rates that
+        # vLLM logs in its SpecDecoding metrics. Answers "is the head accurate at
+        # serve?" directly instead of trusting the offline AUROC.
+        _measure = bool(os.environ.get("DSPARK_MEASURE"))
+        if self._conf_threshold <= 0.0 and not _measure:
+            return
+        conf_head = getattr(getattr(self.model, "model", None), "confidence_head", None)
+        if conf_head is None:
+            return
+        conf_logits = self.model.predict_confidence_step(
+            sample_hidden_states,
+            draft_block,
+            first_prev_token_ids,
+            self.num_speculative_tokens,
+        )  # [B, num_spec] float32
+        probs = conf_logits.sigmoid()
+        if _measure:
+            # accumulate per-position predicted sigmoid (sum + count) and log a
+            # running per-position mean; offline-comparable to realized per-pos accept.
+            pos_sum = getattr(self, "_dspark_pos_sum", None)
+            if pos_sum is None:
+                pos_sum = probs.float().sum(0)
+                self._dspark_pos_cnt = probs.shape[0]
+            else:
+                pos_sum = pos_sum + probs.float().sum(0)
+                self._dspark_pos_cnt += probs.shape[0]
+            self._dspark_pos_sum = pos_sum
+            n = getattr(self, "_dspark_meas_n", 0)
+            if n < 40:
+                self._dspark_meas_n = n + 1
+                mean_pp = (pos_sum / self._dspark_pos_cnt).tolist()
+                logger.info(
+                    "DSPARK measure perpos_pred_accept(mean)=%s (rows=%d)",
+                    [round(x, 3) for x in mean_pp], self._dspark_pos_cnt,
+                )
+        # keep[:, k] = 1 while confident; confident PREFIX length = number of
+        # leading 1s = sum of the running cumprod (zeros out after the first 0).
+        if self._conf_threshold <= 0.0:
+            return  # measure-only: no truncation
+        keep = (probs >= self._conf_threshold).to(torch.int32)
+        lengths = torch.cumprod(keep, dim=1).sum(dim=1)  # [B] in [0, num_spec]
+        if _dbg:
+            n = getattr(self, "_dspark_dbg_n", 0)
+            if n < 8:
+                self._dspark_dbg_n = n + 1
+                p = probs.float()
+                logger.info(
+                    "DSPARK conf: thr=%.3f sigmoid[min=%.3f mean=%.3f max=%.3f] "
+                    "lengths=%s (num_spec=%d)",
+                    self._conf_threshold,
+                    p.min().item(),
+                    p.mean().item(),
+                    p.max().item(),
+                    lengths.to("cpu").tolist(),
+                    self.num_speculative_tokens,
+                )
+        self._last_draft_prefix_lengths = lengths.to(
+            "cpu", dtype=torch.int32
+        ).tolist()
+
+    def propose(
+        self,
+        # [num_tokens]
+        target_token_ids: torch.Tensor,
+        # [num_tokens] or [3, num_tokens] when M-RoPE is enabled
+        target_positions: torch.Tensor,
+        # [num_tokens, hidden_size]
+        target_hidden_states: torch.Tensor,
+        # [batch_size]
+        next_token_ids: torch.Tensor,
+        token_indices_to_sample: torch.Tensor | None,
+        common_attn_metadata: CommonAttentionMetadata,
+        sampling_metadata: SamplingMetadata,
+        mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
+        num_rejected_tokens_gpu: torch.Tensor | None = None,
+        slot_mappings: dict[str, torch.Tensor]
+        | list[dict[str, torch.Tensor]]
+        | None = None,
+    ) -> torch.Tensor:
+        self._last_draft_probs = None
+        self._last_draft_prefix_lengths = None
+        batch_size = common_attn_metadata.batch_size()
+
+        if self.method in ("eagle3", "dflash"):
+            assert isinstance(
+                self.model,
+                (
+                    Eagle3LlamaForCausalLM,
+                    Eagle3DeepseekV2ForCausalLM,
+                    DFlashQwen3ForCausalLM,
+                ),
+            )
+            target_hidden_states = self.model.combine_hidden_states(
+                target_hidden_states
+            )
+            assert target_hidden_states.shape[-1] == self.hidden_size
+
+        num_tokens, token_indices_to_sample, common_attn_metadata = (
+            self.set_inputs_first_pass(
+                target_token_ids=target_token_ids,
+                next_token_ids=next_token_ids,
+                target_positions=target_positions,
+                target_hidden_states=target_hidden_states,
+                token_indices_to_sample=token_indices_to_sample,
+                cad=common_attn_metadata,
+                num_rejected_tokens_gpu=num_rejected_tokens_gpu,
+            )
+        )
+
+        per_group_attn_metadata, per_layer_attn_metadata = (
+            self.build_per_group_and_layer_attn_metadata(common_attn_metadata)
+        )
+
+        cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
+            self._determine_batch_execution_and_padding(num_tokens)
+        )
+
+        model_kwargs, slot_mapping_size = self.build_model_inputs_first_pass(
+            num_tokens, num_input_tokens, mm_embed_inputs
+        )
+        # Step 0 of index_share_for_mtp_iteration: let the MTP layer
+        # compute its own indices (skip_topk=False) so subsequent steps
+        # can reuse them.
+        if self._share_mtp_indices and hasattr(self.model.model, "set_skip_topk"):
+            self.model.model.set_skip_topk(False)
+
+        with set_forward_context(
+            per_layer_attn_metadata,
+            self.vllm_config,
+            num_tokens=num_input_tokens,
+            num_tokens_across_dp=num_tokens_across_dp,
+            cudagraph_runtime_mode=cudagraph_runtime_mode,
+            slot_mapping=self._get_slot_mapping(
+                slot_mapping_size, common_attn_metadata.slot_mapping
+            ),
+        ):
+            ret_hidden_states = self.model(**model_kwargs)
+            if not self.model_returns_tuple():
+                last_hidden_states = ret_hidden_states
+                hidden_states = last_hidden_states
+            else:
+                last_hidden_states, hidden_states = ret_hidden_states
+
+        # After step 0: switch to reuse mode so steps 1+ skip the indexer
+        # and read the indices that step 0 just wrote into the shared buffer.
+        if self._share_mtp_indices and hasattr(self.model.model, "set_skip_topk"):
+            self.model.model.set_skip_topk(True)
+
+        sample_hidden_states = last_hidden_states[token_indices_to_sample]
+
+        # Early exit if there is only one draft token to be generated.
+        if self.num_speculative_tokens == 1 or self.parallel_drafting:
+            # DSpark VanillaMarkov semi-AR greedy shortcut: when the draft model
+            # carries a markov head and the whole batch is greedy, draft the
+            # parallel block LEFT-TO-RIGHT with a per-position markov logit bias
+            # instead of one parallel argmax. LOSSLESS under greedy: the target
+            # verify re-checks every token vs the target argmax, independent of
+            # how the draft was produced. next_token_ids is the verified bonus
+            # token before draft pos 0 (== markov first_prev). This path
+            # intentionally takes priority over the use_local_argmax_reduction
+            # get_top_tokens branch inside _sample_draft_tokens, which would
+            # bypass full logits and drop the markov bias.
+            mk = getattr(getattr(self.model, "model", None), "markov_head", None)
+            if (
+                self.parallel_drafting
+                and self.num_speculative_tokens > 1
+                and mk is not None
+                and sampling_metadata.all_greedy
+            ):
+                # [B, num_spec], already matching the .view(-1, num_spec) contract
+                draft_block = self.model.sample_draft_block_semiar(
+                    sample_hidden_states,
+                    next_token_ids,
+                    self.num_speculative_tokens,
+                )
+                # DSpark confidence-head dynamic-K (PREFIX-ONLY truncate); no-op
+                # unless DSPARK_CONF_THRESHOLD>0. Greedy realized prev == argmax'd
+                # block (semi-AR), so confidence reads the true serve predecessor.
+                self._compute_conf_prefix_lengths(
+                    sample_hidden_states, draft_block, next_token_ids
+                )
+                return draft_block
+            # DSpark VanillaMarkov semi-AR SAMPLING shortcut (temp>0 lossless).
+            # Only taken in the probabilistic regime where draft_probs ARE
+            # consumed by the rejection sampler (self._enable_probabilistic_draft_probs)
+            # AND the batch is not all-greedy (mixed/random). The L->R markov
+            # draw reports the EXACT q_k it sampled from, so standard rejection
+            # sampling stays lossless. CRITICAL GUARD: if the probabilistic gate
+            # is off, draft_probs would be None -- a markov-SAMPLED draft without
+            # a reported q is NON-lossless -- so we must fall through to the
+            # existing greedy-argmax draft (correctly one-hot) instead. The
+            # all_greedy case is handled by the greedy markov branch above.
+            if (
+                self.parallel_drafting
+                and self.num_speculative_tokens > 1
+                and mk is not None
+                and self._enable_probabilistic_draft_probs
+                and not sampling_metadata.all_greedy
+            ):
+                sampled_tokens, corrected_probs = (
+                    self.model.sample_draft_block_semiar_sample(
+                        sample_hidden_states,
+                        next_token_ids,
+                        self.num_speculative_tokens,
+                        sampling_metadata.temperature,
+                        sampling_metadata.all_random,
+                        self.use_fp64_gumbel,
+                    )
+                )
+                # Match the parallel path's _last_draft_probs contract EXACTLY
+                # (shape [B, num_spec, V], float32, contiguous); see the
+                # draft_probs.view(...).contiguous() population below.
+                self._last_draft_probs = corrected_probs.view(
+                    -1, self.num_speculative_tokens, corrected_probs.shape[-1]
+                ).contiguous()
+                # DSpark confidence-head dynamic-K (PREFIX-ONLY truncate); no-op
+                # unless DSPARK_CONF_THRESHOLD>0. Realized prev == the L->R SAMPLED
+                # block, so confidence reads the true serve predecessor. PREFIX
+                # truncation stays lossless: q_0..q_{L-1} in _last_draft_probs are
+                # unchanged and the verifier slices draft_probs[row, :num_draft].
+                self._compute_conf_prefix_lengths(
+                    sample_hidden_states, sampled_tokens, next_token_ids
+                )
+                return sampled_tokens.view(-1, self.num_speculative_tokens)
+            draft_token_ids, draft_probs = self._sample_draft_tokens(
+                sample_hidden_states, sampling_metadata
+            )
+            if draft_probs is not None:
+                self._last_draft_probs = draft_probs.view(
+                    -1, self.num_speculative_tokens, draft_probs.shape[-1]
+                ).contiguous()
+            return draft_token_ids.view(-1, self.num_speculative_tokens)
+
+        if self.uses_mrope:
+            positions = self.mrope_positions[:, token_indices_to_sample]
+        else:
+            positions = self.positions[token_indices_to_sample]
+        hidden_states = hidden_states[token_indices_to_sample]
+
+        if self.constant_draft_positions:
+            # Write the sampling positions into the front of the
+            # positions buffer so that subsequent loop iterations
+            # (which read via _get_positions) use the correct values.
+            self.positions[:batch_size] = positions
+
+        draft_token_ids, draft_probs = self._sample_draft_tokens(
+            sample_hidden_states, sampling_metadata
+        )
+        draft_probs_list = None if draft_probs is None else [draft_probs]
+
+        if self.allowed_attn_types is not None:
+            for group_md in per_group_attn_metadata:
+                if not isinstance(group_md, self.allowed_attn_types):
+                    raise ValueError(
+                        f"Unsupported attention metadata type for speculative "
+                        "decoding with num_speculative_tokens > 1: "
+                        f"{type(group_md)}. Supported types are: "
+                        f"{self.allowed_attn_types}"
+                    )
+
+        # Generate the remaining draft tokens.
+        draft_token_ids_list = [draft_token_ids]
+
+        cudagraph_runtime_mode, input_batch_size, batch_size_across_dp = (
+            self._determine_batch_execution_and_padding(batch_size)
+        )
+
+        common_attn_metadata.num_actual_tokens = batch_size
+        common_attn_metadata.max_query_len = 1
+        common_attn_metadata.query_start_loc = self.arange[: batch_size + 1]
+        common_attn_metadata.query_start_loc_cpu = torch.from_numpy(
+            self.token_arange_np[: batch_size + 1]
+        ).clone()
+
+        # In padded drafter batch, we need to adjust the sequence lengths
+        # to remove the "padding" (i.e. rejected tokens).
+        # Only apply this adjustment when we have rejected tokens
+        # (i.e., not the first proposal).
+        if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None:
+            common_attn_metadata.seq_lens -= num_rejected_tokens_gpu
+            # Invalidate the CPU-side shadows to avoid H<>D sync.
+            common_attn_metadata._seq_lens_cpu = None
+            common_attn_metadata._num_computed_tokens_cpu = None
+
+        block_size = self.block_size
+        assert block_size > 0, "block_size has not been initialized."
+        for token_index in range(self.num_speculative_tokens - 1):
+            # Update the inputs.
+            # cast to int32 is crucial when eagle model is compiled.
+            # tensor.argmax() returns int64 by default.
+            input_ids = draft_token_ids_list[-1].int()
+
+            if not self.constant_draft_positions:
+                positions = self._update_positions_dependent_metadata(
+                    positions,
+                    common_attn_metadata,
+                    batch_size,
+                    input_batch_size,
+                    block_size,
+                )
+
+            # Rebuild attention metadata. When draft positions are constant
+            # (e.g. Gemma4 MTP), common_attn_metadata is invariant across
+            # loop iterations so we build once and reuse.
+            if not self.constant_draft_positions or token_index == 0:
+                _, per_layer_attn_metadata = (
+                    self.build_per_group_and_layer_attn_metadata(
+                        common_attn_metadata, draft_index=token_index + 1
+                    )
+                )
+
+            # copy inputs to buffer for cudagraph
+            self.input_ids[:batch_size] = input_ids
+            self.hidden_states[:batch_size] = hidden_states
+            if self.supports_mm_inputs:
+                self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids)
+
+                input_ids = None
+                inputs_embeds = self.inputs_embeds[:input_batch_size]
+            else:
+                input_ids = self.input_ids[:input_batch_size]
+                inputs_embeds = None
+
+            # Run the model.
+            model_kwargs = {
+                "input_ids": input_ids,
+                "positions": self._get_positions(input_batch_size),
+                "inputs_embeds": inputs_embeds,
+            }
+            if self.pass_hidden_states_to_model:
+                model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size]
+
+            with set_forward_context(
+                per_layer_attn_metadata,
+                self.vllm_config,
+                num_tokens=input_batch_size,
+                num_tokens_across_dp=batch_size_across_dp,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
+                slot_mapping=self._get_slot_mapping(input_batch_size),
+            ):
+                ret_hidden_states = self.model(**model_kwargs)
+                if not self.model_returns_tuple():
+                    last_hidden_states = ret_hidden_states
+                    hidden_states = ret_hidden_states
+                else:
+                    last_hidden_states, hidden_states = ret_hidden_states
+
+            hidden_states = hidden_states[:batch_size]
+            draft_token_ids, draft_probs = self._sample_draft_tokens(
+                last_hidden_states[:batch_size], sampling_metadata
+            )
+            if draft_probs is not None:
+                assert draft_probs_list is not None
+                draft_probs_list.append(draft_probs)
+            draft_token_ids_list.append(draft_token_ids)
+
+        # [batch_size, num_speculative_tokens]
+        draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
+        if draft_probs_list is not None:
+            self._last_draft_probs = torch.stack(draft_probs_list, dim=1).contiguous()
+        return draft_token_ids
+
+    def _update_positions_dependent_metadata(
+        self,
+        positions: torch.Tensor,
+        common_attn_metadata,
+        batch_size: int,
+        input_batch_size: int,
+        block_size: int,
+    ) -> torch.Tensor:
+        """Update positions, slot mappings, and sequence metadata for the
+        next draft step. Returns the updated positions tensor."""
+        positions_1d = positions[0] if self.uses_mrope else positions
+        if self.uses_mrope:
+            out_pos = self.mrope_positions[0, :batch_size]
+        elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
+            out_pos = self.xdrope_positions[0, :batch_size]
+        else:
+            out_pos = self.positions[:batch_size]
+        eagle_step_update_slot_mapping_and_metadata(
+            positions_1d=positions_1d,
+            block_table_tensor=common_attn_metadata.block_table_tensor,
+            seq_lens=common_attn_metadata.seq_lens,
+            block_size=block_size,
+            max_model_len=self.max_model_len,
+            out_clamped_positions=out_pos,
+            out_slot_mapping=self._slot_mapping_buffer[:input_batch_size],
+            input_batch_size=input_batch_size,
+        )
+        common_attn_metadata.slot_mapping = self._slot_mapping_buffer[:batch_size]
+        if self.uses_mrope:
+            self.mrope_positions[1:, :batch_size] = self.mrope_positions[0, :batch_size]
+            positions = self.mrope_positions[:, :batch_size]
+        elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
+            self.xdrope_positions[1:, :batch_size] = self.xdrope_positions[
+                0, :batch_size
+            ]
+            positions = self.xdrope_positions[0, :batch_size]
+        else:
+            positions = self.positions[:batch_size]
+        common_attn_metadata.max_seq_len = min(
+            common_attn_metadata.max_seq_len + 1,
+            self.max_model_len,
+        )
+
+        if common_attn_metadata._seq_lens_cpu is not None:
+            common_attn_metadata._seq_lens_cpu += 1
+        if common_attn_metadata._num_computed_tokens_cpu is not None:
+            common_attn_metadata._num_computed_tokens_cpu += 1
+        if common_attn_metadata.seq_lens_cpu_upper_bound is not None:
+            common_attn_metadata.seq_lens_cpu_upper_bound += 1
+
+        return positions
+
+    def set_inputs_first_pass(
+        self,
+        target_token_ids: torch.Tensor,
+        next_token_ids: torch.Tensor,
+        target_positions: torch.Tensor,
+        target_hidden_states: torch.Tensor,
+        token_indices_to_sample: torch.Tensor | None,
+        cad: CommonAttentionMetadata,
+        num_rejected_tokens_gpu: torch.Tensor | None,
+    ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:
+        if not self.needs_extra_input_slots:
+            # Default EAGLE pathway: no reshaping of input tensors needed.
+            # Simply rotate the input ids and leave the positions unchanged,
+            # Inserting the next token ids at the last slot in each request.
+            if token_indices_to_sample is None:
+                token_indices_to_sample = cad.query_start_loc[1:] - 1
+
+            num_tokens = target_token_ids.shape[0]
+            # Shift the input ids by one token.
+            # E.g., [a1, b1, b2, c1, c2, c3] -> [b1, b2, c1, c2, c3, c3]
+            self.input_ids[: num_tokens - 1] = target_token_ids[1:]
+            # Replace the last token with the next token.
+            # E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4]
+            self.input_ids[token_indices_to_sample] = next_token_ids
+
+            # copy inputs to buffer for cudagraph
+            if self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim == 0:
+                target_positions = target_positions[0]
+            self._set_positions(num_tokens, target_positions)
+
+            self.hidden_states[:num_tokens] = target_hidden_states
+
+            return num_tokens, token_indices_to_sample, cad
+        else:
+            assert self.is_rejected_token_mask is not None
+            assert self.is_masked_token_mask is not None
+            # 1.
+            # Call a custom triton kernel to copy input_ids and positions
+            # into the correct slots in the preallocated buffers self.input_ids,
+            # self.positions.
+            batch_size = cad.batch_size()
+            # Since we might have to copy a lot of data for prefills, we select the
+            # block size based on the max query length and limit to max 256 slots/block.
+            max_num_tokens_per_request = (
+                cad.max_query_len + self.net_num_new_slots_per_request
+            )
+            BLOCK_SIZE_TOKENS = min(256, next_power_of_2(max_num_tokens_per_request))
+            num_blocks = (
+                max_num_tokens_per_request + BLOCK_SIZE_TOKENS - 1
+            ) // BLOCK_SIZE_TOKENS
+            total_num_input_tokens = target_token_ids.shape[0]
+            total_num_output_tokens = total_num_input_tokens + (
+                self.net_num_new_slots_per_request * batch_size
+            )
+
+            token_indices_to_sample = torch.empty(
+                batch_size * self.extra_slots_per_request,
+                dtype=torch.int32,
+                device=self.device,
+            )
+
+            # Destination indices to write target_hidden_states into drafting buffer.
+            out_hidden_state_mapping = torch.empty(
+                total_num_input_tokens, dtype=torch.int32, device=self.device
+            )
+
+            # Kernel grid: one program per request (row)
+            grid = (batch_size, num_blocks)
+            query_start_loc = cad.query_start_loc
+            query_end_loc = cad.query_start_loc[1:] - 1
+            if num_rejected_tokens_gpu is not None:
+                query_end_loc = query_end_loc - num_rejected_tokens_gpu
+
+            copy_and_expand_eagle_inputs_kernel[grid](
+                # (Padded) Inputs from the target model
+                target_token_ids_ptr=target_token_ids,
+                target_positions_ptr=target_positions,
+                next_token_ids_ptr=next_token_ids,  # sampled tokens, one per request
+                # Outputs to the drafting buffers
+                out_input_ids_ptr=self.input_ids,
+                out_positions_ptr=self.positions,  # Doesn't support mrope for now
+                out_is_rejected_token_mask_ptr=self.is_rejected_token_mask,
+                out_is_masked_token_mask_ptr=self.is_masked_token_mask,
+                out_new_token_indices_ptr=token_indices_to_sample,
+                out_hidden_state_mapping_ptr=out_hidden_state_mapping,
+                # Input metadata
+                query_start_loc_ptr=query_start_loc,
+                query_end_loc_ptr=query_end_loc,
+                padding_token_id=0,
+                parallel_drafting_token_id=self.parallel_drafting_token_id,
+                # Sizing info
+                # Note that we can deduce batch_size for free from the grid size
+                total_input_tokens=total_num_input_tokens,
+                num_padding_slots_per_request=self.extra_slots_per_request,
+                shift_input_ids=self.pass_hidden_states_to_model,
+                BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,
+            )
+            if self.pass_hidden_states_to_model:
+                assert self.parallel_drafting_hidden_state_tensor is not None
+                self.hidden_states[out_hidden_state_mapping] = target_hidden_states
+                # Use torch.where to avoid DtoH sync from boolean indexing
+                mask = self.is_masked_token_mask[:total_num_output_tokens]
+                torch.where(
+                    mask.unsqueeze(1),
+                    self.parallel_drafting_hidden_state_tensor,
+                    self.hidden_states[:total_num_output_tokens],
+                    out=self.hidden_states[:total_num_output_tokens],
+                )
+
+            # 2.
+            # Recompute the slot mapping based on the new positions and
+            # rejection mask.
+            assert self.block_size > 0, "block_size has not been initialized."
+            new_slot_mapping = compute_new_slot_mapping(
+                cad=cad,
+                new_positions=self.positions[:total_num_output_tokens],
+                is_rejected_token_mask=self.is_rejected_token_mask[
+                    :total_num_output_tokens
+                ],
+                block_size=self.block_size,
+                num_new_tokens=self.net_num_new_slots_per_request,
+                max_model_len=self.max_model_len,
+            )
+
+            # 3. Update the common attention metadata with the new (meta)data
+            new_cad = extend_all_queries_by_N(
+                cad,
+                N=self.net_num_new_slots_per_request,
+                arange=self.arange,
+                new_slot_mapping=new_slot_mapping,
+            )
+
+            return total_num_output_tokens, token_indices_to_sample, new_cad
+
+    def build_model_inputs_first_pass(
+        self,
+        num_tokens: int,
+        num_input_tokens: int,
+        mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None,
+    ) -> tuple[dict[str, Any], int]:
+        if self.supports_mm_inputs:
+            mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)
+
+            self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(
+                self.input_ids[:num_tokens],
+                multimodal_embeddings=mm_embeds,
+                is_multimodal=is_mm_embed,
+            )
+
+            input_ids = None
+            inputs_embeds = self.inputs_embeds[:num_input_tokens]
+        else:
+            input_ids = self.input_ids[:num_input_tokens]
+            inputs_embeds = None
+
+        model_kwargs = {
+            "input_ids": input_ids,
+            "positions": self._get_positions(num_input_tokens),
+            "inputs_embeds": inputs_embeds,
+        }
+        if self.pass_hidden_states_to_model:
+            model_kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]
+
+        return model_kwargs, num_input_tokens
+
+    def build_per_group_and_layer_attn_metadata(
+        self, common_attn_metadata: CommonAttentionMetadata, draft_index: int = 0
+    ) -> tuple[list[object], dict[str, object]]:
+        per_group_attn_metadata: list[object] = []
+        per_layer_attn_metadata: dict[str, object] = {}
+        for attn_group in self.draft_attn_groups:
+            attn_metadata = attn_group.get_metadata_builder().build_for_drafting(
+                common_attn_metadata=common_attn_metadata, draft_index=draft_index
+            )
+            per_group_attn_metadata.append(attn_metadata)
+            for layer_name in attn_group.layer_names:
+                per_layer_attn_metadata[layer_name] = attn_metadata
+        return per_group_attn_metadata, per_layer_attn_metadata
+
+    def model_returns_tuple(self) -> bool:
+        return self.method not in ("mtp", "draft_model", "dflash")
+
+    def prepare_next_token_ids_cpu(
+        self,
+        sampled_token_ids: list[list[int]],
+        requests: dict[str, CachedRequestState],
+        gpu_input_batch: InputBatch,
+        num_scheduled_tokens: dict[str, int],
+    ) -> torch.Tensor:
+        """
+        This function is used to prepare the inputs for speculative decoding.
+        It calculates the next token ids for each request based on the sampled
+        token ids from the CPU. If a request has no sampled token ids (e.g.,
+        during the initial decoding steps), it falls back to using the request
+        state to get the next token id.
+        """
+        req_ids = gpu_input_batch.req_ids
+        next_token_ids: list[int] = []
+        for i, token_ids in enumerate(sampled_token_ids):
+            if token_ids:
+                # Common case.
+                next_token_id = token_ids[-1]
+            else:
+                # Partial prefill (rare case).
+                # Get the next token id from the request state.
+                req_id = req_ids[i]
+                req_state = requests[req_id]
+                seq_len = req_state.num_computed_tokens + num_scheduled_tokens[req_id]
+                next_token_id = req_state.get_token_id(seq_len)
+            next_token_ids.append(next_token_id)
+        next_token_ids = torch.tensor(
+            next_token_ids, dtype=torch.int32, device=self.input_ids.device
+        )
+        return next_token_ids
+
+    def prepare_next_token_ids_padded(
+        self,
+        sampled_token_ids: torch.Tensor,
+        requests: dict[str, CachedRequestState],
+        gpu_input_batch: InputBatch,
+        discard_request_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        This function is used to prepare the inputs for speculative decoding.
+        It calculates the next token ids and the number of valid sampled tokens
+        for each request, considering the "discarded" requests whose next token
+        is not sampled and comes from `request.get_token_id()` instead. This is denoted
+        the "backup" token id. It also counts rejected tokens via `sampled_token_ids`.
+        """
+        # Precompute backup token IDs for discarded requests.
+        num_reqs = gpu_input_batch.num_reqs
+        for i in range(num_reqs):
+            self.backup_next_token_ids.np[i] = requests[
+                gpu_input_batch.req_ids[i]
+            ].get_token_id(gpu_input_batch.num_tokens_no_spec[i] - 1)
+        self.backup_next_token_ids.copy_to_gpu(num_reqs)
+        backup_tokens_gpu = self.backup_next_token_ids.gpu
+
+        batch_size, num_tokens = sampled_token_ids.shape
+        device = sampled_token_ids.device
+
+        assert discard_request_mask.dtype == torch.bool
+        assert backup_tokens_gpu.dtype == torch.int32
+
+        next_token_ids = torch.empty(batch_size, dtype=torch.int32, device=device)
+        valid_sampled_tokens_count = next_token_ids.new_empty(batch_size)
+
+        # Kernel grid: one program per request (row)
+        grid = (batch_size,)
+
+        # Find the next power of 2 for block sizes
+        BLOCK_SIZE_TOKENS = next_power_of_2(num_tokens)
+        eagle_prepare_next_token_padded_kernel[grid](
+            sampled_token_ids,
+            discard_request_mask,
+            backup_tokens_gpu,
+            next_token_ids,
+            valid_sampled_tokens_count,
+            gpu_input_batch.vocab_size,
+            num_tokens,
+            batch_size,
+            sampled_token_ids.stride(0),
+            BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,
+        )
+
+        return next_token_ids, valid_sampled_tokens_count
+
+    def prepare_inputs_padded(
+        self,
+        common_attn_metadata: CommonAttentionMetadata,
+        spec_decode_metadata: SpecDecodeMetadata,
+        valid_sampled_tokens_count: torch.Tensor,
+    ) -> tuple[CommonAttentionMetadata, torch.Tensor, torch.Tensor]:
+        """
+        This function is used to prepare the inputs for speculative decoding
+        It updates the common_attn_metadata for speculative decoding,
+        but does not consider the rejected tokens. Instead, all tokens
+        are included as inputs to the speculator, with the rejected tokens
+        used as padding and filtered out later by `token_indices_to_sample`.
+        No blocking CPU operations should be introduced in this function.
+        """
+        num_reqs = common_attn_metadata.num_reqs
+        device = valid_sampled_tokens_count.device
+
+        token_indices_to_sample = torch.empty(
+            (num_reqs,), dtype=torch.int32, device=device
+        )
+        num_rejected_tokens_gpu = torch.empty(
+            (num_reqs,), dtype=torch.int32, device=device
+        )
+
+        grid = (num_reqs,)
+        eagle_prepare_inputs_padded_kernel[grid](
+            spec_decode_metadata.cu_num_draft_tokens,
+            valid_sampled_tokens_count,
+            common_attn_metadata.query_start_loc,
+            token_indices_to_sample,
+            num_rejected_tokens_gpu,
+            num_reqs,
+        )
+
+        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
+        new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
+
+        total_num_tokens = query_start_loc_cpu[-1].item()
+
+        spec_common_attn_metadata = CommonAttentionMetadata(
+            query_start_loc=common_attn_metadata.query_start_loc,
+            seq_lens=common_attn_metadata.seq_lens,
+            query_start_loc_cpu=query_start_loc_cpu,
+            _seq_lens_cpu=common_attn_metadata._seq_lens_cpu,
+            _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
+            seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound,
+            num_reqs=common_attn_metadata.num_reqs,
+            num_actual_tokens=total_num_tokens,
+            max_query_len=new_query_len_per_req.max().item(),
+            max_seq_len=common_attn_metadata.max_seq_len,
+            block_table_tensor=common_attn_metadata.block_table_tensor,
+            slot_mapping=common_attn_metadata.slot_mapping[:total_num_tokens],
+            causal=True,
+            dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,
+        )
+
+        return (
+            spec_common_attn_metadata,
+            token_indices_to_sample,
+            num_rejected_tokens_gpu,
+        )
+
+    def prepare_inputs(
+        self,
+        common_attn_metadata: CommonAttentionMetadata,
+        sampled_token_ids: list[list[int]],
+        num_draft_tokens: list[int],
+    ) -> tuple[CommonAttentionMetadata, torch.Tensor]:
+        """
+        This function is used to prepare the inputs for speculative decoding.
+        It updates to the common_attn_metadata to account for the rejected
+        tokens (and newly sampled tokens). It also returns the token indices
+        of the tokens that should be fed to the speculator.
+        """
+        # E.g.
+        #  common_attn_metadata.query_start_loc{_cpu}:
+        #       [0, q1, q1 + q2, q1 + q2 + q3]
+        #  common_attn_metadata.seq_lens{_cpu}: [s1, s2, s3]
+        #  num_rejected_tokens: [n1, n2, n3]
+        # This function computes the intermediate values:
+        #  num_tokens_per_req: [q1 - n1, q2 - n2, q3 - n3]
+        # And returns:
+        #  common_attn_metadata.query_start_loc{_cpu}:
+        #       [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]
+        #  common_attn_metadata.seq_lens{_cpu}:
+        #       [s1 - n1 + 1, s2 - n2 + 1, s3 - n3 + 1]
+        #  token_indices: [0, 1, ..., q1 - n1 - 1,
+        #                 q1, q1 + 1, ..., q1 + q2 - n2 - 1,
+        #                 q1 + q2, q1 + q2 + 1, ..., q1 + q2 + q3 - n3 - 1]
+
+        num_rejected_tokens = [
+            n + 1 - len(sampled_token_ids[i]) if n > 0 else 0
+            for i, n in enumerate(num_draft_tokens)
+        ]
+        num_rejected_tokens = torch.tensor(num_rejected_tokens, dtype=torch.int32)
+
+        device = common_attn_metadata.query_start_loc.device
+        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
+        # upper_bound - rejected = actual post-rejection seq_lens (no D2H sync).
+        assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
+        new_seq_lens_cpu = (
+            common_attn_metadata.seq_lens_cpu_upper_bound - num_rejected_tokens
+        )
+
+        # [0, q1, q1 + q2, q1 + q2 + q3] -> [q1, q2, q3]
+        new_query_len_per_req = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
+        # [q1, q2, q3] -> [q1 - n1, q2 - n2, q3 - n3]
+        new_num_tokens_per_req = new_query_len_per_req - num_rejected_tokens
+        new_num_tokens_per_req_np = new_num_tokens_per_req.numpy()
+
+        # [q1 - n1, q2 - n2, q3 - n3] ->
+        # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]
+        new_query_start_loc_cpu = torch.zeros(
+            query_start_loc_cpu.shape,
+            dtype=torch.int32,
+            pin_memory=is_pin_memory_available(),
+        )
+        new_query_start_loc_np = new_query_start_loc_cpu.numpy()
+        np.cumsum(new_num_tokens_per_req_np, out=new_query_start_loc_np[1:])
+
+        total_num_tokens = new_query_start_loc_np[-1]
+        # Example assuming num_tokens_per_req_np = [2, 4, 3]
+        # this implies that `new_query_start_locs` is:
+        # [0, 2, 6, 9] ->
+        # [0, 0, 2, 2, 2, 2, 6, 6, 6]
+        #  _r1_  ____r2____  ___r3__
+        new_query_start_locs_expanded = np.repeat(
+            new_query_start_loc_np[:-1], new_num_tokens_per_req_np
+        )
+        # [0, 1, 2, 3, 4, 5, 6, 7, 8] ->
+        # [0, 1, 0, 1, 2, 3, 0, 1, 2]
+        #  _r1_  ____r2____  ___r3__
+        token_offsets = (
+            self.token_arange_np[:total_num_tokens] - new_query_start_locs_expanded
+        )
+
+        # Expand starting positions to match token pattern
+        # [0, q1, q1 + q2] ->
+        # [0, 0, q1, q1, q1, q1, q1 + q2, q1 + q2, q1 + q2]
+        #  _r1_  _____r2_______  ___________r3____________
+        old_query_start_locs_expanded = np.repeat(
+            query_start_loc_cpu[:-1].numpy(), new_num_tokens_per_req_np
+        )
+        # Final token indices are:
+        # [0, 1,                                // req 1
+        #  q1 + 0, q1 + 1, q1 + 2, q1 + 3,       // req 2
+        #  q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3
+        token_indices_np = token_offsets + old_query_start_locs_expanded
+        token_indices = torch.from_numpy(token_indices_np).to(device, non_blocking=True)
+
+        spec_common_attn_metadata = CommonAttentionMetadata(
+            query_start_loc=new_query_start_loc_cpu.to(device, non_blocking=True),
+            seq_lens=new_seq_lens_cpu.to(device, non_blocking=True),
+            query_start_loc_cpu=new_query_start_loc_cpu,
+            _seq_lens_cpu=new_seq_lens_cpu,
+            _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
+            seq_lens_cpu_upper_bound=new_seq_lens_cpu,
+            num_reqs=common_attn_metadata.num_reqs,
+            num_actual_tokens=total_num_tokens,
+            max_query_len=new_query_len_per_req.max().item(),
+            max_seq_len=new_seq_lens_cpu.max().item(),
+            block_table_tensor=common_attn_metadata.block_table_tensor,
+            slot_mapping=common_attn_metadata.slot_mapping[token_indices],
+            causal=True,
+            dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,
+        )
+
+        return spec_common_attn_metadata, token_indices
+
+    def get_model_name(self, model: nn.Module) -> str:
+        if hasattr(model, "module"):  # multi-GPU
+            model = model.module
+        return model.__class__.__name__
+
+    def _create_draft_vllm_config(self) -> VllmConfig:
+        """Return a VllmConfig with kernel-level overrides for the proposer.
+        Subclasses may override to apply additional config changes.
+        """
+        spec_cfg = self.speculative_config
+        base = self.vllm_config
+
+        if spec_cfg.moe_backend is not None:
+            base = replace(
+                base,
+                kernel_config=replace(
+                    base.kernel_config,
+                    moe_backend=spec_cfg.moe_backend,
+                ),
+            )
+
+        # Note (matt): Never inherit the attention backend from base, because there are
+        # many opportunities for incompatibility, so we always independently autoselect
+        # unless explicitly specified in the speculative config.
+        base = replace(
+            base,
+            attention_config=replace(
+                base.attention_config,
+                backend=spec_cfg.attention_backend,
+            ),
+        )
+
+        return base
+
+    def _get_model(self) -> nn.Module:
+        """
+        Default method to call get_model(). Can be overridden by subclasses which
+        need to customize model loading.
+        """
+        from vllm.compilation.backends import set_model_tag
+
+        draft_vllm_config = self._create_draft_vllm_config()
+        with set_model_tag("eagle_head"):
+            model = get_model(
+                vllm_config=draft_vllm_config,
+                model_config=self.speculative_config.draft_model_config,
+                load_config=self.speculative_config.draft_load_config,
+            )
+        return model
+
+    def load_model(self, target_model: nn.Module) -> None:
+        target_attn_layer_names = set(
+            get_layers_from_vllm_config(
+                self.vllm_config,
+                AttentionLayerBase,  # type: ignore[type-abstract]
+            ).keys()
+        )
+
+        self.model = self._get_model()
+
+        # Find draft layers (attention layers added by draft model)
+        all_attn_layers = get_layers_from_vllm_config(
+            self.vllm_config,
+            AttentionLayerBase,  # type: ignore[type-abstract]
+        )
+        # Filter to only layers that have KV cache specs.
+        self._draft_attn_layer_names = {
+            name
+            for name in (set(all_attn_layers.keys()) - target_attn_layer_names)
+            if all_attn_layers[name].get_kv_cache_spec(self.vllm_config) is not None
+        }
+
+        if self.supports_mm_inputs:
+            # Even if the target model is multimodal, we can also use
+            # text-only draft models
+            try:
+                dummy_input_ids = torch.tensor([[1]], device=self.input_ids.device)
+                self.model.embed_input_ids(dummy_input_ids, multimodal_embeddings=None)
+            except (NotImplementedError, AttributeError, TypeError):
+                logger.warning(
+                    "Draft model does not support multimodal inputs, "
+                    "falling back to text-only mode"
+                )
+                self.supports_mm_inputs = False
+
+        if supports_multimodal(target_model):
+            # handle multimodality
+            assert hasattr(target_model, "config")
+            if self.get_model_name(target_model) in [
+                "Cohere2VisionForConditionalGeneration",
+                "Exaone4_5_ForConditionalGeneration",
+                "GlmOcrForConditionalGeneration",
+                "HunYuanVLForConditionalGeneration",
+                "InternS2PreviewForConditionalGeneration",
+                "MiMoV2OmniForCausalLM",
+                "Qwen2_5_VLForConditionalGeneration",
+                "Qwen3_5ForConditionalGeneration",
+                "Qwen3_5MoeForConditionalGeneration",
+                "Qwen3VLForConditionalGeneration",
+                "Qwen3VLMoeForConditionalGeneration",
+                "Gemma4ForConditionalGeneration",
+                "Gemma4UnifiedForConditionalGeneration",
+                "Step3p7ForConditionalGeneration",
+            ]:
+                self.model.config.image_token_index = target_model.config.image_token_id
+            elif self.get_model_name(target_model) == "PixtralForConditionalGeneration":
+                self.model.config.image_token_index = (
+                    target_model.config.vision_config.image_token_id
+                )
+            elif self.get_model_name(target_model) == "KimiK25ForConditionalGeneration":
+                self.model.config.image_token_index = (
+                    target_model.config.media_placeholder_token_id
+                )
+            else:
+                self.model.config.image_token_index = (
+                    target_model.config.image_token_index
+                )
+            target_language_model = cast(
+                SupportsMultiModal, target_model
+            ).get_language_model()
+        else:
+            target_language_model = target_model
+
+        self._maybe_share_embeddings(target_language_model)
+        self._maybe_share_lm_head(target_language_model)
+
+        if (
+            self.parallel_drafting
+            and self.pass_hidden_states_to_model
+            and self.parallel_drafting_hidden_state_tensor is not None
+        ):
+            flat_mask = self.model.mask_hidden.view(-1)
+            if self.eagle3_use_aux_hidden_state:
+                # EAGLE3: mask_hidden stores all aux hidden states,
+                # project through combine_hidden_states
+                self.parallel_drafting_hidden_state_tensor.copy_(
+                    self.model.combine_hidden_states(flat_mask)
+                )
+            else:
+                self.parallel_drafting_hidden_state_tensor.copy_(flat_mask)
+
+    def _maybe_share_embeddings(self, target_language_model: nn.Module) -> None:
+        """
+        Some draft models may not have their own embedding layers, and some may
+        have a duplicate copy of the target model's embedding layers. In these cases,
+        we share the target model's embedding layers with the draft model to save
+        memory.
+        """
+        if get_pp_group().world_size == 1:
+            inner_model = getattr(target_language_model, "model", None)
+            if inner_model is None:
+                raise AttributeError("Target model does not have 'model' attribute")
+            if hasattr(inner_model, "embed_tokens"):
+                target_embed_tokens = inner_model.embed_tokens
+            elif hasattr(inner_model, "embedding"):
+                target_embed_tokens = inner_model.embedding
+            else:
+                raise AttributeError(
+                    "Target model does not have 'embed_tokens' or 'embedding' attribute"
+                )
+
+            share_embeddings = False
+            if hasattr(self.model, "has_own_embed_tokens"):
+                # EAGLE model
+                if not self.model.has_own_embed_tokens:
+                    share_embeddings = True
+                    logger.info(
+                        "Detected EAGLE model without its own embed_tokens in the"
+                        " checkpoint. Sharing target model embedding weights with the"
+                        " draft model."
+                    )
+                elif (
+                    isinstance(target_embed_tokens.weight, torch.Tensor)
+                    and isinstance(self.model.model.embed_tokens.weight, torch.Tensor)
+                    # TODO: Offload to CPU for comparison to avoid extra GPU memory
+                    # usage in CI testing environments with limited GPU memory
+                    and torch.equal(
+                        target_embed_tokens.weight.cpu(),
+                        self.model.model.embed_tokens.weight.cpu(),
+                    )
+                ):
+                    share_embeddings = True
+                    logger.info(
+                        "Detected EAGLE model with embed_tokens identical to the target"
+                        " model. Sharing target model embedding weights with the draft"
+                        " model."
+                    )
+                else:
+                    logger.info(
+                        "Detected EAGLE model with distinct embed_tokens weights. "
+                        "Keeping separate embedding weights from the target model."
+                    )
+            else:
+                # MTP model
+                share_embeddings = True
+                logger.info(
+                    "Detected MTP model. "
+                    "Sharing target model embedding weights with the draft model."
+                )
+
+            if share_embeddings:
+                if hasattr(self.model.model, "embed_tokens"):
+                    del self.model.model.embed_tokens
+                self.model.model.embed_tokens = target_embed_tokens
+        else:
+            logger.info(
+                "The draft model's vocab embedding will be loaded separately"
+                " from the target model."
+            )
+
+    def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None:
+        """
+        Some draft models may not have their own LM head, and some may have a
+        duplicate copy of the target model's LM head. In these cases, we share
+        the target model's LM head with the draft model to save memory.
+        """
+        share_lm_head = False
+        if hasattr(self.model, "has_own_lm_head"):
+            # EAGLE model
+            if not self.model.has_own_lm_head:
+                share_lm_head = True
+                logger.info(
+                    "Detected EAGLE model without its own lm_head in the checkpoint. "
+                    "Sharing target model lm_head weights with the draft model."
+                )
+            elif (
+                hasattr(target_language_model, "lm_head")
+                and hasattr(target_language_model.lm_head, "weight")
+                and hasattr(self.model.lm_head, "weight")
+                and isinstance(target_language_model.lm_head.weight, torch.Tensor)
+                and isinstance(self.model.lm_head.weight, torch.Tensor)
+                # TODO: Offload to CPU for comparison to avoid extra GPU memory
+                # usage in CI testing environments with limited GPU memory
+                and torch.equal(
+                    target_language_model.lm_head.weight.cpu(),
+                    self.model.lm_head.weight.cpu(),
+                )
+            ):
+                share_lm_head = True
+                logger.info(
+                    "Detected EAGLE model with lm_head identical to the target model. "
+                    "Sharing target model lm_head weights with the draft model."
+                )
+            else:
+                logger.info(
+                    "Detected EAGLE model with distinct lm_head weights. "
+                    "Keeping separate lm_head weights from the target model."
+                )
+        else:
+            # MTP model
+            share_lm_head = True
+            logger.info(
+                "Detected MTP model. "
+                "Sharing target model lm_head weights with the draft model."
+            )
+
+        if share_lm_head and hasattr(target_language_model, "lm_head"):
+            if hasattr(self.model, "lm_head"):
+                del self.model.lm_head
+            self.model.lm_head = target_language_model.lm_head
+
+            # MTP models call compute_logits via shared_head.head (a
+            # ParallelLMHead inside each MTP layer), not self.model.lm_head.
+            # If the checkpoint omits a copy of the lm_head weights at the
+            # MTP layer path, shared_head.head stays uninitialised and
+            # produces NaN logits. Always share it explicitly.
+            inner = getattr(self.model, "model", None)
+            layers = getattr(inner, "layers", None) if inner else None
+            if layers is not None:
+                items = layers.values() if isinstance(layers, nn.ModuleDict) else layers
+                for layer in items:
+                    sh = getattr(layer, "shared_head", None)
+                    if sh is not None and hasattr(sh, "head"):
+                        del sh.head
+                        sh.head = target_language_model.lm_head
+                        logger.info(
+                            "Shared target model lm_head with MTP shared_head.head."
+                        )
+
+        if hasattr(target_language_model.model, "topk_indices_buffer"):
+            target_buffer = target_language_model.model.topk_indices_buffer
+            if hasattr(self.model.model, "topk_indices_buffer"):
+                del self.model.model.topk_indices_buffer
+            self.model.model.topk_indices_buffer = target_buffer
+            # Also share at per-module level so that the indexer and
+            # sparse-attention backends in each MTP layer read from
+            # the target model's buffer.
+            for _, module in self.model.model.named_modules():
+                if hasattr(module, "topk_indices_buffer"):
+                    module.topk_indices_buffer = target_buffer
+            logger.info(
+                "Detected MTP model with topk_indices_buffer. "
+                "Sharing target model topk_indices_buffer with the draft model."
+            )
+
+        # Detect index_share_for_mtp_iteration: when True, the proposer
+        # toggles skip_topk so step 0 computes MTP's own indices and
+        # steps 1+ reuse them.
+        spec_config = self.vllm_config.speculative_config
+        draft_hf_config = (
+            spec_config.draft_model_config.hf_config
+            if spec_config is not None
+            else None
+        )
+        self._share_mtp_indices = getattr(
+            draft_hf_config, "index_share_for_mtp_iteration", False
+        )
+
+        if self.use_local_argmax_reduction:
+            if not hasattr(self.model, "get_top_tokens"):
+                raise ValueError(
+                    "use_local_argmax_reduction is enabled but draft model "
+                    f"{self.model.__class__.__name__} does not implement "
+                    "get_top_tokens()."
+                )
+            # Warn if draft model has vocab remapping, which forces fallback
+            # to the full-logits path (negating the optimization).
+            if (
+                hasattr(self.model, "draft_id_to_target_id")
+                and self.model.draft_id_to_target_id is not None
+            ):
+                logger.warning(
+                    "use_local_argmax_reduction is enabled but draft model "
+                    "uses draft_id_to_target_id vocab remapping. The "
+                    "optimization will be bypassed (falling back to full "
+                    "logits gather + argmax)."
+                )
+            else:
+                logger.info(
+                    "Using local argmax reduction for draft token generation "
+                    "(communication: O(2*tp_size) vs O(vocab_size))."
+                )
+
+    @torch.inference_mode()
+    def dummy_run(
+        self,
+        num_tokens: int,
+        use_cudagraphs: bool = True,
+        is_graph_capturing: bool = False,
+        slot_mappings: dict[str, torch.Tensor] | None = None,
+    ) -> None:
+        # FIXME: when using tree-based specdec, adjust number of forward-passes
+        # according to the depth of the tree.
+        only_one_forward_pass = is_graph_capturing or self.parallel_drafting
+        for fwd_idx in range(
+            1 if only_one_forward_pass else self.num_speculative_tokens
+        ):
+            if fwd_idx <= 1:
+                cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
+                    self._determine_batch_execution_and_padding(
+                        num_tokens, use_cudagraphs=use_cudagraphs
+                    )
+                )
+
+            # Make sure to use EAGLE's own buffer during cudagraph capture.
+            if (
+                self._draft_attn_layer_names
+                and slot_mappings is not None
+                and next(iter(self._draft_attn_layer_names)) in slot_mappings
+            ):
+                slot_mapping_dict = self._get_slot_mapping(num_input_tokens)
+            else:
+                slot_mapping_dict = slot_mappings or {}
+
+            with set_forward_context(
+                None,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                cudagraph_runtime_mode=cudagraph_runtime_mode,
+                slot_mapping=slot_mapping_dict,
+            ):
+                if self.supports_mm_inputs:
+                    input_ids = None
+                    inputs_embeds = self.inputs_embeds[:num_input_tokens]
+                else:
+                    input_ids = self.input_ids[:num_input_tokens]
+                    inputs_embeds = None
+
+                kwargs = dict(
+                    input_ids=input_ids,
+                    positions=self._get_positions(num_input_tokens),
+                    inputs_embeds=inputs_embeds,
+                )
+                if self.pass_hidden_states_to_model:
+                    kwargs["hidden_states"] = self.hidden_states[:num_input_tokens]
+                self.model(**kwargs)
+
+    def _get_eagle3_use_aux_hidden_state_from_config(self) -> bool:
+        """
+        Some eagle3 heads (e.g., nvidia/gpt-oss-120b-Eagle3-v2) do not use auxiliary
+        hidden states and directly uses the last layer output just like eagle1.
+        They might indicate this by setting "use_aux_hidden_state" to False
+        inside the "eagle_config" dict of their hf_config.
+        """
+        if self.method != "eagle3":
+            return False
+        # Assume that eagle3 heads use aux hidden states by default
+        use_aux_hidden_state = True
+        eagle_config = getattr(self.draft_model_config.hf_config, "eagle_config", None)
+        if eagle_config is not None:
+            use_aux_hidden_state = eagle_config.get("use_aux_hidden_state", True)
+        return use_aux_hidden_state
+
+    def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
+        """
+        Validate that all drafting layers belong to the same KVCacheGroup.
+        Need this assumption to ensure all drafting layers can use the
+        same AttentionMetadata.
+        May extend to multiple AttentionMetadata in the future.
+        """
+        kv_cache_groups: dict[str, int] = {}
+        for id, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
+            for layer_name in kv_cache_group.layer_names:
+                kv_cache_groups[layer_name] = id
+        assert (
+            len(
+                set(
+                    [
+                        kv_cache_groups[layer_name]
+                        for layer_name in self._draft_attn_layer_names
+                    ]
+                )
+            )
+            == 1
+        ), "All drafting layers should belong to the same kv cache group"
+
+    def initialize_attn_backend(
+        self,
+        kv_cache_config: KVCacheConfig,
+        kernel_block_sizes: list[int] | None = None,
+    ) -> None:
+        """
+        Initialize AttentionGroups for draft layers using kv_cache_config.
+        Called from the model runner's initialize_metadata_builders.
+        """
+        all_attn_layers = get_layers_from_vllm_config(
+            self.vllm_config,
+            AttentionLayerBase,  # type: ignore[type-abstract]
+        )
+
+        # Find which kv_cache_group the draft layers belong to
+        self.validate_same_kv_cache_group(kv_cache_config)
+        kv_cache_spec = None
+        for gid, group in enumerate(kv_cache_config.kv_cache_groups):
+            if self._draft_attn_layer_names & set(group.layer_names):
+                self.kv_cache_gid = gid
+                kv_cache_spec = group.kv_cache_spec
+                break
+
+        attention_groups: dict[tuple[str, str], AttentionGroup] = {}
+        if kv_cache_spec is not None:
+            for layer_name in self._draft_attn_layer_names:
+                attn_backend = all_attn_layers[layer_name].get_attn_backend()
+                backend_key = attn_backend.full_cls_name()
+                if backend_key not in attention_groups:
+                    layer_kv_cache_spec = kv_cache_spec
+                    if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
+                        layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[
+                            layer_name
+                        ]
+
+                    kernel_block_size = (
+                        kernel_block_sizes[self.kv_cache_gid]
+                        if kernel_block_sizes is not None
+                        and self.kv_cache_gid < len(kernel_block_sizes)
+                        else None
+                    )
+                    attn_group = AttentionGroup(
+                        backend=attn_backend,
+                        layer_names=[layer_name],
+                        kv_cache_spec=layer_kv_cache_spec,
+                        kv_cache_group_id=self.kv_cache_gid,
+                    )
+                    attn_group.create_metadata_builders(
+                        self.vllm_config,
+                        self.device,
+                        kernel_block_size=kernel_block_size,
+                    )
+                    attention_groups[backend_key] = attn_group
+                else:
+                    attention_groups[backend_key].layer_names.append(layer_name)
+
+        self.draft_attn_groups = list(attention_groups.values())
+        self.block_size = (
+            self.draft_attn_groups[0].get_metadata_builder().kv_cache_spec.block_size
+        )
+        logger.debug("Using block size %d for drafting layers", self.block_size)
+
+    def _determine_batch_execution_and_padding(
+        self,
+        num_tokens: int,
+        use_cudagraphs: bool = True,
+    ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:
+        cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
+            num_tokens,
+            valid_modes=({CUDAGraphMode.NONE} if not use_cudagraphs else None),
+        )
+        num_tokens_padded = batch_desc.num_tokens
+
+        # Extra coordination when running data-parallel since we need to
+        # coordinate across ranks
+        # TODO(Flechman): support DBO ubatching
+        should_ubatch, num_tokens_across_dp = False, None
+        if self.vllm_config.parallel_config.data_parallel_size > 1:
+            should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = (
+                coordinate_batch_across_dp(
+                    num_tokens_unpadded=num_tokens,
+                    parallel_config=self.vllm_config.parallel_config,
+                    allow_microbatching=False,
+                    num_tokens_padded=num_tokens_padded,
+                    cudagraph_mode=cudagraph_mode.value,
+                )
+            )
+            assert not should_ubatch, "DBO ubatching not implemented for EAGLE"
+
+            # Extract DP-synced values
+            if num_tokens_across_dp is not None:
+                dp_rank = self.dp_rank
+                num_tokens_padded = int(num_tokens_across_dp[dp_rank].item())
+                # Re-dispatch with DP padding so we have the correct
+                # batch_descriptor
+                cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
+                    num_tokens_padded,
+                    valid_modes={CUDAGraphMode(synced_cudagraph_mode)},
+                )
+                # Assert to make sure the agreed upon token count is correct
+                # otherwise num_tokens_across_dp will no-longer be valid
+                assert batch_desc.num_tokens == num_tokens_padded
+                num_tokens_across_dp[dp_rank] = num_tokens_padded
+
+        return cudagraph_mode, num_tokens_padded, num_tokens_across_dp
+
+
+# NOTE(woosuk): Currently, the below code is not used and we always use argmax
+# to sample the draft tokens. We will use this after we find a way to manage
+# the draft prob tensor.
+# Refer to https://github.com/vllm-project/vllm/pull/16899 for the details.
+# FIXME(woosuk): The logic here is duplicated with the main sampling code.
+# We should refactor this to reuse the same sampling implementation.
+def compute_probs_and_sample_next_token(
+    logits: torch.Tensor,
+    sampling_metadata: SamplingMetadata,
+    use_fp64_gumbel: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if sampling_metadata.all_greedy:
+        # For greedy requests, draft_probs is not used in rejection sampling.
+        # Therefore, we can just return the logits.
+        probs = logits
+        next_token_ids = logits.argmax(dim=-1)
+        return next_token_ids, probs
+
+    assert sampling_metadata.temperature is not None
+
+    # Use epsilon comparison to detect greedy sampling (temperature ~ 0.0)
+    # consistent with sampler.py's _SAMPLING_EPS threshold
+    temperature = sampling_metadata.temperature
+    # Avoid division by zero if there are greedy requests.
+    if not sampling_metadata.all_random:
+        is_greedy = temperature < _SAMPLING_EPS
+        temperature = torch.where(is_greedy, 1.0, temperature)
+    logits.div_(temperature.view(-1, 1))
+    probs = logits.softmax(dim=-1, dtype=torch.float32)
+
+    # NOTE(woosuk): Currently, we ignore most of the sampling parameters in
+    # generating the draft tokens. We only use the temperature. While this
+    # could degrade the acceptance rate, it does not affect the distribution
+    # of the generated tokens after rejection sampling.
+
+    # TODO(woosuk): Consider seeds.
+    q = empty_exponential_noise_like(probs, use_fp64_gumbel)
+    q.exponential_()
+    # NOTE(woosuk): We shouldn't use `probs.div_(q)` because the draft_probs
+    # will be used later for rejection sampling.
+    next_token_ids = sample_with_exponential_noise(probs.clone(), q)
+    if not sampling_metadata.all_random:
+        greedy_token_ids = probs.argmax(dim=-1)
+        next_token_ids = torch.where(is_greedy, greedy_token_ids, next_token_ids)
+    return next_token_ids, probs

vllm_patches/qwen3_dflash.py

@@ -0,0 +1,1133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+from collections.abc import Iterable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import Qwen3Config
+
+from vllm import _custom_ops as ops
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.multimodal.inputs import NestedTensors
+from vllm.transformers_utils.config import set_default_rope_theta
+from vllm.v1.attention.backend import AttentionType
+
+from .qwen2 import Qwen2MLP as Qwen3MLP
+from .qwen3 import Qwen3ForCausalLM
+from .utils import (
+    AutoWeightsLoader,
+    get_draft_quant_config,
+    maybe_prefix,
+    process_eagle_weight,
+)
+
+logger = init_logger(__name__)
+
+
+class _VanillaMarkov(nn.Module):
+    """Inlined DSpark VanillaMarkov low-rank transition-bias head.
+
+    Mirrors the training repo's markov_head.py VanillaMarkov param names EXACTLY
+    (markov_w1 = nn.Embedding(vocab, rank); markov_w2 = nn.Linear(rank, vocab,
+    bias=False)) so checkpoint tensors markov_head.markov_w1.weight /
+    markov_head.markov_w2.weight line up. Inlined (NOT imported from the repo)
+    because vLLM workers do not have the training repo root on sys.path.
+
+    B(x_{k-1}, :) = W2(W1[x_{k-1}]); the corrected logit for draft position k is
+    base_k + B(x_{k-1}, :), sampled left-to-right (semi-autoregressive).
+    """
+
+    def __init__(self, vocab_size: int, markov_rank: int) -> None:
+        super().__init__()
+        self.markov_w1 = nn.Embedding(vocab_size, markov_rank)
+        self.markov_w2 = nn.Linear(markov_rank, vocab_size, bias=False)
+
+    def compute_step_bias(
+        self, prev_token_ids: torch.Tensor, hidden_states: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        # (B,) prev token ids -> (B, V) additive logit bias
+        # hidden_states ignored for vanilla (memoryless)
+        del hidden_states
+        return self.markov_w2(self.markov_w1(prev_token_ids.long()))
+
+    def compute_step_vec(
+        self, prev_token_ids: torch.Tensor, hidden_states: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """Pre-W2 low-rank step vector v s.t. compute_step_bias == markov_w2(v).
+
+        Returns (B, rank) so a top-N caller can gather only the N needed rows of
+        markov_w2.weight ([V, rank]) instead of reading the full V x rank matrix.
+        compute_step_bias(prev) == self.markov_w2(compute_step_vec(prev)) EXACTLY
+        (same w1 lookup, same gate for the gated subclass); this method is the
+        factored inner half, never a separate approximation.
+        """
+        del hidden_states
+        return self.markov_w1(prev_token_ids.long())
+
+
+class _GatedMarkovHead(_VanillaMarkov):
+    """Inlined DSpark GatedMarkovHead (official DeepSpec).
+
+    Adds a sigmoid gate conditioned on [hidden_state; prev_embedding] to
+    modulate the markov bias. Uses backbone hidden state for adaptive gating.
+    Param names match the training repo: markov_w1, markov_w2, gate_proj.
+    """
+
+    def __init__(self, vocab_size: int, markov_rank: int, hidden_size: int) -> None:
+        super().__init__(vocab_size=vocab_size, markov_rank=markov_rank)
+        self.gate_proj = nn.Linear(hidden_size + markov_rank, markov_rank)
+
+    def compute_step_bias(
+        self, prev_token_ids: torch.Tensor, hidden_states: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        prev_emb = self.markov_w1(prev_token_ids.long())
+        if hidden_states is None:
+            return self.markov_w2(prev_emb)
+        gate = torch.sigmoid(
+            self.gate_proj(torch.cat([hidden_states, prev_emb], dim=-1))
+        ).to(dtype=prev_emb.dtype)
+        return self.markov_w2(gate * prev_emb)
+
+    def compute_step_vec(
+        self, prev_token_ids: torch.Tensor, hidden_states: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """Pre-W2 low-rank step vector for the gated head (see base docstring).
+
+        compute_step_bias(prev, h) == self.markov_w2(compute_step_vec(prev, h))
+        EXACTLY: the gate (or the hidden_states-is-None passthrough) is applied
+        here, identically to compute_step_bias, so the only thing the top-N
+        caller changes is which rows of markov_w2.weight get multiplied in.
+        """
+        prev_emb = self.markov_w1(prev_token_ids.long())
+        if hidden_states is None:
+            return prev_emb
+        gate = torch.sigmoid(
+            self.gate_proj(torch.cat([hidden_states, prev_emb], dim=-1))
+        ).to(dtype=prev_emb.dtype)
+        return gate * prev_emb
+
+
+def _markov_topn_from_env() -> int:
+    """Read DSPARK_MARKOV_TOPN (opt-in). 0 (default) / <=0 / unparsable -> off.
+
+    When >0 the markov transition-bias is computed only for the top-N base-logit
+    candidates per position, replacing the full V x rank markov_w2 GEMV
+    (~127 MB HBM read for V=248320) with an N x rank gather (~1 MB for N=2048).
+    Default 0 keeps the byte-identical full path; this is a pure serve knob.
+    """
+    try:
+        return int(os.environ.get("DSPARK_MARKOV_TOPN", "0") or "0")
+    except (TypeError, ValueError):
+        return 0
+
+
+def _markov_topn_sparse_bias(
+    step_vec: torch.Tensor,  # [B, rank] pre-W2 vector (compute_step_vec output)
+    w2_weight: torch.Tensor,  # [V, rank] markov_w2.weight
+    topn_idx: torch.Tensor,  # [B, N] long, per-row top-N base-logit indices
+) -> torch.Tensor:
+    """Sparse markov bias at only the top-N indices: bias[b, n] = <W2[idx], v>.
+
+    Gathers N rows of markov_w2.weight per row and dots them with the low-rank
+    step vector, so bias[b, n] == compute_step_bias(prev)[b, topn_idx[b, n]] up
+    to fp accumulation order (same sum over `rank`). Returns [B, N] in the head
+    dtype (caller casts to the base-logit dtype), mirroring compute_step_bias's
+    own dtype before its .to(base.dtype) cast. HBM read is B*N*rank vs B*V*rank.
+    """
+    # w2_weight[topn_idx] -> [B, N, rank]; bmm with v[B, rank, 1] -> [B, N, 1].
+    w2_rows = w2_weight[topn_idx]
+    return torch.bmm(w2_rows, step_vec.unsqueeze(-1)).squeeze(-1)
+
+
+def _markov_semiar_sample_block(
+    base_logits: torch.Tensor,  # [B, num_spec, V] base draft logits (model dtype)
+    first_prev_token_ids: torch.Tensor,  # [B] verified token before pos 0
+    compute_step_bias,  # callable: prev_ids[B] long -> bias[B, V]
+    temperature: torch.Tensor,  # [B, num_spec] per-request-per-pos sampling temp
+    all_random: bool,  # sampling_metadata.all_random
+    sample_fn,  # callable: probs[B, V] float32 -> tokens[B] long (random draw)
+    sampling_eps: float = 1e-5,
+    topn: int = 0,  # DSPARK_MARKOV_TOPN; >0 & <V -> sparse top-N bias path
+    compute_step_vec=None,  # callable: prev_ids[B] long -> step_vec[B, rank]
+    w2_weight: torch.Tensor | None = None,  # markov_w2.weight [V, rank]
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pure-torch LEFT-TO-RIGHT semi-AR markov-biased SAMPLING block.
+
+    Self-contained (only torch + injected callables) so it is unit-testable on
+    CPU without vLLM/CUDA. Mirrors compute_probs_and_sample_next_token's math
+    EXACTLY -- per-row temperature division, fp32 softmax, and the mixed-batch
+    greedy-row torch.where override -- but folds in the markov bias and feeds the
+    realized predecessor back at each step.
+
+    At position k the realized previous token biases the logit; the returned
+    probs[:, k] is the EXACT softmax distribution position k sampled from
+    (q_k = softmax((base_k + bias(prev_{k-1})) / T_row)). This identity is what
+    makes standard rejection sampling lossless: accepted dist == target dist iff
+    reported q == the draft's actual sampling dist.
+
+    Returns (tokens [B, num_spec] long, probs [B, num_spec, V] float32).
+    """
+    B, S, V = base_logits.shape
+    if S == 0:
+        return (
+            base_logits.new_empty((B, 0), dtype=torch.long),
+            base_logits.new_empty((B, 0, V), dtype=torch.float32),
+        )
+    out = base_logits.new_empty((B, S), dtype=torch.long)
+    probs_out = base_logits.new_empty((B, S, V), dtype=torch.float32)
+    prev = first_prev_token_ids.long()
+    # Opt-in top-N sparse markov bias. topn>=V would cover the whole vocab, so it
+    # falls back to the exact full path (identical result, no reason to gather V
+    # rows); topn<=0 is off. Only 0<topn<V takes the gather/scatter branch.
+    # PERF GUARD: the gather reads B*topn rows of markov_w2.weight (no cross-row
+    # reuse) vs the full path streaming the [V, rank] weight ONCE for the whole
+    # batch. Break-even is B*topn == V; above it the "sparse" path reads MORE HBM
+    # than full (and gather is less bandwidth-efficient than a streamed GEMM), so
+    # at high concurrency we must fall back to full. Only take top-N when
+    # B*topn < V.
+    use_topn = (
+        topn > 0
+        and topn < V
+        and B * topn < V
+        and compute_step_vec is not None
+        and w2_weight is not None
+    )
+    for k in range(S):
+        if use_topn:
+            base_k = base_logits[:, k, :]
+            # top-N base-logit candidates; only these receive a bias, every other
+            # token keeps its exact base logit (full-V softmax preserved below).
+            _, idx = torch.topk(base_k, topn, dim=-1)  # [B, N]
+            step_vec = compute_step_vec(prev)
+            sparse = _markov_topn_sparse_bias(step_vec, w2_weight, idx)
+            logits_k = base_k.clone()
+            logits_k.scatter_add_(-1, idx, sparse.to(logits_k.dtype))
+        else:
+            logits_k = base_logits[:, k, :] + compute_step_bias(prev).to(
+                base_logits.dtype
+            )
+        t_k = temperature[:, k]
+        is_greedy = None
+        if not all_random:
+            # Mixed batch: greedy rows divide by 1.0 (then argmax-override below),
+            # exactly as compute_probs_and_sample_next_token does.
+            is_greedy = t_k < sampling_eps
+            t_k = torch.where(is_greedy, torch.ones_like(t_k), t_k)
+        probs_k = (logits_k / t_k.unsqueeze(-1)).softmax(dim=-1, dtype=torch.float32)
+        # Record the reported q BEFORE drawing, so an in-place sampler cannot
+        # corrupt it (indexed assignment copies into probs_out's own storage).
+        probs_out[:, k, :] = probs_k
+        tok_k = sample_fn(probs_k).to(torch.long)
+        if is_greedy is not None:
+            greedy_tok = probs_k.argmax(dim=-1)
+            tok_k = torch.where(is_greedy, greedy_tok, tok_k)
+        out[:, k] = tok_k
+        prev = out[:, k]
+    return out, probs_out
+
+
+class DFlashQwen3Attention(nn.Module):
+    """Attention for DFlash speculative decoding.
+
+    Context KVs are pre-inserted into the KV cache before the forward pass.
+    This layer handles only query tokens via standard attention.
+    Adapted from Qwen3Attention."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_parameters: dict,
+        max_position: int = 4096 * 32,
+        head_dim: int | None = None,
+        rms_norm_eps: float = 1e-06,
+        attention_bias: bool = False,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        attn_type: str = AttentionType.DECODER,
+    ) -> None:
+        super().__init__()
+        self.layer_name = prefix
+        self.hidden_size = hidden_size
+        tp_size = get_tensor_model_parallel_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.head_dim = head_dim or hidden_size // self.total_num_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size,
+            self.head_dim,
+            self.total_num_heads,
+            self.total_num_kv_heads,
+            bias=attention_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            self.total_num_heads * self.head_dim,
+            hidden_size,
+            bias=attention_bias,  # DFlash has o_proj bias when using attention bias
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            max_position=max_position,
+            rope_parameters=rope_parameters,
+        )
+        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,
+            prefix=f"{prefix}.attn",
+            attn_type=attn_type,
+        )
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        """DFlash attention assumes that the KV cache is already populated
+        with the context K/V from the target model's hidden states. This forward op
+        computes attention for the query tokens only.
+        See also: precompute_and_store_context_kv"""
+        qkv = F.linear(hidden_states, self.qkv_proj.weight, self.qkv_proj.bias)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        # Per-head RMSNorm
+        q_shape, k_shape = q.shape, k.shape
+        q = self.q_norm(
+            q.view(*q_shape[:-1], q_shape[-1] // self.head_dim, self.head_dim)
+        ).view(q_shape)
+        k = self.k_norm(
+            k.view(*k_shape[:-1], k_shape[-1] // self.head_dim, self.head_dim)
+        ).view(k_shape)
+
+        q, k = self.rotary_emb(positions, q, k)
+
+        attn_output = self.attn(q, k, v)
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class DFlashQwen3DecoderLayer(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        *,
+        config: Qwen3Config,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        set_default_rope_theta(config, default_theta=1000000)
+        attn_type = AttentionType.DECODER
+
+        self.self_attn = DFlashQwen3Attention(
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            max_position=config.max_position_embeddings,
+            num_kv_heads=config.num_key_value_heads,
+            rms_norm_eps=config.rms_norm_eps,
+            attention_bias=getattr(config, "attention_bias", False),
+            head_dim=getattr(config, "head_dim", None),
+            cache_config=cache_config,
+            quant_config=quant_config,
+            rope_parameters=config.rope_parameters,
+            prefix=f"{prefix}.self_attn",
+            attn_type=attn_type,
+        )
+        self.mlp = Qwen3MLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is not None:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+        else:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+@support_torch_compile
+class DFlashQwen3Model(nn.Module):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        start_layer_id: int = 0,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = vllm_config.speculative_config.draft_model_config.hf_config
+        self.vocab_size = self.config.vocab_size
+        self.quant_config = get_draft_quant_config(vllm_config)
+
+        drafter_config = getattr(self.config, "eagle_config", {})
+        drafter_config.update(getattr(self.config, "dflash_config", {}))
+
+        if drafter_config is not None and "use_aux_hidden_state" in drafter_config:
+            self.use_aux_hidden_state = drafter_config["use_aux_hidden_state"]
+        else:
+            self.use_aux_hidden_state = True
+
+        current_vllm_config = get_current_vllm_config()
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.config.vocab_size,
+            self.config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                DFlashQwen3DecoderLayer(
+                    current_vllm_config,
+                    config=self.config,
+                    cache_config=current_vllm_config.cache_config,
+                    quant_config=self.quant_config,
+                    prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"),
+                )
+                for layer_idx in range(self.config.num_hidden_layers)
+            ]
+        )
+        if self.use_aux_hidden_state:
+            num_features_to_use = self.config.num_hidden_layers
+            if "target_layer_ids" in drafter_config:
+                num_features_to_use = len(drafter_config["target_layer_ids"])
+            elif "layer_ids" in drafter_config:
+                num_features_to_use = len(drafter_config["layer_ids"])
+            if hasattr(self.config, "target_hidden_size"):
+                fc_input_size = self.config.target_hidden_size * num_features_to_use
+            else:
+                fc_input_size = self.config.hidden_size * num_features_to_use
+            self.fc = ReplicatedLinear(
+                input_size=fc_input_size,
+                output_size=self.config.hidden_size,
+                bias=False,
+                params_dtype=vllm_config.model_config.dtype,
+                quant_config=self.quant_config,
+                prefix=maybe_prefix(prefix, "fc"),
+                return_bias=False,
+            )
+        self.hidden_norm = RMSNorm(
+            self.config.hidden_size,
+            eps=self.config.rms_norm_eps,
+        )
+        self.norm = RMSNorm(
+            self.config.hidden_size,
+            eps=self.config.rms_norm_eps,
+        )
+        # DSpark VanillaMarkov semi-AR draft head. Built only when the config
+        # declares markov_rank>0; otherwise None, so non-markov checkpoints
+        # (markov_rank absent -> 0) build nothing and stay byte-identical to
+        # before. Submodule path is model.markov_head.*, which the existing
+        # load_weights "model."+name rename loads automatically (no
+        # load_weights change needed).
+        markov_rank = int(getattr(self.config, "markov_rank", 0) or 0)
+        markov_head_type = str(getattr(self.config, "markov_head_type", "vanilla") or "vanilla").lower()
+        if markov_rank > 0:
+            if markov_head_type == "gated":
+                self.markov_head = _GatedMarkovHead(
+                    self.config.vocab_size, markov_rank, self.config.hidden_size
+                )
+            else:
+                self.markov_head = _VanillaMarkov(self.config.vocab_size, markov_rank)
+        else:
+            self.markov_head = None
+        # DSpark confidence head (DeepSpec AcceptRatePredictor, common.py:43-49):
+        # a single Linear predicting the per-position accept-rate LOGIT
+        # ("is one more draft token worth it?"). Built ONLY when the config
+        # declares confidence_head=True; otherwise None, so checkpoints without
+        # it (key absent -> False) build nothing and stay byte-identical to
+        # before. Submodule path model.confidence_head.* is loaded by the same
+        # "model."+name rename in DFlashQwen3ForCausalLM.load_weights that
+        # already handles markov_head (no load_weights change needed) -- this is
+        # exactly what resolves the `KeyError: confidence_head.bias` on serve.
+        # in_dim = hidden_size (+ markov_rank when confidence_head_with_markov),
+        # matching the trained checkpoint (5120 + 256 = 5376). NO explicit dtype:
+        # the head inherits the draft-model init dtype (bf16) just like markov_head
+        # and the official DeepSpec AcceptRatePredictor (common.py:43-49). This is
+        # deliberate -- the head was TRAINED with a bf16 GEMM (train_head.py:594-616:
+        # feats=draft_hidden(bf16), conf_head(feats), output .float()), so a bf16
+        # serve GEMM reproduces the exact calibration the threshold was tuned for.
+        if bool(getattr(self.config, "confidence_head", False)):
+            conf_in = self.config.hidden_size + (
+                markov_rank
+                if bool(getattr(self.config, "confidence_head_with_markov", False))
+                else 0
+            )
+            self.confidence_head = nn.Linear(conf_in, 1, bias=True)
+        else:
+            self.confidence_head = None
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def _build_fused_kv_buffers(self) -> None:
+        """Build fused weight buffers for precompute_and_store_context_kv.
+
+        Must be called after weights are loaded. Stacks the KV-projection
+        weights, K-norm weights, and RoPE parameters from every attention
+        layer so that precompute_and_store_context_kv can run one fused
+        GEMM for all layers at once. Also aliases the weight of the hidden_norm.
+        """
+        layers_attn = [layer.self_attn for layer in self.layers]
+        attn0 = layers_attn[0]
+        has_bias = attn0.qkv_proj.bias is not None
+
+        self._hidden_norm_weight = self.hidden_norm.weight.data
+
+        # KV projection weights: [num_layers * 2 * kv_size, hidden_size]
+        kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn]
+        self._fused_kv_weight = torch.cat(kv_weights, dim=0)
+        if has_bias:
+            kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn]
+            self._fused_kv_bias: torch.Tensor | None = torch.cat(kv_biases, dim=0)
+        else:
+            self._fused_kv_bias = None
+
+        # K-norm weights: list of [head_dim] tensors, one per layer.
+        self._k_norm_weights = [a.k_norm.weight.data for a in layers_attn]
+
+        # RoPE parameters
+        self._rope_head_size = attn0.rotary_emb.head_size
+        self._rope_cos_sin_cache = attn0.rotary_emb.cos_sin_cache
+        self._rope_is_neox = attn0.rotary_emb.is_neox_style
+        # Validation that RoPE params are the same across all layers
+        for attn in layers_attn[1:]:
+            assert (
+                attn.rotary_emb.head_size == self._rope_head_size
+                and attn.rotary_emb.is_neox_style == self._rope_is_neox
+            ), "All layers must have the same RoPE parameters for DFlash precomputation"
+
+        # Layer metadata
+        self._num_attn_layers = len(layers_attn)
+        self._kv_size = attn0.kv_size
+        self._head_dim = attn0.head_dim
+        self._num_kv_heads = attn0.num_kv_heads
+        self._rms_norm_eps = attn0.q_norm.variance_epsilon
+        # Validation that all layers have the same attention config
+        for attn in layers_attn[1:]:
+            assert (
+                attn.kv_size == self._kv_size
+                and attn.head_dim == self._head_dim
+                and attn.num_kv_heads == self._num_kv_heads
+                and attn.q_norm.variance_epsilon == self._rms_norm_eps
+            ), "All layers must have the same attn config for DFlash precomputation"
+
+        # References to inner Attention layers for direct cache writes
+        self._attn_layers = [layer.self_attn.attn for layer in self.layers]
+
+    def precompute_and_store_context_kv(
+        self,
+        context_states: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mapping: torch.Tensor | None = None,
+    ) -> None:
+        """Precompute K/V for context states write them into each layer's KV cache.
+
+        Input context states are projected to K/V, normed, and have RoPE applied.
+        Since the context shape is different than the query shape, we can't rely on the
+        regular forward pass to apply torch.compile and CUDA graphs to this section.
+        As such, this function is optimized to minimize the number of torch ops present:
+        we use fused vLLM kernels for RMSNorm and RoPE, fuse the GEMM into one
+        large projection, and avoid cloning buffers (with .contiguous()) where possible.
+
+        When context_slot_mapping is None (e.g. during dummy_run) only
+        the computation runs, and no K/V is written to cache.
+        """
+        if not hasattr(self, "_num_attn_layers"):
+            logger.warning_once(
+                "DFlash buffer initialization was skipped. If dummy weights are not "
+                "in use, this may indicate an error in weight loading."
+            )
+            self._build_fused_kv_buffers()
+
+        num_ctx = context_states.shape[0]
+        L = self._num_attn_layers
+        kv = self._kv_size
+        hd = self._head_dim
+        nkv = self._num_kv_heads
+
+        # --- Fused KV projection (one GEMM for all layers) ---
+        normed_context_states = torch.empty_like(context_states)
+        ops.rms_norm(
+            normed_context_states,
+            context_states,
+            self._hidden_norm_weight,
+            self._rms_norm_eps,
+        )
+        all_kv_flat = F.linear(
+            normed_context_states, self._fused_kv_weight, self._fused_kv_bias
+        )
+        # Single contiguous copy that separates K/V and transposes to
+        # layer-major layout.  Result: [2, L, num_ctx, nkv, hd] contiguous.
+        # Indexing dim-0 gives contiguous [L, num_ctx, nkv, hd] for K and V.
+        all_kv = (
+            all_kv_flat.view(num_ctx, L, 2, nkv, hd).permute(2, 1, 0, 3, 4).contiguous()
+        )
+        all_k = all_kv[0]  # [L, num_ctx, nkv, hd], contiguous
+        all_v = all_kv[1]  # [L, num_ctx, nkv, hd], contiguous
+
+        # --- Per-layer RMSNorm K (3D: [num_ctx, nkv, hd] per layer) ---
+        all_k_normed = torch.empty_like(all_k)
+        for i in range(L):
+            ops.rms_norm(
+                all_k_normed[i],
+                all_k[i],
+                self._k_norm_weights[i],
+                self._rms_norm_eps,
+            )
+
+        # --- Fused RoPE across all layers ---
+        # View as [L * num_ctx, kv] so RoPE sees one big batch (no copy).
+        # In-place RoPE: pass K as the "query" arg with key=None.
+        all_k_flat = all_k_normed.view(L * num_ctx, kv)
+        positions_repeated = context_positions.repeat(L)
+        cos_sin_cache = self._rope_cos_sin_cache
+        if cos_sin_cache.dtype != all_k_flat.dtype:
+            cos_sin_cache = cos_sin_cache.to(dtype=all_k_flat.dtype)
+        ops.rotary_embedding(
+            positions_repeated,
+            all_k_flat,
+            None,
+            self._rope_head_size,
+            cos_sin_cache,
+            self._rope_is_neox,
+        )
+
+        if context_slot_mapping is None:
+            return
+
+        # --- Per-layer cache insert ---
+        all_k_final = all_k_flat.view(L, num_ctx, nkv, hd)
+        for i in range(L):
+            attn = self._attn_layers[i]
+            kv_cache = attn.kv_cache
+            attn.impl.do_kv_cache_update(
+                attn,
+                all_k_final[i],
+                all_v[i],
+                kv_cache,
+                context_slot_mapping,
+            )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        input_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if input_embeds is None:
+            input_embeds = self.embed_input_ids(input_ids)
+
+        hidden_states = input_embeds
+
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(
+                positions=positions,
+                hidden_states=hidden_states,
+                residual=residual,
+            )
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if "midlayer." in name:
+                name = name.replace("midlayer.", "layers.0.")
+            if "scale" in name:
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
+
+class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        nn.Module.__init__(self)
+        self.config = vllm_config.speculative_config.draft_model_config.hf_config
+        if getattr(self.config, "draft_vocab_size", None) is None:
+            self.config.draft_vocab_size = getattr(self.config, "vocab_size", None)
+        target_layer_num = vllm_config.model_config.get_num_layers(
+            vllm_config.parallel_config
+        )
+        self.model = DFlashQwen3Model(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+            start_layer_id=target_layer_num,
+        )
+
+        logit_scale = getattr(self.config, "logit_scale", 1.0)
+        self.lm_head = ParallelLMHead(
+            self.config.draft_vocab_size,
+            self.config.hidden_size,
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+        self.logits_processor = LogitsProcessor(
+            self.config.draft_vocab_size, scale=logit_scale
+        )
+        target_vocab_size = vllm_config.model_config.get_vocab_size()
+        if self.config.draft_vocab_size != target_vocab_size:
+            self.draft_id_to_target_id = nn.Parameter(
+                torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
+                requires_grad=False,
+            )
+        else:
+            self.draft_id_to_target_id = None
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: NestedTensors | None = None,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.model(input_ids, positions, inputs_embeds)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        if self.draft_id_to_target_id is None:
+            return logits
+
+        base = torch.arange(self.config.draft_vocab_size, device=logits.device)
+        targets = base + self.draft_id_to_target_id
+        logits_new = logits.new_full(
+            (logits.shape[0], self.config.vocab_size),
+            float("-inf"),
+        )
+        logits_new[:, targets] = logits
+        return logits_new
+
+    @torch.no_grad()
+    def sample_draft_block_semiar(
+        self,
+        sample_hidden_states: torch.Tensor,
+        first_prev_token_ids: torch.Tensor,
+        num_spec: int,
+    ) -> torch.Tensor:
+        """LEFT-TO-RIGHT (semi-autoregressive) markov-biased greedy draft block.
+
+        Produces the parallel-drafted block one position at a time: position k's
+        argmax is taken over base_k + markov_bias(prev_token), where prev_token
+        is the token argmax'd at position k-1 (first_prev = the verified bonus
+        token before draft pos 0). LOSSLESS under greedy because the target
+        verify re-checks every drafted token against the target argmax,
+        independent of how the draft was produced.
+
+        Returns [B, num_spec] long tensor, matching the .view(-1, num_spec)
+        contract of the parallel-drafting early-exit path.
+        """
+        mk = self.model.markov_head
+        assert mk is not None, (
+            "sample_draft_block_semiar called without a markov_head"
+        )
+        # The markov head is trained in the (draft == target) vocab space. The
+        # d2t remap path is intentionally NOT supported here: with d2t != None
+        # the base logits live in draft-vocab space while the markov bias lives
+        # in target-vocab space (shape mismatch), and the L->R feedback loop
+        # would index markov_w1 with draft-vocab ids. Fail loudly rather than
+        # silently dropping the bias. (Follow-up could remap out at the end via
+        # out = self.draft_id_to_target_id[out] AND project the bias/feedback
+        # into draft space, but AEON has draft_vocab == vocab so d2t is None.)
+        assert self.draft_id_to_target_id is None, (
+            "VanillaMarkov semi-AR draft requires draft_vocab_size == target "
+            "vocab_size (draft_id_to_target_id is None); the d2t remap path is "
+            "not implemented."
+        )
+        # Base draft logits ONCE, draft-vocab space, mirroring compute_logits
+        # but WITHOUT the d2t scatter (guarded None above). Stay on-device:
+        # no .item()/D2H inside the loop.
+        base = self.logits_processor(self.lm_head, sample_hidden_states)
+        B = sample_hidden_states.shape[0] // num_spec
+        V = base.shape[-1]
+        assert V == mk.markov_w2.weight.shape[0], (
+            f"draft logits vocab {V} != markov bias vocab "
+            f"{mk.markov_w2.weight.shape[0]}"
+        )
+        assert first_prev_token_ids.shape[0] == B, (
+            f"first_prev_token_ids batch {first_prev_token_ids.shape[0]} != "
+            f"derived batch {B} (rows={sample_hidden_states.shape[0]}, "
+            f"num_spec={num_spec})"
+        )
+        base = base.view(B, num_spec, V)
+        # Fused semi-AR: use mk.compute_step_bias but with pre-allocated buffers
+        # and in-place add+argmax to reduce kernel launches per iteration.
+        # Original: 8 × (embed+linear+add+argmax) = 32 launches
+        # Optimized: 8 × (compute_step_bias+add+argmax) = 24, no intermediate allocs
+        # GatedMarkovHead: passes hidden_states per position for gate conditioning.
+        out = base.new_empty((B, num_spec), dtype=torch.long)
+        prev = first_prev_token_ids.long()
+        # Pre-allocate bias buffer to avoid per-step allocation
+        bias_buf = base.new_empty((B, V))
+        # Reshape sample_hidden_states to [B, num_spec, H] for gated head access
+        hs_per_pos = sample_hidden_states.view(B, num_spec, -1)
+        # Opt-in top-N sparse markov bias (DSPARK_MARKOV_TOPN). Every non-top-N
+        # token keeps its base logit (its bias is defined to be 0), so the top-N
+        # argmax equals the TRUE full-base+bias argmax ONLY WHEN the full-bias
+        # winner lies inside base-top-N; if a token outside base-top-N would have
+        # won under the full bias, top-N misses it -> an accept-affecting
+        # TRUNCATION (an approximation of the full-bias argmax), NOT an identity.
+        # This is safe: losslessness is guaranteed by verify (the target re-checks
+        # every drafted token); top-N can at most shift which token is proposed
+        # (accept rate), never correctness. topn>=V falls back to the full path.
+        # PERF GUARD (same as the block path): only worthwhile when B*topn < V,
+        # else the per-row gather reads more HBM than the streamed full weight.
+        topn = _markov_topn_from_env()
+        use_topn = 0 < topn < V and B * topn < V
+        w2w = mk.markov_w2.weight if use_topn else None
+        for k in range(num_spec):
+            hs_k = hs_per_pos[:, k, :] if hasattr(mk, 'gate_proj') else None
+            if use_topn:
+                base_k = base[:, k, :]
+                _, idx = torch.topk(base_k, topn, dim=-1)  # [B, N]
+                step_vec = mk.compute_step_vec(prev, hidden_states=hs_k)
+                sparse = _markov_topn_sparse_bias(step_vec, w2w, idx).to(base.dtype)
+                bias_buf.copy_(base_k)
+                bias_buf.scatter_add_(-1, idx, sparse)
+            else:
+                bias = mk.compute_step_bias(prev, hidden_states=hs_k).to(base.dtype)
+                # Fused add + argmax via in-place add then argmax
+                torch.add(base[:, k, :], bias, out=bias_buf)
+            tok = bias_buf.argmax(dim=-1)
+            out[:, k] = tok
+            prev = tok
+        return out  # [B, num_spec]
+
+    @torch.no_grad()
+    def sample_draft_block_semiar_sample(
+        self,
+        sample_hidden_states: torch.Tensor,
+        first_prev_token_ids: torch.Tensor,
+        num_spec: int,
+        temperature: torch.Tensor,
+        all_random: bool,
+        use_fp64_gumbel: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """LEFT-TO-RIGHT (semi-AR) markov-biased SAMPLING draft block (temp>0).
+
+        LOSSLESS counterpart of sample_draft_block_semiar (the greedy variant):
+        each position k is SAMPLED -- not argmax'd -- from
+            q_k = softmax((base_k + markov_bias(prev_{k-1})) / T_row)
+        and that EXACT q_k is returned as the corrected proposal probs. Standard
+        rejection sampling is lossless iff the reported q equals the distribution
+        the draft actually sampled from; here they are the same tensor, computed
+        from the same logits, same per-row temperature, and the realized
+        left-to-right predecessor -- so the contract holds by construction.
+
+        Math mirrors compute_probs_and_sample_next_token EXACTLY (per-row temp
+        division, fp32 softmax, exponential-noise / Gumbel-max draw, and the
+        greedy-row argmax override for mixed batches), applied position-by-
+        position with the markov bias folded in and prev fed back.
+
+        temperature is sampling_metadata.temperature, length B*num_spec, laid out
+        request-major / position-minor (same layout as sample_hidden_states and
+        the parallel path's draft_probs.view(-1, num_spec, V)).
+
+        Returns (sampled_tokens [B, num_spec] long, corrected_probs
+        [B, num_spec, V] float32). corrected_probs matches the parallel path's
+        _last_draft_probs layout (request-major, position, vocab; float32;
+        contiguous).
+        """
+        from vllm.v1.sample.ops.topk_topp_sampler import (
+            empty_exponential_noise_like,
+            sample_with_exponential_noise,
+        )
+        from vllm.v1.sample.sampler import _SAMPLING_EPS
+
+        mk = self.model.markov_head
+        assert mk is not None, (
+            "sample_draft_block_semiar_sample called without a markov_head"
+        )
+        # Same d2t guard as the greedy path: the markov head lives in the
+        # draft==target vocab space; the remap path is not implemented.
+        assert self.draft_id_to_target_id is None, (
+            "VanillaMarkov semi-AR SAMPLING draft requires draft_vocab_size == "
+            "target vocab_size (draft_id_to_target_id is None); the d2t remap "
+            "path is not implemented."
+        )
+        base = self.logits_processor(self.lm_head, sample_hidden_states)
+        B = sample_hidden_states.shape[0] // num_spec
+        V = base.shape[-1]
+        assert V == mk.markov_w2.weight.shape[0], (
+            f"draft logits vocab {V} != markov bias vocab "
+            f"{mk.markov_w2.weight.shape[0]}"
+        )
+        assert first_prev_token_ids.shape[0] == B, (
+            f"first_prev_token_ids batch {first_prev_token_ids.shape[0]} != "
+            f"derived batch {B} (rows={sample_hidden_states.shape[0]}, "
+            f"num_spec={num_spec})"
+        )
+        assert temperature is not None, (
+            "sample_draft_block_semiar_sample requires a temperature tensor "
+            "(only reached when sampling_metadata is not all_greedy)."
+        )
+        base = base.view(B, num_spec, V)
+        # temperature may be per-request [B] (vLLM sampling_metadata.temperature),
+        # per-position [B*num_spec], or a scalar -> broadcast to [B, num_spec].
+        _t = temperature.reshape(-1)
+        if _t.numel() == B * num_spec:
+            temp = _t.view(B, num_spec)
+        elif _t.numel() == B:
+            temp = _t.view(B, 1).expand(B, num_spec)
+        else:
+            temp = _t.reshape(1, 1).expand(B, num_spec)
+
+        def _gumbel_sample(probs: torch.Tensor) -> torch.Tensor:
+            # Mirror compute_probs_and_sample_next_token's draw exactly:
+            # exponential (Gumbel-max) noise, sampled from probs.clone() so the
+            # returned probs (used as rejection-sampling q) stay intact.
+            noise = empty_exponential_noise_like(probs, use_fp64_gumbel)
+            noise.exponential_()
+            return sample_with_exponential_noise(probs.clone(), noise)
+
+        # Opt-in top-N sparse markov bias (DSPARK_MARKOV_TOPN). The sampling path
+        # keeps the full-V softmax; only the bias term is sparsified (top-N base
+        # candidates get the bias, all other tokens keep their exact base logit).
+        # The resulting q_k is a TRUNCATION of the full-bias distribution (tokens
+        # whose bias would have mattered but sit outside base-top-N are omitted) --
+        # an approximation that shifts accept, NOT correctness. Losslessness still
+        # holds because rejection sampling only needs the reported q_k to equal the
+        # distribution actually sampled from, and here they are the SAME tensor.
+        # NOTE: like the full path here, the gated head is fed prev only
+        # (hidden=None) -> compute_step_vec matches compute_step_bias(prev)
+        # bit-for-bit at the inner (pre-W2) vector.
+        # PERF: the block applies the same B*topn<V guard internally.
+        topn = _markov_topn_from_env()
+        return _markov_semiar_sample_block(
+            base,
+            first_prev_token_ids,
+            mk.compute_step_bias,
+            temp,
+            all_random,
+            _gumbel_sample,
+            sampling_eps=_SAMPLING_EPS,
+            topn=topn,
+            compute_step_vec=mk.compute_step_vec,
+            w2_weight=mk.markov_w2.weight,
+        )
+
+    @torch.no_grad()
+    def predict_confidence_step(
+        self,
+        sample_hidden_states: torch.Tensor,
+        draft_token_ids: torch.Tensor,
+        first_prev_token_ids: torch.Tensor,
+        num_spec: int,
+    ) -> torch.Tensor:
+        """Per-position confidence (accept-rate) LOGITS for the realized draft block.
+
+        Mirrors the official DeepSpec DSparkModel.predict_confidence_step
+        (qwen3/modeling.py:293-308) + draft_ops._predict_confidence_logits EXACTLY:
+
+          prev_token_ids = cat([first_prev, realized_block[:, :-1]])   # SERVE prev
+          prev_emb       = markov_head.get_prev_embeddings(prev).to(hidden.dtype)
+          features       = cat([draft_hidden, prev_emb], dim=-1)
+          logit          = confidence_head(features).float()
+
+        draft_hidden == sample_hidden_states (the post-norm hidden fed to lm_head;
+        same tensor used as draft_hidden in train_head.py:567-568). The predecessor
+        is the SERVE-PATH realized token (sampled/argmax'd at k-1, first_prev for
+        k=0), NOT the teacher-forced prev used in training -- this is the intended
+        serve feature (design CONFIDENCE_HEAD_DESIGN.md §9). This logit is consumed
+        ONLY for the dynamic-K prefix-length decision; it never changes which tokens
+        are proposed, so it cannot affect losslessness.
+
+        Args:
+          sample_hidden_states: [B*num_spec, H] request-major / position-minor.
+          draft_token_ids:      [B, num_spec] the realized drafted block.
+          first_prev_token_ids: [B] the verified bonus token before draft pos 0.
+          num_spec:             block size (num_speculative_tokens).
+        Returns:
+          [B, num_spec] float32 accept-rate logits.
+        """
+        conf = self.model.confidence_head
+        assert conf is not None, (
+            "predict_confidence_step called without a confidence_head"
+        )
+        H = sample_hidden_states.shape[-1]
+        B = sample_hidden_states.shape[0] // num_spec
+        hidden = sample_hidden_states.view(B, num_spec, H)
+        if bool(getattr(self.config, "confidence_head_with_markov", False)):
+            mk = self.model.markov_head
+            assert mk is not None, (
+                "confidence_head_with_markov=True but markov_head is None"
+            )
+            # prev[:,0]=first_prev, prev[:,k]=realized draft token at k-1.
+            prev = torch.cat(
+                [
+                    first_prev_token_ids.view(B, 1).long(),
+                    draft_token_ids[:, :-1].long(),
+                ],
+                dim=1,
+            )  # [B, num_spec]
+            # get_prev_embeddings == markov_w1(prev) (markov_head.py:73-74).
+            prev_emb = mk.markov_w1(prev).to(dtype=hidden.dtype)  # [B, num_spec, r]
+            features = torch.cat([hidden, prev_emb], dim=-1)
+        else:
+            features = hidden
+        # Match the head's own param dtype for the GEMM (bf16 to mirror training);
+        # cast guards against a float32 head built from an older checkpoint.
+        features = features.to(conf.weight.dtype)
+        return conf(features).squeeze(-1).float()  # [B, num_spec]
+
+    def precompute_and_store_context_kv(
+        self,
+        context_states: torch.Tensor,
+        context_positions: torch.Tensor,
+        context_slot_mapping: torch.Tensor | None = None,
+    ) -> None:
+        """Precompute projected + RoPE'd K/V and write to cache."""
+        self.model.precompute_and_store_context_kv(
+            context_states, context_positions, context_slot_mapping
+        )
+
+    def combine_hidden_states(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        if not self.model.use_aux_hidden_state:
+            return hidden_states
+        needs_squeeze = hidden_states.dim() == 1
+        if needs_squeeze:
+            hidden_states = hidden_states.unsqueeze(0)
+        result = self.model.fc(hidden_states)
+        if needs_squeeze:
+            result = result.squeeze(0)
+        return result
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        model_weights = {}
+        includes_draft_id_mapping = False
+        includes_embed_tokens = False
+        for name, loaded_weight in weights:
+            assert "mask_hidden" not in name, (
+                "DFlash should use mask_token_id to embed the padding hidden state"
+            )
+            if "t2d" in name:
+                continue
+            if "d2t" in name:
+                name = name.replace("d2t", "draft_id_to_target_id")
+                includes_draft_id_mapping = True
+            elif "lm_head" not in name:
+                name = "model." + name
+            if "embed_tokens" in name:
+                includes_embed_tokens = True
+            model_weights[name] = loaded_weight
+            process_eagle_weight(self, name)
+
+        skip_substrs = []
+        if not includes_draft_id_mapping:
+            skip_substrs.append("draft_id_to_target_id")
+        if not includes_embed_tokens:
+            skip_substrs.append("embed_tokens")
+        if not self.model.use_aux_hidden_state:
+            skip_substrs.append("fc.")
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=None,
+            skip_substrs=skip_substrs,
+        )
+        loader.load_weights(model_weights.items())
+        self.model._build_fused_kv_buffers()