[Feat] spec: Support CUDA attention DP for DFlash and DFlash2

docs/docs/advanced_features/speculative_decoding.mdx

@@ -101,7 +101,7 @@ SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3,
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>DFlash draft model (linear block verification)</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Yes</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>--speculative-algorithm DFLASH</code> + <code>--speculative-draft-model-path ...</code></td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>No DP attention (<code>--attn-dp-size</code> &gt; 1); <code>pp_size == 1</code>; disables overlap scheduler &amp; mixed chunked prefill</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>CUDA DP attention requires <code>--enable-dp-lm-head</code> and <code>attn_cp_size == 1</code>; <code>pp_size == 1</code></td>
     </tr>
 <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>STANDALONE</td>
@@ -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.

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(

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

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

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(

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