diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 0c24ae32e..427acdd38 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -147,8 +147,10 @@ def handle_speculative_decoding(server_args: ServerArgs) -> None: def _handle_dflash(server_args: ServerArgs) -> None: from sglang.srt.arg_groups.overrides import resolved_view - if not server_args.device.startswith("cuda"): - raise ValueError("DFLASH speculative decoding only supports CUDA device.") + if not (server_args.device.startswith("cuda") or server_args.device == "npu"): + raise ValueError( + "DFLASH speculative decoding only supports CUDA and NPU devices." + ) if resolved_view(server_args).enable_dp_attention: raise ValueError( diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 6b9e3e50c..30761f228 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -27,7 +27,7 @@ from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_flags -from sglang.srt.speculative.spec_info import SpecInput +from sglang.srt.speculative.spec_info import SpecInput, SpecInputType from sglang.srt.utils import get_bool_env_var, get_current_device_stream_fast if TYPE_CHECKING: @@ -44,6 +44,13 @@ logger = logging.getLogger(__name__) FULL_ATTENTION_WINDOW = 2147483647 +def _is_dflash_verify(spec_info: Optional[SpecInput]) -> bool: + return ( + spec_info is not None + and spec_info.spec_input_type == SpecInputType.DFLASH_VERIFY + ) + + def _expand_dsa_sparse_indices(topk_indices: torch.Tensor) -> torch.Tensor: """Expand [T, K] to [T, 1, K] for NPU sparse attention.""" if topk_indices.dim() == 2: @@ -480,7 +487,9 @@ class AscendAttnBackend(AttentionBackend): seq_lens_list_cumsum = np.cumsum(forward_batch.extend_seq_lens_cpu) self.forward_metadata.seq_lens_list_cumsum = seq_lens_list_cumsum - if forward_batch.forward_mode.is_target_verify(): + if forward_batch.forward_mode.is_target_verify() and not _is_dflash_verify( + forward_batch.spec_info + ): self.forward_metadata.seq_lens_cpu_int += self.speculative_num_draft_tokens elif ( forward_batch.forward_mode.is_decode_or_idle() @@ -667,7 +676,7 @@ class AscendAttnBackend(AttentionBackend): self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc) ) max_len = seq_lens_cpu[:bs].max().item() - if forward_mode.is_target_verify(): + if forward_mode.is_target_verify() and not _is_dflash_verify(spec_info): max_len += self.speculative_num_draft_tokens elif forward_mode.is_decode_or_idle() and spec_info is not None: max_len += self.speculative_step_id + 1