[NPU] Support DFlash speculative decoding for MiMo-V2.5-Pro (mxfp4) (#37565)

python/sglang/srt/arg_groups/speculative_hook.py

@@ -206,9 +206,11 @@ def _handle_dflash(server_args: ServerArgs) -> None:
             "DFLASH speculative decoding only supports CUDA, NPU and XPU devices."
         )
 
-    if resolved_view(server_args).enable_dp_attention:
+    # DFLASH + dp attention is validated on NPU only.
+    if cfg.enable_dp_attention and not cfg.device == "npu":
         raise ValueError(
-            "Currently DFLASH speculative decoding does not support dp attention."
+            "Currently DFLASH speculative decoding does not support dp "
+            "attention on non-NPU devices."
         )
 
     if cfg.pp_size != 1:

python/sglang/srt/configs/model_config.py

@@ -510,6 +510,15 @@ class ModelConfig:
         if self.is_fp4_experts:
             logger.info("Detected mixed checkpoint layout: routed experts are MXFP4.")
 
+        # MiMo-V2 mxfp4 ckpts declare the routed-expert layout via store_dtype.
+        if (
+            not self.is_fp4_experts
+            and _hf_arch(self.hf_config) in MIMO_V2_MODEL_ARCHS
+            and str(quantization_config.get("store_dtype") or "").lower() == "mxfp4"
+        ):
+            self.is_fp4_experts = True
+            logger.info("Detected MiMo-V2 mxfp4 routed-expert layout.")
+
         # DSV4 mxfp4 layout applies only when the ckpt does not opt in above.
         if is_deepseek_v4(self.hf_config) and routed_experts_quant_method is None:
             self.is_fp4_experts = envs.SGLANG_DSV4_FP4_EXPERTS.get()

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -476,12 +476,23 @@ class AscendAttnBackend(AttentionBackend):
         self.forward_metadata = ForwardMetadata()
         seq_lens_max = forward_batch.seq_lens.max()
         if forward_batch.forward_mode.is_target_verify():
+            if (
+                forward_batch.spec_algorithm is not None
+                and forward_batch.spec_algorithm.is_dflash()
+            ):
+                # dflash_worker_v2 already publishes seq_lens_cpu as prefix +
+                # one verify block, which already covers the draft block.
+                seq_lens_max = forward_batch.seq_lens_cpu.max().item()
+            else:
+                # Overlap scheduling can publish the CPU sequence length one
+                # step ahead of the device tensor. FIA consumes seq_lens_cpu
+                # below, so derive the block-table width from the same source.
+                # Otherwise a page-aligned request can expose KV_S=N while
+                # asking FIA for N+1.
                 spec_tokens_per_req = int(forward_batch.spec_info.draft_token_num)
-            # Overlap scheduling can publish the CPU sequence length one step
+                seq_lens_max = (
-            # ahead of the device tensor. FIA consumes seq_lens_cpu below, so
+                    forward_batch.seq_lens_cpu.max().item() + spec_tokens_per_req
-            # derive the block-table width from the same source. Otherwise a
+                )
-            # page-aligned request can expose KV_S=N while asking FIA for N+1.
-            seq_lens_max = forward_batch.seq_lens_cpu.max().item() + spec_tokens_per_req
         elif (
             forward_batch.forward_mode.is_decode_or_idle()
             and forward_batch.spec_info is not None
@@ -522,6 +533,10 @@ class AscendAttnBackend(AttentionBackend):
             ).int()
 
         self.forward_metadata.seq_lens_cpu_int = forward_batch.seq_lens_cpu.int()
+        # In graph mode (see _init_cuda_graph_metadata) seq_lens_cpu_int stays
+        # None so forward_mtp binds seq_lens_cpu_list instead: graph.update can
+        # only rebind the Host-side IntArray when captured as a Python list.
+
         if (
             not forward_batch.forward_mode.is_draft_extend_v2()
             and not forward_batch.forward_mode.is_target_verify()
@@ -531,8 +546,10 @@ class AscendAttnBackend(AttentionBackend):
 
         if forward_batch.forward_mode.is_target_verify():
             spec_algorithm = forward_batch.spec_algorithm
-            if spec_algorithm is None or not spec_algorithm.is_dspark():
+            if spec_algorithm is None or not spec_algorithm.is_dflash_family():
-                self.forward_metadata.seq_lens_cpu_int += spec_tokens_per_req
+                self.forward_metadata.seq_lens_cpu_int += int(
+                    forward_batch.spec_info.draft_token_num
+                )
         elif (
             forward_batch.forward_mode.is_decode_or_idle()
             and forward_batch.spec_info is not None
@@ -680,6 +697,8 @@ class AscendAttnBackend(AttentionBackend):
             metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[:num_tokens]
         metadata.seq_lens_cpu_list = seq_lens.cpu().int().tolist()
         metadata.seq_lens = seq_lens
+        # Deliberately no seq_lens_cpu_int here: forward_mtp must bind
+        # seq_lens_cpu_list so graph.update can rebind it.
         if forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2():
             metadata.actual_seq_lengths_q = torch.arange(
                 self.speculative_num_draft_tokens,
@@ -775,7 +794,16 @@ class AscendAttnBackend(AttentionBackend):
             metadata.block_tables_swa[:bs, max_seq_pages:].fill_(0)
             metadata.block_tables_swa[bs:, :].fill_(0)
 
-            # Update SWA mask: True = masked out (don't attend), False = attend
+            # Update SWA mask: True = masked out (don't attend), False = attend.
+            # DFlash verify seq_lens is prefix-only, so use seq_lens_cpu
+            # (= prefix + block_size) to keep draft KV inside the mask window.
+            if (
+                forward_mode.is_target_verify()
+                and _is_dflash_verify(spec_info)
+                and seq_lens_cpu is not None
+            ):
+                seq_lens_int = seq_lens_cpu[:bs].int()
+            else:
                 seq_lens_int = seq_lens[:bs].int()
             starts = torch.clamp(seq_lens_int - self.sliding_window_size, min=0)
             indices = self.graph_metadata["swa_indices"]
@@ -796,6 +824,13 @@ class AscendAttnBackend(AttentionBackend):
 
         if forward_mode.is_target_verify():
             seq_lens = seq_lens + self.speculative_num_draft_tokens
+            # For DFlash, seq_lens_cpu (= prefix + block_size) is the true KV
+            # length; other spec algorithms already added the draft tokens above.
+            if _is_dflash_verify(spec_info) and seq_lens_cpu is not None:
+                kv_lens = seq_lens_cpu[:bs]
+            else:
+                kv_lens = seq_lens[:bs]
+            metadata.seq_lens_cpu_list = kv_lens.cpu().int().tolist()
         elif forward_mode.is_decode_or_idle() and spec_info is not None:
             seq_lens = seq_lens + self.speculative_step_offset_npu
         metadata.seq_lens[:bs].copy_(seq_lens[:bs])
@@ -2111,13 +2146,26 @@ class AscendAttnBackend(AttentionBackend):
             if not self.graph_mode:
                 num_token_padding = query.shape[0]
                 query = query[: forward_batch.global_num_token_non_padded_cpu]
+                # Trim DP padding rows so actualSeqLengthsKv matches the
+                # operator's batchSize (TND layout); only target_verify has a
+                # uniform per-request width.
+                if forward_batch.forward_mode.is_target_verify():
+                    real_bs = query.shape[0] // self.speculative_num_draft_tokens
 
             if self.forward_metadata.seq_lens_cpu_int is None:
+                # Graph mode: bind the Python list, which graph.update can
+                # rebind (a captured CPU tensor would be baked as constant).
                 actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list
             else:
                 actual_seq_lengths_kv = (
                     self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
                 )
+            if (
+                not self.graph_mode
+                and forward_batch.forward_mode.is_target_verify()
+                and len(actual_seq_lengths_kv) > real_bs
+            ):
+                actual_seq_lengths_kv = actual_seq_lengths_kv[:real_bs]
 
             if forward_batch.forward_mode.is_draft_extend_v2():
                 extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
@@ -2127,6 +2175,7 @@ class AscendAttnBackend(AttentionBackend):
                     ]
                 actual_seq_lengths = np.array(extend_seq_lens_cpu).cumsum().tolist()
             else:
+                # Static across replays ([spec_draft, 2*spec_draft, ...]).
                 actual_seq_lengths = np.arange(
                     self.speculative_num_draft_tokens,
                     self.speculative_num_draft_tokens + query.shape[0],
@@ -2146,6 +2195,13 @@ class AscendAttnBackend(AttentionBackend):
                 block_table = self.forward_metadata.block_tables_swa
             else:
                 block_table = self.forward_metadata.block_tables
+            if (
+                not self.graph_mode
+                and forward_batch.forward_mode.is_target_verify()
+                and block_table.shape[0] > real_bs
+            ):
+                # Drop DP padding rows (see real_bs comment above).
+                block_table = block_table[:real_bs]
 
             if layer.attn_type == AttentionType.ENCODER_ONLY:
                 mask = None
@@ -2182,7 +2238,7 @@ class AscendAttnBackend(AttentionBackend):
                     query,
                     k_cache,
                     v_cache,
-                    block_table=self.forward_metadata.block_tables,
+                    block_table=block_table,
                     block_size=self.page_size,
                     num_heads=layer.tp_q_head_num,
                     num_key_value_heads=layer.tp_k_head_num,

python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py

@@ -40,9 +40,14 @@ from sglang.srt.configs.model_config import (
     is_deepseek_dsa,
     is_deepseek_v4,
 )
-from sglang.srt.distributed.parallel_state import GroupCoordinator
+from sglang.srt.distributed.parallel_state import (
+    GroupCoordinator,
+)
 from sglang.srt.environ import envs
 from sglang.srt.model_executor.runner import DecodeCudaGraphRunner
+from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
+    build_replay_fb_view,
+)
 from sglang.srt.utils import (
     empty_context,
     get_bool_env_var,
@@ -62,7 +67,12 @@ if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
 
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_executor.forward_batch_info import (
+    ForwardBatch,
+    PPProxyTensors,
+    compute_local_num_token_non_padded_cpu,
+    enable_num_token_non_padded,
+)
 
 
 @contextmanager
@@ -109,6 +119,11 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
         self.update_attr_name = None
         self.update_attr_type = None
         self.model_runner = model_runner
+        # DFLASH verify under dp attention replays through the generic DP
+        # graph machinery: the scheduler-level DP vote keeps all DP ranks on
+        # the same graph/eager decision, and load_batch pads every rank to
+        # the same global max bucket so the captured dp-gather geometry stays
+        # valid when per-rank batch sizes diverge.
         self._init_arch_map()
         self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
         self.if_use_v2 = any(
@@ -133,6 +148,8 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
         self.attr_type: Dict[str, Union[list, torch.Tensor]] = {
             AttentionArch.MLA: [],
             AttentionArch.MHA: torch.Tensor(),
+            # TARGET_VERIFY must use a Python list: graph.update can only
+            # rebind the Host-side IntArray when captured as a list.
             "TARGET_VERIFY": [],
         }
 
@@ -215,6 +232,46 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
             self.load_batch(forward_batch, pp_proxy_tensors)
         else:
             # In speculative decoding, these two fields are still needed.
+            # NPU skips the DFLASH verify pre-planning, so load_batch may
+            # never have recorded the padded batch shapes; recompute them on
+            # every batch (the verify batch size varies with concurrency).
+            raw_bs = forward_batch.batch_size
+            if self.require_mlp_tp_gather:
+                bs = self._pad_to_bucket(
+                    self._max_dp_batch_size(forward_batch), self.capture_bs
+                )
+            else:
+                bs = self._pad_to_bucket(raw_bs, self.capture_bs)
+            self.raw_bs = raw_bs
+            self.raw_num_token = raw_bs * self.captured_req_width
+            self.bs = bs
+            # Restore the DeepEP dispatch mode recorded at capture time
+            # (mirrors load_batch); an interleaved eager extend may have
+            # switched it.
+            self.deepep_adapter.replay()
+            # Refresh the static DP token buffers bound by the captured
+            # graph (stale values misalign dp-gather segments across ranks);
+            # mirror the capture-side uniform [padded_num_tokens] * dp_size.
+            if self.require_mlp_tp_gather:
+                _padded_num_tokens = bs * self.captured_req_width
+                self.buffers.global_num_tokens_gpu.fill_(_padded_num_tokens)
+                self.buffers.global_num_tokens_for_logprob_gpu.fill_(_padded_num_tokens)
+            if (
+                enable_num_token_non_padded()
+                and self.require_gathered_buffer
+                and not self.enable_prefill_cp
+            ):
+                self.buffers.num_token_non_padded.fill_(
+                    compute_local_num_token_non_padded_cpu(
+                        global_num_token_non_padded=(
+                            forward_batch.global_num_token_non_padded_cpu
+                        ),
+                        num_tokens_per_dp=bs * self.captured_req_width,
+                        sharded=self.model_runner.attn_tp_sequence_sharded(
+                            bs * self.captured_req_width
+                        ),
+                    )
+                )
             self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
             self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
             if (
@@ -233,6 +290,42 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
                     forward_batch.mrope_positions
                 )
 
+            # The pre-planned path skipped init_forward_metadata_out_graph;
+            # refresh attention metadata so replay reads correct KV pages.
+            self.buffers.seq_lens[: self.raw_bs].copy_(
+                forward_batch.seq_lens_cpu[: self.raw_bs]
+            )
+            self.buffers.seq_lens[self.raw_bs : self.bs].fill_(self.seq_len_fill_value)
+            self.buffers.seq_lens_cpu[: self.raw_bs].copy_(
+                forward_batch.seq_lens_cpu[: self.raw_bs]
+            )
+            self.buffers.seq_lens_cpu[self.raw_bs : self.bs].fill_(
+                self.seq_len_fill_value
+            )
+            self.buffers.req_pool_indices[: self.raw_bs].copy_(
+                forward_batch.req_pool_indices[: self.raw_bs]
+            )
+            self.buffers.req_pool_indices[self.raw_bs : self.bs].fill_(0)
+            # Refresh the static out_cache_loc bound by the captured graph
+            # for full-pool KV writes in save_kv_cache (replay would
+            # otherwise write verify KV to stale capture-time slots).
+            if forward_batch.out_cache_loc is not None:
+                _padded_num_token = self.bs * self.captured_req_width
+                _n = min(self.raw_num_token, forward_batch.out_cache_loc.shape[0])
+                self.buffers.out_cache_loc[:_n].copy_(forward_batch.out_cache_loc[:_n])
+                self.buffers.out_cache_loc[_n:_padded_num_token].zero_()
+            fb_view = build_replay_fb_view(
+                forward_batch=forward_batch,
+                buffers=self.buffers,
+                bs=self.bs,
+                raw_bs=self.raw_bs,
+                num_tokens=self.bs * self.captured_req_width,
+                seq_len_fill_value=self.seq_len_fill_value,
+                capture_forward_mode=self.capture_forward_mode,
+                is_encoder_decoder=self.is_encoder_decoder,
+            )
+            self._replay_attn_backend().init_forward_metadata_out_graph(fb_view)
+
         graph_key = self._make_graph_key(self.bs)
 
         if not (
@@ -240,7 +333,22 @@ class NPUGraphRunner(DecodeCudaGraphRunner):
             or is_deepseek_v4(self.model_runner.model_config.hf_config)
         ):
             if forward_batch.forward_mode.is_target_verify():
-                seq_lens_cpu = forward_batch.seq_lens.cpu() + self.captured_req_width
+                _attn = self._replay_attn_backend()
+                _meta = getattr(_attn, "forward_metadata", None)
+                _meta_list = getattr(_meta, "seq_lens_cpu_list", None)
+                if _meta_list is not None:
+                    # graph.update must carry the exact KV length already
+                    # computed in forward_metadata.seq_lens_cpu_list (it
+                    # already includes the draft block for DFlash); do not
+                    # recompute and double-add here.
+                    seq_lens = list(_meta_list)
+                else:
+                    # Wrapper backends (e.g. hybrid linear attention) keep
+                    # forward_metadata only on their children, so it stays
+                    # None here; fall back to the pre-DFlash computation.
+                    seq_lens_cpu = (
+                        forward_batch.seq_lens.cpu() + self.captured_req_width
+                    )
                     seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
             else:
                 seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (

python/sglang/srt/layers/quantization/fp8.py

@@ -261,6 +261,7 @@ class Fp8Config(QuantizationConfig):
         # model_loader from ModelConfig. Default False off the DSV4 path.
         self.is_fp4_experts = is_fp4_experts
         self.dequant_fp4_to_fp8 = False
+        self.is_dsv4_fp4_experts = False
         self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
         if is_checkpoint_fp8_serialized:
             log_info_on_rank0(logger, "Detected fp8 checkpoint.")
@@ -409,7 +410,14 @@ class Fp8Config(QuantizationConfig):
                 )
                 return fp8_method
 
-            if self.is_fp4_experts and is_npu_arch35():
+            if (
+                self.is_fp4_experts
+                and is_npu_arch35()
+                # NPUW4A4Fp4MoEMethod is DSV4-specific (deepep dispatch,
+                # swiglu_limit); other FP4-expert checkpoints fall through
+                # to Fp8MoEMethod.
+                and self.is_dsv4_fp4_experts
+            ):
                 from sglang.srt.hardware_backend.npu.quantization.fp4_moe_methods import (
                     NPUW4A4Fp4MoEMethod,
                 )
@@ -505,6 +513,15 @@ class Fp8LinearMethod(LinearMethodBase):
             self.w8a8_mxfp8_linear = dispatch_w8a8_mxfp8_linear()
         else:
             self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear()
+            if _is_npu and is_npu_arch35() and self.quant_config.scale_fmt != "ue8m0":
+                # The A5 backend expects the ue8m0 weight layout installed by
+                # the arch35 load path; keep plain block-FP8 checkpoints on
+                # the generic triton backend.
+                from sglang.srt.layers.quantization.fp8_utils import (
+                    triton_w8a8_block_fp8_linear,
+                )
+
+                self.w8a8_block_fp8_linear = triton_w8a8_block_fp8_linear
         self.is_checkpoint_fp8_serialized = (
             self.quant_config.is_checkpoint_fp8_serialized
         )
@@ -708,7 +725,7 @@ class Fp8LinearMethod(LinearMethodBase):
             layer.weight_scale_inv.format_ue8m0 = True
             self._process_mxfp8_linear_weight_scale(layer)
             return
-        elif _is_npu and is_npu_arch35():
+        elif _is_npu and is_npu_arch35() and self.quant_config.scale_fmt == "ue8m0":
             from sglang.srt.hardware_backend.npu.quantization.w8a8_mxfp8 import (
                 process_npu_arch35_mxfp8_linear_weights,
             )
@@ -1767,6 +1784,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                 logger.warning_once("Dequantized FP4 expert weights to FP8.")
 
             if self.is_fp4_expert:
+                if _is_npu:
+                    self._process_npu_fp4_expert_weights(layer)
+                    return
+
                 if get_moe_runner_backend().is_marlin():
                     layer.w13_weight.data = layer.w13_weight.data.view(torch.int8)
                     layer.w2_weight.data = layer.w2_weight.data.view(torch.int8)
@@ -2454,6 +2475,20 @@ class Fp8MoEMethod(FusedMoEMethodBase):
     ):
         self._owns_moe_runner = False
         self.moe_runner_config = moe_runner_config
+
+        # MXFP4 experts on NPU run through the ASCEND runner, which
+        # quantizes activations internally.
+        if _is_npu and self.is_fp4_expert:
+            from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
+                NPUW4A8MXFP4MoEMethod,
+            )
+
+            layer.w13_kernel = NPUW4A8MXFP4MoEMethod()
+            layer.w2_kernel = NPUW4A8MXFP4MoEMethod()
+            moe_runner_config.layer = layer
+            self.runner = MoeRunner(MoeRunnerBackend.ASCEND, moe_runner_config)
+            return
+
         moe_runner_backend = get_moe_runner_backend()
 
         if moe_runner_backend.is_auto():
@@ -2505,6 +2540,53 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             block_shape=self.weight_block_size,
         )
 
+    def _process_npu_fp4_expert_weights(self, layer: torch.nn.Module) -> None:
+        """Convert HF MXFP4 experts to the NPU W4A8 kernel layout.
+
+        HF stores packed FP4 weights as int8 and block scales as float32; the
+        NPU kernel expects the weight in FRACTAL_NZ (transposed) and the scale
+        as e8m0 (uint8) reshaped to [E, K//64, N, 2].
+        """
+        from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
+            _get_float4_e2m1fn_x2_dtype,
+        )
+        from sglang.srt.hardware_backend.npu.utils import npu_format_cast
+
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+        if fp4_dtype is None:
+            raise RuntimeError("NPU W4A8 MXFP MoE requires float4 support.")
+
+        for prefix in ("w13", "w2"):
+            weight = getattr(layer, f"{prefix}_weight")
+            weight.data = npu_format_cast(
+                weight.data.view(torch.uint8),
+                customize_dtype=torch.float8_e4m3fn,
+                input_dtype=fp4_dtype,
+            ).transpose(-1, -2)
+
+            # Two e8m0 storage conventions in the fp32 scales: (a) the value
+            # IS the e8m0 byte (int in 0..255, MiMo-V2.5-Pro); (b) the value
+            # is 2^(e-127) -> recover e from the exponent field. Never
+            # re-encode with round(log2)+127: it re-biases the exponent.
+            scale_inv = getattr(layer, f"{prefix}_weight_scale_inv")
+            scale = scale_inv.data.to(torch.float32)
+            s_flat = scale.detach().float().flatten()
+            is_int_like = bool(
+                torch.all(s_flat >= 0)
+                and torch.all(s_flat <= 255)
+                and torch.allclose(s_flat, s_flat.round())
+            )
+            if is_int_like:
+                e8m0 = scale.to(torch.uint8)
+            else:
+                e8m0 = (scale.view(torch.int32) >> 23 & 0xFF).to(torch.uint8)
+            scale_inv.data = e8m0.reshape(
+                scale.shape[0],
+                scale.shape[1],
+                scale.shape[2] // 2,
+                2,
+            ).transpose(1, 2)
+
     def apply(
         self,
         layer: torch.nn.Module,
@@ -2516,6 +2598,17 @@ class Fp8MoEMethod(FusedMoEMethodBase):
         x = dispatch_output.hidden_states
         moe_runner_config = self.moe_runner_config
 
+        if _is_npu and self.is_fp4_expert:
+            from sglang.srt.layers.moe.moe_runner.ascend import AscendQuantInfo
+
+            quant_info = AscendQuantInfo(
+                w13_weight=layer.w13_weight,
+                w2_weight=layer.w2_weight,
+                w13_weight_scale=layer.w13_weight_scale_inv,
+                w2_weight_scale=layer.w2_weight_scale_inv,
+            )
+            return self.runner.run(dispatch_output, quant_info)
+
         if use_intel_amx_backend(layer):
             from sglang.srt.layers.moe.topk import apply_topk_weights_cpu
 

python/sglang/srt/model_loader/loader.py

@@ -214,6 +214,9 @@ def _get_quantization_config(
 
         if isinstance(quant_config, Fp8Config):
             quant_config.is_fp4_experts = model_config.is_fp4_experts
+            from sglang.srt.configs.model_config import is_deepseek_v4
+
+            quant_config.is_dsv4_fp4_experts = is_deepseek_v4(model_config.hf_config)
             quant_config.dequant_fp4_to_fp8 = envs.SGLANG_DSV4_FP4_DEQUANT.get()
             # Handle hybrid NVFP4 moe (nvidia/DeepSeek-V4-Pro-NVFP4)
             nvfp4_meta = model_config.nvfp4_moe_meta

python/sglang/srt/models/dflash.py

@@ -212,6 +212,7 @@ class DFlashAttention(nn.Module):
             base=rope_theta,
             rope_scaling=rope_scaling,
             is_neox_style=rope_is_neox_style,
+            partial_rotary_factor=float(getattr(config, "partial_rotary_factor", 1.0)),
         )
 
         self.scaling = head_dim**-0.5
@@ -225,6 +226,8 @@ class DFlashAttention(nn.Module):
         self.sliding_window_size, self.attn_type = _get_dflash_layer_attention_params(
             config, layer_id
         )
+        draft_cfg = parse_dflash_draft_config(draft_hf_config=config)
+        self.v_scale: Optional[float] = draft_cfg.attention_value_scale
         self.attention_sink_bias = None
         if is_nemotron_35_draft_config(config) and bool(
             getattr(config, "attention_sink_bias", False)
@@ -243,6 +246,16 @@ class DFlashAttention(nn.Module):
                 self.attention_sink_bias,
                 {"weight_loader": sharded_weight_loader(0)},
             )
+        elif draft_cfg.attention_sink_bias:
+            # Per-head sink bias; each TP rank owns its slice of the
+            # all-heads checkpoint tensor.
+            self.attention_sink_bias = nn.Parameter(
+                torch.empty(self.num_heads, dtype=torch.float32), requires_grad=False
+            )
+            set_weight_attrs(
+                self.attention_sink_bias,
+                {"weight_loader": sharded_weight_loader(0)},
+            )
         self.attn = RadixAttention(
             num_heads=self.num_heads,
             head_dim=head_dim,
@@ -301,6 +314,8 @@ class DFlashAttention(nn.Module):
             q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
             q, k = apply_qk_norm(q, k, self.q_norm, self.k_norm, self.head_dim)
             q, k = self.rotary_emb(positions, q, k)
+        if self.v_scale is not None:
+            v = v * self.v_scale
         if self.attention_sink_bias is None:
             attn_output = self.attn(q, k, v, forward_batch)
         else:
@@ -334,11 +349,14 @@ class DFlashAttention(nn.Module):
             )
             kv = F.linear(hidden_states, weight, bias)
             k, v = kv.split([self.kv_size, self.kv_size], dim=-1)
-            return k, v
+        else:
-
             # Fallback: compute full QKV and discard Q (keeps compatibility with quantized weights).
             qkv, _ = self.qkv_proj(hidden_states)
             _, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        # Keep V scaling consistent with forward() so ctx and self-generated V
+        # align in the draft cache.
+        if self.v_scale is not None:
+            v = v * self.v_scale
         return k, v
 
     def apply_k_norm(self, k: torch.Tensor) -> torch.Tensor:

python/sglang/srt/models/mimo_v2.py

@@ -31,6 +31,10 @@ from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_r
 from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
 from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
 from sglang.srt.layers.activation import SiluAndMul
+from sglang.srt.layers.aux_hidden_states import (
+    AuxHiddenStateAccumulator,
+    AuxHiddenStatePacker,
+)
 from sglang.srt.layers.communicator import (
     LayerCommunicator,
     LayerScatterModes,
@@ -873,10 +877,16 @@ class MiMoV2DecoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
         residual: Optional[torch.Tensor],
+        captured_last_layer_outputs: Optional[AuxHiddenStateAccumulator] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         # Self Attention
-        hidden_states, residual = self.layer_communicator.prepare_attn(
+        hidden_states, residual = (
-            hidden_states, residual, forward_batch
+            self.layer_communicator.prepare_attn_and_capture_last_layer_outputs(
+                hidden_states,
+                residual,
+                forward_batch,
+                captured_last_layer_outputs=captured_last_layer_outputs,
+            )
         )
 
         if hidden_states.shape[0] != 0:
@@ -994,6 +1004,7 @@ class MiMoV2Model(nn.Module):
         self.padding_idx = getattr(config, "pad_token_id", None)
         self.vocab_size = config.vocab_size
         self.pp_group = get_pp_group()
+        self.layers_to_capture = []
 
         if self.pp_group.is_first_rank:
             self.embed_tokens = VocabParallelEmbedding(
@@ -1053,7 +1064,8 @@ class MiMoV2Model(nn.Module):
             hidden_states = pp_proxy_tensors["hidden_states"]
             residual = pp_proxy_tensors["residual"]
 
-        if forward_batch.can_run_tbo:
+        aux_hidden_states = AuxHiddenStatePacker(len(self.layers_to_capture))
+        if forward_batch.can_run_tbo and not self.layers_to_capture:
             tbo_start_layer = self.start_layer
             tbo_end_layer = self.end_layer
 
@@ -1088,6 +1100,20 @@ class MiMoV2Model(nn.Module):
                     hidden_states,
                     forward_batch,
                     residual,
+                    captured_last_layer_outputs=(
+                        aux_hidden_states if i in self.layers_to_capture else None
+                    ),
+                )
+
+        # A draft targeting the final layer ("after layer
+        # num_hidden_layers-1") maps to capture index num_hidden_layers,
+        # past the layer loop; capture the pre-norm output here instead.
+        if (
+            self.pp_group.is_last_rank
+            and self.config.num_hidden_layers in self.layers_to_capture
+        ):
+            aux_hidden_states.append(
+                hidden_states if residual is None else hidden_states + residual
             )
 
         hidden_states_before_norm = None
@@ -1109,7 +1135,9 @@ class MiMoV2Model(nn.Module):
                 else:
                     hidden_states, _ = self.norm(hidden_states, residual)
 
+        if len(aux_hidden_states) == 0:
             return hidden_states, hidden_states_before_norm
+        return hidden_states, hidden_states_before_norm, aux_hidden_states.finalize()
 
     # If this function is called, it should always initialize KV cache scale
     # factors (or else raise an exception). Thus, handled exceptions should
@@ -1196,6 +1224,7 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin):
         self.logits_processor = (
             LogitsProcessor(config) if not self.config.encoder_only else None
         )
+        self.capture_aux_hidden_states = False
 
         vision_config = getattr(config, "vision_config", None)
         audio_config = getattr(config, "audio_config", None)
@@ -1353,7 +1382,16 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin):
             "forward() should not be called in encoder_only mode"
         )
 
-        if self._is_multimodal:
+        aux_hidden_states = None
+        if self.capture_aux_hidden_states:
+            hidden_states, hidden_states_before_norm, aux_hidden_states = self.model(
+                input_ids,
+                positions,
+                forward_batch,
+                input_embeds,
+                pp_proxy_tensors=pp_proxy_tensors,
+            )
+        elif self._is_multimodal:
             hidden_states, hidden_states_before_norm = general_mm_embed_routine(
                 input_ids=input_ids,
                 forward_batch=forward_batch,
@@ -1378,6 +1416,7 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin):
                 self.lm_head,
                 forward_batch,
                 hidden_states_before_norm=hidden_states_before_norm,
+                aux_hidden_states=aux_hidden_states,
             )
         else:
             return hidden_states
@@ -1390,6 +1429,20 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin):
     def end_layer(self):
         return self.model.end_layer if self.model is not None else 0
 
+    def set_dflash_layers_to_capture(self, layer_ids: List[int]):
+        if not self.pp_group.is_last_rank:
+            return
+
+        if layer_ids is None:
+            raise ValueError(
+                "DFLASH requires explicit layer_ids for aux hidden capture."
+            )
+
+        self.capture_aux_hidden_states = True
+        # target_layer_ids are "after layer X" ids; capture before layer X+1,
+        # matching the draft's extract_context_feature (offset=1).
+        self.model.layers_to_capture = [val + 1 for val in layer_ids]
+
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
@@ -1582,6 +1635,10 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin):
                     if weight_name not in name:
                         continue
                     name = name.replace(weight_name, param_name)
+                    # mxfp4 ckpts store expert scales without the `_inv` suffix,
+                    # while Fp8MoEMethod registers them as *_weight_scale_inv.
+                    if name.endswith("weight_scale") and (name + "_inv" in params_dict):
+                        name = name + "_inv"
                     param = params_dict[name]
                     weight_loader = param.weight_loader
                     weight_loader(

python/sglang/srt/speculative/dflash_disaggregation.py

@@ -13,7 +13,7 @@ if TYPE_CHECKING:
     from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
 
 
-def build_dflash_family_disagg_draft_input(
+def build_dflash_disagg_draft_input(
     batch: ScheduleBatch,
     last_tokens_tensor: torch.Tensor,
     future_map: FutureMap,

python/sglang/srt/speculative/dflash_utils.py

@@ -544,6 +544,8 @@ class DFlashDraftConfig:
     pure_draft_prefix_len: Optional[int]
     gru_hidden_dim: Optional[int]
     emb_dim: Optional[int]
+    attention_sink_bias: bool = False
+    attention_value_scale: Optional[float] = None
 
     @property
     def is_domino(self) -> bool:
@@ -708,6 +710,29 @@ def parse_dflash_draft_config(*, draft_hf_config: Any) -> DFlashDraftConfig:
                 f"got {mask_token_id}."
             )
 
+    # MiMo DFlash draft extras: per-head attention sink bias and V value scale.
+    raw_attention_sink_bias = dflash_cfg.get("attention_sink_bias", False)
+    if not isinstance(raw_attention_sink_bias, bool):
+        raise ValueError(
+            "DFLASH dflash_config.attention_sink_bias must be a bool, "
+            f"got {raw_attention_sink_bias!r} (type={type(raw_attention_sink_bias).__name__})."
+        )
+    attention_sink_bias = bool(raw_attention_sink_bias)
+
+    raw_attention_value_scale = dflash_cfg.get("attention_value_scale", None)
+    if raw_attention_value_scale is None:
+        attention_value_scale: Optional[float] = None
+    else:
+        if isinstance(raw_attention_value_scale, bool) or not isinstance(
+            raw_attention_value_scale, (int, float)
+        ):
+            raise ValueError(
+                "DFLASH dflash_config.attention_value_scale must be int|float|None, "
+                f"got {raw_attention_value_scale!r} "
+                f"(type={type(raw_attention_value_scale).__name__})."
+            )
+        attention_value_scale = float(raw_attention_value_scale)
+
     projector_type = dflash_cfg.get(
         "projector_type", _cfg_get(draft_hf_config, "projector_type", None)
     )
@@ -792,6 +817,8 @@ def parse_dflash_draft_config(*, draft_hf_config: Any) -> DFlashDraftConfig:
         pure_draft_prefix_len=pure_draft_prefix_len,
         gru_hidden_dim=gru_hidden_dim,
         emb_dim=emb_dim,
+        attention_sink_bias=attention_sink_bias,
+        attention_value_scale=attention_value_scale,
     )
 
 

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -1,9 +1,11 @@
 import logging
 import math
+import os
 from dataclasses import replace
 from typing import List, Optional, Tuple
 
 import torch
+import torch.distributed as dist
 
 from sglang.kernels.ops.speculative.cache_locs import (
     assign_extend_cache_locs_func,
@@ -32,6 +34,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     ForwardBatch,
     ForwardMode,
     compute_position,
+    enable_num_token_non_padded,
 )
 from sglang.srt.model_executor.model_runner import SamplingPrewarmResult
 from sglang.srt.model_executor.runner_utils.pool import (
@@ -40,6 +43,7 @@ from sglang.srt.model_executor.runner_utils.pool import (
 )
 from sglang.srt.runtime_context import (
     get_exec,
+    get_parallel,
     get_schedule,
     get_spec,
     mamba_track_grid,
@@ -81,8 +85,10 @@ from sglang.srt.speculative.spec_utils import (
     GrammarTree,
     assign_req_to_token_pool_func,
     build_grammar_vocab_mask,
+    draft_tp_context,
 )
 from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
+from sglang.srt.utils.common import empty_context
 
 _is_npu = is_npu()
 
@@ -379,8 +385,25 @@ class DFlashWorkerV2(BaseSpecWorker):
         self._warned_sampling_fallback = False
         self._draft_probs_buf = None
         self._logged_first_verify = False
-        self._tp_sync = SpecTpSync(get_tp_group())
+        self._full_embed_gpu: Optional[torch.Tensor] = None
+        # Under dp attention, peer DP ranks run different (idle) paths, so
+        # spec broadcasts must stay within the attn-TP group.
+        self._tp_sync = SpecTpSync(
+            get_parallel().attn_tp_group
+            if get_parallel().enable_dp_attention
+            else get_tp_group()
+        )
 
+        # Under dp attention, the draft worker runs on the per-DP attn-TP
+        # group, independent of idle peer DP ranks.
+        self.draft_tp_context = (
+            draft_tp_context if get_parallel().enable_dp_attention else empty_context
+        )
+        if get_parallel().enable_dp_attention:
+            draft_init_ctx = draft_tp_context(get_parallel().attn_tp_group)
+        else:
+            draft_init_ctx = empty_context()
+        with draft_init_ctx:
             bundle = build_draft_tp_worker(
                 server_args=server_args,
                 gpu_id=gpu_id,
@@ -466,6 +489,8 @@ class DFlashWorkerV2(BaseSpecWorker):
             if hasattr(target_model, "get_dflash_noise_embedding_scale")
             else 1.0
         )
+        self._maybe_merge_trained_mask_embedding()
+        self._cache_full_embed_weight()
         if self.ps.tp_rank == 0:
             logger.info(
                 "Initialized DFLASH draft runner. attention_backend=%s, model=%s, block_size=%s, draft_window_size=%s, compact_cache=%s",
@@ -574,6 +599,7 @@ class DFlashWorkerV2(BaseSpecWorker):
         )
 
     def init_attention_backends(self):
+        with self.draft_tp_context(self.draft_model_runner.tp_group):
             self._draft_worker.init_attention_backends()
         self._need_mamba_verify_commit = mambaish_config(
             self.model_runner.model_config
@@ -583,9 +609,20 @@ class DFlashWorkerV2(BaseSpecWorker):
         )
 
     def init_cuda_graphs(self):
+        with self.draft_tp_context(self.draft_model_runner.tp_group):
             capture_decode_cuda_graph = (
                 get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED
             )
+            if get_parallel().enable_dp_attention 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.
+                capture_decode_cuda_graph = False
+                if self.ps.tp_rank == 0:
+                    logger.warning(
+                        "Disable DFLASH draft cuda graph because dp attention "
+                        "is enabled (draft runs eager)."
+                    )
             if is_cuda() and capture_decode_cuda_graph:
                 available_mem = self._tp_sync.available_memory_gb(
                     SpecTpSyncSite.DFLASH_MEM,
@@ -1143,6 +1180,101 @@ class DFlashWorkerV2(BaseSpecWorker):
                 bs,
             )
 
+    def _maybe_merge_trained_mask_embedding(self) -> None:
+        """Merge a trained mask embedding (mask_embedding.pt) into the target table.
+
+        The target's own mask-token row is typically an untrained added-vocab
+        slot; overwrite it with the learned vector. VocabParallelEmbedding-aware:
+        each rank only updates the row if the token falls within its shard.
+        """
+        from sglang.srt.arg_groups.overrides import resolving_view
+
+        draft_model_path = resolving_view(self.server_args).speculative_draft_model_path
+        if draft_model_path is None:
+            return
+
+        mask_emb_path = os.path.join(draft_model_path, "mask_embedding.pt")
+        if not os.path.exists(mask_emb_path):
+            return
+
+        saved = torch.load(mask_emb_path, map_location=self.device, weights_only=True)
+        embedding_tensor = saved["embedding"]
+        saved_token_id = int(saved["mask_token_id"])
+
+        if saved_token_id != self._mask_token_id:
+            raise ValueError(
+                f"DFLASH mask_embedding.pt was trained with mask_token_id={saved_token_id}, "
+                f"but the current resolved mask_token_id={self._mask_token_id}. "
+                "These must match."
+            )
+
+        target_model = self._target_worker.model_runner.model
+        embed_module = target_model.get_input_embeddings()
+
+        if saved.get("per_position"):
+            raise NotImplementedError(
+                "DFLASH per-position mask embeddings are not supported in "
+                "inference yet (mask_embedding.pt per_position=True)."
+            )
+
+        token_id = self._mask_token_id
+        shard_indices = getattr(embed_module, "shard_indices", None)
+        with torch.no_grad():
+            if shard_indices is not None:
+                # VocabParallelEmbedding: only update if this token is in our shard.
+                start = shard_indices.org_vocab_start_index
+                end = shard_indices.org_vocab_end_index
+                if start <= token_id < end:
+                    local_idx = token_id - start
+                    embed_module.weight[local_idx].copy_(
+                        embedding_tensor.to(embed_module.weight.dtype)
+                    )
+            else:
+                embed_module.weight[token_id].copy_(
+                    embedding_tensor.to(embed_module.weight.dtype)
+                )
+
+        if self.ps.tp_rank == 0:
+            logger.info(
+                "Merged trained mask embedding into target model "
+                "(mask_token_id=%s, source=%s)",
+                self._mask_token_id,
+                mask_emb_path,
+            )
+
+    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().enable_dp_attention:
+            return
+
+        tp_group = get_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 self.ps.tp_rank == 0:
+            logger.info(
+                "DFLASH cached full embed on GPU for dp attention: shape=%s",
+                list(self._full_embed_gpu.shape),
+            )
+
     def _resolve_mask_token_id(
         self, *, mask_token: str, mask_token_id: Optional[int] = None
     ) -> int:
@@ -1601,6 +1733,22 @@ class DFlashWorkerV2(BaseSpecWorker):
                 f"DFLASH positions must be 1D, got shape={tuple(positions.shape)}."
             )
         num_tokens = int(target_hidden.shape[0])
+        # Drop trailing alignment-padding rows (from TP/EP-size rounding)
+        # before materializing into the draft KV; they would corrupt other
+        # requests' cache slots.
+        expected_tokens = int(cache_loc.numel())
+        if num_tokens > expected_tokens:
+            if not getattr(self, "_logged_padding_trim", False):
+                logger.warning(
+                    "DFLASH target_hidden has %d trailing padding row(s); trimming "
+                    "to cache_loc length=%d (target_hidden=%d). Logged once per worker.",
+                    num_tokens - expected_tokens,
+                    expected_tokens,
+                    num_tokens,
+                )
+                self._logged_padding_trim = True
+            target_hidden = target_hidden[:expected_tokens]
+            num_tokens = expected_tokens
         if int(cache_loc.numel()) != num_tokens:
             raise ValueError(
                 "DFLASH cache_loc length mismatch: "
@@ -1650,7 +1798,10 @@ class DFlashWorkerV2(BaseSpecWorker):
             if commit_lens.dtype != torch.int32:
                 commit_lens = commit_lens.to(torch.int32)
 
-        with torch.inference_mode():
+        with (
+            torch.inference_mode(),
+            self.draft_tp_context(self.draft_model_runner.tp_group),
+        ):
             ctx_hidden = self.draft_model.project_target_hidden(target_hidden)
 
             if cache_loc_2d is not None:
@@ -1737,6 +1888,9 @@ class DFlashWorkerV2(BaseSpecWorker):
             )
             if _is_npu:
                 _, k, v = attn.forward_prepare_npu(ctx_positions, layer_ctx_hidden)
+                # Keep V scaling consistent with forward().
+                if attn.v_scale is not None:
+                    v = v * attn.v_scale
             else:
                 k, v = attn.kv_proj_only(layer_ctx_hidden)
                 k = attn.apply_k_norm(k)
@@ -2054,6 +2208,15 @@ class DFlashWorkerV2(BaseSpecWorker):
             if on_publish is not None:
                 on_publish(batch_output.new_seq_lens)
 
+            # An idle DP rank runs the empty target prefill above to stay in
+            # the DP collective, but must skip the draft KV materialization,
+            # which needs per-request extend info.
+            if batch.forward_mode.is_idle():
+                batch_output.next_draft_input = DFlashDraftInputV2.create_idle_input(
+                    device=self.device
+                )
+                return batch_output
+
             if logits_output.hidden_states is None:
                 raise RuntimeError(
                     "DFLASH requires target aux hidden capture for prefill, but got None. "
@@ -2110,6 +2273,32 @@ class DFlashWorkerV2(BaseSpecWorker):
             )
 
         if batch.forward_mode.is_idle():
+            # Under dp attention an idle DP rank must still run the target
+            # verify forward (IDLE mode) so its cross-DP collectives stay in
+            # lockstep with the active DP group; the draft block's
+            # collectives are within-rank and skipped.
+            if get_parallel().enable_dp_attention:
+                idle_verify_input = DFlashVerifyInput(
+                    draft_token=torch.empty((0,), dtype=torch.long, device=self.device),
+                    positions=torch.empty((0,), dtype=torch.int64, device=self.device),
+                    draft_token_num=int(self.block_size),
+                    custom_mask=None,
+                    capture_hidden_mode=CaptureHiddenMode.FULL,
+                )
+                idle_verify_forward_batch, _ = idle_verify_input.prepare_for_verify(
+                    batch, self._target_worker
+                )
+                # Force eager: the active ranks run eager verify when any DP
+                # rank is idle (see the idle-guard below); a graph-replaying
+                # idle rank would disagree with them on DP-gather segment
+                # offsets (padded bucket vs raw counts).
+                idle_verify_forward_batch.can_run_decode_cuda_graph = False
+                self._target_worker.forward_batch_generation(
+                    batch=None,
+                    forward_batch=idle_verify_forward_batch,
+                    is_verify=True,
+                    skip_attn_backend_init=True,
+                )
             empty_ids = torch.empty((0,), dtype=torch.int64, device=self.device)
             empty_lens = torch.empty((0,), dtype=torch.int32, device=self.device)
             next_draft_input = self._make_next_draft_input_decode(
@@ -2221,6 +2410,13 @@ class DFlashWorkerV2(BaseSpecWorker):
             )
             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.
+            noise_embedding = torch.nn.functional.embedding(
+                block_ids, self._full_embed_gpu
+            )
+        else:
             noise_embedding = embed_module(block_ids)
         if self._noise_embed_scale != 1.0:
             noise_embedding = noise_embedding * self._noise_embed_scale
@@ -2281,6 +2477,12 @@ class DFlashWorkerV2(BaseSpecWorker):
             spec_algorithm=SpeculativeAlgorithm.DFLASH,
             spec_info=self._draft_block_spec_info,
             capture_hidden_mode=CaptureHiddenMode.NULL,
+            num_token_non_padded=(
+                torch.tensor(bs * block_size, dtype=torch.int32, device=device)
+                if enable_num_token_non_padded()
+                else None
+            ),
+            global_num_token_non_padded_cpu=bs * block_size,
         )
 
         if self.selector is not None:
@@ -2291,7 +2493,10 @@ class DFlashWorkerV2(BaseSpecWorker):
                     bs=bs, sampling_info=batch.sampling_info
                 )
 
-        with torch.inference_mode():
+        with (
+            torch.inference_mode(),
+            self.draft_tp_context(self.draft_model_runner.tp_group),
+        ):
             draft_out = self.draft_model_runner.forward(forward_batch)
         draft_logits_output = draft_out.logits_output
 
@@ -2349,6 +2554,7 @@ class DFlashWorkerV2(BaseSpecWorker):
                     self._draft_sampler.q_out[:bs],
                 )
         elif self.selector is not None:
+            with self.draft_tp_context(self.draft_model_runner.tp_group):
                 draft_next = self._propose_selector_block(
                     draft_logits_output=draft_logits_output,
                     bs=bs,
@@ -2361,6 +2567,7 @@ class DFlashWorkerV2(BaseSpecWorker):
             if draft_hidden is None:
                 raise RuntimeError("DFLASH draft model returned no hidden states.")
             draft_hidden = draft_hidden.view(bs, int(self.block_size), -1)
+            with self.draft_tp_context(self.draft_model_runner.tp_group):
                 draft_next = self._greedy_sample_from_vocab_parallel_head(
                     hidden_states=draft_hidden[:, 1:, :].reshape(
                         -1, draft_hidden.shape[-1]
@@ -2413,6 +2620,23 @@ class DFlashWorkerV2(BaseSpecWorker):
         batch.seq_lens_cpu = seq_lens_cpu_backup
         batch.seq_lens_sum = seq_lens_sum_backup
 
+        # Idle-DP guard: an idle rank's eager verify contributes raw scheduler
+        # counts while a graph-replaying rank assumes the padded bucket, so
+        # their DP-gather segment offsets disagree. Fall back to eager verify
+        # whenever any DP rank is idle this round.
+        if (
+            get_parallel().enable_dp_attention
+            and verify_forward_batch.original_global_num_tokens_cpu is not None
+            and min(verify_forward_batch.original_global_num_tokens_cpu) == 0
+        ):
+            verify_forward_batch.can_run_decode_cuda_graph = False
+
+        # Mixed-round guard: an extend rank's raw token counts disagree with
+        # the verify batch's spec-scaled counts on the DP-gather layout; run
+        # eager, symmetric with the idle guard above.
+        if get_parallel().enable_dp_attention and batch.is_extend_in_batch:
+            verify_forward_batch.can_run_decode_cuda_graph = False
+
         target_out = self.target_worker.forward_batch_generation(
             batch=None,
             forward_batch=verify_forward_batch,

python/sglang/srt/speculative/spec_info.py

@@ -211,6 +211,14 @@ class SpeculativeAlgorithm(Enum):
             return build_dspark_disagg_draft_input(
                 batch, last_tokens_tensor, future_map
             )
+        if self.is_dflash():
+            from sglang.srt.speculative.dflash_disaggregation import (
+                build_dflash_disagg_draft_input,
+            )
+
+            return build_dflash_disagg_draft_input(
+                batch, last_tokens_tensor, future_map
+            )
         return None
 
     def need_topk(self) -> bool:

python/sglang/test/ascend/test_ascend_utils.py

@@ -149,6 +149,12 @@ META_LLAMA_3_1_8B_INSTRUCT = os.path.join(
 )
 MIMO_7B_RL_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "XiaomiMiMo/MiMo-7B-RL")
 MIMO_V2_FLASH_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "XiaomiMiMo/MiMo-V2-Flash")
+MIMO_V2_5_PRO_FP4_DFLASH_WEIGHTS_PATH = os.path.join(
+    MODEL_WEIGHTS_DIR, "XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash"
+)
+MIMO_V2_5_PRO_FP4_DFLASH_DRAFT_WEIGHTS_PATH = os.path.join(
+    MIMO_V2_5_PRO_FP4_DFLASH_WEIGHTS_PATH, "dflash"
+)
 MIMO_V2_5_W8A8_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "solinliu/MiMo-V2.5-W8A8")
 MINICPM3_4B_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "OpenBMB/MiniCPM3-4B")
 MISTRAL_7B_INSTRUCT_V0_2_WEIGHTS_PATH = os.path.join(

test/manual/ascend/llm_models/test_npu_mimo_v2_5_pro_fp4_dflash.py

@@ -0,0 +1,47 @@
+import unittest
+
+from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin
+from sglang.test.ascend.test_ascend_utils import (
+    MIMO_V2_5_PRO_FP4_DFLASH_DRAFT_WEIGHTS_PATH,
+    MIMO_V2_5_PRO_FP4_DFLASH_WEIGHTS_PATH,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestMiMoV25ProFP4GraphWithDFlash(GSM8KAscendMixin, CustomTestCase):
+    """Testcase: Verify the inference accuracy of MiMo-V2.5-Pro-FP4 on GSM8K with npu graph and DFlash speculative decoding.
+
+    [Test Category] Model
+    [Test Target] MiMo-V2.5-Pro-FP4-DFlash
+    [Test Config] Prefill+Decode, npu graph enabled, DFLASH speculative decoding, mxfp4 quantization, dp attention
+    """
+
+    model = MIMO_V2_5_PRO_FP4_DFLASH_WEIGHTS_PATH
+    accuracy = 0.9
+    other_args = [
+        "--trust-remote-code",
+        "--mem-fraction-static",
+        "0.87",
+        "--attention-backend",
+        "ascend",
+        "--tp-size",
+        "8",
+        "--moe-a2a-backend",
+        "deepep",
+        "--deepep-mode",
+        "auto",
+        "--speculative-algorithm",
+        "DFLASH",
+        "--speculative-draft-model-path",
+        MIMO_V2_5_PRO_FP4_DFLASH_DRAFT_WEIGHTS_PATH,
+        "--speculative-num-draft-tokens",
+        "8",
+        "--dp-size",
+        "2",
+        "--enable-dp-attention",
+        "--enable-dp-lm-head",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/spec/dspark/test_dspark_stacked_ctx_kv_parity.py

@@ -58,6 +58,9 @@ def _make_attn(rope, *, eps=EPS, has_bias=False, quantized=False, g=None):
         k_norm.weight.copy_(torch.randn(HEAD_DIM, device=DEVICE, generator=g))
     attn.k_norm = k_norm
     attn.rotary_emb = rope
+    # Real DFlashAttention always sets v_scale in __init__; the mock must match
+    # so kv_proj_only's v_scale read does not AttributeError.
+    attn.v_scale = None
     for name in ("kv_proj_only", "apply_k_norm", "apply_k_rope"):
         setattr(attn, name, types.MethodType(getattr(DFlashAttention, name), attn))
     return attn