# coding=utf-8
# DFlash2 draft model — self-contained (single file) on purpose.
#
# Sources, cross-verified against each other:
#   1. drafttrain/dflash/_vendored/dflash_model.py — the SpecForge DFlash1 backbone this
#      repo already trains and serves through SGLang's DFLASH algorithm. The backbone
#      code below is copied VERBATIM from it (attention / decoder layer / model /
#      spec_generate) so train<->serve behavior matches the proven DFlash1 path.
#   2. z-lab/dflash dflash/model.py (Apache-2.0) — the official DFlash2DraftModel
#      reference: GroupedDynamicCausalConv (conv_kernel_size=2, conv_group_size=16) and
#      CandidateSelector (selector_rank=256, selector_top_k=16). Parameter names are
#      kept identical (layers.{i}.attention_conv/mlp_conv.{base_kernel,kernel_projection},
#      candidate_selector.{predecessor_codebook,successor_codebook,hidden_projection})
#      so exported weights load under SGLang's DFlash2 support and under the z-lab
#      reference implementation unchanged.
#
# Two deliberate deviations from z-lab's reference:
#   - The conv is BLOCK-LOCAL during training: when the noise stream is a whole number
#     of blocks (the training layout: N blocks of [anchor, mask*bs-1] concatenated),
#     the predecessor tap zero-pads at each block start instead of reading across
#     blocks — exactly matching inference, where the draft sees one block at a time.
#     With a single block (<= block_size, the inference shape) the behavior is
#     bit-identical to z-lab's F.pad reference.
#   - conv_identity_init (default True): base kernel starts as [1, 0] (identity) and
#     kernel_projection at zero, so a warm-started DFlash1 checkpoint behaves
#     identically at step 0 and the convs grow in smoothly during training.
#
# This file must stay importable stand-alone: export_draft.py copies it into the served
# draft directory as ``dflash.py`` (auto_map -> "dflash.DFlash2DraftModel"), where no
# ``drafttrain`` package exists.

from typing import Callable, ClassVar, Optional

import torch
import torch.nn.functional as F
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


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 _sampling_probs(
    logits: torch.Tensor,
    temperature: float,
    top_p: float = 1.0,
    top_k: int = 0,
) -> torch.Tensor:
    """Softmax over (optionally top-k/top-p filtered) logits, scattered back to full vocab."""
    scores = logits.float() / temperature
    vocab_size = scores.shape[-1]
    if 0 < top_k < vocab_size:
        scores, indices = torch.topk(scores, top_k, dim=-1)
    else:
        indices = None

    probs = torch.softmax(scores, dim=-1)
    if top_p < 1.0:
        sorted_probs, order = probs.sort(dim=-1, descending=True)
        keep = sorted_probs.cumsum(dim=-1) - sorted_probs < top_p
        sorted_probs = sorted_probs * keep
        probs = torch.zeros_like(probs).scatter(-1, order, sorted_probs)
        probs = probs / probs.sum(dim=-1, keepdim=True)

    if indices is not None:
        probs = torch.zeros_like(logits, dtype=probs.dtype).scatter(-1, indices, probs)
    return probs


def _sample_probs(probs: torch.Tensor) -> torch.Tensor:
    shape = probs.shape[:-1]
    return torch.multinomial(probs.view(-1, probs.shape[-1]), 1).view(shape)


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 _dflash_config(config) -> dict:
    return getattr(config, "dflash_config", {}) or {}


def _draft_value(config, name, default=None):
    return _dflash_config(config).get(name, getattr(config, name, default))


class Qwen3DFlashAttention(nn.Module):
    """Dual-stream attention (copied from the vendored DFlash1 model, unchanged).

    K/V are computed from BOTH the captured target hidden states (context stream) and
    the draft's own noise-stream hidden states, concatenated.
    """

    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
        self.is_causal = 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]
        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)
        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_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


# --------------------------------------------------------------------------- #
# DFlash2 addition 1: two-tap grouped dynamic causal convolution
# --------------------------------------------------------------------------- #

def _grouped_dynamic_convolve(hidden, dynamic, base, group_size):
    """z-lab reference convolve: out[t] = sum_offset (base[offset] + dynamic[t, offset]) * x[t-offset].

    ``hidden`` (B', L', H), ``dynamic`` (B', L', kernel_size, groups), ``base``
    (kernel_size, H). Zero left-pad: x[t-offset] = 0 for t < offset.
    """
    batch, length, hidden_size = hidden.shape
    groups = hidden_size // group_size
    blocks = hidden.view(batch, length, groups, group_size)
    dynamic = dynamic.view(batch, length, base.shape[0], groups, 1)
    output = torch.zeros_like(blocks)
    for offset in range(base.shape[0]):
        values = blocks if offset == 0 else F.pad(blocks[:, :-offset], (0, 0, 0, 0, offset, 0))
        kernel = base[offset].view(1, 1, groups, group_size).to(hidden.dtype)
        output = output + kernel * values
        output = torch.addcmul(output, dynamic[:, :, offset], values)
    return output.view_as(hidden)


class GroupedDynamicCausalConv(nn.Module):
    """Two-tap grouped dynamic causal convolution (DFlash2).

    Parameter layout is identical to z-lab/dflash's GroupedDynamicCausalConv:
      - ``base_kernel``: (2, kernel_size, hidden). Index 0 holds the static taps for
        the PRE-sublayer conv (``prepare``), index 1 for the POST-sublayer conv
        (``finish``). The two dynamic kernels are computed once from the pre-sublayer
        (normed) hidden state and reused by both.
      - ``kernel_projection``: Linear(hidden, 2 * kernel_size * groups, bias=False) —
        per-position dynamic tap corrections; every ``group_size`` channels share one.

    Block-local extension (training): when the input length is a whole number of
    blocks AND longer than one block, the sequence is processed as concatenated
    independent blocks — the predecessor tap zero-pads at each block start instead of
    reading the previous block's tail. This exactly matches inference, where the draft
    forward sees a single block ([anchor, mask, ...]) at a time. With a single block
    (length <= block_size) the computation is bit-identical to the z-lab reference.
    """

    def __init__(self, hidden_size: int, kernel_size: int, group_size: int, block_size: int):
        super().__init__()
        if hidden_size % group_size != 0:
            raise ValueError(
                f"GroupedDynamicCausalConv requires group_size to divide hidden_size; "
                f"got hidden_size={hidden_size}, group_size={group_size}"
            )
        if kernel_size < 1:
            raise ValueError(f"kernel_size must be >= 1, got {kernel_size}")
        if block_size < 1:
            raise ValueError(f"block_size must be >= 1, got {block_size}")
        self.kernel_size = kernel_size
        self.group_size = group_size
        self.block_size = block_size
        self.num_groups = hidden_size // group_size
        self.base_kernel = nn.Parameter(torch.empty(2, kernel_size, hidden_size))
        self.kernel_projection = nn.Linear(
            hidden_size, 2 * kernel_size * self.num_groups, bias=False
        )

    def _split_blocks(self, hidden: torch.Tensor):
        """(B, L, ...) -> ((B*N, bs, ...), bsz, n). Whole-block sequences (L = N*bs > bs,
        the training noise layout) become N independent blocks; anything else (a single
        full or partial block, the inference layout) is returned as one block."""
        bsz, seq_len = hidden.shape[0], hidden.shape[1]
        if seq_len > self.block_size and seq_len % self.block_size == 0:
            n = seq_len // self.block_size
            return hidden.reshape(bsz * n, self.block_size, *hidden.shape[2:]), bsz, n
        return hidden, bsz, 1

    def prepare(self, hidden: torch.Tensor):
        """Pre-sublayer: convolve the (normed) input; stash the finish-step taps.

        Returns (conv_out (B, L, H), dynamic (B, L, kernel_size, groups)).
        """
        bsz, seq_len = hidden.shape[0], hidden.shape[1]
        blocked, _, _ = self._split_blocks(hidden)
        dynamic = self.kernel_projection(blocked).view(
            *blocked.shape[:-1], 2, self.kernel_size, self.num_groups
        )
        out = _grouped_dynamic_convolve(
            blocked, dynamic[..., 0, :, :], self.base_kernel[0], self.group_size
        )
        return (
            out.reshape(bsz, seq_len, hidden.shape[-1]),
            dynamic[..., 1, :, :].reshape(
                bsz, seq_len, self.kernel_size, self.num_groups
            ),
        )

    def finish(self, hidden: torch.Tensor, dynamic: torch.Tensor) -> torch.Tensor:
        """Post-sublayer: convolve the sublayer output with the stashed taps."""
        bsz, seq_len = hidden.shape[0], hidden.shape[1]
        blocked, _, _ = self._split_blocks(hidden)
        dyn_blocked = dynamic.reshape(
            blocked.shape[0], blocked.shape[1], self.kernel_size, self.num_groups
        )
        out = _grouped_dynamic_convolve(
            blocked, dyn_blocked, self.base_kernel[1], self.group_size
        )
        return out.reshape(bsz, seq_len, hidden.shape[-1])


# --------------------------------------------------------------------------- #
# Decoder layer (DFlash1 layer + optional DFlash2 conv hooks, z-lab pattern)
# --------------------------------------------------------------------------- #

class Qwen3DFlashDecoderLayer(GradientCheckpointingLayer):
    """DFlash decoder layer with optional DFlash2 conv hooks.

    With ``attention_conv``/``mlp_conv`` None (DFlash1) the forward is identical to
    the vendored SpecForge layer. When set (DFlash2), each sublayer is wrapped as:
    prepare(normed input) -> sublayer -> finish(sublayer output), exactly z-lab's
    placement. The convs act ONLY on the noise stream (``hidden_states``); the context
    stream (``target_hidden``) feeds k_ctx/v_ctx directly, unconverted.
    """

    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
        )
        self.attention_conv: Optional[GroupedDynamicCausalConv] = None
        self.mlp_conv: Optional[GroupedDynamicCausalConv] = None

    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)
        attention_kernel = None
        if self.attention_conv is not None:
            hidden_states, attention_kernel = self.attention_conv.prepare(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_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
            **kwargs,
        )[0]
        if attention_kernel is not None:
            hidden_states = self.attention_conv.finish(hidden_states, attention_kernel)
        hidden_states = residual + hidden_states
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        mlp_kernel = None
        if self.mlp_conv is not None:
            hidden_states, mlp_kernel = self.mlp_conv.prepare(hidden_states)
        hidden_states = self.mlp(hidden_states)
        if mlp_kernel is not None:
            hidden_states = self.mlp_conv.finish(hidden_states, mlp_kernel)
        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


# --------------------------------------------------------------------------- #
# DFlash2 addition 2: candidate path selector
# --------------------------------------------------------------------------- #

class CandidateSelector(nn.Module):
    """Top-k candidate path selector (z-lab-compatible parameter layout).

    Scores adjacent candidate pairs with a gated low-rank bilinear form::

        S_t(a, b) = U_t(b) + < A(a) ⊙ H(h_t), B(b) >

    where ``U_t`` is the draft's own logit for candidate ``b`` (how much the drafter
    liked it on its own), ``A``/``B`` are compact per-token codebooks and ``H(h_t)`` is
    a context gate projected from the draft hidden state deciding which parts of the
    predecessor/successor match count.
    """

    def __init__(self, config):
        super().__init__()
        rank = int(_draft_value(config, "selector_rank", 256))
        top_k = int(_draft_value(config, "selector_top_k", 16))
        if rank <= 0:
            raise ValueError(f"selector_rank must be > 0, got {rank}")
        if top_k <= 0:
            raise ValueError(f"selector_top_k must be > 0, got {top_k}")
        self.rank = rank
        self.top_k = top_k
        self.predecessor_codebook = nn.Embedding(config.vocab_size, rank)
        self.successor_codebook = nn.Embedding(config.vocab_size, rank)
        self.hidden_projection = nn.Linear(config.hidden_size, rank, bias=False)

    def pairwise_scores(
        self,
        hidden: torch.Tensor,
        logits: torch.Tensor,
        prev_ids: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Teacher-forced pairwise scores for TRAINING.

        hidden (B, L, H), logits (B, L, V), prev_ids (B, L) true predecessor token ids.
        Returns (candidates (B, L, k), scores (B, L, k)) with
        ``scores[b, t, j] = logits[b, t, cand_j] + <A(prev) ⊙ H(h_t), B(cand_j)>``.
        Fully parallel — no sequential walk (teacher forcing supplies the predecessor).
        """
        unary, candidates = torch.topk(logits, self.top_k, dim=-1, sorted=False)
        gate = self.hidden_projection(hidden)                     # (B, L, R)
        pred = self.predecessor_codebook(prev_ids.long())         # (B, L, R)
        succ = self.successor_codebook(candidates)                # (B, L, k, R)
        scores = unary + torch.einsum("blr,blr,blkr->blk", pred, gate, succ)
        return candidates, scores

    def select(
        self,
        hidden: torch.Tensor,
        logits: torch.Tensor,
        anchor_ids: torch.Tensor,
        temperature: float,
    ):
        """Inference-time path walk (z-lak reference): greedy at T=0, else sampling
        from softmax over the k candidate scores (also returned for lossless
        rejection sampling). ``anchor_ids`` is the last verified token."""
        unary, candidates = torch.topk(logits, self.top_k, dim=-1, sorted=False)
        hidden = self.hidden_projection(hidden)
        # Accept (B,) or (B, 1) anchor ids (the walk keeps a flat (B,) predecessor).
        predecessor = anchor_ids.reshape(anchor_ids.shape[0], -1).squeeze(-1)
        path, q_rows = [], []
        for position in range(hidden.shape[1]):
            scores = unary[:, position] + torch.einsum(
                "br,bkr->bk",
                self.predecessor_codebook(predecessor) * hidden[:, position],
                self.successor_codebook(candidates[:, position]),
            )
            if temperature > 0:
                q = _sampling_probs(scores[:, None], temperature)[:, 0]
                index = _sample_probs(q)
                q_rows.append(q)
            else:
                index = torch.argmax(scores, dim=-1)
            predecessor = candidates[:, position].gather(-1, index[:, None])[:, 0]
            path.append(predecessor)
        return (
            torch.stack(path, dim=1),
            candidates,
            torch.stack(q_rows, dim=1) if q_rows else None,
        )


# --------------------------------------------------------------------------- #
# DFlash (1) backbone — copied verbatim from drafttrain/dflash/_vendored/dflash_model.py
# --------------------------------------------------------------------------- #

class DFlashDraftModel(Qwen3PreTrainedModel):
    config_class = Qwen3Config
    _no_split_modules: ClassVar[list[str]] = ["Qwen3DFlashDecoderLayer"]

    def __init__(self, config) -> None:
        super().__init__(config)
        self.config = config
        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)
        self.projector_type = dflash_config.get("projector_type", None)
        self.pure_draft_prefix_len = dflash_config.get("pure_draft_prefix_len", 0)
        self.shift_label = dflash_config.get("shift_label", False)

        if self.projector_type == "domino":
            self.emb_dim = dflash_config["emb_dim"]
            self.gru_hidden_dim = dflash_config["gru_hidden_dim"]
            self.prefix_gru = nn.GRU(
                input_size=config.hidden_size,
                hidden_size=self.gru_hidden_dim,
                num_layers=1,
                batch_first=True,
                bias=False,
            )
            in_dim = config.hidden_size + self.gru_hidden_dim
            self.embed_proj = nn.Sequential(
                nn.Linear(in_dim, self.emb_dim, bias=False),
                nn.SiLU(),
                nn.Linear(self.emb_dim, config.vocab_size, bias=False),
            )
        elif self.projector_type is not None:
            raise ValueError(f"Unknown draft projector_type: {self.projector_type}")
        self.post_init()

    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 = 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


# --------------------------------------------------------------------------- #
# DFlash2 draft model
# --------------------------------------------------------------------------- #

class DFlash2DraftModel(DFlashDraftModel):
    """DFlash backbone with the two DFlash2 additions.

    - Every decoder layer gets ``attention_conv`` + ``mlp_conv``
      (GroupedDynamicCausalConv, kernel_size=2 / group_size=16 by default).
    - A ``candidate_selector`` (CandidateSelector, rank=256 / top_k=16 by default).

    Config keys (in ``dflash_config``): ``conv_kernel_size``, ``conv_group_size``,
    ``conv_identity_init``, ``selector_rank``, ``selector_top_k`` — mirroring the
    released z-lab/Qwen3.8-27B-DFlash2 config.json (which carries the first two and
    the last two; ``conv_identity_init`` is a training-only knob).

    Weight layout is z-lab-compatible, so exported checkpoints load under SGLang's
    DFlash2 support and under z-lab/dflash's reference implementation.
    """

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        # The SERVED checkpoint stores the selector codebooks under bare keys
        # (no ".weight" — z-lab/SGLang format; see export_draft.py). Map them onto
        # this model's nn.Embedding parameters when loading via HF transformers.
        kwargs.setdefault(
            "key_mapping",
            {
                f"candidate_selector.{name}": f"candidate_selector.{name}.weight"
                for name in ("predecessor_codebook", "successor_codebook")
            },
        )
        return super().from_pretrained(*args, **kwargs)

    def __init__(self, config) -> None:
        super().__init__(config)
        dflash_config = _dflash_config(config)
        kernel_size = int(dflash_config.get("conv_kernel_size", 2))
        group_size = int(dflash_config.get("conv_group_size", 16))
        self.conv_identity_init = bool(dflash_config.get("conv_identity_init", True))
        for layer in self.layers:
            layer.attention_conv = GroupedDynamicCausalConv(
                config.hidden_size, kernel_size, group_size, self.block_size
            )
            layer.mlp_conv = GroupedDynamicCausalConv(
                config.hidden_size, kernel_size, group_size, self.block_size
            )
        self.candidate_selector = CandidateSelector(config)
        # post_init() initializes the newly added Linear/Embedding modules (already
        # initialized parent modules are skipped); the raw base_kernel Parameters and
        # the identity pattern are then set explicitly so initialization does not
        # depend on transformers' double-post_init semantics.
        self.post_init()
        self._init_dflash2_weights()

    def _init_dflash2_weights(self) -> None:
        std = float(getattr(self.config, "initializer_range", 0.02))
        with torch.no_grad():
            for layer in self.layers:
                for conv in (layer.attention_conv, layer.mlp_conv):
                    if conv is None:  # pragma: no cover (always set for DFlash2)
                        continue
                    # Static taps start as identity: tap-0 (self) = 1, tap-1
                    # (predecessor) = 0. Anything else destroys the residual stream at
                    # init (a zero tap-0 zeroes the sublayer input).
                    conv.base_kernel.zero_()
                    conv.base_kernel[:, 0, :].fill_(1.0)
                    if self.conv_identity_init:
                        conv.kernel_projection.weight.zero_()
                    else:
                        conv.kernel_projection.weight.normal_(mean=0.0, std=std)
            sel = self.candidate_selector
            sel.predecessor_codebook.weight.normal_(mean=0.0, std=std)
            sel.successor_codebook.weight.normal_(mean=0.0, std=std)
            sel.hidden_projection.weight.normal_(mean=0.0, std=std)

    def propose(
        self,
        hidden: torch.Tensor,
        anchor_ids: torch.Tensor,
        output_head: nn.Module,
        temperature: float,
    ):
        """Select a draft path through the top-k candidates (SGLang/spec entry point)."""
        logits = output_head(hidden)
        return self.candidate_selector.select(hidden, logits, anchor_ids, temperature)

    @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,
    ):
        """DFlash1 spec_generate with the candidate-selector walk replacing the
        independent per-position argmax/sample (the only behavioral change)."""
        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_hidden = 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)
            draft_logits = target.lm_head(draft_hidden)
            # DFlash2: walk one coherent path through the top-k candidates instead of
            # taking each position's top pick independently.
            draft_tokens, _, _ = self.candidate_selector.select(
                draft_hidden,
                draft_logits,
                block_output_ids[:, 0],
                temperature,
            )
            block_output_ids[:, 1:] = draft_tokens

            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