Inkling-Small DSpark speculator preview (epoch_0_step_7250)

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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

README.md

@@ -0,0 +1,76 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- speculative-decoding
+- dspark
+- dflash
+- specforge
+- sglang
+inference: false
+---
+
+# Inkling-Small DSpark speculator
+
+## Overview
+
+A DSpark speculator for the Inkling-Small NVFP4 target, enabling faster inference through speculative decoding. DSpark extends the DFlash parallel-draft backbone with a Markov logit-bias head and a per-position confidence head. This checkpoint was trained with [SpecForge](https://github.com/sgl-project/SpecForge/) using hidden states from a live SGLang target engine.
+
+**Preview:** Preview: This is a research preview. Official training and evaluation are in progress.
+
+## Model Specifications
+
+- **Base model:** `thinkingmachines/Inkling-Small-NVFP4`; the exact matching target checkpoint and tokenizer are required.
+- **Format:** Safetensors (single-file BF16, 0.90B trainable parameters; draft weights only).
+- **Draft:** 5 layers (Qwen3-style GQA), hidden 4096, 64 heads / 16 KV heads, head_dim 64, FFN 8192, rope_theta 8000000, `block_size=15`.
+- **RoPE:** YaRN, factor 128, `original_max_position_embeddings=8192`, `max_position_embeddings=1048576` (matches the target addressable range); config carries both `rope_scaling` and `rope_parameters` schemas.
+- **Vocabulary:** 200,058 tokenizer entries and 201,024 padded weight rows; `mask_token_id=200064`.
+- **DSpark heads:** Markov rank 256 (vanilla) and confidence head (with-Markov).
+- **Aux hidden-state layers:** `[5, 11, 23, 29, 35]`.
+- **Trained context:** 8,192 native; long-context adaptation with sequences up to 65,536.
+- **Target weights:** target embedding and unembedding weights are not included in this checkpoint.
+
+## Evaluation Results
+
+Acceptance length (`acc_len`) at temperature 0 (128 prompts per task, DSPARK block size 7, thinking effort 0.99):
+
+| Dataset | acc_len |
+|:--|--:|
+| GSM8K | 4.7871 |
+| MATH500 | 4.1426 |
+| MBPP | 3.4394 |
+| HumanEval | 3.3491 |
+| MT-Bench | 3.1140 |
+| LiveCodeBench | 2.9288 |
+| AIME25 | 2.8938 |
+| Alpaca | 2.7823 |
+| Arena-Hard-v2 | 2.6982 |
+| **Mean** | **3.3484** |
+
+## Serving with SGLang
+
+Requires a SGLang build with DSpark support
+
+```bash
+SGLANG_ENABLE_UNIFIED_RADIX_TREE=1 \
+sglang serve --trust-remote-code \
+  --model-path thinkingmachines/Inkling-Small-NVFP4 --tp 8 \
+  --quantization modelopt_fp4 --attention-backend fa4 --page-size 128 \
+  --fp4-gemm-backend flashinfer_trtllm --moe-runner-backend flashinfer_trtllm_routed \
+  --enable-torch-symm-mem --mamba-radix-cache-strategy extra_buffer \
+  --mem-fraction-static 0.68 --swa-full-tokens-ratio 0.1 --mamba-full-memory-ratio 0.1 \
+  --max-running-requests 68 --reasoning-parser inkling --tool-call-parser inkling \
+  --skip-server-warmup --speculative-algorithm DSPARK \
+  --speculative-draft-model-path skx618/InkLing-Small-DSpark-Preview \
+  --speculative-draft-model-quantization unquant \
+  --speculative-dspark-block-size 7 \
+  --chunked-prefill-size 8192 --cuda-graph-max-bs-prefill 8192 \
+  --disable-flashinfer-autotune --host 0.0.0.0 --port 30000
+```
+
+## Training Details
+
+- **Framework:** SpecForge online distillation with hidden states captured from a frozen Inkling-Small target served by a colocated per-node TP4 SGLang engine.
+- **Stage 1 (original RoPE, 8K):** 444,264 Inkling-Small self-regenerations from `open-perfectblend` (temperature 1.0, top-p 0.95, reasoning effort 0.99), 2 epochs, AdamW, global batch 512, gradient clipping 1.0; peak LR `6e-4` (cosine, 4% warmup).
+- **Stage 2 (YaRN-128 adaptation — this checkpoint):** weights-only warm start from stage 1; 250,000 chat rows (192,055 previously unseen regenerations + 57,945 replayed) plus 1,092 real 16K-64K agentic trajectories upsampled x8; `max_length=65536`; constant LR `5e-4` with a 64-step re-warm.
+- **Loss:** `0.1 CE + 0.9 L1 distillation + 1.0 confidence BCE`, 512 sampled anchors per sequence, `block_size=15`, within-block decay gamma `28/3`.

config.json

@@ -0,0 +1,76 @@
+{
+  "architectures": [
+    "DSparkDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "dspark.DSparkDraftModel"
+  },
+  "block_size": 15,
+  "bos_token_id": null,
+  "confidence_head_with_markov": true,
+  "dflash_config": {
+    "mask_token_id": 200064,
+    "target_layer_ids": [
+      5,
+      11,
+      23,
+      29,
+      35
+    ]
+  },
+  "dtype": "bfloat16",
+  "enable_confidence_head": true,
+  "eos_token_id": [
+    200006,
+    200000,
+    200002,
+    200003
+  ],
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "markov_head_type": "vanilla",
+  "markov_rank": 256,
+  "max_position_embeddings": 1048576,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "num_attention_heads": 64,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 16,
+  "num_target_layers": 42,
+  "pad_token_id": 200006,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 128.0,
+    "original_max_position_embeddings": 8192,
+    "rope_theta": 8000000,
+    "rope_type": "yarn"
+  },
+  "rope_scaling": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 128.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "yarn"
+  },
+  "rope_theta": 8000000,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.12.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 201024
+}

dflash.py

@@ -0,0 +1,470 @@
+import os
+from functools import partial
+from typing import Callable, Optional
+
+import torch
+from torch import nn
+from torch.nn.attention.flex_attention import BlockMask, flex_attention
+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
+
+# FlashAttention-4 flex backend, opt-in via env SPECFORGE_DRAFT_FLEX_BACKEND=fa4.
+# flex_attention with kernel_options={"BACKEND": "FLASH"} runs the FA4 kernel instead
+# of the Triton flex kernel (ref: meta-pytorch/attention-gym flex_flash_attention.py).
+#
+# STATUS on this stack (torch 2.11 / GB300 sm_10.3): the FA4 kernel works for a small
+# head_dim (64/128) BUT ONLY for a mask_mod that captures no tensors (it needs the
+# asymmetric block sparsity BLOCK_SIZE=(q=256, kv=128); see core/dflash.py). The DSpark
+# dual-source mask (`create_dflash_block_mask`) MUST capture per-sample `anchor_positions`
+# / `block_keep_mask` tensors, and the FA4 CuteDSL template fails on any captured-tensor
+# mask_mod ("CuteDSL template failed"), both dynamic=True and False. => FA4 is currently
+# NOT usable for the DSpark drafter; the default Triton flex backend (used when this env
+# is unset) handles the captured-tensor mask correctly and is the supported path.
+# (The DeepSeek-V4 DSpark draft was likewise FA4-ruled-out, there for head_dim 512.)
+# The code path is kept, gated + off by default, for a future stack / a captured-tensor-
+# free mask formulation. dynamic=True: draft Q/context lengths vary per batch.
+_FLEX_FA4_COMPILED = None
+
+
+def _flex_fa4():
+    global _FLEX_FA4_COMPILED
+    if _FLEX_FA4_COMPILED is None:
+        _FLEX_FA4_COMPILED = torch.compile(
+            partial(flex_attention, kernel_options={"BACKEND": "FLASH"}),
+            dynamic=True,
+        )
+    return _FLEX_FA4_COMPILED
+
+
+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
+
+
+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
+        )
+        num_attention_heads = int(config.num_attention_heads)
+        num_key_value_heads = int(config.num_key_value_heads)
+        if (
+            num_attention_heads <= 0
+            or num_key_value_heads <= 0
+            or num_attention_heads % num_key_value_heads != 0
+        ):
+            raise ValueError(
+                "Qwen3DFlashAttention requires positive attention head counts and "
+                "num_attention_heads divisible by num_key_value_heads, got "
+                f"num_attention_heads={num_attention_heads}, "
+                f"num_key_value_heads={num_key_value_heads}."
+            )
+        self.num_key_value_groups = num_attention_heads // 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
+        )
+
+    # Keep the dual-source attention OUT of any outer torch.compile region: the
+    # flex_attention block-mask HOP fails inductor lowering when nested inside a
+    # larger dynamic-shape graph on this stack (CantSplit / "unsupported operand &"),
+    # even though it compiles fine on its own (HF's flex integration compiles it
+    # separately). Marking the attention compiler-disabled lets SPECFORGE_COMPILE_DRAFT
+    # fuse the rest of the block (RoPE, RMSNorm, MLP, residual) while the flex
+    # attention keeps its own (block-sparse, separately-compiled) fast path.
+    @torch.compiler.disable
+    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)
+        # FA4 flex path (opt-in): call flex_attention with the FLASH kernel directly on
+        # the prebuilt dual-source BlockMask, bypassing HF's Triton-flex wrapper. q/k/v
+        # are already [B, H, S, D]; flex returns [B, H, S, D] -> transpose to [B, S, H, D]
+        # to match the reshape below. GQA (num_kv_heads < num_heads) via enable_gqa.
+        if (
+            self.config._attn_implementation in ("flex_attention", "flex")
+            and os.environ.get("SPECFORGE_DRAFT_FLEX_BACKEND") == "fa4"
+            and isinstance(attention_mask, BlockMask)
+        ):
+            attn_output = _flex_fa4()(
+                q,
+                k,
+                v,
+                block_mask=attention_mask,
+                scale=self.scaling,
+                enable_gqa=(self.num_key_value_groups > 1),
+            )
+            attn_output = attn_output.transpose(1, 2).contiguous()
+            attn_weights = None
+        else:
+            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
+        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

dspark.py

@@ -0,0 +1,160 @@
+# coding=utf-8
+"""DSpark draft model: DFlash backbone + EAGLE-style Markov and confidence heads.
+
+DSpark shares SpecForge's DFlash block-diffusion drafter (dual-source KV
+injection via :class:`DFlashDraftModel`, anchor sampling, MASK-token noise
+stream) and adds two heads on top:
+
+  - Markov head: a low-rank learned bigram bias added to the draft logits,
+    conditioned on the (teacher-forced) previous token. Improves the per-token
+    distribution without touching the backbone.
+  - Confidence head (AcceptRatePredictor): predicts a per-draft-position
+    acceptance probability, trained against the empirical draft-vs-target
+    accept rate (used at inference time for adaptive block length).
+
+Ported from TorchSpec PR #129 (``torchspec/models/draft/dspark.py``). The Markov
+/ confidence modeling code is adapted from DeepSeek's DeepSpec
+(``deepspec/modeling/dspark/{markov_head,common}.py``, MIT License).
+
+SpecForge differences vs TorchSpec (load-bearing):
+  - There is no ``DFlashConfig``; SpecForge's :class:`DFlashDraftModel` uses a
+    plain ``Qwen3Config`` plus a ``config.dflash_config`` dict. So
+    :class:`DSparkConfig` subclasses ``Qwen3Config`` and declares the DSpark
+    fields as top-level attributes; DFlash-carried fields (``block_size``,
+    ``num_target_layers``, ``dflash_config``) stay as before.
+  - The draft model has no ``embed_tokens`` of its own (the embedding lives on
+    the target and is passed into the online wrapper), and the context
+    projection is ``self.fc`` (not ``context_proj``). The heads only depend on
+    ``config.hidden_size`` / ``config.vocab_size``, so this does not matter for
+    construction.
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from transformers.models.qwen3.modeling_qwen3 import Qwen3Config
+
+from specforge.modeling.draft.dflash import DFlashDraftModel
+
+
+class DSparkConfig(Qwen3Config):
+    """Configuration for the DSpark draft model.
+
+    Extends ``Qwen3Config`` (SpecForge's DFlash draft is config-light and reads a
+    plain ``Qwen3Config``). DSpark-specific fields are declared here; the
+    DFlash-carried fields (``block_size``, ``num_target_layers``, and the nested
+    ``dflash_config`` dict holding ``target_layer_ids`` / ``mask_token_id``) are
+    consumed by the :class:`DFlashDraftModel` base ``__init__`` and must be
+    present on the config object before constructing the model.
+    """
+
+    model_type = "dspark"
+
+    def __init__(
+        self,
+        markov_rank: int = 256,
+        markov_head_type: str = "vanilla",
+        enable_confidence_head: bool = True,
+        confidence_head_with_markov: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.markov_rank = markov_rank
+        self.markov_head_type = markov_head_type
+        self.enable_confidence_head = enable_confidence_head
+        self.confidence_head_with_markov = confidence_head_with_markov
+
+
+class VanillaMarkov(nn.Module):
+    """Low-rank learned bigram bias added to the draft logits.
+
+    Adapted from DeepSpec's ``deepspec/modeling/dspark/markov_head.py``.
+    """
+
+    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) -> torch.Tensor:
+        return self.project_bias(self.get_prev_embeddings(token_ids))
+
+    def apply_block_logits(
+        self,
+        base_logits: torch.Tensor,
+        *,
+        token_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        if base_logits.size(2) == 0:
+            return base_logits
+        return base_logits + self.compute_step_bias(token_ids)
+
+
+class AcceptRatePredictor(nn.Module):
+    """Per-position acceptance-probability predictor (a single linear head).
+
+    Adapted from DeepSpec's ``deepspec/modeling/dspark/common.py``.
+    """
+
+    def __init__(self, input_dim: int):
+        super().__init__()
+        self.proj = nn.Linear(int(input_dim), 1)
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        return self.proj(features).squeeze(-1)
+
+
+def build_markov_head(config) -> Optional[nn.Module]:
+    markov_rank = int(getattr(config, "markov_rank", 0))
+    assert markov_rank >= 0, f"markov_rank must be >= 0, got {markov_rank}"
+    if markov_rank == 0:
+        return None
+
+    markov_head_type = str(getattr(config, "markov_head_type", "vanilla")).lower()
+    if markov_head_type == "vanilla":
+        return VanillaMarkov(vocab_size=config.vocab_size, markov_rank=markov_rank)
+    raise NotImplementedError(
+        f"markov_head_type={markov_head_type!r} is not supported yet; only 'vanilla' "
+        "is implemented as it is recommended by the authors."
+    )
+
+
+class DSparkDraftModel(DFlashDraftModel):
+    """DSpark draft network: DFlash backbone + Markov / confidence heads."""
+
+    config_class = DSparkConfig
+
+    def __init__(self, config) -> None:
+        super().__init__(config)
+
+        self.markov_rank = int(getattr(config, "markov_rank", 0))
+        self.confidence_head_with_markov = bool(
+            getattr(config, "confidence_head_with_markov", True)
+        )
+
+        self.markov_head = build_markov_head(config)
+
+        self.confidence_head: Optional[nn.Module] = None
+        if getattr(config, "enable_confidence_head", False):
+            conf_input_dim = config.hidden_size
+            if self.confidence_head_with_markov:
+                if self.markov_head is None:
+                    raise ValueError(
+                        "confidence_head_with_markov=True requires a Markov head "
+                        "(markov_rank > 0)."
+                    )
+                conf_input_dim += self.markov_rank
+            self.confidence_head = AcceptRatePredictor(conf_input_dim)

model.safetensors

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