feat: add DFlash2 training support

Add the grouped dynamic causal convolution and predecessor-conditioned candidate selector from the public DFlash2 architecture. Match the public serving implementations for transformed top-K unary logits, bilinear transition scores, verified-anchor alignment, and checkpoint schema.

Reuse the shared DFlash training path, configuration dispatch, optimizer sharding, and strict serving export hooks. Public sources do not publish selector supervision, so define a TorchSpec-specific teacher-forced cross-entropy objective that inserts a missing gold successor into the training candidate set.

Add focused configuration, loss, gradient, masking, inference-lattice parity, distributed export, and artifact tests.

Signed-off-by: Aaron Batilo <AaronBatilo@gmail.com>

Author: abatilo

README.md

@@ -229,6 +229,8 @@ python tools/convert_to_hf.py --input-dir ./outputs/my_experiment/iter_0010000/
 
 Vocabulary pruning, which reduces the draft model `lm_head` to a smaller token set and emits `d2t` and `t2d` mappings, can be applied either during training or at conversion time.
 
+DFlash2 requires the full target vocabulary. Do not set `draft_vocab_size` or pass `--prune-vocab` for a DFlash2 model.
+
 - **Pre-pruning**: set `draft_vocab_size` in your training config. The checkpoint already contains the pruned `lm_head` and `d2t`/`t2d` buffers, so the basic conversion command is enough.
 - **Post-pruning**: train with the full vocabulary, then pass `--prune-vocab` at conversion time together with a representative dataset to compute token frequencies.
 

THIRDPARTYNOTICES

@@ -204,6 +204,31 @@ Notice for THUDM/slime
    See the License for the specific language governing permissions and
    limitations under the License.
 
+Notice for z-lab/dflash
+-------------------------------
+
+MIT License
+
+Copyright (c) 2026 Z Lab
+
+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.
+
 Notice for sgl-project/SpecForge
 -------------------------------
 

configs/sglang_qwen3_8b_dflash2.yaml

@@ -0,0 +1,68 @@
+model:
+  target_model_path: Qwen/Qwen3-8B
+  trust_remote_code: true
+  draft_model_config: torchspec/config/dflash2_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.01
+  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
+  dflash_block_size: 8
+  dflash_num_anchors: 512
+  dflash_loss_decay_gamma: 7.0
+  dflash_num_target_layers: 5
+  dflash2_selector_loss_alpha: 1.0
+
+inference:
+  inference_engine_type: sgl
+  store_last_hidden_states: false
+  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-dflash2
+cache_dir: ./cache/qwen3-8b-dflash2
+model_download_dir: null
+
+debug:
+  save_debug_train_data: null