[Refactor] Promote DFly to a first-class DFlash-family drafter and add an end-to-end multi-step TV loss

LICENSE

@@ -1,8 +1,8 @@
-Tencent is pleased to support the open source community by making AngelSpec available. 
+Tencent is pleased to support the open source community by making AngelSpec available.
 
-Copyright (C) 2026 Tencent.  All rights reserved. 
+Copyright (C) 2026 Tencent.  All rights reserved.
 
-The open-source software included in this distribution may have been modified by Tencent (“Tencent Modifications”). All Tencent Modifications are Copyright (C) Tencent. 
+The open-source software included in this distribution may have been modified by Tencent (“Tencent Modifications”). All Tencent Modifications are Copyright (C) Tencent.
 
 AngelSpec is licensed under the Apache-2.0 except for the third-party components listed below.
 

angelspec/config/dfly_hy3_draft_config.json

@@ -1,7 +1,6 @@
 {
-    "architectures": ["DSparkDraftModel"],
-    "model_type": "dspark",
-    "model_arch": "dfly",
+    "architectures": ["Qwen3DFlyModel"],
+    "model_type": "qwen3",
     "hidden_size": 4096,
     "intermediate_size": 13312,
     "head_dim": 128,
@@ -18,9 +17,6 @@
     "target_layer_ids": [1, 20, 39, 58, 77],
     "mask_token_id": 120009,
     "tie_word_embeddings": false,
-    "markov_rank": 0,
-    "enable_confidence_head": false,
-    "confidence_head_with_markov": false,
     "enable_hidden_correction": true,
     "hidden_correction_intermediate_size": null
 }

angelspec/config/dfly_qwen3_4b_draft_config.json

@@ -1,7 +1,6 @@
 {
-    "architectures": ["DSparkDraftModel"],
-    "model_type": "dspark",
-    "model_arch": "dfly",
+    "architectures": ["Qwen3DFlyModel"],
+    "model_type": "qwen3",
     "hidden_size": 2560,
     "intermediate_size": 9728,
     "head_dim": 128,
@@ -18,9 +17,6 @@
     "target_layer_ids": [1, 9, 17, 25, 33],
     "mask_token_id": 151669,
     "tie_word_embeddings": false,
-    "markov_rank": 0,
-    "enable_confidence_head": false,
-    "confidence_head_with_markov": false,
     "enable_hidden_correction": true,
     "hidden_correction_intermediate_size": null
 }

angelspec/config/dfly_qwen3_8b_draft_config.json

@@ -1,7 +1,6 @@
 {
-    "architectures": ["DSparkDraftModel"],
-    "model_type": "dspark",
-    "model_arch": "dfly",
+    "architectures": ["Qwen3DFlyModel"],
+    "model_type": "qwen3",
     "hidden_size": 4096,
     "intermediate_size": 12288,
     "head_dim": 128,
@@ -18,9 +17,6 @@
     "target_layer_ids": [1, 9, 17, 25, 33],
     "mask_token_id": 151669,
     "tie_word_embeddings": false,
-    "markov_rank": 0,
-    "enable_confidence_head": false,
-    "confidence_head_with_markov": false,
     "enable_hidden_correction": true,
     "hidden_correction_intermediate_size": null
 }

angelspec/config/train_config.py

@@ -49,6 +49,7 @@ class DatasetConfig:
     prompt_key: str = "conversations"
     shuffle_dataset: bool = True
     train_data_path: str = ""
+    num_proc: int = 64
 
 
 @dataclass
@@ -224,13 +225,14 @@ class TrainingConfig:
     dflash_ce_loss_alpha: float = 1.0
     dflash_l1_loss_alpha: float = 0.0
     dflash_kl_loss_weight: float = 0.0
-    dflash_kl_temperature: float = 1.0
     dflash_kl_topk: int = 10
-    dflash_kl_topk_renormalize: bool = True
     dflash_lk_loss_weight: float = 0.0
     dflash_lk_loss_type: str = "hybrid"  # "alpha" or "hybrid"
     dflash_lk_eta: float = 3.0
-    dflash_lk_temperature: float = 1.0
+    # End-to-end multi-step TV loss (γ-step MTP; γ=block_size). Independent term
+    # added on top of the total loss (not mutually exclusive with KL/LK). Needs
+    # target last_hidden_states. 0 disables.
+    dflash_e2e_tv_loss_weight: float = 0.0
     # Gated-sum layer-selection run only (fusion_type=gated_sum in the draft config).
     # Optional sparsity penalty weight on the layer gate: adds
     # `weight * H(softmax(gate))` to the loss to push the gate toward a peakier

angelspec/config/utils.py

@@ -100,7 +100,8 @@ def generate_draft_model_config(
         warnings.warn(
             "No template config provided for draft model. "
             "Auto-generating config entirely from target model. "
-            "Consider providing a template via draft_model_config for full control."
+            "Consider providing a template via draft_model_config for full control.",
+            stacklevel=2,
         )
         draft_config = {
             "architectures": ["LlamaForCausalLMEagle3"],

angelspec/controller/eval.py

@@ -30,7 +30,10 @@
 import wandb
 from tqdm import tqdm
 
-from angelspec.training.checkpoint import _read_checkpoint_metadata, _write_checkpoint_metadata
+from angelspec.training.checkpoint import (
+    _read_checkpoint_metadata,
+    _write_checkpoint_metadata,
+)
 from angelspec.utils.logging import logger
 
 EVAL_CACHE_IDLE_TIMEOUT = 300.0

angelspec/controller/inference_manager.py

@@ -421,9 +421,9 @@ def _prepare_engine_inputs(self, entries: list[InferenceInput]) -> dict:
         if self._defer_tokenization:
             input_ids_ref = None
             packed_loss_mask_list = None
-            assert all(e.formatted_prompt is not None for e in entries), (
-                "formatted_prompt is required when defer_tokenization is True"
-            )
+            assert all(
+                e.formatted_prompt is not None for e in entries
+            ), "formatted_prompt is required when defer_tokenization is True"
             formatted_prompts = [e.formatted_prompt for e in entries]
         else:
             input_ids_ref = ray.put([e.input_ids for e in entries])

angelspec/controller/training_controller.py

@@ -477,9 +477,9 @@ def compute_vocab_mapping(self, target_vocab_size: int, draft_vocab_size: int) -
         from angelspec.data.preprocessing import generate_vocab_mapping
 
         assert self._stored_dataset is not None, "No stored dataset for vocab mapping"
-        assert "input_ids" in self._stored_dataset[0], (
-            "compute_vocab_mapping requires input_ids in dataset. Set defer_tokenization=False to enable tokenization."
-        )
+        assert (
+            "input_ids" in self._stored_dataset[0]
+        ), "compute_vocab_mapping requires input_ids in dataset. Set defer_tokenization=False to enable tokenization."
         return generate_vocab_mapping(
             prompts=self._stored_dataset,
             target_vocab_size=target_vocab_size,

angelspec/data/dataset.py

@@ -27,7 +27,10 @@
 from tqdm import tqdm
 
 from angelspec.data.parse import create_parser, has_thinking_content
-from angelspec.data.preprocessing import _normalize_conversation, preprocess_conversations
+from angelspec.data.preprocessing import (
+    _normalize_conversation,
+    preprocess_conversations,
+)
 from angelspec.data.template import TEMPLATE_REGISTRY
 from angelspec.data.utils import (
     estimate_row_count,

angelspec/data/parse.py

@@ -178,7 +178,8 @@ def format(self, conversation: "Conversation", **kwargs) -> str:
 
         if conversation[0]["role"] == "system":
             warnings.warn(
-                "The first message is from system, we will use the system prompt from the data and ignore the system prompt from the template"
+                "The first message is from system, we will use the system prompt from the data and ignore the system prompt from the template",
+                stacklevel=2,
             )
             messages.append({"role": "system", "content": conversation[0]["content"]})
             conversation = conversation[1:]
@@ -191,19 +192,22 @@ def format(self, conversation: "Conversation", **kwargs) -> str:
             if j == 0:
                 if role != "user":
                     warnings.warn(
-                        f"Conversation must start with a 'user' role, but found '{role}'. Conversation truncated."
+                        f"Conversation must start with a 'user' role, but found '{role}'. Conversation truncated.",
+                        stacklevel=2,
                     )
                     break
             else:
                 prev_role = conversation[j - 1]["role"]
                 if role == "tool" and prev_role not in ["assistant", "tool"]:
                     warnings.warn(
-                        f"A 'tool' message must follow an 'assistant' or 'tool' message, but was preceded by '{prev_role}'. Conversation truncated."
+                        f"A 'tool' message must follow an 'assistant' or 'tool' message, but was preceded by '{prev_role}'. Conversation truncated.",
+                        stacklevel=2,
                     )
                     break
                 if role == "assistant" and prev_role not in ["user", "tool"]:
                     warnings.warn(
-                        f"An 'assistant' message must follow a 'user' or 'tool' message, but was preceded by '{prev_role}'. Conversation truncated."
+                        f"An 'assistant' message must follow a 'user' or 'tool' message, but was preceded by '{prev_role}'. Conversation truncated.",
+                        stacklevel=2,
                     )
                     break
             messages.append(sentence)
@@ -216,7 +220,9 @@ def format(self, conversation: "Conversation", **kwargs) -> str:
         try:
             return self._apply_chat_template(messages, **kwargs)
         except (ValueError, TypeError):
-            warnings.warn("Tokenizer does not have a chat_template, using fallback rendering.")
+            warnings.warn(
+                "Tokenizer does not have a chat_template, using fallback rendering.", stacklevel=2
+            )
             add_generation_prompt = kwargs.get("add_generation_prompt", False)
             parts = []
             bos_token = getattr(self.tokenizer, "bos_token", None)

angelspec/data/preprocessing.py

@@ -283,7 +283,8 @@ def process_token_dict_to_mappings(
     else:
         warnings.warn(
             f"Unique tokens ({len(token_dict)}) exceed draft vocab size ({draft_vocab_size}). "
-            f"{len(token_dict) - draft_vocab_size} tokens will be dropped from the vocab mapping."
+            f"{len(token_dict) - draft_vocab_size} tokens will be dropped from the vocab mapping.",
+            stacklevel=2,
         )
     total_frequency = sum(token_dict.values())
     top_N = token_dict.most_common(draft_vocab_size)

angelspec/data/template.py

@@ -41,9 +41,9 @@ def __init__(self):
         self.templates = {}
 
     def register(self, name: str, template: ChatTemplate, override: bool = False):
-        assert override or name not in self.templates, (
-            f"Chat template for the model type {name} has already been registered"
-        )
+        assert (
+            override or name not in self.templates
+        ), f"Chat template for the model type {name} has already been registered"
         self.templates[name] = template
 
     def get(self, name: str) -> ChatTemplate:

angelspec/inference/engine/mooncake_hidden_states_connector.py

@@ -153,9 +153,9 @@ def __init__(
         self.cache_layers: list[str] = []
         self._cache_layer_group_id: int = self._find_cache_layer_group_id(kv_cache_config)
 
-        assert self._vllm_config.speculative_config is not None, (
-            "MooncakeHiddenStatesConnector requires 'extract_hidden_states' speculative method"
-        )
+        assert (
+            self._vllm_config.speculative_config is not None
+        ), "MooncakeHiddenStatesConnector requires 'extract_hidden_states' speculative method"
         spec_config = self._vllm_config.speculative_config.draft_model_config.hf_config
         self._layer_ids = list(getattr(spec_config, "eagle_aux_hidden_state_layer_ids", []))
         self.num_hidden_states = len(self._layer_ids)
@@ -271,9 +271,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             self._vllm_config, CacheOnlyAttentionLayer, list(kv_caches.keys())
         )
         self.cache_layers = list(layers.keys())
-        assert len(self.cache_layers) == 1, (
-            f"Expected 1 CacheOnlyAttentionLayer, got {len(self.cache_layers)}"
-        )
+        assert (
+            len(self.cache_layers) == 1
+        ), f"Expected 1 CacheOnlyAttentionLayer, got {len(self.cache_layers)}"
 
     def save_kv_layer(
         self,

angelspec/inference/engine/score_worker_ext.py

@@ -121,7 +121,9 @@ def score_packed(
         from vllm.config import CUDAGraphMode
         from vllm.forward_context import set_forward_context
 
-        from angelspec.models.ops.flex_attention import compile_friendly_create_block_mask
+        from angelspec.models.ops.flex_attention import (
+            compile_friendly_create_block_mask,
+        )
         from angelspec.models.ops.tree_mask import create_tree_mask_mod
 
         layers = self._score_attention_layers()  # fail-closed guard

angelspec/inference/engine/sgl_engine.py

@@ -290,9 +290,9 @@ def init(
         else:
             engine_kwargs["disable_cuda_graph"] = True
 
-        assert pre_allocated_port is not None, (
-            f"SglEngine rank {self.rank}: pre_allocated_port is required (ports must be pre-allocated by the factory)"
-        )
+        assert (
+            pre_allocated_port is not None
+        ), f"SglEngine rank {self.rank}: pre_allocated_port is required (ports must be pre-allocated by the factory)"
         engine_kwargs["port"] = pre_allocated_port
         engine_kwargs["nccl_port"] = pre_allocated_port + 1
 

angelspec/models/__init__.py

@@ -19,6 +19,7 @@
 # SOFTWARE.
 
 from angelspec.models.dflash import DFlashModel
+from angelspec.models.dfly import DFlyModel
 from angelspec.models.dspark import DSparkModel
 from angelspec.models.eagle3 import (
     Eagle3Model,
@@ -32,6 +33,7 @@
 __all__ = [
     "Eagle3Model",
     "DFlashModel",
+    "DFlyModel",
     "DSparkModel",
     "MTPModel",
     "compute_lazy_target_padded",