[WIP] Add EDSD: Entropy-Driven Design for Faster Speculative Decoding

README.md

@@ -29,6 +29,7 @@ Check out [**our documentation**](https://docs.sglang.ai/SpecForge/) to get star
 | Method | Description | Training | Example | Optimization |
 | --- | --- | --- | --- | --- |
 | **[EAGLE3](https://arxiv.org/abs/2503.01840)** | Feature-based autoregressive drafting | [`scripts/train_eagle3.py`](./scripts/train_eagle3.py) | [`examples/run_qwen3_8b_eagle3_online.sh`](./examples/run_qwen3_8b_eagle3_online.sh) | [LK loss](https://arxiv.org/pdf/2602.23881)
+| **[EDSD](https://aclanthology.org/2026.acl-long.2145.pdf)** | EAGLE3 with entropy-driven design | [`scripts/train_edsd.py`](./scripts/train_edsd.py) | [`examples/run_qwen3_8b_edsd_online.sh`](./examples/run_qwen3_8b_edsd_online.sh) | 
 | **[DFlash](https://arxiv.org/abs/2602.06036)** | Block-parallel drafting | [`scripts/train_dflash.py`](./scripts/train_dflash.py) | [`examples/run_qwen3_8b_dflash_online.sh`](./examples/run_qwen3_8b_dflash_online.sh) | [D-PACE](https://arxiv.org/abs/2605.18810)
 | **[Domino](https://arxiv.org/html/2605.29707v1)** | DFlash with GRU logit correction | [`scripts/train_domino.py`](./scripts/train_domino.py) | [`examples/run_qwen3_8b_domino_online.sh`](./examples/run_qwen3_8b_domino_online.sh) |
 
@@ -47,6 +48,7 @@ SpecBundle is a collection of production-grade speculative decoding models that
 
 ## 🎉 News
 
+- [2026-07] 🔥 Added EDSD online training for EAGLE3 draft models.
 - [2026-06] 🔥 Added D-PACE as an optional loss for DFlash training.
 - [2026-06] 🔥 Added Domino online training for DFlash draft models.
 - [2026-01] 🔥 Added DFlash block-parallel online training with SGLang serving support.

configs/qwen3-8b-edsd.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "EdsdDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "max_position_embeddings": 40960,
+  "max_window_layers": 36,
+  "model_type": "llama",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 1,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936,
+  "draft_vocab_size": 32000,
+  "target_layer_ids": [23, 34]
+}

examples/run_qwen3_8b_edsd_online.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+ROOT_DIR=$(dirname $SCRIPT_DIR)
+export TORCHINDUCTOR_CACHE_DIR=$ROOT_DIR/cache/compiled_kernels
+
+# support tp8 train eagle3 for Qwen3-4B/8B/32B up to tp_size = 8
+NUM_GPUS=${1:-4}
+TP_SIZE=${4:-1}
+BUILD_DATASET_NUM_PROC=${BUILD_DATASET_NUM_PROC:-64}
+
+torchrun \
+    --standalone \
+    --nproc_per_node $NUM_GPUS \
+    $ROOT_DIR/scripts/train_edsd.py \
+    --target-model-path /root/spec-train/ckpt/Qwen3-8B \
+    --draft-model-config $ROOT_DIR/configs/qwen3-8b-edsd.json \
+    --train-data-path /root/spec-train/data/qwen3-8b-test.jsonl \
+    --build-dataset-num-proc $BUILD_DATASET_NUM_PROC \
+    --output-dir $ROOT_DIR/outputs/qwen3-8b-edsd-sharegpt \
+    --num-epochs 10 \
+    --batch-size 4 \
+    --learning-rate 1e-4 \
+    --max-length 4096 \
+    --chat-template qwen \
+    --cache-dir $ROOT_DIR/cache \
+    --embedding-key model.embed_tokens.weight \
+    --tp-size $TP_SIZE \
+    --target-model-backend sglang \
+    --drop-ratio-scale 0.03 \
+    --report-to tensorboard