From 36afd442c7a96c4997f14fbc30058616581810b0 Mon Sep 17 00:00:00 2001 From: wenxuewuhd Date: Thu, 30 Jul 2026 14:13:14 +0800 Subject: [PATCH] [DLLM] vectorized joint/low-confidence decoding and skip redundant attn init (#21094) Co-authored-by: ronnie_zheng --- python/sglang/srt/dllm/algorithm/base.py | 8 + .../srt/dllm/algorithm/joint_threshold.py | 199 ++++++++++++++++++ .../dllm/test_npu_llada2_mini.py | 1 - 3 files changed, 207 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/dllm/algorithm/base.py b/python/sglang/srt/dllm/algorithm/base.py index 5272c518b..0400c0540 100644 --- a/python/sglang/srt/dllm/algorithm/base.py +++ b/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): diff --git a/python/sglang/srt/dllm/algorithm/joint_threshold.py b/python/sglang/srt/dllm/algorithm/joint_threshold.py index 6866996c4..00769a924 100644 --- a/python/sglang/srt/dllm/algorithm/joint_threshold.py +++ b/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] = [] diff --git a/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py b/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py index 7550fd88d..4d55a335a 100644 --- a/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py +++ b/test/registered/ascend/basic_function/dllm/test_npu_llada2_mini.py @@ -19,7 +19,6 @@ class TestLLaDA2Mini(GSM8KAscendMixin, CustomTestCase): other_args = [ "--trust-remote-code", - "--disable-radix-cache", "--mem-fraction-static", "0.9", "--max-running-requests",