diff --git a/README.md b/README.md
index 044aacf8..6193382d 100644
--- a/README.md
+++ b/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,6 +96,7 @@ 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)
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 |
@@ -102,6 +104,18 @@ drafter backend you use.
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):
diff --git a/ci/run_example_test.sh b/ci/run_example_test.sh
index 531562a8..76b9c4f0 100644
--- a/ci/run_example_test.sh
+++ b/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}"
diff --git a/examples/run_qwen3-8b_drafter_dflash2_vllm.sh b/examples/run_qwen3-8b_drafter_dflash2_vllm.sh
new file mode 100644
index 00000000..36d87db8
--- /dev/null
+++ b/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 $@
diff --git a/pyproject.toml b/pyproject.toml
index ca8a9c7b..e4755cb2 100644
--- a/pyproject.toml
+++ b/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__" }
diff --git a/tests/integration/test_dflash2_backend_contract.py b/tests/integration/test_dflash2_backend_contract.py
index 576f7f05..82a93d38 100644
--- a/tests/integration/test_dflash2_backend_contract.py
+++ b/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, "", 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:
diff --git a/tests/integration/test_vllm_dflash2_runtime_contract.py b/tests/integration/test_vllm_dflash2_runtime_contract.py
new file mode 100644
index 00000000..9bc94024
--- /dev/null
+++ b/tests/integration/test_vllm_dflash2_runtime_contract.py
@@ -0,0 +1,298 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""vLLM co-training contract for the DFlash2 drafter (served through DFlash)."""
+
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from verl_speco.integration import vllm_runtime
+from verl_speco.integration.vllm_runtime import (
+ _assert_vllm_supports_dflash2,
+ _dflash2_engine_param_name,
+ _draft_param_name_candidates,
+ _normalize_dflash2_runtime_aliases,
+ _validate_vllm_dflash2_block_size,
+ _validate_vllm_dflash_drafter_config,
+ build_vllm_speculative_config_from_drafter,
+)
+
+_DFLASH2_CONFIG = {
+ "architectures": ["DFlash2DraftModel"],
+ "model_type": "qwen3",
+ "dflash_config": {
+ "block_size": 8,
+ "conv_kernel_size": 2,
+ "conv_group_size": 16,
+ "selector_rank": 256,
+ "selector_top_k": 16,
+ "mask_token_id": 151669,
+ "target_layer_ids": [0, 8, 16, 24, 32],
+ },
+}
+
+
+def _write_drafter(tmp_path, config, name="dflash2-drafter"):
+ model_path = tmp_path / name
+ model_path.mkdir()
+ (model_path / "config.json").write_text(json.dumps(config), encoding="utf-8")
+ return model_path
+
+
+def _drafter(model_path, **overrides):
+ config = {
+ "enable": True,
+ "enable_drafter_training": True,
+ "speculative_algorithm": "DFLASH2",
+ "model_path": str(model_path),
+ "rollout": {"spec_steps": 1, "spec_verify_tokens": 7},
+ "training": {"dflash2_block_size": 8},
+ "vllm": {},
+ }
+ config.update(overrides)
+ return config
+
+
+@pytest.fixture
+def vllm_has_dflash2(monkeypatch):
+ monkeypatch.setattr(vllm_runtime, "_vllm_supports_dflash2", lambda: True)
+
+
+def test_dflash2_speculative_config_is_the_dflash_contract(
+ tmp_path, vllm_has_dflash2
+) -> None:
+ model_path = _write_drafter(tmp_path, _DFLASH2_CONFIG)
+
+ config = build_vllm_speculative_config_from_drafter(_drafter(model_path))
+
+ assert config == {
+ "draft_sample_method": "greedy",
+ "method": "dflash",
+ "model": str(model_path),
+ "num_speculative_tokens": 7,
+ }
+
+
+def test_dflash2_speculative_config_needs_spec_verify_tokens(
+ tmp_path, vllm_has_dflash2
+) -> None:
+ model_path = _write_drafter(tmp_path, _DFLASH2_CONFIG)
+
+ with pytest.raises(ValueError, match="spec_verify_tokens"):
+ build_vllm_speculative_config_from_drafter(
+ _drafter(model_path, rollout={"spec_steps": 3})
+ )
+
+
+def test_dflash2_speculative_config_refuses_a_vllm_without_dflash2(
+ tmp_path, monkeypatch
+) -> None:
+ model_path = _write_drafter(tmp_path, _DFLASH2_CONFIG)
+ monkeypatch.setattr(vllm_runtime, "_vllm_supports_dflash2", lambda: False)
+
+ with pytest.raises(ValueError, match="vLLM 0.28.0"):
+ build_vllm_speculative_config_from_drafter(_drafter(model_path))
+
+
+def test_dflash2_capability_probe_is_skipped_without_vllm(monkeypatch) -> None:
+ monkeypatch.setattr(vllm_runtime, "_vllm_supports_dflash2", lambda: None)
+ _assert_vllm_supports_dflash2()
+
+
+def test_dflash2_capability_probe_reads_the_installed_vllm(monkeypatch) -> None:
+ import importlib.util
+
+ real_find_spec = importlib.util.find_spec
+
+ def fake_find_spec(name, *args, **kwargs):
+ if name == "vllm":
+ return object()
+ if name == vllm_runtime._VLLM_DFLASH2_MODULE:
+ return None
+ return real_find_spec(name, *args, **kwargs)
+
+ monkeypatch.setattr(importlib.util, "find_spec", fake_find_spec)
+ assert vllm_runtime._vllm_supports_dflash2() is False
+
+ monkeypatch.setattr(importlib.util, "find_spec", lambda name, *a, **k: None)
+ assert vllm_runtime._vllm_supports_dflash2() is None
+
+
+def test_dflash2_validator_rejects_a_plain_dflash_checkpoint(tmp_path) -> None:
+ """A DFlash checkpoint would load and silently serve without the DFlash2 modules."""
+ model_path = _write_drafter(
+ tmp_path,
+ {"architectures": ["DFlashDraftModel"], **_DFLASH2_CONFIG["dflash_config"]},
+ )
+
+ with pytest.raises(ValueError, match="DFlash2 drafter checkpoint"):
+ _validate_vllm_dflash_drafter_config(model_path, algorithm="DFLASH2")
+
+
+def test_dflash2_validator_requires_the_runtime_hyperparameters(tmp_path) -> None:
+ config = json.loads(json.dumps(_DFLASH2_CONFIG))
+ del config["dflash_config"]["selector_top_k"]
+ del config["dflash_config"]["conv_group_size"]
+ model_path = _write_drafter(tmp_path, config)
+
+ with pytest.raises(
+ ValueError, match=r"missing \['conv_group_size', 'selector_top_k'\]"
+ ):
+ _validate_vllm_dflash_drafter_config(model_path, algorithm="DFLASH2")
+
+
+def test_dflash2_validator_accepts_top_level_hyperparameters(tmp_path) -> None:
+ """Configs saved by the trainer keep the knobs flat; the alias patch nests them."""
+ config = {
+ "architectures": ["Qwen3DFlash2Model"],
+ **_DFLASH2_CONFIG["dflash_config"],
+ }
+ model_path = _write_drafter(tmp_path, config)
+
+ _validate_vllm_dflash_drafter_config(model_path, algorithm="DFLASH2")
+
+
+def test_dflash2_validator_skips_a_missing_checkpoint(tmp_path) -> None:
+ _validate_vllm_dflash_drafter_config(tmp_path / "absent", algorithm="DFLASH2")
+ _validate_vllm_dflash_drafter_config(None, algorithm="DFLASH2")
+
+
+def test_dflash2_block_size_must_match_the_engine_block(
+ tmp_path, vllm_has_dflash2
+) -> None:
+ model_path = _write_drafter(tmp_path, _DFLASH2_CONFIG)
+
+ with pytest.raises(ValueError, match="spec_verify_tokens=16 but block_size=8"):
+ build_vllm_speculative_config_from_drafter(
+ _drafter(model_path, rollout={"spec_verify_tokens": 16})
+ )
+
+
+def test_dflash2_block_size_falls_back_to_the_checkpoint() -> None:
+ config = json.loads(json.dumps(_DFLASH2_CONFIG))
+ _validate_vllm_dflash2_block_size(config, {"training": {}}, 7)
+ with pytest.raises(ValueError, match="block_size=8"):
+ _validate_vllm_dflash2_block_size(config, {"training": {}}, 3)
+ # Training config wins over the checkpoint.
+ _validate_vllm_dflash2_block_size(
+ config, {"training": {"dflash2_block_size": 4}}, 3
+ )
+ # Nothing to compare against: accept.
+ _validate_vllm_dflash2_block_size(None, {"training": {}}, 3)
+ _validate_vllm_dflash2_block_size({"architectures": ["DFlash2DraftModel"]}, {}, 3)
+
+
+def test_dflash2_codebooks_publish_under_the_engine_parameter_names() -> None:
+ """Trainer nn.Embedding ``.weight`` -> vLLM bare ``nn.Parameter``."""
+ candidates = _draft_param_name_candidates(
+ "draft_model.candidate_selector.predecessor_codebook.weight"
+ )
+ assert "model.candidate_selector.predecessor_codebook" in candidates
+ assert "candidate_selector.predecessor_codebook" in candidates
+ # The spelled-out name is still tried first for engines that keep the module.
+ assert candidates.index(
+ "candidate_selector.predecessor_codebook.weight"
+ ) < candidates.index("candidate_selector.predecessor_codebook")
+
+ successor = _draft_param_name_candidates(
+ "module.draft_model.candidate_selector.successor_codebook.weight"
+ )
+ assert "model.candidate_selector.successor_codebook" in successor
+
+ # Every other DFlash2 parameter already matches the engine spelling.
+ projection = _draft_param_name_candidates(
+ "draft_model.candidate_selector.hidden_projection.weight"
+ )
+ assert "model.candidate_selector.hidden_projection.weight" in projection
+ assert not any(name.endswith("hidden_projection") for name in projection)
+ conv = _draft_param_name_candidates(
+ "draft_model.layers.0.attention_conv.base_kernel"
+ )
+ assert "model.layers.0.attention_conv.base_kernel" in conv
+
+
+def test_dflash2_runtime_aliases_nest_flat_hyperparameters() -> None:
+ config = {
+ "architectures": ["DFlash2DraftModel"],
+ "block_size": 8,
+ "conv_kernel_size": 2,
+ "conv_group_size": 16,
+ "selector_rank": 256,
+ "selector_top_k": 16,
+ "dflash_config": {"target_layer_ids": [0, 8]},
+ }
+
+ assert _normalize_dflash2_runtime_aliases(config) is True
+ assert config["dflash_config"] == {
+ "target_layer_ids": [0, 8],
+ "block_size": 8,
+ "conv_kernel_size": 2,
+ "conv_group_size": 16,
+ "selector_rank": 256,
+ "selector_top_k": 16,
+ }
+ # Idempotent once nested.
+ assert _normalize_dflash2_runtime_aliases(config) is False
+
+
+def test_dflash2_runtime_aliases_create_the_nested_block_and_ignore_others() -> None:
+ config = {"architectures": ["Qwen3DFlash2Model"], "selector_top_k": 4}
+ assert _normalize_dflash2_runtime_aliases(config) is True
+ assert config["dflash_config"] == {"selector_top_k": 4}
+
+ dflash = {"architectures": ["DFlashDraftModel"], "selector_top_k": 4}
+ assert _normalize_dflash2_runtime_aliases(dflash) is False
+ assert "dflash_config" not in dflash
+
+ assert (
+ _normalize_dflash2_runtime_aliases({"architectures": "DFlash2DraftModel"})
+ is False
+ )
+
+
+def test_dflash2_runtime_aliases_refuse_conflicts() -> None:
+ config = {
+ "architectures": ["DFlash2DraftModel"],
+ "selector_top_k": 16,
+ "dflash_config": {"selector_top_k": 8},
+ }
+ with pytest.raises(ValueError, match="selector_top_k conflicts"):
+ _normalize_dflash2_runtime_aliases(config)
+
+ with pytest.raises(TypeError, match="must be a mapping"):
+ _normalize_dflash2_runtime_aliases(
+ {"architectures": ["DFlash2DraftModel"], "dflash_config": "bad"}
+ )
+
+
+def test_dflash2_engine_param_name_renames_only_the_codebooks() -> None:
+ """The IPC receiver hands vLLM's load_weights the engine spelling."""
+ assert (
+ _dflash2_engine_param_name("candidate_selector.predecessor_codebook.weight")
+ == "candidate_selector.predecessor_codebook"
+ )
+ assert (
+ _dflash2_engine_param_name("candidate_selector.successor_codebook.weight")
+ == "candidate_selector.successor_codebook"
+ )
+ for untouched in (
+ "candidate_selector.hidden_projection.weight",
+ "candidate_selector.predecessor_codebook",
+ "layers.0.attention_conv.base_kernel",
+ "layers.0.mlp_conv.kernel_projection.weight",
+ "fc.weight",
+ ):
+ assert _dflash2_engine_param_name(untouched) == untouched
diff --git a/tests/integration/test_vllm_draft_update_allocator_contract.py b/tests/integration/test_vllm_draft_update_allocator_contract.py
new file mode 100644
index 00000000..1124c0b6
--- /dev/null
+++ b/tests/integration/test_vllm_draft_update_allocator_contract.py
@@ -0,0 +1,64 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""The vLLM draft publish stages its CUDA-IPC buckets in non-expandable segments."""
+
+from __future__ import annotations
+
+import sys
+import types
+
+import pytest
+
+from verl_speco.integration.vllm_runtime import _ipc_safe_allocator
+
+
+@pytest.fixture
+def fake_verl_device(monkeypatch):
+ calls: list[bool] = []
+ # The CPU CI runs without verl installed; stub the parents so the import
+ # machinery reaches the fake leaf module.
+ for parent in ("verl", "verl.utils"):
+ if parent not in sys.modules:
+ monkeypatch.setitem(sys.modules, parent, types.ModuleType(parent))
+ module = types.ModuleType("verl.utils.device")
+ module.set_expandable_segments = calls.append # type: ignore[attr-defined]
+ monkeypatch.setitem(sys.modules, "verl.utils.device", module)
+ return calls
+
+
+def test_ipc_send_disables_expandable_segments_and_restores_them(
+ fake_verl_device,
+) -> None:
+ with _ipc_safe_allocator(True):
+ assert fake_verl_device == [False]
+ assert fake_verl_device == [False, True]
+
+
+def test_ipc_send_restores_expandable_segments_on_failure(fake_verl_device) -> None:
+ with pytest.raises(RuntimeError, match="send failed"):
+ with _ipc_safe_allocator(True):
+ raise RuntimeError("send failed")
+ assert fake_verl_device == [False, True]
+
+
+def test_shm_transport_leaves_the_allocator_alone(fake_verl_device) -> None:
+ with _ipc_safe_allocator(False):
+ pass
+ assert fake_verl_device == []
+
+
+def test_verl_without_the_helper_is_a_no_op(monkeypatch) -> None:
+ monkeypatch.setitem(sys.modules, "verl.utils.device", None)
+ with _ipc_safe_allocator(True):
+ pass
diff --git a/tests/unit/test_convert_speculators_dflash2.py b/tests/unit/test_convert_speculators_dflash2.py
new file mode 100644
index 00000000..c0e235c8
--- /dev/null
+++ b/tests/unit/test_convert_speculators_dflash2.py
@@ -0,0 +1,215 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import json
+import os
+
+import pytest
+
+from verl_speco.convert_speculators_dflash2 import (
+ convert_speculators_dflash2_checkpoint,
+ convert_speculators_dflash2_config,
+ main,
+)
+
+
+def _speculators_config(**overrides):
+ config = {
+ "architectures": ["DFlash2DraftModel"],
+ "aux_hidden_state_layer_ids": [1, 9, 17, 25, 33],
+ "block_size": 8,
+ "conv_group_size": 16,
+ "conv_kernel_size": 2,
+ "draft_vocab_size": 151936,
+ "dtype": "bfloat16",
+ "mask_token_id": 151669,
+ "selector_rank": 256,
+ "selector_top_k": 16,
+ "speculators_config": {"algorithm": "dflash2"},
+ "speculators_model_type": "dflash2",
+ "transformer_layer_config": {
+ "attention_bias": False,
+ "head_dim": 128,
+ "hidden_act": "silu",
+ "hidden_size": 2560,
+ "intermediate_size": 9728,
+ "layer_types": ["sliding_attention"] * 5,
+ "max_position_embeddings": 40960,
+ "max_window_layers": 28,
+ "model_type": "qwen3",
+ "num_attention_heads": 32,
+ "num_hidden_layers": 5,
+ "num_key_value_heads": 8,
+ "rms_norm_eps": 1e-6,
+ "rope_parameters": {"rope_theta": 1000000, "rope_type": "default"},
+ "sliding_window": 2048,
+ "use_sliding_window": True,
+ "vocab_size": 151936,
+ },
+ }
+ config.update(overrides)
+ return config
+
+
+_TARGET_CONFIG = {
+ "num_hidden_layers": 36,
+ "bos_token_id": 151643,
+ "eos_token_id": 151645,
+ "vocab_size": 151936,
+}
+
+
+def test_convert_emits_the_zlab_layout() -> None:
+ converted = convert_speculators_dflash2_config(
+ _speculators_config(), _TARGET_CONFIG
+ )
+
+ assert converted["architectures"] == ["DFlash2DraftModel"]
+ assert converted["model_type"] == "qwen3"
+ assert converted["is_causal"] is False
+ assert converted["num_target_layers"] == 36
+ assert converted["hidden_size"] == 2560
+ assert converted["num_hidden_layers"] == 5
+ assert converted["rope_theta"] == 1000000
+ assert converted["rope_parameters"] == {
+ "rope_theta": 1000000,
+ "rope_type": "default",
+ }
+ assert converted["eos_token_id"] == 151645
+ assert converted["bos_token_id"] == 151643
+ assert converted["dflash_config"] == {
+ "block_size": 8,
+ "conv_group_size": 16,
+ "conv_kernel_size": 2,
+ "mask_token_id": 151669,
+ "selector_rank": 256,
+ "selector_top_k": 16,
+ # hidden_states indices minus one: decoder-layer indices.
+ "target_layer_ids": [0, 8, 16, 24, 32],
+ }
+ # Speculators keys must not leak into the runtime contract.
+ assert "speculators_config" not in converted
+ assert "transformer_layer_config" not in converted
+ assert "aux_hidden_state_layer_ids" not in converted
+
+
+def test_convert_defaults_to_full_attention_but_can_keep_the_window() -> None:
+ full = convert_speculators_dflash2_config(_speculators_config(), _TARGET_CONFIG)
+ assert full["use_sliding_window"] is False
+ assert full["sliding_window"] is None
+ assert full["layer_types"] == ["full_attention"] * 5
+
+ windowed = convert_speculators_dflash2_config(
+ _speculators_config(), _TARGET_CONFIG, keep_sliding_window=True
+ )
+ assert windowed["use_sliding_window"] is True
+ assert windowed["sliding_window"] == 2048
+ assert windowed["layer_types"] == ["sliding_attention"] * 5
+
+
+def test_convert_prefers_explicit_target_layer_ids() -> None:
+ converted = convert_speculators_dflash2_config(
+ _speculators_config(target_layer_ids=[2, 4, 6, 8, 10]), _TARGET_CONFIG
+ )
+ assert converted["dflash_config"]["target_layer_ids"] == [2, 4, 6, 8, 10]
+
+
+@pytest.mark.parametrize(
+ ("overrides", "match"),
+ [
+ (
+ {"speculators_model_type": "eagle3", "speculators_config": {}},
+ "not a speculators DFlash2",
+ ),
+ ({"transformer_layer_config": None}, "transformer_layer_config"),
+ ({"aux_hidden_state_layer_ids": [1, 9, 17, 25, 40]}, "exceed the target"),
+ ({"aux_hidden_state_layer_ids": [0, 9, 17, 25, 33]}, "must be >= 1"),
+ ({"draft_vocab_size": 32000}, "share the target vocabulary"),
+ ({"mask_token_id": None}, "lacks mask_token_id"),
+ ({"selector_rank": None}, "lacks \\['selector_rank'\\]"),
+ ],
+)
+def test_convert_rejects_malformed_configs(overrides, match) -> None:
+ with pytest.raises(ValueError, match=match):
+ convert_speculators_dflash2_config(
+ _speculators_config(**overrides), _TARGET_CONFIG
+ )
+
+
+def test_convert_rejects_missing_layer_ids() -> None:
+ config = _speculators_config()
+ del config["aux_hidden_state_layer_ids"]
+ with pytest.raises(ValueError, match="neither target_layer_ids"):
+ convert_speculators_dflash2_config(config, _TARGET_CONFIG)
+
+
+def _write_dirs(tmp_path):
+ source = tmp_path / "speculators"
+ target = tmp_path / "target"
+ source.mkdir()
+ target.mkdir()
+ (source / "config.json").write_text(
+ json.dumps(_speculators_config()), encoding="utf-8"
+ )
+ (source / "model.safetensors").write_bytes(b"weights")
+ (target / "config.json").write_text(json.dumps(_TARGET_CONFIG), encoding="utf-8")
+ return source, target
+
+
+def test_convert_checkpoint_writes_config_and_copies_weights(tmp_path) -> None:
+ source, target = _write_dirs(tmp_path)
+ output = tmp_path / "converted"
+
+ converted = convert_speculators_dflash2_checkpoint(
+ str(source), str(target), str(output)
+ )
+
+ written = json.loads((output / "config.json").read_text(encoding="utf-8"))
+ assert written == converted
+ assert (output / "model.safetensors").read_bytes() == b"weights"
+ assert not os.path.islink(output / "model.safetensors")
+
+
+def test_convert_checkpoint_can_symlink_weights_and_overwrites(tmp_path) -> None:
+ source, target = _write_dirs(tmp_path)
+ output = tmp_path / "converted"
+ output.mkdir()
+ (output / "model.safetensors").write_bytes(b"stale")
+
+ convert_speculators_dflash2_checkpoint(
+ str(source), str(target), str(output), link_weights=True
+ )
+
+ assert os.path.islink(output / "model.safetensors")
+ assert (output / "model.safetensors").read_bytes() == b"weights"
+
+
+def test_convert_checkpoint_requires_weights(tmp_path) -> None:
+ source, target = _write_dirs(tmp_path)
+ (source / "model.safetensors").unlink()
+ with pytest.raises(FileNotFoundError, match="no safetensors"):
+ convert_speculators_dflash2_checkpoint(
+ str(source), str(target), str(tmp_path / "out")
+ )
+
+
+def test_main_converts_and_reports(tmp_path, capsys) -> None:
+ source, target = _write_dirs(tmp_path)
+ output = tmp_path / "converted"
+
+ main(["--input", str(source), "--target", str(target), "--output", str(output)])
+
+ assert (output / "config.json").exists()
+ assert "DFlash2DraftModel" in capsys.readouterr().out
diff --git a/verl_speco/backends/dflash2_trainer_backend.py b/verl_speco/backends/dflash2_trainer_backend.py
index 7434ed24..8382568a 100644
--- a/verl_speco/backends/dflash2_trainer_backend.py
+++ b/verl_speco/backends/dflash2_trainer_backend.py
@@ -161,6 +161,8 @@ class DFlash2TrainerBackend(DFlashTrainerBackend):
# ``nn.Parameter`` tensors, while this overlay holds them in ``nn.Embedding``
# modules, whose state dict spells the same tensor with a trailing
# ``.weight``. Every other DFlash2 parameter name matches upstream exactly.
+ # Publish-side inverse: vllm_runtime._dflash2_engine_param_name — keep the
+ # two in sync when adding aliases here.
_CHECKPOINT_KEY_ALIASES = {
"candidate_selector.predecessor_codebook": (
"candidate_selector.predecessor_codebook.weight"
diff --git a/verl_speco/convert_speculators_dflash2.py b/verl_speco/convert_speculators_dflash2.py
new file mode 100644
index 00000000..ed88bed9
--- /dev/null
+++ b/verl_speco/convert_speculators_dflash2.py
@@ -0,0 +1,278 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Convert a speculators-format DFlash2 drafter into the z-lab checkpoint layout.
+
+Public DFlash2 drafters trained with the ``speculators`` library (for example
+``mgoin/Qwen3-4B-speculator.dflash2``) ship a ``config.json`` that nests the
+transformer hyperparameters under ``transformer_layer_config`` and describes the
+target through ``speculators_config``. Both this overlay's trainer
+(``DFlash2Config.from_dflash2_pretrained``) and vLLM releases up to 0.28.0 read
+the released z-lab layout instead: a flat Qwen3-style config with the DFlash2
+knobs under ``dflash_config``. The weights use the same parameter names in both
+layouts, so only ``config.json`` is rewritten; the safetensors are copied as-is.
+
+Run::
+
+ python -m verl_speco.convert_speculators_dflash2 \
+ --input /path/to/speculators-dflash2 \
+ --target /path/to/target-model \
+ --output /path/to/dflash2-drafter
+"""
+
+from __future__ import annotations
+
+import argparse
+import glob
+import json
+import os
+import shutil
+from typing import Any
+
+# Transformer hyperparameters copied verbatim from ``transformer_layer_config``.
+_TRANSFORMER_KEYS = (
+ "attention_bias",
+ "attention_dropout",
+ "head_dim",
+ "hidden_act",
+ "hidden_size",
+ "initializer_range",
+ "intermediate_size",
+ "max_position_embeddings",
+ "num_attention_heads",
+ "num_hidden_layers",
+ "num_key_value_heads",
+ "rms_norm_eps",
+ "rope_scaling",
+ "rope_theta",
+ "rope_parameters",
+ "vocab_size",
+)
+# Sliding-window attention knobs. The trainer's DFlash backbone attends over the
+# full anchored context, so the converted drafter defaults to full attention to
+# keep training and rollout on the same model.
+_SLIDING_WINDOW_KEYS = (
+ "layer_types",
+ "max_window_layers",
+ "sliding_window",
+ "use_sliding_window",
+)
+_DFLASH2_KEYS = (
+ "block_size",
+ "conv_kernel_size",
+ "conv_group_size",
+ "selector_rank",
+ "selector_top_k",
+)
+_WEIGHT_PATTERNS = ("*.safetensors", "*.safetensors.index.json")
+
+
+def _load_json(path: str) -> dict[str, Any]:
+ with open(path, "r", encoding="utf-8") as f:
+ loaded = json.load(f)
+ if not isinstance(loaded, dict):
+ raise ValueError(f"{path} does not hold a JSON object")
+ return loaded
+
+
+def _target_layer_ids(speculators_config: dict[str, Any]) -> list[int]:
+ """The DFlash context layers, as decoder-layer indices.
+
+ speculators records ``aux_hidden_state_layer_ids`` as indices into
+ ``output_hidden_states`` (entry 0 is the embedding output), while the z-lab
+ layout's ``target_layer_ids`` count decoder layers, hence the ``- 1``; this
+ matches the ``eagle_aux_hidden_state_layer_ids = target_layer_ids + 1`` alias
+ vLLM applies in the other direction.
+ """
+ if "target_layer_ids" in speculators_config:
+ return [int(layer_id) for layer_id in speculators_config["target_layer_ids"]]
+ aux_layer_ids = speculators_config.get("aux_hidden_state_layer_ids")
+ if not aux_layer_ids:
+ raise ValueError(
+ "speculators DFlash2 config has neither target_layer_ids nor "
+ "aux_hidden_state_layer_ids"
+ )
+ layer_ids = [int(layer_id) - 1 for layer_id in aux_layer_ids]
+ if min(layer_ids) < 0:
+ raise ValueError(
+ f"aux_hidden_state_layer_ids={aux_layer_ids!r} must be >= 1 (hidden_states indices)"
+ )
+ return layer_ids
+
+
+def convert_speculators_dflash2_config(
+ speculators_config: dict[str, Any],
+ target_config: dict[str, Any],
+ *,
+ keep_sliding_window: bool = False,
+) -> dict[str, Any]:
+ """Build the z-lab layout ``config.json`` for a speculators DFlash2 drafter."""
+ model_type = str(speculators_config.get("speculators_model_type") or "").lower()
+ algorithm = str(
+ (speculators_config.get("speculators_config") or {}).get("algorithm") or ""
+ ).lower()
+ if "dflash2" not in (model_type, algorithm):
+ raise ValueError(
+ "not a speculators DFlash2 config: speculators_model_type="
+ f"{model_type!r} algorithm={algorithm!r}"
+ )
+ transformer = speculators_config.get("transformer_layer_config")
+ if not isinstance(transformer, dict):
+ raise ValueError("speculators DFlash2 config lacks transformer_layer_config")
+
+ target_text = target_config.get("text_config") or target_config
+ target_num_hidden_layers = int(target_text["num_hidden_layers"])
+ target_layer_ids = _target_layer_ids(speculators_config)
+ if max(target_layer_ids) >= target_num_hidden_layers:
+ raise ValueError(
+ f"target_layer_ids={target_layer_ids} exceed the target's "
+ f"num_hidden_layers={target_num_hidden_layers}"
+ )
+
+ converted: dict[str, Any] = {
+ "architectures": ["DFlash2DraftModel"],
+ "model_type": str(transformer.get("model_type") or "qwen3"),
+ "is_causal": False,
+ "tie_word_embeddings": False,
+ "use_cache": True,
+ "dtype": str(speculators_config.get("dtype") or "bfloat16"),
+ }
+ for key in _TRANSFORMER_KEYS:
+ if transformer.get(key) is not None:
+ converted[key] = transformer[key]
+ rope_parameters = transformer.get("rope_parameters")
+ # transformers < 5 reads rope_theta at the top level.
+ if (
+ "rope_theta" not in converted
+ and isinstance(rope_parameters, dict)
+ and rope_parameters.get("rope_theta") is not None
+ ):
+ converted["rope_theta"] = rope_parameters["rope_theta"]
+ if keep_sliding_window:
+ for key in _SLIDING_WINDOW_KEYS:
+ if key in transformer:
+ converted[key] = transformer[key]
+ else:
+ converted["use_sliding_window"] = False
+ converted["sliding_window"] = None
+ converted["layer_types"] = ["full_attention"] * int(
+ converted["num_hidden_layers"]
+ )
+ for key in ("bos_token_id", "eos_token_id", "pad_token_id"):
+ value = transformer.get(key)
+ if value is None:
+ value = target_text.get(key)
+ converted[key] = value
+
+ draft_vocab_size = speculators_config.get("draft_vocab_size")
+ if draft_vocab_size is not None and int(draft_vocab_size) != int(
+ converted["vocab_size"]
+ ):
+ raise ValueError(
+ "DFlash2 drafts share the target vocabulary; "
+ f"draft_vocab_size={draft_vocab_size} != vocab_size={converted['vocab_size']}"
+ )
+
+ mask_token_id = speculators_config.get("mask_token_id")
+ if mask_token_id is None:
+ raise ValueError("speculators DFlash2 config lacks mask_token_id")
+ dflash_config = {
+ key: speculators_config[key]
+ for key in _DFLASH2_KEYS
+ if speculators_config.get(key) is not None
+ }
+ missing = sorted(set(_DFLASH2_KEYS) - set(dflash_config))
+ if missing:
+ raise ValueError(f"speculators DFlash2 config lacks {missing}")
+ dflash_config["mask_token_id"] = int(mask_token_id)
+ dflash_config["target_layer_ids"] = target_layer_ids
+ converted["dflash_config"] = dflash_config
+ converted["num_target_layers"] = target_num_hidden_layers
+ return converted
+
+
+def convert_speculators_dflash2_checkpoint(
+ input_dir: str,
+ target_dir: str,
+ output_dir: str,
+ *,
+ keep_sliding_window: bool = False,
+ link_weights: bool = False,
+) -> dict[str, Any]:
+ """Write the converted drafter to ``output_dir`` and return its config."""
+ speculators_config = _load_json(os.path.join(input_dir, "config.json"))
+ target_config = _load_json(os.path.join(target_dir, "config.json"))
+ converted = convert_speculators_dflash2_config(
+ speculators_config, target_config, keep_sliding_window=keep_sliding_window
+ )
+
+ weight_files = sorted(
+ path
+ for pattern in _WEIGHT_PATTERNS
+ for path in glob.glob(os.path.join(input_dir, pattern))
+ )
+ if not weight_files:
+ raise FileNotFoundError(f"no safetensors weights under {input_dir}")
+
+ os.makedirs(output_dir, exist_ok=True)
+ with open(os.path.join(output_dir, "config.json"), "w", encoding="utf-8") as f:
+ json.dump(converted, f, indent=2, sort_keys=True)
+ f.write("\n")
+ for source in weight_files:
+ destination = os.path.join(output_dir, os.path.basename(source))
+ if os.path.lexists(destination):
+ os.remove(destination)
+ if link_weights:
+ os.symlink(os.path.abspath(source), destination)
+ else:
+ shutil.copy2(source, destination)
+ return converted
+
+
+def main(argv: list[str] | None = None) -> None:
+ parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+ parser.add_argument(
+ "--input", required=True, help="speculators-format DFlash2 drafter directory"
+ )
+ parser.add_argument(
+ "--target",
+ required=True,
+ help="target model directory (its config.json supplies the layer count and special tokens)",
+ )
+ parser.add_argument("--output", required=True, help="converted drafter directory")
+ parser.add_argument(
+ "--keep-sliding-window",
+ action="store_true",
+ help="keep the speculators sliding-window attention settings instead of full attention",
+ )
+ parser.add_argument(
+ "--link-weights",
+ action="store_true",
+ help="symlink the safetensors into the output instead of copying them",
+ )
+ args = parser.parse_args(argv)
+ converted = convert_speculators_dflash2_checkpoint(
+ args.input,
+ args.target,
+ args.output,
+ keep_sliding_window=args.keep_sliding_window,
+ link_weights=args.link_weights,
+ )
+ print(
+ f"wrote {args.output}: architectures={converted['architectures']} "
+ f"dflash_config={json.dumps(converted['dflash_config'], sort_keys=True)}"
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/verl_speco/integration/rollout_publish.py b/verl_speco/integration/rollout_publish.py
index fdedb93b..104f0d98 100644
--- a/verl_speco/integration/rollout_publish.py
+++ b/verl_speco/integration/rollout_publish.py
@@ -745,8 +745,10 @@ def export_actor_lm_head_weight(
)
normalized_row_indices = _normalize_lm_head_row_indices(row_indices)
+ # Block drafters that train against the target's own lm_head rows.
is_dflash = drafter_speculative_algorithm(getattr(worker, "config", None)) in {
"DFLASH",
+ "DFLASH2",
"DSPARK",
}
if (
diff --git a/verl_speco/integration/sglang_runtime.py b/verl_speco/integration/sglang_runtime.py
index 6a2edffa..3b629120 100644
--- a/verl_speco/integration/sglang_runtime.py
+++ b/verl_speco/integration/sglang_runtime.py
@@ -641,11 +641,14 @@ def _server_args_overrides_from_drafter(
# (dynamic convolutions + candidate selector) ride in the checkpoint's
# dflash_config, not a distinct engine-level method. DFLASH2 is never a
# valid SGLang ServerArgs algorithm, so fail loud instead of forwarding
- # the raw string.
+ # the raw string. The vLLM overlay maps DFLASH2 onto its DFlash path
+ # (vllm_runtime._speculative_method_from_drafter); the SGLang co-train
+ # path is not validated for DFlash2 yet.
raise ValueError(
"DFLASH2 is not an engine-level speculative algorithm; DFlash2 is served as a DFlash "
- "checkpoint. Keep DFLASH2 for drafter training (which this overlay runs offline) and "
- "set actor_rollout_ref.rollout.drafter.speculative_algorithm=DFLASH to serve a trained "
+ "checkpoint. SGLang co-training of a DFlash2 drafter is not wired in this overlay: use "
+ "actor_rollout_ref.rollout.name=vllm (vLLM >= 0.28.0) for DFLASH2 co-training, or set "
+ "actor_rollout_ref.rollout.drafter.speculative_algorithm=DFLASH to serve a trained "
"DFlash2 checkpoint as a frozen rollout drafter; its dflash_config carries the DFlash2 "
"convolution and selector hyperparameters."
)
diff --git a/verl_speco/integration/vllm_runtime.py b/verl_speco/integration/vllm_runtime.py
index e47273b6..60a9f310 100644
--- a/verl_speco/integration/vllm_runtime.py
+++ b/verl_speco/integration/vllm_runtime.py
@@ -814,31 +814,195 @@ def _drafter_algorithm(drafter_cfg: dict[str, Any]) -> str:
# Draft architectures vLLM can serve through its DFlash speculative path.
# Keep in sync with the alias sets in verl_speco/models/auto.py.
+_DFLASH2_SERVABLE_ARCHITECTURES = frozenset({"DFlash2DraftModel", "Qwen3DFlash2Model"})
_DFLASH_SERVABLE_ARCHITECTURES = frozenset(
- {"DFlashDraftModel", "DFlash2DraftModel", "Qwen3DFlash2Model"}
+ {"DFlashDraftModel", *_DFLASH2_SERVABLE_ARCHITECTURES}
+)
+# Hyperparameters vLLM's DFlash2 draft (qwen3_dflash2.py) indexes straight out of
+# ``dflash_config`` when it builds the model; a missing key is a KeyError deep in
+# engine startup, so the checkpoint contract is checked up front instead.
+# Superset lists of the nested-key contract live in
+# verl_speco/models/dflash2/configuration_dflash2.py (_NESTED_DFLASH_KEYS) and
+# verl_speco/convert_speculators_dflash2.py (_DFLASH2_KEYS); they cannot be
+# imported here because this module must stay importable without transformers.
+_DFLASH2_RUNTIME_KEYS = (
+ "conv_kernel_size",
+ "conv_group_size",
+ "selector_rank",
+ "selector_top_k",
+)
+# vLLM module that carries the DFlash2 draft class; its presence is the
+# capability probe for DFlash2 rollout (vllm-project/vllm#52816, v0.28.0+).
+_VLLM_DFLASH2_MODULE = "vllm.model_executor.models.qwen3_dflash2"
+# Trainer-side spelling of the DFlash2 selector codebooks. Load-side inverse:
+# DFlash2TrainerBackend._CHECKPOINT_KEY_ALIASES — keep the two in sync.
+_DFLASH2_CODEBOOK_WEIGHT_SUFFIXES = (
+ "candidate_selector.predecessor_codebook.weight",
+ "candidate_selector.successor_codebook.weight",
)
-def _validate_vllm_dflash_drafter_config(
- spec_model_path: Any, algorithm: str = "DFLASH"
+def _dflash2_engine_param_name(name: str) -> str:
+ """Spell a published DFlash2 parameter the way vLLM's draft names it.
+
+ The trainer keeps the selector codebooks in ``nn.Embedding`` modules
+ (``..._codebook.weight``) while vLLM's ``CandidateSelector`` holds them as
+ bare parameters (``..._codebook``), like the released z-lab checkpoints;
+ every other DFlash2 parameter already matches. vLLM's ``load_weights``
+ refuses the trainer spelling ("Attempted to load nested weight ... into a
+ single parameter"), so the rename has to happen before the weights reach
+ the engine.
+ """
+ if name.endswith(_DFLASH2_CODEBOOK_WEIGHT_SUFFIXES):
+ return name[: -len(".weight")]
+ return name
+
+
+def _normalize_dflash2_runtime_aliases(config: Any) -> bool:
+ """Mirror DFlash2 hyperparameters into ``dflash_config`` for vLLM.
+
+ vLLM's DFlash2 draft reads the convolution and selector knobs strictly from
+ ``dflash_config`` (the z-lab checkpoint layout), while a checkpoint this
+ overlay saved carries them at the top level. Copy missing keys down so both
+ layouts serve; a conflicting pair is a corrupt checkpoint and fails loud.
+ """
+ architectures = _get_nested(config, ("architectures",), None) or []
+ if isinstance(architectures, str):
+ architectures = [architectures]
+ if _DFLASH2_SERVABLE_ARCHITECTURES.isdisjoint(str(name) for name in architectures):
+ return False
+
+ dflash_config = _get_nested(config, ("dflash_config",), None)
+ if dflash_config is not None and not hasattr(dflash_config, "get"):
+ raise TypeError("DFlash2 dflash_config must be a mapping when provided")
+ changed = False
+ for key in (*_DFLASH2_RUNTIME_KEYS, "block_size"):
+ top_level = _get_nested(config, (key,), None)
+ nested = _get_nested(dflash_config, (key,), None)
+ if top_level is None or top_level == nested:
+ continue
+ if nested is not None:
+ raise ValueError(
+ f"DFlash2 {key} conflicts with dflash_config.{key}: {top_level!r} != {nested!r}"
+ )
+ if dflash_config is None:
+ dflash_config = {}
+ _set_child(config, "dflash_config", dflash_config)
+ _set_child(dflash_config, key, top_level)
+ changed = True
+ return changed
+
+
+def _vllm_supports_dflash2() -> bool | None:
+ """Whether the installed vLLM ships the DFlash2 draft; ``None`` without vLLM."""
+ import importlib.util
+
+ try:
+ if importlib.util.find_spec("vllm") is None:
+ return None
+ return importlib.util.find_spec(_VLLM_DFLASH2_MODULE) is not None
+ except (ImportError, ValueError):
+ return None
+
+
+def _assert_vllm_supports_dflash2() -> None:
+ if _vllm_supports_dflash2() is False:
+ raise ValueError(
+ "DFLASH2 rollout needs a vLLM that ships the DFlash2 draft model "
+ f"({_VLLM_DFLASH2_MODULE}, vllm-project/vllm#52816, released in vLLM 0.28.0); "
+ "the installed vLLM does not have it. Upgrade vLLM, or keep "
+ "actor_rollout_ref.rollout.drafter.enable=false and train the DFlash2 drafter offline."
+ )
+
+
+def _dflash2_config_value(config: dict[str, Any], key: str) -> Any:
+ """Read a DFlash2 knob from the checkpoint config, top level first."""
+ return _first_present(
+ _get_nested(config, (key,), None),
+ _get_nested(config, ("dflash_config", key), None),
+ )
+
+
+def _validate_vllm_dflash2_block_size(
+ config: dict[str, Any] | None,
+ drafter_cfg: dict[str, Any],
+ num_speculative_tokens: int,
) -> None:
- if not spec_model_path:
+ """Pin the DFlash2 block to the engine's ``1 + num_speculative_tokens``.
+
+ The dynamic convolutions are causal *within* a block: vLLM sizes that block
+ as the bonus token plus ``num_speculative_tokens`` mask tokens, while the
+ trainer folds the drafted sequence by ``dflash2_block_size``. If the two
+ disagree, the served conv reads across positions the trained conv never saw
+ (or vice versa), which shows up as a silently weak drafter rather than an
+ error, so refuse the mismatch here.
+ """
+ training_cfg = drafter_cfg.get("training") or {}
+ block_size = _positive_int_or_none(training_cfg.get("dflash2_block_size"))
+ if block_size is None and config is not None:
+ block_size = _positive_int_or_none(_dflash2_config_value(config, "block_size"))
+ if block_size is None:
return
+ if int(num_speculative_tokens) + 1 != int(block_size):
+ raise ValueError(
+ "DFLASH2 rollout requires actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens "
+ "== block_size - 1 so the served convolution block matches the trained one: got "
+ f"spec_verify_tokens={num_speculative_tokens} but block_size={block_size} "
+ "(from drafter.training.dflash2_block_size or the checkpoint's dflash_config)."
+ )
+
+def _load_vllm_dflash_drafter_config(spec_model_path: Any) -> dict[str, Any] | None:
+ if not spec_model_path:
+ return None
config_path = os.path.join(os.fspath(spec_model_path), "config.json")
if not os.path.exists(config_path):
- return
-
+ return None
try:
with open(config_path, "r", encoding="utf-8") as f:
- config = json.load(f)
+ return json.load(f)
except json.JSONDecodeError as exc:
raise ValueError(
f"Invalid DFlash drafter config.json at {config_path}: {exc}"
) from exc
+
+def _validate_vllm_dflash_drafter_config(
+ spec_model_path: Any,
+ algorithm: str = "DFLASH",
+ config: dict[str, Any] | None = None,
+) -> None:
+ if config is None:
+ config = _load_vllm_dflash_drafter_config(spec_model_path)
+ if config is None:
+ return
+ config_path = os.path.join(os.fspath(spec_model_path), "config.json")
+
architectures = config.get("architectures") or []
algorithm = str(algorithm or "DFLASH").strip().upper()
+ if algorithm == "DFLASH2":
+ # vLLM dispatches on the architecture: only the DFlash2 names reach the
+ # draft class with the convolutions and the selector. A plain DFlash
+ # checkpoint would load fine and silently serve without them.
+ if _DFLASH2_SERVABLE_ARCHITECTURES.isdisjoint(architectures):
+ raise ValueError(
+ "vLLM DFLASH2 requires actor_rollout_ref.rollout.drafter.model_path to point "
+ "to a DFlash2 drafter checkpoint with architectures in "
+ f"{sorted(_DFLASH2_SERVABLE_ARCHITECTURES)}; got architectures={architectures!r} "
+ f"from {config_path}. Use speculative_algorithm=DFLASH for a plain DFlash drafter."
+ )
+ missing = [
+ key
+ for key in _DFLASH2_RUNTIME_KEYS
+ if _dflash2_config_value(config, key) is None
+ ]
+ if missing:
+ raise ValueError(
+ "vLLM DFLASH2 requires the drafter config.json to carry the DFlash2 "
+ f"hyperparameters {list(_DFLASH2_RUNTIME_KEYS)} (top level or under dflash_config); "
+ f"missing {missing} in {config_path}."
+ )
+ return
if algorithm == "DSPARK":
if not _is_dspark_config(config):
raise ValueError(
@@ -922,18 +1086,15 @@ def _speculative_method_from_drafter(drafter_cfg: dict[str, Any]) -> str:
"drafter training."
)
if algorithm == "DFLASH2":
- # Same story as Domino: DFlash2 is a DFlash variant whose extra modules
- # (dynamic convolutions + candidate selector) ride in the checkpoint's
- # dflash_config, not a distinct engine-level method. DFLASH2 is never a
- # valid vLLM method, so fail loud instead of forwarding the raw string,
- # mirroring sglang_runtime._server_args_overrides_from_drafter.
- raise ValueError(
- "DFLASH2 is not an engine-level speculative algorithm; DFlash2 is served as a DFlash "
- "checkpoint. Keep DFLASH2 for drafter training (which this overlay runs offline) and "
- "set actor_rollout_ref.rollout.drafter.speculative_algorithm=DFLASH to serve a trained "
- "DFlash2 checkpoint as a frozen rollout drafter; its dflash_config carries the DFlash2 "
- "convolution and selector hyperparameters."
- )
+ # DFlash2 is a DFlash variant, not an engine-level method: vLLM runs it
+ # through the DFlash proposer (method="dflash") and picks the DFlash2 draft
+ # class (dynamic convolutions + candidate selector, vllm-project/vllm#52816,
+ # first released in v0.28.0) from the checkpoint's ``DFlash2DraftModel``
+ # architecture. DFLASH2 itself is never a valid vLLM method string, so map
+ # it here; the checkpoint contract is enforced by
+ # ``_validate_vllm_dflash_drafter_config`` and the engine capability by
+ # ``_assert_vllm_supports_dflash2``.
+ return "dflash"
if algorithm == "DSPARK":
return "dflash" if _is_vllm_ascend_runtime_hint() else "dspark"
@@ -1060,8 +1221,11 @@ def build_vllm_speculative_config_from_drafter(
rollout_drafter_cfg = drafter_cfg.get("rollout") or {}
if method in ("dflash", "dspark"):
+ drafter_checkpoint_config = _load_vllm_dflash_drafter_config(spec_model_path)
if method == "dflash" or algorithm == "DSPARK":
- _validate_vllm_dflash_drafter_config(spec_model_path, algorithm=algorithm)
+ _validate_vllm_dflash_drafter_config(
+ spec_model_path, algorithm=algorithm, config=drafter_checkpoint_config
+ )
num_speculative_tokens = _positive_int_or_none(
rollout_drafter_cfg.get("spec_verify_tokens")
)
@@ -1070,6 +1234,13 @@ def build_vllm_speculative_config_from_drafter(
"actor_rollout_ref.rollout.drafter.rollout.spec_verify_tokens "
f"must be positive for vLLM {method.upper()} speculative decoding"
)
+ if algorithm == "DFLASH2":
+ _assert_vllm_supports_dflash2()
+ _validate_vllm_dflash2_block_size(
+ drafter_checkpoint_config,
+ drafter_cfg,
+ num_speculative_tokens,
+ )
else:
num_speculative_tokens = _positive_int_or_none(
_first_present(
@@ -1728,6 +1899,7 @@ def patched_eagle_config_init(self, *args, **kwargs):
current(self, *args, **kwargs)
if str(method or "eagle").strip().lower() == "dflash":
_normalize_dflash_target_layer_aliases(self)
+ _normalize_dflash2_runtime_aliases(self)
if _is_dspark_hf_config(self):
_set_child(self, "architectures", ["DFlashDraftModel"])
@@ -2222,6 +2394,9 @@ def _draft_param_name_candidates(name: str) -> list[str]:
candidates.append(candidate)
if "midlayer." in candidate:
candidates.append(candidate.replace("midlayer.", "model.layers.0."))
+ engine_name = _dflash2_engine_param_name(candidate)
+ if engine_name != candidate:
+ candidates.append(engine_name)
for candidate in list(candidates):
if not candidate.startswith("model."):
candidates.append(f"model.{candidate}")
@@ -2285,6 +2460,36 @@ async def _maybe_call_vllm_server_method(
return await method.remote(*args, **kwargs)
+@contextmanager
+def _ipc_safe_allocator(enabled: bool):
+ """Stage the IPC buckets in non-expandable CUDA segments.
+
+ CUDA tensors shared over IPC out of an expandable segment carry an fd-based
+ handle that the receiver can only import through ``pidfd_getfd`` (Linux >=
+ 5.6); on older kernels the vLLM worker fails the whole draft update with
+ "does not support the pidfd_getfd syscall". verl's own actor->rollout sync
+ flips expandable segments off around its send for the same reason and turns
+ them back on afterwards, so the draft publish mirrors that. Restoring
+ ``True`` unconditionally (torch has no public getter for the prior state)
+ matches the state verl's own per-step sync leaves behind on every verl that
+ ships the helper; a verl without it never enabled expandable segments, and
+ the ImportError guard then leaves the allocator untouched.
+ """
+ if not enabled:
+ yield
+ return
+ try:
+ from verl.utils.device import set_expandable_segments
+ except ImportError:
+ yield
+ return
+ set_expandable_segments(False)
+ try:
+ yield
+ finally:
+ set_expandable_segments(True)
+
+
async def speco_vllm_update_draft_weights(
self, weights: Any, *args, global_steps: int | None = None, **kwargs
):
@@ -2341,12 +2546,15 @@ async def speco_vllm_update_draft_weights(
kwargs={**kwargs, "use_shm": use_shm},
)
- sender = BucketedWeightSender(
- zmq_handle=_draft_zmq_handle_from_base(self.zmq_handle),
- bucket_size_mb=int(bucket_mb),
- use_shm=use_shm,
- )
- await sender.async_send_weights(_named_weight_iter(weights))
+ # Only the CUDA-IPC transport shares device allocations; shm stages
+ # through host memory and is unaffected by the allocator mode.
+ with _ipc_safe_allocator(not use_shm):
+ sender = BucketedWeightSender(
+ zmq_handle=_draft_zmq_handle_from_base(self.zmq_handle),
+ bucket_size_mb=int(bucket_mb),
+ use_shm=use_shm,
+ )
+ await sender.async_send_weights(_named_weight_iter(weights))
if future is not None:
await future
@@ -2810,6 +3018,10 @@ def on_bucket_received(bucket_weights):
changed = True
if "midlayer." in n:
n = n.replace("midlayer.", "layers.0.")
+ if is_dflash:
+ # DFlash2 selector codebooks: trainer ``.weight`` -> engine bare
+ # parameter; vLLM's load_weights rejects the trainer spelling.
+ n = _dflash2_engine_param_name(n)
if (
is_eagle3
and n != "lm_head.weight"
diff --git a/verl_speco/models/dflash2/configuration_dflash2.py b/verl_speco/models/dflash2/configuration_dflash2.py
index 6dd95b3f..2c0bce9c 100644
--- a/verl_speco/models/dflash2/configuration_dflash2.py
+++ b/verl_speco/models/dflash2/configuration_dflash2.py
@@ -75,6 +75,24 @@ def __init__(
self.selector_top_k = int(selector_top_k)
self.selector_loss_weight = float(selector_loss_weight)
+ def to_dict(self):
+ """Serialize with the z-lab ``dflash_config`` block alongside the flat keys.
+
+ vLLM's DFlash2 draft reads the convolution and selector knobs strictly
+ from ``dflash_config``, so a checkpoint saved by the trainer has to carry
+ that block to be servable as a rollout drafter. The flat keys stay for
+ this overlay's own loaders; ``from_dflash2_pretrained`` lets a top-level
+ value win over the nested one, so the two never disagree.
+ """
+ output = super().to_dict()
+ nested = dict(output.get("dflash_config") or {})
+ for key in _NESTED_DFLASH_KEYS:
+ if output.get(key) is not None:
+ nested[key] = output[key]
+ if nested:
+ output["dflash_config"] = nested
+ return output
+
@classmethod
def from_dflash2_pretrained(cls, model_path: str):
config_path = os.path.join(model_path, "config.json")