diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f33d5125d..5a933dc67 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -208,6 +208,17 @@ ATTENTION_BACKEND_CHOICES = [ "intel_xpu", ] +# trtllm_mha is valid for decode-only dense-MQA drafts. DFLASH rejects it +# earlier when its per-layer attention requirements are not met. +DRAFT_ATTENTION_BACKEND_CHOICES = [ + "flashinfer", + "fa3", + "fa4", + "triton", + "ascend", + "trtllm_mha", +] + # Attention backends whose kernels read the chunked prefix-cache layout. # Out-of-tree platforms may extend this list (via # add_chunked_prefix_cache_attention_backend) before ServerArgs construction; @@ -406,6 +417,10 @@ def add_attention_backend_choices(choices): ATTENTION_BACKEND_CHOICES.extend(choices) +def add_draft_attention_backend_choices(choices): + DRAFT_ATTENTION_BACKEND_CHOICES.extend(choices) + + def add_chunked_prefix_cache_attention_backend(backend_name): CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS.append(backend_name) diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index 058f0f782..c7c1c14fe 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/python/sglang/srt/speculative/draft_worker_common.py @@ -9,7 +9,7 @@ import torch from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode -from sglang.srt.server_args import ServerArgs +from sglang.srt.server_args import DRAFT_ATTENTION_BACKEND_CHOICES, ServerArgs from sglang.srt.speculative.dflash_info import DFlashVerifyInput from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 @@ -20,18 +20,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -# trtllm_mha: decode-only dense-MQA drafts (dspark). DFLASH excludes it -# earlier, at arg resolution (speculative_hook.py) -- its draft path needs -# per-layer DFlash attention -- so it never reaches this gate with it. -_SUPPORTED_DRAFT_BACKENDS = ( - "flashinfer", - "fa3", - "fa4", - "triton", - "ascend", - "trtllm_mha", -) - class DraftWorkerBundle(msgspec.Struct, frozen=True): draft_worker: TpModelWorker @@ -48,13 +36,13 @@ def _resolve_draft_attention_backend_fallback( draft_backend, _ = server_args.get_attention_backends() if draft_backend is None: return "triton" if torch.version.hip else "flashinfer" - if draft_backend not in _SUPPORTED_DRAFT_BACKENDS: + if draft_backend not in DRAFT_ATTENTION_BACKEND_CHOICES: fallback = "triton" if torch.version.hip else "flashinfer" logger.warning( "%s draft worker only supports attention_backend in %s for now, " "but got %r. Falling back to '%s'.", algo_label, - _SUPPORTED_DRAFT_BACKENDS, + DRAFT_ATTENTION_BACKEND_CHOICES, draft_backend, fallback, )