[Spec Decode] Enable async scheduling for DFlash (#58065)
Jeff (Junze) Ma committed
Sep 22, 2026 at 01:05 UTC
0961bbae2894d574be790d219651824eb199318e
2 files changed
+69
-2
tests/test_config.py
+65
@@ -1568,6 +1568,71 @@ def test_draft_model_enables_async_scheduling_by_default():
1568
assert cfg.scheduler_config.async_scheduling is True
1569
1570
1571
+def test_dflash_allows_async_scheduling(tmp_path: Path):
1572
+ """DFlash must stay async-schedulable: it was the only parallel-drafting
1573
+ method excluded from the allowlist."""
1574
+ from transformers import LlamaConfig
1575
+
1576
+ common = dict(
1577
+ hidden_size=128,
1578
+ intermediate_size=256,
1579
+ num_hidden_layers=2,
1580
+ num_attention_heads=4,
1581
+ num_key_value_heads=2,
1582
+ vocab_size=256,
1583
+ max_position_embeddings=2048,
1584
+ )
1585
+ target_path = tmp_path / "target"
1586
+ draft_path = tmp_path / "draft"
1587
+ _write_json(
1588
+ target_path / "config.json",
1589
+ LlamaConfig(architectures=["LlamaForCausalLM"], **common).to_dict(),
1590
+ )
1591
+ _write_json(
1592
+ draft_path / "config.json",
1593
+ dict(
1594
+ common,
1595
+ architectures=["DFlash2DraftModel"],
1596
+ model_type="qwen3",
1597
+ num_hidden_layers=1,
1598
+ layer_types=["sliding_attention"],
1599
+ sliding_window=128,
1600
+ dflash_config=dict(
1601
+ block_size=4,
1602
+ mask_token_id=255,
1603
+ target_layer_ids=[0],
1604
+ conv_group_size=2,
1605
+ conv_kernel_size=2,
1606
+ selector_rank=8,
1607
+ selector_top_k=4,
1608
+ ),
1609
+ ),
1610
+ )
1611
+ parallel_config = ParallelConfig(distributed_executor_backend="uni")
1612
+ model_config = ModelConfig(
1613
+ model=str(target_path), tokenizer_mode="skip", max_model_len=2048
1614
+ )
1615
+ speculative_config = SpeculativeConfig(
1616
+ method="dflash",
1617
+ model=str(draft_path),
1618
+ num_speculative_tokens=3,
1619
+ target_model_config=model_config,
1620
+ target_parallel_config=parallel_config,
1621
+ )
1622
+ cfg = VllmConfig(
1623
+ model_config=model_config,
1624
+ scheduler_config=SchedulerConfig(
1625
+ max_model_len=2048,
1626
+ is_encoder_decoder=False,
1627
+ ),
1628
+ parallel_config=parallel_config,
1629
+ speculative_config=speculative_config,
1630
+ )
1631
+
1632
+ assert speculative_config.method == "dflash"
1633
+ assert cfg.scheduler_config.async_scheduling is True
1634
+
1635
+
1636
@pytest.mark.skip_global_cleanup
1637
@pytest.mark.parametrize("tp_size", [1, 2])
1638
@pytest.mark.parametrize("target_ep", [False, True], ids=["ep-off", "ep-on"])
vllm/config/vllm.py
+4
-2
@@ -1457,11 +1457,12 @@ class VllmConfig:
1457
and self.speculative_config.method not in get_args(NgramGPUTypes)
1458
and self.speculative_config.method != "draft_model"
1459
and self.speculative_config.method != "dspark"
1460
+ and self.speculative_config.method != "dflash"
1461
):
1462
raise ValueError(
1463
"Currently, async scheduling is only supported "
1463
- "with EAGLE/MTP/Draft Model/NGram GPU/DSpark kind of "
1464
- "speculative decoding"
1464
+ "with EAGLE/MTP/Draft Model/NGram GPU/DSpark/DFlash "
1465
+ "kind of speculative decoding"
1466
)
1467
if self.speculative_config.disable_padded_drafter_batch:
1468
raise ValueError(
@@ -1490,6 +1491,7 @@ class VllmConfig:
1491
and self.speculative_config.method not in get_args(NgramGPUTypes)
1492
and self.speculative_config.method != "draft_model"
1493
and self.speculative_config.method != "dspark"
1494
+ and self.speculative_config.method != "dflash"
1495
):
1496
logger.warning_once(
1497
"Async scheduling not supported with %s-based "