[DLLM] vectorized joint/low-confidence decoding and skip redundant attn init (#21094)

Co-authored-by: ronnie_zheng <zl19940307@163.com>

python/sglang/srt/dllm/algorithm/base.py

@@ -10,6 +10,9 @@ from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import is_npu
+
+_is_npu = is_npu()
 
 DllmRunOutput = Tuple[
     Union[LogitsProcessorOutput, torch.Tensor],
@@ -81,6 +84,11 @@ class DllmAlgorithm:
             return out.logits_output, [], None, None, out.can_run_graph
 
         states = self.init_step_state(forward_batch)
+        # NPU: attention metadata is stable across a block's denoise steps (the
+        # first forward above already planned it), so mark it ready once and let
+        # every later forward skip re-planning.
+        if _is_npu:
+            forward_batch.mark_forward_metadata_ready()
         for _ in range(self.max_steps(self.block_size)):
             done = self.step(forward_batch, out.logits_output.full_logits, states)
             if all(done):

python/sglang/srt/dllm/algorithm/joint_threshold.py

@@ -7,6 +7,103 @@ import torch.nn.functional as F
 from sglang.srt.dllm.algorithm.base import DllmAlgorithm
 from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.utils import is_npu
+
+_is_npu = is_npu()
+
+
+def joint_threshold_update_step_vectorized(
+    input_ids_1d: torch.Tensor,  # [B*blk]
+    full_logits_2d: torch.Tensor,  # [B*blk, V]
+    prompt_masks: torch.Tensor,  # [B, blk]
+    finished: torch.Tensor,  # [B]
+    post_edit_steps: torch.Tensor,  # [B]
+    mask_id: int,
+    blk: int,
+    threshold: float,
+    edit_threshold: float,
+    max_post_edit_steps: int,
+    penalty_lambda: float,
+):
+    """Batched single denoise step for joint-threshold decoding.
+
+    Advances ``input_ids_1d`` / ``finished`` / ``post_edit_steps`` in place,
+    processing every block at once (no per-row Python loop or ``.item()`` sync).
+    Semantics mirror the per-row path in ``JointThreshold.step``.
+    """
+    B = input_ids_1d.shape[0] // blk
+    V = full_logits_2d.shape[1]
+
+    input_ids = input_ids_1d.view(B, blk)
+    logits = full_logits_2d.view(B, blk, V)
+
+    active = ~finished
+
+    # ---------- penalty ----------
+    if penalty_lambda > 0:
+        prev_ids = input_ids[:, :-1]
+        logits[:, 1:, :].scatter_(
+            dim=2,
+            index=prev_ids.unsqueeze(-1),
+            src=torch.full_like(
+                prev_ids.unsqueeze(-1), -penalty_lambda, dtype=logits.dtype
+            ),
+            reduce="add",
+        )
+
+    # ---------- argmax + confidence ----------
+    # Same ops as the per-row path (argmax over logits, then gather the softmax
+    # probability), just batched: keeps decisions bitwise-aligned with it. On
+    # NPU this also beats a log-domain max+logsumexp variant (fused softmax).
+    x = torch.argmax(logits, dim=-1)
+    p = torch.gather(F.softmax(logits, dim=-1), dim=-1, index=x.unsqueeze(-1)).squeeze(
+        -1
+    )
+
+    mask_pos = input_ids.eq(mask_id)
+    has_mask = mask_pos.any(dim=1)
+
+    # ---------- post-edit ----------
+    no_mask_active = active & (~has_mask)
+    post_edit_steps.add_(no_mask_active.to(post_edit_steps.dtype))
+    exceeded = post_edit_steps > max_post_edit_steps
+    finished |= no_mask_active & exceeded
+
+    # eligible rows (match original semantics)
+    eligible = active & (~(no_mask_active & exceeded))
+
+    # ---------- M2T ----------
+    neg_inf = torch.full_like(p, float("-inf"))
+    conf_m2t = torch.where(mask_pos, p, neg_inf)
+
+    m2t = (conf_m2t > threshold) & (eligible & has_mask).view(B, 1)
+
+    # force-one if needed
+    hit_any = m2t.any(dim=1)
+    need_force = (eligible & has_mask) & (~hit_any)
+
+    # topk (not argmax): the per-row fallback picks its forced position with
+    # torch.topk, and the two ops can break exact-confidence ties differently.
+    best_idx = torch.topk(conf_m2t, k=1, dim=1).indices.squeeze(1)
+    rows = torch.arange(B, device=input_ids.device)
+
+    m2t[rows, best_idx] |= need_force
+
+    # ---------- T2T ----------
+    edit_mask = (~mask_pos) & (~prompt_masks)
+    t2t = (p > edit_threshold) & (input_ids != x) & edit_mask
+    t2t = t2t & eligible.view(B, 1)
+
+    # ---------- combine ----------
+    transfer = m2t | t2t
+    any_transfer_row = transfer.any(dim=1)
+
+    finished |= eligible & (~any_transfer_row)
+
+    # apply update
+    input_ids.copy_(torch.where(transfer, x, input_ids))
+
+    return any_transfer_row.any()
 
 
 class JointThreshold(DllmAlgorithm):
@@ -23,6 +120,14 @@ class JointThreshold(DllmAlgorithm):
             "max_post_edit_steps", 16
         )
         self.penalty_lambda = config.algorithm_config.get("penalty_lambda", 0)
+        # NPU defaults to the batched (vectorized) path; other platforms keep the
+        # upstream per-row path unless explicitly overridden via algorithm_config.
+        self.vectorized_decoding = config.algorithm_config.get(
+            "vectorized_decoding", _is_npu
+        )
+        # The sync loop advances one shared batched state in place across steps;
+        # FDFO must carry state per request, so it gathers/scatters each round.
+        self._use_shared_state = self.vectorized_decoding and not self.fdfo
 
     def max_steps(self, block_size: int) -> int:
         return block_size + self.max_post_edit_steps + 1
@@ -33,6 +138,18 @@ class JointThreshold(DllmAlgorithm):
         # Built once as a GPU tensor and reused across steps (no per-step
         # host/device transfer); the FDFO carry keeps it in-process.
         prompt_mask = input_ids != self.mask_id
+        if self._use_shared_state:
+            # One shared batched state, advanced in place across every step of the
+            # synchronous loop.
+            device = forward_batch.input_ids.device
+            shared = {
+                "prompt_masks": prompt_mask,  # [B, blk]
+                "finished": torch.zeros(batch_size, dtype=torch.bool, device=device),
+                "post_edit_steps": torch.zeros(
+                    batch_size, dtype=torch.int32, device=device
+                ),
+            }
+            return [shared] * batch_size
         return [
             {
                 "post_edit_steps": 0,
@@ -47,6 +164,88 @@ class JointThreshold(DllmAlgorithm):
         forward_batch: ForwardBatch,
         full_logits: torch.Tensor,
         states: List[Any],
+    ) -> List[bool]:
+        if self._use_shared_state:
+            return self._step_vectorized_shared(
+                forward_batch=forward_batch, full_logits=full_logits, states=states
+            )
+        if self.vectorized_decoding:
+            return self._step_vectorized_fdfo(
+                forward_batch=forward_batch, full_logits=full_logits, states=states
+            )
+        return self._step_per_row(
+            forward_batch=forward_batch, full_logits=full_logits, states=states
+        )
+
+    def _step_vectorized_shared(
+        self,
+        forward_batch: ForwardBatch,
+        full_logits: torch.Tensor,
+        states: List[Any],
+    ) -> List[bool]:
+        shared = states[0]
+        joint_threshold_update_step_vectorized(
+            input_ids_1d=forward_batch.input_ids,
+            full_logits_2d=full_logits,
+            prompt_masks=shared["prompt_masks"],
+            finished=shared["finished"],
+            post_edit_steps=shared["post_edit_steps"],
+            mask_id=self.mask_id,
+            blk=self.block_size,
+            threshold=self.threshold,
+            edit_threshold=self.edit_threshold,
+            max_post_edit_steps=self.max_post_edit_steps,
+            penalty_lambda=self.penalty_lambda,
+        )
+        return shared["finished"].tolist()
+
+    def _step_vectorized_fdfo(
+        self,
+        forward_batch: ForwardBatch,
+        full_logits: torch.Tensor,
+        states: List[Any],
+    ) -> List[bool]:
+        # FDFO carries per-request dict states across rounds (stashed on the
+        # request, re-mixed with fresh rows each round), so gather them into
+        # batched tensors for this round's single step, then scatter the results
+        # back onto the per-request dicts.
+        device = forward_batch.input_ids.device
+        prompt_masks = torch.stack([state["prompt_mask"] for state in states])
+        finished = torch.tensor(
+            [state["finished"] for state in states], dtype=torch.bool, device=device
+        )
+        post_edit_steps = torch.tensor(
+            [state["post_edit_steps"] for state in states],
+            dtype=torch.int32,
+            device=device,
+        )
+
+        joint_threshold_update_step_vectorized(
+            input_ids_1d=forward_batch.input_ids,
+            full_logits_2d=full_logits,
+            prompt_masks=prompt_masks,
+            finished=finished,
+            post_edit_steps=post_edit_steps,
+            mask_id=self.mask_id,
+            blk=self.block_size,
+            threshold=self.threshold,
+            edit_threshold=self.edit_threshold,
+            max_post_edit_steps=self.max_post_edit_steps,
+            penalty_lambda=self.penalty_lambda,
+        )
+
+        done = finished.tolist()
+        new_post_edit_steps = post_edit_steps.tolist()
+        for i, state in enumerate(states):
+            state["finished"] = done[i]
+            state["post_edit_steps"] = new_post_edit_steps[i]
+        return done
+
+    def _step_per_row(
+        self,
+        forward_batch: ForwardBatch,
+        full_logits: torch.Tensor,
+        states: List[Any],
     ) -> List[bool]:
         batch_size = forward_batch.batch_size
         done: List[bool] = []