feat: DeepSpec - DSpark trainer support

configs/sglang_qwen3_8b_dspark.yaml

@@ -0,0 +1,89 @@
+# DSpark training config for Qwen3-8B target model
+#
+# DSpark = DFlash block-diffusion drafter + EAGLE-style Markov & confidence
+# heads, trained with cross-entropy + L1 distribution distillation + confidence
+# BCE. The L1 / confidence terms need the target's final hidden state, so
+# inference.store_last_hidden_states MUST be true (DFlash leaves it false).
+#
+# GPU allocation (8x GPU):
+#   - 4 GPUs for inference (SGLang engine, tp_size=1, duplicate mode)
+#   - 4 GPUs for training (FSDP FULL_SHARD)
+#
+# Usage:
+#   python -m torchspec.train_entry --config configs/sglang_qwen3_8b_dspark.yaml
+#   ./examples/qwen3-8b-single-node/run.sh configs/sglang_qwen3_8b_dspark.yaml
+
+model:
+  target_model_path: Qwen/Qwen3-8B
+  trust_remote_code: true
+  draft_model_config: torchspec/config/dspark_draft_config.json
+
+dataset:
+  train_data_path: ../examples/data/sample_conversations.jsonl
+  eval_data_path: null
+  eval_interval: 100
+  chat_template: qwen
+  prompt_key: conversations
+  min_loss_tokens: 32
+
+training:
+  attention_backend: flex_attention
+  micro_batch_size: 1
+  draft_accumulation_steps: 2
+  learning_rate: 6e-4
+  min_lr: 6e-5
+  weight_decay: 0.0
+  max_concurrent_batches: 1
+  max_grad_norm: 1.0
+  max_seq_length: 2048
+  num_epochs: 3
+  seed: 42
+  training_num_gpus_per_node: 4
+  training_num_nodes: 1
+  ttt_length: 7
+  fsdp_strategy: FULL_SHARD
+  fsdp_reduce_dtype: bfloat16
+  prefetch_depth: 8
+  save_interval: 1000
+  save_per_epoch: true
+  max_checkpoints: 2
+  warmup_ratio: 0.04
+
+  # DSpark-specific parameters
+  dflash_block_size: 7
+  dspark_num_anchors: 512
+  dspark_num_target_layers: 5
+  dspark_loss_decay_gamma: 4.0
+  dspark_ce_loss_alpha: 0.1
+  dspark_l1_loss_alpha: 0.9
+  dspark_confidence_head_alpha: 1.0
+
+inference:
+  inference_engine_type: sgl
+  store_last_hidden_states: true
+  inference_num_gpus: 4
+  inference_num_gpus_per_engine: 1
+  inference_num_gpus_per_node: 4
+  max_sample_pool_size: 64
+  inference_buffer_threshold: 32
+  inference_batch_size: 8
+  sglang:
+    tp_size: 1
+    mem_fraction_static: 0.7
+
+mooncake:
+  master_server_address: null
+  metadata_server: null
+  protocol: tcp
+  global_segment_size: 16GB
+  local_buffer_size: 4GB
+  enable_hard_pin: true
+
+output_dir: ./outputs/qwen3-8b-dspark
+cache_dir: ./cache/qwen3-8b-dspark
+model_download_dir: null
+
+debug:
+  save_debug_train_data: null
+  debug_train_only: false
+  debug_inference_only: false

torchspec/__init__.py

@@ -24,12 +24,17 @@
 from torchspec.models.dflash import DFlashModel
 from torchspec.models.draft import AutoDraftModelConfig, AutoEagle3DraftModel
 from torchspec.models.draft.dflash import DFlashConfig, DFlashDraftModel
+from torchspec.models.draft.dspark import DSparkConfig, DSparkDraftModel
+from torchspec.models.dspark import DSparkModel
 
 __all__ = [
     "Eagle3Model",
     "DFlashModel",
     "DFlashConfig",
     "DFlashDraftModel",
+    "DSparkModel",
+    "DSparkConfig",
+    "DSparkDraftModel",
     "AutoDraftModelConfig",
     "AutoEagle3DraftModel",
 ]

torchspec/config/dspark_draft_config.json

@@ -0,0 +1,23 @@
+{
+    "architectures": ["DSparkDraftModel"],
+    "model_type": "dspark",
+    "hidden_size": 4096,
+    "intermediate_size": 12288,
+    "num_hidden_layers": 5,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 8,
+    "vocab_size": 151936,
+    "rms_norm_eps": 1e-6,
+    "max_position_embeddings": 40960,
+    "rope_theta": 1000000.0,
+    "num_target_layers": 5,
+    "target_hidden_size": 4096,
+    "target_num_hidden_layers": 36,
+    "target_layer_ids": [1, 9, 17, 25, 33],
+    "mask_token_id": 151669,
+    "markov_rank": 256,
+    "markov_head_type": "vanilla",
+    "enable_confidence_head": true,
+    "confidence_head_with_markov": true,
+    "tie_word_embeddings": false
+}

torchspec/config/train_config.py

@@ -154,6 +154,14 @@ class TrainingConfig:
     dflash_num_anchors: int = 512
     dflash_num_target_layers: int = 5
 
+    # DSpark-specific parameters (used by DSpark trainer only)
+    dspark_num_anchors: int = 512
+    dspark_num_target_layers: int = 5
+    dspark_loss_decay_gamma: float = 4.0
+    dspark_ce_loss_alpha: float = 0.1
+    dspark_l1_loss_alpha: float = 0.9
+    dspark_confidence_head_alpha: float = 1.0
+
 
 @dataclass
 class DecodeConfig:

torchspec/models/__init__.py

@@ -19,6 +19,7 @@
 # SOFTWARE.
 
 from torchspec.models.dflash import DFlashModel
+from torchspec.models.dspark import DSparkModel
 from torchspec.models.eagle3 import (
     Eagle3Model,
     compute_lazy_target_padded,
@@ -30,6 +31,7 @@
 __all__ = [
     "Eagle3Model",
     "DFlashModel",
+    "DSparkModel",
     "compute_lazy_target_padded",
     "compute_target_p_padded",
     "compiled_forward_kl_loss",

torchspec/models/draft/__init__.py

@@ -22,6 +22,7 @@
 from torchspec.models.draft.base import Eagle3DraftModel
 from torchspec.models.draft.deepseek_eagle import Eagle3DeepseekV2ForCausalLM
 from torchspec.models.draft.dflash import DFlashConfig, DFlashDraftModel
+from torchspec.models.draft.dspark import DSparkConfig, DSparkDraftModel
 from torchspec.models.draft.llama3_eagle import LlamaForCausalLMEagle3
 
 __all__ = [
@@ -31,5 +32,7 @@
     "Eagle3DraftModel",
     "DFlashConfig",
     "DFlashDraftModel",
+    "DSparkConfig",
+    "DSparkDraftModel",
     "LlamaForCausalLMEagle3",
 ]

torchspec/models/draft/auto.py

@@ -28,6 +28,7 @@
 
 from torchspec.models.draft.deepseek_eagle import Eagle3DeepseekV2ForCausalLM
 from torchspec.models.draft.dflash import DFlashConfig, DFlashDraftModel
+from torchspec.models.draft.dspark import DSparkConfig, DSparkDraftModel
 from torchspec.models.draft.llama3_eagle import LlamaForCausalLMEagle3
 from torchspec.utils.logging import logger
 
@@ -37,6 +38,7 @@ class AutoEagle3DraftModel(AutoModelForCausalLMBase):
         LlamaConfig: LlamaForCausalLMEagle3,
         DeepseekV3Config: Eagle3DeepseekV2ForCausalLM,
         DFlashConfig: DFlashDraftModel,
+        DSparkConfig: DSparkDraftModel,
     }
 
     @classmethod
@@ -77,6 +79,7 @@ class AutoDraftModelConfig:
         "LlamaForCausalLMEagle3": LlamaConfig,
         "Eagle3DeepseekV2ForCausalLM": DeepseekV3Config,
         "DFlashDraftModel": DFlashConfig,
+        "DSparkDraftModel": DSparkConfig,
     }
 
     @classmethod

torchspec/models/draft/dspark.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""
+DSpark draft model: DFlash backbone + EAGLE-style Markov and confidence heads.
+
+DSpark shares DFlash's block-diffusion drafter (dual-source KV injection, 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).
+
+Markov / confidence modeling code is adapted from DeepSeek's DeepSpec
+(deepspec/modeling/dspark/{markov_head,common}.py, MIT License).
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from torchspec.models.draft.dflash import DFlashConfig, DFlashDraftModel
+
+
+class DSparkConfig(DFlashConfig):
+    """
+    Configuration for the DSpark draft model. Extends :class:`DFlashConfig`.
+    """
+
+    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):
+    """
+    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):
+    """
+    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 in TorchSpec as it is recommended by the authors."
+    )
+
+
+class DSparkDraftModel(DFlashDraftModel):
+    config_class = DSparkConfig
+
+    def __init__(self, config: DSparkConfig):
+        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 = self.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)

torchspec/train_entry.py

@@ -250,22 +250,27 @@ def _validate_and_configure_dflash(args, draft_model_config) -> None:
     Called before dataset loading to fail fast on misconfigurations.
     """
     from torchspec.models.draft.dflash import DFlashConfig
+    from torchspec.models.draft.dspark import DSparkConfig
 
     if not isinstance(draft_model_config, DFlashConfig):
         return
 
+    # DSparkConfig subclasses DFlashConfig
+    is_dspark = isinstance(draft_model_config, DSparkConfig)
+    algo = "DSpark" if is_dspark else "DFlash"
+
     engine_type = getattr(args, "inference_engine_type", "hf")
     if engine_type not in ("vllm", "sgl"):
         raise NotImplementedError(
-            f"DFlash supports inference_engine_type in ('vllm', 'sgl'), got '{engine_type}'."
+            f"{algo} supports inference_engine_type in ('vllm', 'sgl'), got '{engine_type}'."
         )
     if getattr(args, "defer_tokenization", False):
         raise NotImplementedError("DFlash does not support defer_tokenization=True.")
     block_size = getattr(args, "dflash_block_size", 16)
     min_loss = getattr(args, "min_loss_tokens", 0)
     if min_loss < 2 * block_size:
         raise ValueError(
-            f"DFlash requires dataset.min_loss_tokens >= 2 * training.dflash_block_size "
+            f"{algo} requires dataset.min_loss_tokens >= 2 * training.dflash_block_size "
             f"({min_loss} < {2 * block_size}). Set dataset.min_loss_tokens={2 * block_size}."
         )