Make draft attention backends extensible (#35932)

Co-authored-by: Yichao Fu <yichaofu@meta.com>

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)
 

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,
         )