From 99a44c88d4177f2e3b08bbccb49b33278bcc945a Mon Sep 17 00:00:00 2001 From: Yuhan Chen Date: Sat, 19 Sep 2026 23:54:46 -0700 Subject: [PATCH] Add out-of-tree DFlash extension points (#38740) Co-authored-by: Yuhan Chen Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- .../kernels/ops/speculative/cache_locs.py | 22 ++++ .../sglang/srt/arg_groups/speculative_hook.py | 73 ++++++++--- python/sglang/srt/platforms/interface.py | 24 +++- .../srt/speculative/dflash_worker_v2.py | 12 ++ .../unit/platforms/test_platform_interface.py | 9 ++ .../unit/spec/test_oot_dflash_hooks.py | 119 ++++++++++++++++++ 6 files changed, 241 insertions(+), 18 deletions(-) create mode 100644 test/registered/unit/spec/test_oot_dflash_hooks.py diff --git a/python/sglang/kernels/ops/speculative/cache_locs.py b/python/sglang/kernels/ops/speculative/cache_locs.py index c8aad1ab0..af887c303 100644 --- a/python/sglang/kernels/ops/speculative/cache_locs.py +++ b/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,), diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 390c1f94f..fd50cc8f0 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/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 ( - cfg.device.startswith("cuda") or cfg.device == "npu" or cfg.device == "xpu" - ): + 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. - fallback_backend = ( - "triton" if (get_platform().is_xpu or get_platform().is_hip) else "flashinfer" - ) + + 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 = ( + 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( diff --git a/python/sglang/srt/platforms/interface.py b/python/sglang/srt/platforms/interface.py index b7941002a..7ec272723 100644 --- a/python/sglang/srt/platforms/interface.py +++ b/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. diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index ec94e1ddb..339c8f908 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/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 diff --git a/test/registered/unit/platforms/test_platform_interface.py b/test/registered/unit/platforms/test_platform_interface.py index cd9548672..1e36b9b43 100644 --- a/test/registered/unit/platforms/test_platform_interface.py +++ b/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.""" diff --git a/test/registered/unit/spec/test_oot_dflash_hooks.py b/test/registered/unit/spec/test_oot_dflash_hooks.py new file mode 100644 index 000000000..af7f5eda2 --- /dev/null +++ b/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()