From 17705abae15c2226fb3dc808f67a72e9d8934c3b Mon Sep 17 00:00:00 2001 From: yctseng0211 Date: Wed, 5 Aug 2026 05:08:41 -0500 Subject: [PATCH 1/2] Gate DFLASH non-greedy verify on the target-only kernel being registered --- python/sglang/srt/speculative/dflash_utils.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index d7de9bf49196..32f3042120d8 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -54,7 +54,17 @@ top_p_renorm_probs_triton as top_p_renorm_prob, ) - _DFLASH_SAMPLING_VERIFY_AVAILABLE = True + try: + from sgl_kernel import tree_speculative_sampling_target_only + + # The ROCm sgl-kernel build omits csrc/speculative/speculative_sampling.cu + # while still shipping the Python wrapper, so importing the wrapper says + # nothing about callability; probe the registered op table instead. + _DFLASH_SAMPLING_VERIFY_AVAILABLE = hasattr( + torch.ops.sgl_kernel, "tree_speculative_sampling_target_only" + ) + except Exception: + tree_speculative_sampling_target_only = None else: top_k_renorm_prob = None top_p_renorm_prob = None From f58218f2e36d7fd78bf646bb74579df3a9c50017 Mon Sep 17 00:00:00 2001 From: HaiShaw Date: Thu, 6 Aug 2026 00:06:25 -0700 Subject: [PATCH 2/2] refactor --- python/sglang/srt/speculative/dflash_utils.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 32f3042120d8..7fff7556eead 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -53,18 +53,7 @@ from sglang.kernels.ops.sampling.renorm_triton import ( top_p_renorm_probs_triton as top_p_renorm_prob, ) - - try: - from sgl_kernel import tree_speculative_sampling_target_only - - # The ROCm sgl-kernel build omits csrc/speculative/speculative_sampling.cu - # while still shipping the Python wrapper, so importing the wrapper says - # nothing about callability; probe the registered op table instead. - _DFLASH_SAMPLING_VERIFY_AVAILABLE = hasattr( - torch.ops.sgl_kernel, "tree_speculative_sampling_target_only" - ) - except Exception: - tree_speculative_sampling_target_only = None + tree_speculative_sampling_target_only = None else: top_k_renorm_prob = None top_p_renorm_prob = None