diff --git a/docs/docs/advanced_features/speculative_decoding.mdx b/docs/docs/advanced_features/speculative_decoding.mdx
index a9b1386a52be..3dd7a2e5a868 100644
--- a/docs/docs/advanced_features/speculative_decoding.mdx
+++ b/docs/docs/advanced_features/speculative_decoding.mdx
@@ -101,7 +101,7 @@ SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3,
| DFlash draft model (linear block verification) |
Yes |
--speculative-algorithm DFLASH + --speculative-draft-model-path ... |
- No DP attention (--attn-dp-size > 1); pp_size == 1; disables overlap scheduler & mixed chunked prefill |
+ CUDA DP attention requires --enable-dp-lm-head and attn_cp_size == 1; pp_size == 1 |
| STANDALONE |
@@ -620,6 +620,37 @@ response = client.chat.completions.create(
print(response.choices[0].message.content)
```
+### DFlash and DFlash2 with DP attention
+
+On NVIDIA CUDA, you can use attention DP with plain DFlash drafts and DFlash2
+candidate-selector checkpoints. Both use `--speculative-algorithm DFLASH`.
+Set `--enable-dp-lm-head` so the target head has the same vocabulary shards as
+its local draft replica. Pipeline parallelism and DP combined with context
+parallelism are not supported. Domino and LiLiCorr heads are not supported
+with CUDA DP attention.
+
+For example, on four NVIDIA GPUs, use TP4 with attention DP2 / TP2:
+
+```bash Command
+python3 -m sglang.launch_server \
+ --model-path Qwen/Qwen3-4B \
+ --speculative-algorithm DFLASH \
+ --speculative-draft-model-path z-lab/Qwen3-4B-DFlash-b16 \
+ --tp-size 4 \
+ --attn-dp-size 2 \
+ --enable-dp-lm-head \
+ --attention-backend flashinfer \
+ --speculative-draft-attention-backend flashinfer \
+ --cuda-graph-backend-prefill disabled
+```
+
+For DFlash2, use a matching target and selector checkpoint, such as
+`Qwen/Qwen3.8-27B` and `incoai/Qwen3.8-27B-DFlash2`.
+Each attention-DP group runs its draft within the local attention-TP group.
+Draft execution stays eager; eligible target verify batches can use CUDA graphs.
+The scheduler separates prefill and verify rounds across DP ranks. Keep the
+experimental `SGLANG_ENABLE_DP_SPEC_PREFILL_COORDINATION` setting disabled.
+
### LiLiCorr: reranked DFlash drafts
A DFlash drafter is trained on per-position marginals rather than on the joint block distribution, so its drafted tokens are individually plausible yet jointly incoherent. **LiLiCorr** ([arXiv:2608.20530](https://arxiv.org/abs/2608.20530)) keeps the top-`k` candidates per block position, scores the whole lattice jointly with a small two-layer transformer, and commits a path through it greedily instead of taking the per-slot argmax. Verify is unchanged, so outputs remain distributionally identical to the target model's and only acceptance length moves.
diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py
index 18ceaa1b581c..f4b03a662f1f 100644
--- a/python/sglang/srt/arg_groups/speculative_hook.py
+++ b/python/sglang/srt/arg_groups/speculative_hook.py
@@ -299,12 +299,21 @@ def _handle_dflash(server_args: ServerArgs) -> None:
f"{type(current_platform).__name__} on device {cfg.device!r}."
)
- # DFLASH + dp attention is validated on NPU only.
- if attn_dp_enabled_of(cfg) and not cfg.device == "npu":
- raise ValueError(
- "Currently DFLASH speculative decoding does not support dp "
- "attention on non-NPU devices."
- )
+ if attn_dp_enabled_of(cfg) and cfg.device != "npu":
+ if not cfg.device.startswith("cuda"):
+ raise ValueError("DFLASH with dp attention requires CUDA or NPU.")
+ # Both greedy projection and the selector borrow the target head.
+ if not cfg.enable_dp_lm_head:
+ raise ValueError("DFLASH with dp attention requires --enable-dp-lm-head.")
+ if cfg.attn_cp_size != 1:
+ raise ValueError(
+ "DFLASH with dp attention does not support context parallel."
+ )
+ if envs.SGLANG_ENABLE_DP_SPEC_PREFILL_COORDINATION.get():
+ raise ValueError(
+ "DFLASH with dp attention requires "
+ "SGLANG_ENABLE_DP_SPEC_PREFILL_COORDINATION=0."
+ )
if cfg.pp_size != 1:
raise ValueError(
diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py
index 41ca963bfee4..2d06f2e5b4e2 100644
--- a/python/sglang/srt/model_executor/forward_batch_info.py
+++ b/python/sglang/srt/model_executor/forward_batch_info.py
@@ -1169,6 +1169,9 @@ def init_new(
if ret.forward_mode.is_idle():
ret.positions = torch.empty((0,), dtype=torch.int64, device=device)
+ if model_runner.model_config.model_is_mrope:
+ # DP padding can turn an idle rank into a target verify batch.
+ ret.mrope_positions = ret.positions.unsqueeze(0).expand(3, -1)
if model_runner.lora_manager is not None:
model_runner.lora_manager.reset_lora_batch()
return ret
diff --git a/python/sglang/srt/speculative/dflash_dp.py b/python/sglang/srt/speculative/dflash_dp.py
new file mode 100644
index 000000000000..8e4905a91e8e
--- /dev/null
+++ b/python/sglang/srt/speculative/dflash_dp.py
@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod
+from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
+
+
+@torch.no_grad()
+def gather_draft_embedding(embedding: nn.Module) -> torch.Tensor | None:
+ # Other modules may scale their output or own a checkpoint embedding.
+ if not isinstance(embedding, VocabParallelEmbedding):
+ return None
+ if isinstance(embedding.quant_method, UnquantizedEmbeddingMethod):
+ weight = embedding.weight
+ else:
+ # Packed weights are not embedding rows. Materialize the table through
+ # the same quantization method as an ordinary embedding lookup.
+ ids = torch.arange(
+ embedding.num_embeddings_per_partition,
+ device=next(embedding.parameters()).device,
+ )
+ weight = embedding.quant_method.embedding(embedding, ids)
+ if embedding.output_dtype is not None:
+ weight = weight.to(embedding.output_dtype)
+ if embedding.tp_size == 1 and embedding.num_embeddings == embedding.org_vocab_size:
+ return weight[: embedding.num_embeddings]
+
+ shards = (
+ embedding.tp_group.all_gather(weight, dim=0).view(
+ embedding.tp_size, *weight.shape
+ )
+ if embedding.tp_size > 1
+ else weight.unsqueeze(0)
+ )
+ full_weight = weight.new_empty((embedding.num_embeddings, weight.shape[1]))
+ for rank, shard in enumerate(shards):
+ indices = embedding._get_indices(
+ embedding.num_embeddings_padded,
+ embedding.org_vocab_size_padded,
+ embedding.num_embeddings,
+ embedding.org_vocab_size,
+ rank,
+ embedding.tp_size,
+ )
+ full_weight[indices.org_vocab_start_index : indices.org_vocab_end_index].copy_(
+ shard[: indices.num_org_elements]
+ )
+ offset = indices.num_org_elements_padded
+ full_weight[
+ indices.added_vocab_start_index : indices.added_vocab_end_index
+ ].copy_(shard[offset : offset + indices.num_added_elements])
+ return full_weight
diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py
index 02becd4a228f..5826777d87d3 100644
--- a/python/sglang/srt/speculative/dflash_info.py
+++ b/python/sglang/srt/speculative/dflash_info.py
@@ -67,13 +67,14 @@ def prepare_for_verify(
self,
batch: ScheduleBatch,
target_worker: TpModelWorker,
+ *,
+ defer_metadata_init: bool = False,
) -> tuple[ForwardBatch, bool]:
"""Prepare a DFLASH verify forward batch for overlap scheduling.
The caller computes and stores `batch.out_cache_loc` before this
- method is called. GPU keeps the original pre-planning path. NPU leaves
- attention/graph metadata initialization to ModelRunner because DP/EP
- padding can still change the compressor's runtime shapes.
+ method is called. Deferred callers leave attention/graph planning to
+ ModelRunner, after DP/EP padding has finalized the runtime shapes.
"""
from sglang.srt.speculative.spec_utils import prepare_mamba_track_for_verify
@@ -116,11 +117,9 @@ def prepare_for_verify(
verify_forward_batch
)
)
- if _is_npu:
- # Do not pre-plan target verify on NPU. DP/EP padding can change
- # the compressor's logical batch without changing ForwardBatch's
- # stale-plan shape fields. Let ModelRunner select graph/eager and
- # initialize metadata after final batch preparation.
+ if _is_npu or defer_metadata_init:
+ # DP padding can change the live rows and graph eligibility.
+ # ModelRunner owns planning after final batch preparation.
return verify_forward_batch, can_run_cuda_graph
elif can_run_cuda_graph:
target_worker.model_runner.decode_cuda_graph_runner.load_batch(
diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py
index 71050f712978..fb15b49f62f1 100644
--- a/python/sglang/srt/speculative/dflash_worker_v2.py
+++ b/python/sglang/srt/speculative/dflash_worker_v2.py
@@ -4,7 +4,6 @@
from typing import TYPE_CHECKING, List, Optional, Tuple
import torch
-import torch.distributed as dist
from sglang.kernels.ops.speculative.cache_locs import (
assign_extend_cache_locs_func,
@@ -23,6 +22,7 @@
from sglang.srt.layers.dp_attention import get_dp_tp_group
from sglang.srt.layers.logits_processor import (
LogitsProcessor,
+ LogitsProcessorOutput,
should_apply_lm_head_quant_method,
)
from sglang.srt.layers.logprob_processor import compute_spec_logprobs
@@ -56,6 +56,7 @@
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
+from sglang.srt.speculative.dflash_dp import gather_draft_embedding
from sglang.srt.speculative.dflash_info import DFlashVerifyInput
from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
from sglang.srt.speculative.dflash_utils import (
@@ -451,6 +452,12 @@ def __init__(
draft_hf_config=self.draft_model_runner.model_config.hf_config
)
self._is_domino = draft_config.is_domino
+ if get_parallel().attn_dp_enabled and not _is_npu:
+ if self._is_domino or self.lilicorr is not None:
+ raise ValueError(
+ "DFLASH with dp attention supports plain DFlash and DFlash2 "
+ "drafts; Domino and LiLiCorr heads are not supported."
+ )
self.domino_candidate_pool_size = int(
get_spec().speculative_domino_candidate_pool_size
)
@@ -689,6 +696,7 @@ def init_attention_backends(self):
)
def init_cuda_graphs(self):
+ target_dp_enabled = get_parallel().attn_dp_enabled
with (
draft_pp_context(),
draft_tp_context(self.draft_parallel_layout),
@@ -707,10 +715,9 @@ def init_cuda_graphs(self):
"device graph capture.",
type(current_platform).__name__,
)
- if get_parallel().attn_dp_enabled and capture_decode_cuda_graph:
- # Idle DP ranks skip the draft step, so they cannot join a
- # shared graph capture/replay; keep the draft eager under dp
- # attention.
+ if target_dp_enabled and capture_decode_cuda_graph:
+ # Keep the eager-draft policy for idle/unequal DP batches.
+ # The narrowed draft scope hides the target's DP flag.
capture_decode_cuda_graph = False
if self._target_tp_rank == 0:
logger.warning(
@@ -741,6 +748,25 @@ def init_cuda_graphs(self):
self._draft_worker.init_cuda_graphs(
capture_decode_cuda_graph=capture_decode_cuda_graph
)
+ if target_dp_enabled and not _is_npu and self.selector is not None:
+ self._warm_eager_selector()
+
+ def _warm_eager_selector(self) -> None:
+ # Build selector kernels before serving rather than under the watchdog.
+ hidden = torch.zeros(
+ (int(self.block_size), int(self.draft_model.config.hidden_size)),
+ dtype=self.draft_model_runner.dtype,
+ device=self.device,
+ )
+ self._propose_selector_block(
+ draft_logits_output=LogitsProcessorOutput(
+ next_token_logits=None, hidden_states=hidden
+ ),
+ bs=1,
+ lm_head=unwrap_lora_layer(self._target_worker.model_runner.model.lm_head),
+ anchor_token_ids=torch.zeros((1,), dtype=torch.int64, device=self.device),
+ sampling_info=None,
+ )
def _prewarm_batch_size(self, block_size: int) -> int:
"""Largest batch the non-greedy verify path can see in one step."""
@@ -1352,37 +1378,15 @@ def _maybe_merge_trained_mask_embedding(self) -> None:
)
def _cache_full_embed_weight(self) -> None:
- """Cache the full target embedding replicated on GPU during init.
-
- Under dp attention, idle DP ranks skip the draft step, so the attn-TP
- all_reduce inside VocabParallelEmbedding gets mismatched calls. Gather
- the full embedding once during init and keep it replicated, making
- the draft block-id lookup a collective-free on-device F.embedding.
- """
if not get_parallel().attn_dp_enabled:
return
-
- tp_group = get_parallel().tp_group
- tp_size = int(tp_group.world_size)
- if tp_size <= 1:
- return
-
- target_model = self._target_worker.model_runner.model
- embed_module = target_model.get_input_embeddings()
- local_w = embed_module.weight.data
- shard = getattr(embed_module, "shard_indices", None)
- num_org = int(shard.num_org_elements) if shard else local_w.shape[0]
- vocab_size = int(self._target_worker.model_runner.model_config.vocab_size)
-
- shard_t = local_w[:num_org].contiguous()
- parts = [torch.empty_like(shard_t) for _ in range(tp_size)]
- dist.all_gather(parts, shard_t, group=tp_group.device_group)
- self._full_embed_gpu = torch.cat(parts, dim=0)[:vocab_size]
- if get_parallel().tp_rank == 0:
- logger.info(
- "DFLASH cached full embed on GPU for dp attention: shape=%s",
- list(self._full_embed_gpu.shape),
+ embedding = unwrap_lora_layer(
+ _resolve_dflash_embedding_module(
+ self.draft_model, self._target_worker.model_runner.model
)
+ )
+ # Reconstruct on the table's own group before DP ranks can go idle.
+ self._full_embed_gpu = gather_draft_embedding(embedding)
def _resolve_mask_token_id(
self, *, mask_token: str, mask_token_id: Optional[int] = None
@@ -2441,7 +2445,7 @@ def forward_batch_generation(
capture_hidden_mode=CaptureHiddenMode.FULL,
)
idle_verify_forward_batch, _ = idle_verify_input.prepare_for_verify(
- batch, self._target_worker
+ batch, self._target_worker, defer_metadata_init=True
)
# Force eager: the active ranks run eager verify when any DP
# rank is idle (see the idle-guard below); a graph-replaying
@@ -2452,7 +2456,7 @@ def forward_batch_generation(
batch=None,
forward_batch=idle_verify_forward_batch,
is_verify=True,
- skip_attn_backend_init=True,
+ skip_attn_backend_init=True if _is_npu else None,
)
empty_ids = torch.empty((0,), dtype=torch.int64, device=self.device)
empty_lens = torch.empty((0,), dtype=torch.int32, device=self.device)
@@ -2566,8 +2570,8 @@ def forward_batch_generation(
verify_out_cache_loc_2d.copy_(verify_out_cache_loc.view(bs, block_size))
if self._full_embed_gpu is not None:
- # Replicated lookup avoids the mismatched attn-TP all_reduce
- # inside VocabParallelEmbedding under dp attention.
+ # A borrowed embedding may use target-wide collectives, while
+ # other DP replicas are idle or looking up different token IDs.
noise_embedding = torch.nn.functional.embedding(
block_ids, self._full_embed_gpu
)
@@ -2843,7 +2847,9 @@ def forward_batch_generation(
batch.seq_lens_sum = int(draft_input.nxt_kv_lens_sum)
verify_forward_batch, _ = verify_input.prepare_for_verify(
- batch, self.target_worker
+ batch,
+ self.target_worker,
+ defer_metadata_init=get_parallel().attn_dp_enabled,
)
batch.seq_lens_cpu = seq_lens_cpu_backup
batch.seq_lens_sum = seq_lens_sum_backup
@@ -2869,7 +2875,9 @@ def forward_batch_generation(
batch=None,
forward_batch=verify_forward_batch,
is_verify=True,
- skip_attn_backend_init=True if not _is_npu else None,
+ skip_attn_backend_init=(
+ None if _is_npu or get_parallel().attn_dp_enabled else True
+ ),
)
logits_output = target_out.logits_output
can_run_cuda_graph = target_out.can_run_cuda_graph
diff --git a/test/registered/e2e/speculative/test_dflash_dp_attention.py b/test/registered/e2e/speculative/test_dflash_dp_attention.py
new file mode 100644
index 000000000000..200d0e484456
--- /dev/null
+++ b/test/registered/e2e/speculative/test_dflash_dp_attention.py
@@ -0,0 +1,149 @@
+"""DFLASH DP replicas must preserve target outputs with sharded vocabulary tables."""
+
+import unittest
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+
+import requests
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+ DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+ DEFAULT_URL_FOR_TEST,
+ CustomTestCase,
+ popen_launch_server,
+ terminate_and_kill_process_tree,
+)
+
+# TP4 / DP2 exercises both local attention TP and idle attention-DP peers.
+register_cuda_ci(est_time=480, stage="nightly", runner_config="4-gpu-h100")
+
+
+class TestDFlashDPAttention(CustomTestCase):
+ @contextmanager
+ def _server(self, model, revision, draft=None, draft_revision=None):
+ args = [
+ "--revision",
+ revision,
+ "--tp-size",
+ "4",
+ "--attn-dp-size",
+ "2",
+ "--enable-dp-lm-head",
+ "--attention-backend",
+ "fa3",
+ "--enable-deterministic-inference",
+ "--mem-fraction-static",
+ "0.6",
+ "--context-length",
+ "4096",
+ "--max-total-tokens",
+ "16384",
+ "--max-running-requests",
+ "8",
+ "--chunked-prefill-size",
+ "1024",
+ "--cuda-graph-backend-prefill",
+ "disabled",
+ "--cuda-graph-max-bs-decode",
+ "4",
+ "--random-seed",
+ "42",
+ ]
+ if draft is not None:
+ args += [
+ "--speculative-algorithm",
+ "DFLASH",
+ "--speculative-draft-model-path",
+ draft,
+ "--speculative-draft-model-revision",
+ draft_revision,
+ "--speculative-draft-attention-backend",
+ "flashinfer",
+ ]
+ process = popen_launch_server(
+ model,
+ DEFAULT_URL_FOR_TEST,
+ timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+ other_args=args,
+ env={"SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY": "1"},
+ )
+ try:
+ yield
+ finally:
+ terminate_and_kill_process_tree(process)
+
+ def _generate(self, prompt, rank, length, temperature=0):
+ response = requests.post(
+ DEFAULT_URL_FOR_TEST + "/generate",
+ json={
+ "text": prompt,
+ "routed_dp_rank": rank,
+ "sampling_params": {
+ "temperature": temperature,
+ "max_new_tokens": length,
+ "ignore_eos": True,
+ "sampling_seed": 42,
+ },
+ },
+ timeout=180,
+ )
+ response.raise_for_status()
+ result = response.json()
+ self.assertEqual(len(result["output_ids"]), length)
+ return result
+
+ def _check_model(self, model, revision, draft, draft_revision):
+ prompts = [
+ ("The capital of France is", 16),
+ (
+ "Explain how a compiler translates source code into machine instructions. "
+ * 12,
+ 64,
+ ),
+ ]
+ with self._server(model, revision):
+ reference = [
+ self._generate(text, 0, n)["output_ids"] for text, n in prompts
+ ]
+ with self._server(model, revision, draft, draft_revision):
+ # Isolated requests exercise idle peers, in both rank directions.
+ for rank in (0, 1):
+ for (text, n), expected in zip(prompts, reference):
+ result = self._generate(text, rank, n)
+ self.assertEqual(result["output_ids"], expected)
+ self.assertGreater(result["meta_info"]["spec_verify_ct"], 0)
+ # Unequal lengths transition between active and idle peers, with
+ # target graphs enabled and cached prefixes on both replicas.
+ for rank in (0, 1):
+ with ThreadPoolExecutor(2) as pool:
+ futures = [
+ pool.submit(self._generate, text, (rank + i) % 2, n)
+ for i, (text, n) in enumerate(prompts)
+ ]
+ for future, expected in zip(futures, reference):
+ self.assertEqual(future.result()["output_ids"], expected)
+ sampled = self._generate(
+ "Write a short story about an astronomer.", 1, 64, 0.7
+ )
+ self.assertGreater(sampled["meta_info"]["spec_verify_ct"], 0)
+
+ def test_dflash(self):
+ self._check_model(
+ "Qwen/Qwen3-4B",
+ "1cfa9a7208912126459214e8b04321603b3df60c",
+ "z-lab/Qwen3-4B-DFlash-b16",
+ "b74e3a329c4d963783143b1e970d95b002be72bd",
+ )
+
+ def test_dflash2_hybrid_target(self):
+ self._check_model(
+ "Qwen/Qwen3.8-27B",
+ "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0",
+ "incoai/Qwen3.8-27B-DFlash2",
+ "015e795645c74b1a0eeef3b570031fb62e769bc5",
+ )
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/test/registered/unit/spec/test_dflash_dp_attention.py b/test/registered/unit/spec/test_dflash_dp_attention.py
new file mode 100644
index 000000000000..0c4372aee3cb
--- /dev/null
+++ b/test/registered/unit/spec/test_dflash_dp_attention.py
@@ -0,0 +1,225 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+from torch import nn
+
+from sglang.srt.arg_groups.overrides import resolution_result
+from sglang.srt.arg_groups.speculative_hook import _handle_dflash
+from sglang.srt.environ import envs
+from sglang.srt.layers import linear, vocab_parallel_embedding
+from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative import dflash_info
+from sglang.srt.speculative.dflash_dp import gather_draft_embedding
+from sglang.srt.speculative.dflash_info import DFlashVerifyInput
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=8, suite="base-a-test-cpu")
+
+
+class TestDFlashDPAdmission(CustomTestCase):
+ def _args(self, **overrides):
+ fields = dict(
+ model_path="dummy",
+ device="cuda",
+ tp_size=4,
+ attn_dp_size=2,
+ attn_cp_size=1,
+ enable_dp_lm_head=True,
+ speculative_algorithm="DFLASH",
+ speculative_draft_model_path="draft",
+ speculative_num_draft_tokens=8,
+ speculative_draft_attention_backend="flashinfer",
+ )
+ return ServerArgs(**(fields | overrides))
+
+ def test_cuda_dp_keeps_explicit_placement_and_raw_input(self):
+ args = self._args()
+ _handle_dflash(args)
+ self.assertTrue(args.enable_dp_lm_head)
+ self.assertEqual(resolution_result(args, "attn_dp_size"), 2)
+ self.assertIsNone(args.speculative_num_steps)
+ self.assertEqual(resolution_result(args, "speculative_num_steps"), 1)
+
+ def test_incompatible_cuda_dp_layouts_fail(self):
+ for fields, message in [
+ ({"enable_dp_lm_head": False}, "enable-dp-lm-head"),
+ ({"attn_cp_size": 2}, "context parallel"),
+ ({"pp_size": 2}, "pp_size"),
+ ({"device": "xpu"}, "CUDA or NPU"),
+ ]:
+ with (
+ self.subTest(fields=fields),
+ self.assertRaisesRegex(ValueError, message),
+ ):
+ _handle_dflash(self._args(**fields))
+
+ def test_unimplemented_mixed_prefill_coordination_fails_at_startup(self):
+ with (
+ envs.SGLANG_ENABLE_DP_SPEC_PREFILL_COORDINATION.override(True),
+ self.assertRaisesRegex(ValueError, "DP_SPEC_PREFILL_COORDINATION=0"),
+ ):
+ _handle_dflash(self._args())
+
+ def test_non_dp_does_not_require_dp_head(self):
+ args = self._args(attn_dp_size=1, enable_dp_lm_head=False)
+ _handle_dflash(args)
+ self.assertFalse(resolution_result(args, "enable_dp_lm_head"))
+
+
+class _TableGroup:
+ def __init__(self, shards):
+ self.shards = shards
+
+ def all_gather(self, weight, dim):
+ if not any(torch.equal(weight, shard) for shard in self.shards):
+ raise AssertionError("embedding used a different table's group")
+ return torch.cat(self.shards, dim=dim)
+
+
+def _embedding_shards(org_vocab, added_vocab, tp_size):
+ embeddings = []
+ for rank in range(tp_size):
+ parallel = SimpleNamespace(tp_size=tp_size, tp_rank=rank, tp_group=None)
+ with (
+ patch.object(
+ vocab_parallel_embedding, "get_parallel", return_value=parallel
+ ),
+ patch.object(linear, "get_parallel", return_value=parallel),
+ ):
+ embedding = VocabParallelEmbedding(
+ num_embeddings=org_vocab + added_vocab,
+ org_num_embeddings=org_vocab,
+ embedding_dim=2,
+ padding_size=8,
+ )
+ indices = embedding.shard_indices
+ with torch.no_grad():
+ embedding.weight.fill_(-999)
+ for token in range(
+ indices.org_vocab_start_index, indices.org_vocab_end_index
+ ):
+ embedding.weight[token - indices.org_vocab_start_index].fill_(token)
+ for token in range(
+ indices.added_vocab_start_index, indices.added_vocab_end_index
+ ):
+ offset = (
+ indices.num_org_elements_padded
+ + token
+ - indices.added_vocab_start_index
+ )
+ embedding.weight[offset].fill_(token)
+ embeddings.append(embedding)
+ group = _TableGroup([embedding.weight for embedding in embeddings])
+ for embedding in embeddings:
+ embedding.tp_group = group
+ return embeddings
+
+
+class TestDFlashDPEmbedding(CustomTestCase):
+ def test_padding_and_added_vocabulary_keep_logical_token_ids(self):
+ """Uneven and empty vocabulary shards must not shift token IDs after gathering."""
+ for org_vocab, added_vocab, tp_size in [(5, 0, 4), (9, 3, 4), (9, 3, 1)]:
+ with self.subTest(
+ org_vocab=org_vocab, added_vocab=added_vocab, tp_size=tp_size
+ ):
+ expected = (
+ torch.arange(org_vocab + added_vocab)
+ .float()
+ .unsqueeze(1)
+ .expand(-1, 2)
+ )
+ for embedding in _embedding_shards(org_vocab, added_vocab, tp_size):
+ # The table's retained group wins over any ambient draft/target scope.
+ actual = gather_draft_embedding(embedding)
+ torch.testing.assert_close(actual, expected)
+
+ def test_replicated_table_keeps_storage(self):
+ embedding = _embedding_shards(9, 0, 1)[0]
+ actual = gather_draft_embedding(embedding)
+ self.assertEqual(actual.data_ptr(), embedding.weight.data_ptr())
+ torch.testing.assert_close(actual[:, 0], torch.arange(9).float())
+
+ def test_materialized_embedding_preserves_quant_method_and_output_dtype(self):
+ embeddings = _embedding_shards(9, 3, 4)
+ group = _TableGroup([e.weight.double() * 2 for e in embeddings])
+ for embedding in embeddings:
+ embedding.quant_method = SimpleNamespace(
+ embedding=lambda layer, ids: layer.weight[ids].double() * 2
+ )
+ embedding.output_dtype = torch.float64
+ embedding.tp_group = group
+ actual = gather_draft_embedding(embedding)
+ expected = (torch.arange(12).double() * 2).unsqueeze(1).expand(-1, 2)
+ torch.testing.assert_close(actual, expected)
+
+ def test_other_embedding_keeps_its_forward_semantics(self):
+ class ScaledEmbedding(nn.Embedding):
+ def forward(self, ids):
+ return super().forward(ids) * 2
+
+ embedding = ScaledEmbedding(8, 2)
+ self.assertIsNone(gather_draft_embedding(embedding))
+
+
+class TestDFlashDPMetadata(CustomTestCase):
+ def test_deferred_verify_preserves_live_inputs_until_final_planning(self):
+ """Early graph loading must not replace an idle or unequal DP batch's live rows."""
+ for mode in (ForwardMode.IDLE, ForwardMode.DECODE):
+ for graph_eligible in (False, True):
+ with self.subTest(mode=mode, graph_eligible=graph_eligible):
+ n = 0 if mode.is_idle() else 4
+ ids = torch.arange(n)
+ forward = ForwardBatch(
+ forward_mode=mode,
+ batch_size=0 if mode.is_idle() else 1,
+ input_ids=ids,
+ req_pool_indices=torch.zeros(
+ 0 if mode.is_idle() else 1, dtype=torch.long
+ ),
+ seq_lens=torch.ones(
+ 0 if mode.is_idle() else 1, dtype=torch.long
+ ),
+ out_cache_loc=ids,
+ seq_lens_sum=0 if mode.is_idle() else 1,
+ positions=ids,
+ )
+
+ def preplan(batch):
+ batch.input_ids = torch.full((8,), -1)
+ batch.mark_forward_metadata_ready()
+
+ runner = SimpleNamespace(
+ decode_cuda_graph_runner=SimpleNamespace(
+ can_run_graph=lambda batch: graph_eligible,
+ load_batch=preplan,
+ ),
+ attn_backend=SimpleNamespace(init_forward_metadata=preplan),
+ )
+ batch = SimpleNamespace(forward_mode=mode)
+ spec = DFlashVerifyInput(
+ draft_token=ids, positions=ids, draft_token_num=4
+ )
+ with (
+ patch.object(dflash_info, "_is_npu", False),
+ patch.object(ForwardBatch, "init_new", return_value=forward),
+ patch(
+ "sglang.srt.speculative.spec_utils.prepare_mamba_track_for_verify"
+ ),
+ ):
+ result, can_graph = spec.prepare_for_verify(
+ batch,
+ SimpleNamespace(model_runner=runner),
+ defer_metadata_init=True,
+ )
+ torch.testing.assert_close(result.input_ids, ids)
+ self.assertTrue(result.needs_forward_metadata_init())
+ self.assertEqual(can_graph, graph_eligible)
+
+
+if __name__ == "__main__":
+ unittest.main()