feat(dflash2): co-train the DFlash2 drafter against vLLM rollout

README.md

@@ -16,7 +16,7 @@ training, and hot-update logic through `verl_speco`.
 - **Multiple drafter backends**: includes EAGLE-1, EAGLE-2, EAGLE3, DFlash,
   DSpark, Domino, and P-EAGLE trainer backends under `verl_speco.backends`.
 - **vLLM and SGLang integration**: supports EAGLE-1, EAGLE-2, EAGLE3, DFlash,
-  and DSpark speculative decoding on vLLM, plus EAGLE3 and DFlash on SGLang,
+  DFlash2, and DSpark speculative decoding on vLLM, plus EAGLE3 and DFlash on SGLang,
   with drafter collection and hot-update logic integrated through the rollout
   engine.
 - **GPU and NPU examples**: provides example scripts for vLLM, SGLang, and
@@ -63,6 +63,7 @@ faster end-to-end training without accuracy regression.
 | EAGLE-2 | vLLM | FSDP | Available |
 | EAGLE3 | vLLM, SGLang | FSDP | Available |
 | DFlash | vLLM, SGLang | FSDP | Available |
+| DFlash2 | vLLM via DFlash | FSDP | Available |
 | DSpark | vLLM | FSDP | Available |
 | Domino | vLLM, SGLang via DFlash | FSDP | Available |
 | P-EAGLE | Not wired in this overlay | FSDP | Training only |
@@ -95,13 +96,26 @@ drafter backend you use.
 | EAGLE-1 / EAGLE-2 | Engine version with native EAGLE support | Runtime-specific | - |
 | EAGLE3 | >= 0.18.0 | >= 0.18.0 | >= 0.5.10 |
 | DFlash | >= 0.20.2 | >= 0.20.2 | >= 0.5.12 |
+| DFlash2 | >= 0.28.0 (served as DFlash) | - | - |
 | DSpark | GPU: [main](https://github.com/vllm-project/vllm/tree/main)<br>NPU: [`dc68bd8`](https://github.com/vllm-project/vllm/tree/dc68bd8c4199b00631fe71eb37313f406cc66ac1) | NPU: [`8214d19`](https://github.com/vllm-project/vllm-ascend/tree/8214d19f8b505484b839469444887b404db2e3a8) | - |
 | Domino | DFlash-compatible runtime with Domino projector support | Runtime-specific | Runtime-specific |
 | P-EAGLE | Not wired | Not wired | Not wired |
 
 For vLLM DFlash, the drafter checkpoint must use the DFlash draft model config
 expected by the runtime.
 
+For vLLM DFlash2, keep `speculative_algorithm=DFLASH2`: the overlay maps it onto
+vLLM's DFlash method and the engine picks the DFlash2 draft (dynamic
+convolutions plus candidate selector) from the checkpoint's `DFlash2DraftModel`
+architecture, so both the drafter training loop and the rollout drafter run
+DFlash2. The checkpoint must use the z-lab layout with the DFlash2 knobs under
+`dflash_config`; `python -m verl_speco.convert_speculators_dflash2` rewrites a
+speculators-format drafter (for example `mgoin/Qwen3-4B-speculator.dflash2`) into
+it. vLLM sizes the convolution block as the bonus token plus
+`rollout.spec_verify_tokens`, so set `spec_verify_tokens = dflash2_block_size - 1`
+(see `examples/run_qwen3-8b_drafter_dflash2_vllm.sh`). SGLang co-training of
+DFlash2 is not wired yet.
+
 For vLLM DSpark on GPU, use vLLM main. For vLLM DSpark on NPU, follow the
 version pairing documented by
 [vLLM-Ascend PR #11153](https://github.com/vllm-project/vllm-ascend/pull/11153):

ci/run_example_test.sh

@@ -15,6 +15,9 @@ case "${platform}/${backend}/${drafter}" in
   gpu/vllm/dspark)
     example="examples/run_qwen3-8b_drafter_dspark_vllm.sh"
     ;;
+  gpu/vllm/dflash2)
+    example="examples/run_qwen3-8b_drafter_dflash2_vllm.sh"
+    ;;
   gpu/vllm/peagle|gpu/vllm/domino)
     example="examples/run_qwen3-8b_drafter_domino_peagle_separate_training.sh"
     ;;
@@ -43,7 +46,7 @@ case "${platform}/${backend}/${drafter}" in
     example="examples/run_qwen3-8b_drafter_dflash_sglang.sh"
     ;;
   *)
-    echo "usage: $0 {gpu|npu} {vllm|sglang} {eagle3|megatron-eagle3|dflash|dspark|peagle|domino}" >&2
+    echo "usage: $0 {gpu|npu} {vllm|sglang} {eagle3|megatron-eagle3|dflash|dflash2|dspark|peagle|domino}" >&2
     exit 2
     ;;
 esac
@@ -83,6 +86,10 @@ case "${drafter}" in
     draft_model="${SPECO_DSPARK_DRAFT_MODEL:-}"
     draft_algorithm="DSPARK"
     ;;
+  dflash2)
+    draft_model="${SPECO_DFLASH2_DRAFT_MODEL:-}"
+    draft_algorithm="DFLASH2"
+    ;;
   peagle)
     draft_model="${SPECO_EAGLE3_DRAFT_MODEL:-}"
     draft_algorithm="EAGLE3"
@@ -228,6 +235,18 @@ if [[ "${drafter}" == "dflash" ]]; then
   )
 fi
 
+if [[ "${drafter}" == "dflash2" ]]; then
+  overrides+=(
+    # vLLM sizes the DFlash2 convolution block as 1 + spec_verify_tokens, and
+    # the trainer folds by dflash2_block_size (default 8), so the two must agree.
+    "actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens=${SPECO_DFLASH2_SPEC_VERIFY_TOKENS:-7}"
+    "actor_rollout_ref.rollout.drafter.training.dflash2_block_size=${SPECO_DFLASH2_BLOCK_SIZE:-8}"
+    "actor_rollout_ref.rollout.drafter.training.dflash2_num_anchors=${SPECO_DFLASH2_NUM_ANCHORS:-8}"
+    "actor_rollout_ref.rollout.drafter.training.dflash2_loss_decay_gamma=${SPECO_DFLASH2_LOSS_DECAY_GAMMA:-7}"
+    "actor_rollout_ref.rollout.drafter.training.dflash_max_window=${SPECO_DFLASH_MAX_WINDOW:-64}"
+  )
+fi
+
 if [[ "${drafter}" == "dspark" ]]; then
   overrides+=(
     "actor_rollout_ref.rollout.drafter.rollout.spec_steps=${SPECO_DSPARK_SPEC_STEPS:-1}"

examples/run_qwen3-8b_drafter_dflash2_vllm.sh

@@ -0,0 +1,110 @@
+set -x
+
+# GPU example for a DFlash2 drafter co-trained against vLLM's DFlash proposer.
+#
+# DFlash2 (dynamic convolutions + candidate selector) is served by vLLM >= 0.28.0
+# through the DFlash speculative method; the engine picks the DFlash2 draft from
+# the checkpoint's "DFlash2DraftModel" architecture. The drafter checkpoint must
+# use the z-lab layout (DFlash2 knobs under dflash_config); convert a
+# speculators-format drafter with
+#   python -m verl_speco.convert_speculators_dflash2 --input ... --target ... --output ...
+# The rollout block is the bonus token plus spec_verify_tokens mask tokens, so
+# spec_verify_tokens must equal dflash2_block_size - 1 (8 -> 7 below).
+project_name='verl_grpo_example_dflash2_drafter'
+exp_name='qwen3_8b_dflash2_drafter_vllm_gpu'
+
+gen_tp=2
+train_sp=1
+ray_num_cpus=${SPECO_RAY_NUM_CPUS:-64}
+ray_worker_soft_limit=${SPECO_RAY_WORKER_SOFT_LIMIT:-8}
+
+MODEL_PATH=/path/to/model
+CKPTS_DIR=/path/to/checkpoint
+TRAIN_FILE=/path/to/train_file
+TEST_FILE=/path/to/test_file
+DRAFTER_PATH=/path/to/vllm-compatible-dflash2-drafter
+
+
+PYTHONUNBUFFERED=1 python3 -m verl_speco.main \
+    algorithm.adv_estimator=grpo \
+    ray_kwargs.ray_init.num_cpus=${ray_num_cpus} \
+    +ray_kwargs.ray_init._system_config.prestart_worker_first_driver=false \
+    +ray_kwargs.ray_init._system_config.num_workers_soft_limit=${ray_worker_soft_limit} \
+    data.train_files=${TRAIN_FILE} \
+    data.val_files=${TEST_FILE} \
+    data.train_batch_size=16 \
+    data.max_prompt_length=512 \
+    data.max_response_length=8192 \
+    data.filter_overlong_prompts=True \
+    data.filter_overlong_prompts_workers=256 \
+    data.truncation='error' \
+    actor_rollout_ref.rollout.temperature=0.6 \
+    actor_rollout_ref.model.path=${MODEL_PATH} \
+    actor_rollout_ref.actor.optim.lr=1e-6 \
+    actor_rollout_ref.model.use_remove_padding=True \
+    actor_rollout_ref.actor.ppo_mini_batch_size=16 \
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \
+    actor_rollout_ref.actor.use_kl_loss=True \
+    actor_rollout_ref.actor.kl_loss_coef=0.001 \
+    actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+    actor_rollout_ref.actor.entropy_coeff=0 \
+    actor_rollout_ref.actor.calculate_entropy=False \
+    actor_rollout_ref.model.enable_gradient_checkpointing=True \
+    actor_rollout_ref.actor.fsdp_config.param_offload=True \
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
+    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${train_sp} \
+    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${train_sp} \
+    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \
+    actor_rollout_ref.actor.use_dynamic_bsz=True \
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \
+    actor_rollout_ref.rollout.name=vllm \
+    actor_rollout_ref.rollout.enforce_eager=False \
+    actor_rollout_ref.rollout.enable_chunked_prefill=True \
+    actor_rollout_ref.rollout.enable_prefix_caching=True \
+    actor_rollout_ref.rollout.max_num_seqs=256 \
+    actor_rollout_ref.rollout.max_num_batched_tokens=12288 \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \
+    actor_rollout_ref.rollout.n=5 \
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \
+    actor_rollout_ref.ref.fsdp_config.param_offload=True \
+    actor_rollout_ref.rollout.drafter.enable=True \
+    actor_rollout_ref.rollout.drafter.enable_drafter_training=True \
+    actor_rollout_ref.rollout.drafter.model_path=${DRAFTER_PATH} \
+    actor_rollout_ref.rollout.drafter.speculative_algorithm=DFLASH2 \
+    actor_rollout_ref.rollout.drafter.training.collect_hidden_states_from_sgl=False \
+    actor_rollout_ref.rollout.drafter.training.collect_hidden_states_from_old_logprob=True \
+    actor_rollout_ref.rollout.drafter.training.old_logprob_hidden_capture_impl=forward_hook \
+    actor_rollout_ref.rollout.drafter.training.dflash2_block_size=8 \
+    actor_rollout_ref.rollout.drafter.training.dflash2_num_anchors=64 \
+    actor_rollout_ref.rollout.drafter.training.dflash2_loss_decay_gamma=7 \
+    actor_rollout_ref.rollout.drafter.training.dflash2_selector_loss_weight=1.0 \
+    actor_rollout_ref.rollout.drafter.training.dflash_max_window=512 \
+    actor_rollout_ref.rollout.drafter.rollout.spec_steps=1 \
+    actor_rollout_ref.rollout.drafter.rollout.spec_topk=1 \
+    actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens=7 \
+    actor_rollout_ref.rollout.drafter.training.step=10 \
+    actor_rollout_ref.rollout.drafter.training.collect_interval_steps=5 \
+    actor_rollout_ref.rollout.drafter.training.training_interval_steps=5 \
+    actor_rollout_ref.rollout.drafter.training.publish_async=True \
+    actor_rollout_ref.rollout.drafter.training.publish_dtype=bf16 \
+    actor_rollout_ref.rollout.drafter.training.draft_update_weights_bucket_megabytes=512 \
+    actor_rollout_ref.rollout.drafter.training.draft_update_pause_generation=True \
+    actor_rollout_ref.rollout.drafter.training.draft_update_flush_before=False \
+    actor_rollout_ref.rollout.drafter.training.draft_update_flush_after=True \
+    actor_rollout_ref.rollout.load_format="auto" \
+    actor_rollout_ref.actor.strategy=fsdp2 \
+    algorithm.use_kl_in_reward=False \
+    trainer.val_before_train=False \
+    trainer.critic_warmup=0 \
+    trainer.logger='["console", "wandb"]' \
+    trainer.project_name=${project_name} \
+    trainer.experiment_name=${exp_name} \
+    trainer.n_gpus_per_node=16 \
+    trainer.nnodes=1 \
+    trainer.default_local_dir=${CKPTS_DIR} \
+    trainer.total_training_steps=100 \
+    trainer.save_freq=20 \
+    trainer.test_freq=5 \
+    trainer.total_epochs=6 $@

pyproject.toml

@@ -23,6 +23,7 @@ Repository = "https://github.com/verl-project/verl-SpeCo"
 verl-speco = "verl_speco.main:main"
 verl-speco-draft-train = "verl_speco.draft_train_launcher:main"
 verl-speco-inspect-features = "verl_speco.inspect_feature_store:main"
+verl-speco-convert-speculators-dflash2 = "verl_speco.convert_speculators_dflash2:main"
 
 [tool.setuptools.dynamic]
 version = { attr = "verl_speco.__version__" }

tests/integration/test_dflash2_backend_contract.py

@@ -22,6 +22,8 @@
 
 from __future__ import annotations
 
+import json
+
 import pytest
 
 
@@ -775,11 +777,47 @@ def test_plain_dflash_checkpoint_warm_starts_without_the_dflash2_modules() -> No
     backend._load_draft_checkpoint(model, "<test>", normalized_state=backbone_only)
 
 
-def test_dflash2_rejected_by_vllm_config_builder() -> None:
+def test_dflash2_maps_to_the_vllm_dflash_method() -> None:
+    """vLLM serves DFlash2 through its DFlash proposer, dispatching on the architecture."""
     from verl_speco.integration.vllm_runtime import _speculative_method_from_drafter
 
-    with pytest.raises(ValueError, match="not an engine-level speculative algorithm"):
-        _speculative_method_from_drafter({"speculative_algorithm": "DFLASH2"})
+    assert _speculative_method_from_drafter({"speculative_algorithm": "DFLASH2"}) == "dflash"
+
+
+def test_dflash2_config_serializes_the_runtime_dflash_config_block() -> None:
+    """A trainer-saved DFlash2 config must stay servable by vLLM (nested knobs)."""
+    pytest.importorskip("torch")
+    pytest.importorskip("transformers")
+    from verl_speco.models.dflash2 import DFlash2Config
+
+    config = DFlash2Config(
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        vocab_size=256,
+        target_layer_ids=[1, 3],
+        mask_token_id=255,
+        block_size=4,
+        selector_top_k=8,
+    )
+
+    serialized = json.loads(config.to_json_string())
+    nested = serialized["dflash_config"]
+    assert nested["block_size"] == 4
+    assert nested["conv_kernel_size"] == 2
+    assert nested["conv_group_size"] == 16
+    assert nested["selector_rank"] == 256
+    assert nested["selector_top_k"] == 8
+    assert nested["mask_token_id"] == 255
+    assert nested["target_layer_ids"] == [1, 3]
+    assert serialized["architectures"] == ["DFlash2DraftModel"]
+
+    # The nested block must round-trip through the overlay's own loader too.
+    reloaded = DFlash2Config.from_dict(serialized)
+    assert reloaded.selector_top_k == 8
+    assert reloaded.block_size == 4
 
 
 def test_dflash2_rejected_by_sglang_config_builder() -> None: