Add out-of-tree DFlash extension points (#38740)

Co-authored-by: Yuhan Chen <yuhanc@fb.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

python/sglang/kernels/ops/speculative/cache_locs.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from collections.abc import Callable
+
 import torch
 import triton
 import triton.language as tl
@@ -25,6 +27,14 @@ if _is_cpu:
     from sgl_kernel import assign_extend_cache_locs_cpu
 
 
+def _get_oot_speculative_cache_locs_fn() -> Callable[..., torch.Tensor] | None:
+    from sglang.srt.platforms import current_platform
+
+    if not current_platform.is_out_of_tree():
+        return None
+    return current_platform.get_speculative_cache_locs_fn()
+
+
 @triton.jit
 def assign_draft_cache_locs_contiguous(
     req_pool_indices,
@@ -489,6 +499,18 @@ def assign_extend_cache_locs_func(
     draft_token_num: int,
     device,
 ) -> torch.Tensor:
+    platform_fn = _get_oot_speculative_cache_locs_fn()
+    if platform_fn is not None:
+        return platform_fn(
+            req_pool_indices=req_pool_indices,
+            req_to_token=req_to_token,
+            start_offset=start_offset,
+            end_offset=end_offset,
+            batch_size=batch_size,
+            draft_token_num=draft_token_num,
+            device=device,
+        )
+
     if _is_cuda or _is_hip or _is_musa or _is_xpu:
         out_cache_loc = torch.empty(
             (batch_size * draft_token_num,),

python/sglang/srt/arg_groups/speculative_hook.py

@@ -16,6 +16,7 @@ from sglang.srt.arg_groups.overrides import (
     resolving_view,
     run_post_process_pass,
 )
+from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_platform
 
 if TYPE_CHECKING:
@@ -230,11 +231,17 @@ def handle_speculative_decoding(server_args: ServerArgs) -> None:
 def _handle_dflash(server_args: ServerArgs) -> None:
     cfg = resolving_view(server_args)
 
-    if not (
+    algorithm = "DFLASH"
+    if current_platform.is_out_of_tree():
+        is_supported = current_platform.supports_speculative_algorithm(algorithm)
+    else:
+        is_supported = (
             cfg.device.startswith("cuda") or cfg.device == "npu" or cfg.device == "xpu"
-    ):
+        )
+    if not is_supported:
         raise ValueError(
-            "DFLASH speculative decoding only supports CUDA, NPU and XPU devices."
+            f"{algorithm} speculative decoding is not supported by "
+            f"{type(current_platform).__name__} on device {cfg.device!r}."
         )
 
     # DFLASH + dp attention is validated on NPU only.
@@ -759,15 +766,55 @@ def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None:
     cfg = resolving_view(server_args)
 
     supported_draft_backends = DRAFT_ATTENTION_BACKEND_CHOICES
-    # FlashInfer is CUDA-only; fall back to triton on XPU and ROCm.
+
+    def is_supported_backend(backend: str) -> bool:
+        if current_platform.is_out_of_tree():
+            return current_platform.supports_speculative_draft_attention_backend(
+                "DFLASH", backend
+            )
+        return backend in supported_draft_backends
+
+    def get_fallback_backend() -> str:
+        if current_platform.is_out_of_tree():
+            try:
                 fallback_backend = (
-        "triton" if (get_platform().is_xpu or get_platform().is_hip) else "flashinfer"
+                    current_platform.get_default_speculative_draft_attention_backend(
+                        "DFLASH"
+                    )
+                )
+            except NotImplementedError as error:
+                raise ValueError(
+                    f"{type(current_platform).__name__} must implement "
+                    "get_default_speculative_draft_attention_backend() to use DFLASH."
+                ) from error
+            if not is_supported_backend(fallback_backend):
+                raise ValueError(
+                    f"{type(current_platform).__name__} returned unsupported DFLASH "
+                    f"draft attention backend {fallback_backend!r} from "
+                    "get_default_speculative_draft_attention_backend()."
+                )
+            return fallback_backend
+        # FlashInfer is CUDA-only; fall back to triton on XPU and ROCm.
+        return (
+            "triton"
+            if (get_platform().is_xpu or get_platform().is_hip)
+            else "flashinfer"
         )
 
     draft_backend = cfg.speculative_draft_attention_backend
     if draft_backend is None:
         draft_backend, _ = attention_backends_of(resolved_view(server_args))
     if draft_backend is None:
+        draft_backend = get_fallback_backend()
+    elif not is_supported_backend(draft_backend):
+        fallback_backend = get_fallback_backend()
+        logger.warning(
+            "DFLASH draft worker does not support attention_backend %r on %s. "
+            "Falling back to '%s'.",
+            draft_backend,
+            type(current_platform).__name__,
+            fallback_backend,
+        )
         draft_backend = fallback_backend
     elif draft_backend == "trtllm_mha":
         from sglang.srt.speculative.dflash_utils import get_dflash_layer_types
@@ -791,6 +838,7 @@ def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None:
         )
         all_causal = getattr(draft_text_config, "is_causal", False) is True
         if not (all_sliding or all_causal):
+            fallback_backend = get_fallback_backend()
             logger.warning(
                 "DFLASH only enables 'trtllm_mha' when all layers use sliding "
                 "attention or the draft is explicitly causal; got "
@@ -801,15 +849,6 @@ def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None:
                 fallback_backend,
             )
             draft_backend = fallback_backend
-    elif draft_backend not in supported_draft_backends:
-        logger.warning(
-            "DFLASH draft worker only supports attention_backend in %s for now, "
-            "but got %r. Falling back to '%s'.",
-            supported_draft_backends,
-            draft_backend,
-            fallback_backend,
-        )
-        draft_backend = fallback_backend
     # FIXME: avoid overriding server args directly; pass the resolved draft
     # backend to the draft worker explicitly instead.
     declare_resolution(

python/sglang/srt/platforms/interface.py

@@ -12,11 +12,13 @@ Out-of-tree platforms register via setuptools entry_points under the
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Optional, Type
+from typing import TYPE_CHECKING, Any, Callable, Optional, Type
 
 from sglang.srt.platforms.device_mixin import DeviceMixin, PlatformEnum
 
 if TYPE_CHECKING:
+    import torch
+
     from sglang.srt.layers.quantization.base_config import QuantizationConfig
 
 # Re-export for convenience
@@ -56,6 +58,10 @@ class SRTPlatform(DeviceMixin):
         """Return the default attention backend name for this platform."""
         raise NotImplementedError
 
+    def get_default_speculative_draft_attention_backend(self, algorithm: str) -> str:
+        """Return the default draft attention backend for an algorithm."""
+        raise NotImplementedError
+
     def get_graph_runner_cls(self) -> type:
         """Return the graph runner class for this platform."""
         raise NotImplementedError
@@ -91,6 +97,12 @@ class SRTPlatform(DeviceMixin):
         """Return the piecewise compilation backend class for this platform."""
         raise NotImplementedError
 
+    def get_speculative_cache_locs_fn(
+        self,
+    ) -> Optional[Callable[..., torch.Tensor]]:
+        """Return a platform implementation for speculative KV-cache locations."""
+        return None
+
     def get_quantization_config(
         self, quantization: str
     ) -> Optional[Type[QuantizationConfig]]:
@@ -107,6 +119,16 @@ class SRTPlatform(DeviceMixin):
         """Whether this platform supports FP8 quantization."""
         return False
 
+    def supports_speculative_algorithm(self, algorithm: str) -> bool:
+        """Whether this platform supports the named speculative algorithm."""
+        return False
+
+    def supports_speculative_draft_attention_backend(
+        self, algorithm: str, backend: str
+    ) -> bool:
+        """Whether this platform supports a draft backend for an algorithm."""
+        return False
+
     def support_cuda_graph(self) -> bool:
         """Whether this platform supports device graph capture and replay.
         Controls CUDA graph (CudaGraphRunner) for the decode path.

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -40,6 +40,7 @@ from sglang.srt.model_executor.runner_utils.pool import (
     disable_graph_pool_borrow,
     graph_pool_borrow_enabled,
 )
+from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import (
     get_exec,
     get_parallel,
@@ -619,6 +620,17 @@ class DFlashWorkerV2(BaseSpecWorker):
             capture_decode_cuda_graph = (
                 get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED
             )
+            if (
+                capture_decode_cuda_graph
+                and current_platform.is_out_of_tree()
+                and not current_platform.support_cuda_graph()
+            ):
+                capture_decode_cuda_graph = False
+                logger.warning(
+                    "Disable DFLASH draft cuda graph because %s does not support "
+                    "device graph capture.",
+                    type(current_platform).__name__,
+                )
             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

test/registered/unit/platforms/test_platform_interface.py

@@ -210,6 +210,15 @@ class TestSRTPlatform(CustomTestCase):
         self.assertFalse(base.is_pin_memory_available())
         self.assertFalse(base.is_pin_memory_available(device="cpu"))
 
+    def test_base_speculative_capability_defaults_are_conservative(self):
+        base = SRTPlatform()
+        self.assertFalse(base.supports_speculative_algorithm("DFLASH"))
+        self.assertFalse(
+            base.supports_speculative_draft_attention_backend(
+                "DFLASH", "custom_backend"
+            )
+        )
+
 
 class TestCudaDeviceMixin(CustomTestCase):
     """Tests for CUDA device operation defaults."""

test/registered/unit/spec/test_oot_dflash_hooks.py

@@ -0,0 +1,119 @@
+import unittest
+from unittest.mock import Mock, patch, sentinel
+
+from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func
+from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding
+from sglang.srt.platforms.interface import SRTPlatform
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+HOOK_MODULE = "sglang.srt.arg_groups.speculative_hook"
+
+
+def _platform(
+    default_backend: str = "custom_backend",
+    supported_backends: set[str] | None = None,
+    cache_result=None,
+) -> Mock:
+    supported_backends = supported_backends or {"custom_backend"}
+    platform = Mock(spec=SRTPlatform)
+    platform.is_out_of_tree.return_value = True
+    platform.supports_speculative_algorithm.return_value = True
+    platform.supports_speculative_draft_attention_backend.side_effect = (
+        lambda algorithm, backend: (
+            algorithm == "DFLASH" and backend in supported_backends
+        )
+    )
+    platform.get_default_speculative_draft_attention_backend.return_value = (
+        default_backend
+    )
+    platform.get_speculative_cache_locs_fn.return_value = (
+        None if cache_result is None else lambda **_: cache_result
+    )
+    return platform
+
+
+def _make_dflash_args(draft_backend: str | None) -> ServerArgs:
+    return ServerArgs(
+        model_path="dummy",
+        device="custom",
+        speculative_algorithm="DFLASH",
+        speculative_draft_model_path="draft",
+        speculative_num_draft_tokens=4,
+        speculative_draft_attention_backend=draft_backend,
+    )
+
+
+class TestOOTDFlashHooks(CustomTestCase):
+    def _resolve_backend(self, draft_backend: str | None, platform: Mock) -> str | None:
+        args = _make_dflash_args(draft_backend)
+        with (
+            patch(f"{HOOK_MODULE}.current_platform", platform),
+            patch(f"{HOOK_MODULE}.attention_backends_of", return_value=(None, None)),
+        ):
+            handle_speculative_decoding(args)
+        return args.speculative_draft_attention_backend
+
+    def test_explicit_backends_follow_platform_capabilities(self):
+        cases = (
+            ("custom_backend", {"custom_backend"}, "custom_backend"),
+            ("flashinfer", {"custom_backend"}, "custom_backend"),
+            ("flashinfer", {"flashinfer"}, "flashinfer"),
+            ("trtllm_mha", {"custom_backend"}, "custom_backend"),
+        )
+        for draft_backend, supported_backends, expected in cases:
+            with self.subTest(
+                draft_backend=draft_backend, supported_backends=supported_backends
+            ):
+                self.assertEqual(
+                    self._resolve_backend(
+                        draft_backend,
+                        _platform(supported_backends=supported_backends),
+                    ),
+                    expected,
+                )
+
+    def test_unknown_backend_warns_and_falls_back(self):
+        with self.assertLogs(
+            "sglang.srt.arg_groups.speculative_hook", "WARNING"
+        ) as logs:
+            resolved = self._resolve_backend("typo", _platform())
+
+        self.assertEqual(resolved, "custom_backend")
+        self.assertIn("attention_backend 'typo'", "\n".join(logs.output))
+
+    def test_invalid_platform_defaults_raise_actionable_errors(self):
+        missing_default = _platform()
+        missing_default.get_default_speculative_draft_attention_backend.side_effect = (
+            NotImplementedError
+        )
+        for platform, error in (
+            (_platform(default_backend="flashinfer"), "returned unsupported"),
+            (
+                missing_default,
+                "get_default_speculative_draft_attention_backend",
+            ),
+        ):
+            with (
+                self.subTest(error=error),
+                self.assertRaisesRegex(ValueError, error),
+            ):
+                self._resolve_backend(None, platform)
+
+    def test_cache_location_dispatch_tracks_platform_changes(self):
+        for expected in (sentinel.first, sentinel.second):
+            with patch(
+                "sglang.srt.platforms.current_platform",
+                _platform(cache_result=expected),
+            ):
+                result = assign_extend_cache_locs_func(
+                    None, None, None, None, 0, 0, None
+                )
+            self.assertIs(result, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()