From 5c415e3985026bc9b1daf120613e753b44adae55 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 20 Oct 2026 20:13:46 -0700 Subject: [PATCH 01/10] take post-verify kv move targets from the verify window; add test --- python/sglang/srt/speculative/spec_utils.py | 47 ++++------ .../spec/test_move_accept_tokens_targets.py | 89 +++++++++++++++++++ 2 files changed, 106 insertions(+), 30 deletions(-) create mode 100644 test/registered/unit/spec/test_move_accept_tokens_targets.py diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index bb3eed548611..f8e2cc3be59e 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -62,13 +62,11 @@ max_speculative_num_draft_tokens, ) from sglang.srt.utils import ( - is_cpu, is_cuda, is_hip, is_musa, is_npu, is_xpu, - next_power_of_2, ) from sglang.srt.utils.async_probe import maybe_detect_oob from sglang.srt.utils.common import fast_topk @@ -79,7 +77,6 @@ _is_npu = is_npu() _is_musa = is_musa() _is_xpu = is_xpu() -_is_cpu = is_cpu() if TYPE_CHECKING: from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject @@ -89,10 +86,6 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo -if _is_cpu: - from sgl_kernel import assign_extend_cache_locs_cpu - - logger = logging.getLogger(__name__) @@ -817,31 +810,25 @@ def move_accept_tokens_to_target_kvcache( "spec v2 move_accept_tokens accept_index", ) - tgt_cache_loc = torch.zeros( + # The targets lead each request's verify window. Take them from out_cache_loc, + # not req_to_token: after verify's read-done the scheduler may hand a finished + # request's row to a new request. + window = batch.out_cache_loc.size(0) // bs + cols = torch.arange(accept_index.shape[1], device=device) + front_index = torch.where( + cols < (num_correct_drafts + 1).unsqueeze(1), + torch.arange(bs, device=device).unsqueeze(1) * window + cols, + -1, + ).to(accept_index.dtype) + + tgt_cache_loc = torch.zeros(size, dtype=torch.int64, device=device) + accept_out_cache_loc = torch.zeros(size, dtype=torch.int64, device=device) + fill_accept_out_cache_loc_func( + front_index, + batch.out_cache_loc, + tgt_cache_loc, size, - dtype=torch.int64, - device=device, ) - accept_out_cache_loc = torch.zeros(size, dtype=torch.int64, device=device) - if _is_cpu: - assign_extend_cache_locs_cpu( - batch.req_pool_indices, - batch.req_to_token_pool.req_to_token, - batch.seq_lens, - batch.seq_lens + num_correct_drafts + 1, - tgt_cache_loc, - batch.req_to_token_pool.req_to_token.shape[1], - ) - else: - assign_extend_cache_locs[(bs,)]( - batch.req_pool_indices, - batch.req_to_token_pool.req_to_token, - batch.seq_lens, - batch.seq_lens + num_correct_drafts + 1, - tgt_cache_loc, - batch.req_to_token_pool.req_to_token.shape[1], - next_power_of_2(bs), - ) fill_accept_out_cache_loc_func( accept_index, batch.out_cache_loc, diff --git a/test/registered/unit/spec/test_move_accept_tokens_targets.py b/test/registered/unit/spec/test_move_accept_tokens_targets.py new file mode 100644 index 000000000000..dd7eda231136 --- /dev/null +++ b/test/registered/unit/spec/test_move_accept_tokens_targets.py @@ -0,0 +1,89 @@ +"""The post-verify KV move takes its targets from the verify window, not from +req_to_token: the scheduler may rewrite a finished request's row once verify's +read-done event fires, and a reread would move KV into the new owner's slots.""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.speculative.spec_utils import move_accept_tokens_to_target_kvcache +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") + + +class _RecordingKVCache: + def move_kv_cache(self, tgt_loc, src_loc): + self.tgt_loc, self.src_loc = tgt_loc.clone(), src_loc.clone() + + +@unittest.skipUnless(torch.cuda.is_available(), "the move kernels are Triton") +class TestMoveAcceptTokensTargets(CustomTestCase): + def _run(self, window, accept_rows): + device = "cuda" + rows = torch.tensor([5, 2, 7], device=device) + seq_lens = torch.tensor([10, 20, 30], device=device) + req_to_token = torch.arange(512, dtype=torch.int32, device=device).view(8, 64) + out_cache_loc = torch.cat( + [req_to_token[r, s : s + window] for r, s in zip(rows, seq_lens)] + ).to(torch.int64) + accept_index = torch.tensor(accept_rows, dtype=torch.int32, device=device) + num_accepted = (accept_index != -1).sum(dim=1) + + expected_tgt = torch.zeros(accept_index.numel(), dtype=torch.int64) + packed = torch.cat( + [ + req_to_token[r, s : s + n] + for r, s, n in zip(rows, seq_lens, num_accepted) + ] + ) + expected_tgt[: packed.numel()] = packed.cpu() + valid = accept_index[accept_index != -1].long() + expected_src = torch.zeros(accept_index.numel(), dtype=torch.int64) + expected_src[: valid.numel()] = out_cache_loc[valid].cpu() + + # The finished requests' rows go to new requests after read-done. + req_to_token[rows] = -7 + + kv_cache = _RecordingKVCache() + move_accept_tokens_to_target_kvcache( + SimpleNamespace( + seq_lens=seq_lens, + req_pool_indices=rows, + req_to_token_pool=SimpleNamespace(req_to_token=req_to_token), + out_cache_loc=out_cache_loc, + ), + accept_index, + num_accepted - 1, + SimpleNamespace(get_kvcache=lambda: kv_cache), + ) + self.assertTrue(torch.equal(kv_cache.tgt_loc.cpu(), expected_tgt)) + self.assertTrue(torch.equal(kv_cache.src_loc.cpu(), expected_src)) + + def test_chain_targets_ignore_rewritten_rows(self): + # NGRAM: accept_index spans the whole window. + self._run( + window=6, + accept_rows=[ + [0, 1, 2, -1, -1, -1], + [6, -1, -1, -1, -1, -1], + [12, 13, 14, 15, 16, 17], + ], + ) + + def test_tree_targets_ignore_rewritten_rows(self): + # topk > 1: the accepted path is shorter than the window and not a prefix of it. + self._run( + window=8, + accept_rows=[ + [0, 2, 5, -1], + [8, -1, -1, -1], + [16, 19, 20, 23], + ], + ) + + +if __name__ == "__main__": + unittest.main() From 450eb79adb348d56728cdfc5fe7d09b1da4809e4 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 20 Oct 2026 20:29:37 -0700 Subject: [PATCH 02/10] reuse assign_extend_cache_locs over the verify window --- python/sglang/srt/speculative/spec_utils.py | 53 ++++++++++++++------- 1 file changed, 36 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index f8e2cc3be59e..9fa3336047cc 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -62,11 +62,13 @@ max_speculative_num_draft_tokens, ) from sglang.srt.utils import ( + is_cpu, is_cuda, is_hip, is_musa, is_npu, is_xpu, + next_power_of_2, ) from sglang.srt.utils.async_probe import maybe_detect_oob from sglang.srt.utils.common import fast_topk @@ -77,6 +79,7 @@ _is_npu = is_npu() _is_musa = is_musa() _is_xpu = is_xpu() +_is_cpu = is_cpu() if TYPE_CHECKING: from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject @@ -86,6 +89,10 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo +if _is_cpu: + from sgl_kernel import assign_extend_cache_locs_cpu + + logger = logging.getLogger(__name__) @@ -810,25 +817,37 @@ def move_accept_tokens_to_target_kvcache( "spec v2 move_accept_tokens accept_index", ) - # The targets lead each request's verify window. Take them from out_cache_loc, - # not req_to_token: after verify's read-done the scheduler may hand a finished - # request's row to a new request. - window = batch.out_cache_loc.size(0) // bs - cols = torch.arange(accept_index.shape[1], device=device) - front_index = torch.where( - cols < (num_correct_drafts + 1).unsqueeze(1), - torch.arange(bs, device=device).unsqueeze(1) * window + cols, - -1, - ).to(accept_index.dtype) - - tgt_cache_loc = torch.zeros(size, dtype=torch.int64, device=device) - accept_out_cache_loc = torch.zeros(size, dtype=torch.int64, device=device) - fill_accept_out_cache_loc_func( - front_index, - batch.out_cache_loc, - tgt_cache_loc, + tgt_cache_loc = torch.zeros( size, + dtype=torch.int64, + device=device, ) + accept_out_cache_loc = torch.zeros(size, dtype=torch.int64, device=device) + # The targets lead each request's verify window. Read them from that window, + # not req_to_token: after verify's read-done the scheduler may hand a finished + # request's row to a new request. + window = batch.out_cache_loc.view(bs, -1) + window_rows = torch.arange(bs, device=device) + window_starts = torch.zeros_like(num_correct_drafts) + if _is_cpu: + assign_extend_cache_locs_cpu( + window_rows, + window.int(), + window_starts, + num_correct_drafts + 1, + tgt_cache_loc, + window.shape[1], + ) + else: + assign_extend_cache_locs[(bs,)]( + window_rows, + window, + window_starts, + num_correct_drafts + 1, + tgt_cache_loc, + window.shape[1], + next_power_of_2(bs), + ) fill_accept_out_cache_loc_func( accept_index, batch.out_cache_loc, From 0f6175e0c96899e67f4cfe4ee3a9e7d12be675b5 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 20 Oct 2026 20:51:23 -0700 Subject: [PATCH 03/10] add fill_front_out_cache_loc_func for the move targets --- .../kernels/ops/speculative/__init__.py | 1 + .../kernels/ops/speculative/cache_locs.py | 31 ++++++++++++++++ python/sglang/srt/speculative/spec_utils.py | 37 ++++--------------- 3 files changed, 40 insertions(+), 29 deletions(-) diff --git a/python/sglang/kernels/ops/speculative/__init__.py b/python/sglang/kernels/ops/speculative/__init__.py index c662f80b207f..f7b3c1079680 100644 --- a/python/sglang/kernels/ops/speculative/__init__.py +++ b/python/sglang/kernels/ops/speculative/__init__.py @@ -12,6 +12,7 @@ _TRITON_KERNELS = [ ("cache_locs", "assign_extend_cache_locs_func"), ("cache_locs", "assign_extend_cache_locs_uniform_func"), + ("cache_locs", "fill_front_out_cache_loc_func"), ("cache_locs", "generate_draft_decode_kv_indices"), ("cache_locs", "generate_draft_decode_window_kv_indices"), ("eagle", "fill_bonus_tokens"), diff --git a/python/sglang/kernels/ops/speculative/cache_locs.py b/python/sglang/kernels/ops/speculative/cache_locs.py index a93975e392ee..2059d7fd687a 100644 --- a/python/sglang/kernels/ops/speculative/cache_locs.py +++ b/python/sglang/kernels/ops/speculative/cache_locs.py @@ -800,3 +800,34 @@ def assign_extend_cache_locs_func( ) return out_cache_loc + + +def fill_front_out_cache_loc_func( + num_front: torch.Tensor, + out_cache_loc: torch.Tensor, + front_out_cache_loc: torch.Tensor, # mutable +): + """Pack the first ``num_front[i]`` slots of request i's window in ``out_cache_loc``.""" + bs = num_front.numel() + window = out_cache_loc.view(bs, -1) + rows = torch.arange(bs, device=window.device) + starts = torch.zeros_like(num_front) + if _is_cpu: + assign_extend_cache_locs_cpu( + rows, + window.int(), + starts, + num_front, + front_out_cache_loc, + window.shape[1], + ) + return + assign_extend_cache_locs[(bs,)]( + rows, + window, + starts, + num_front, + front_out_cache_loc, + window.shape[1], + next_power_of_2(bs), + ) diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 9fa3336047cc..0432a7fe0473 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -17,6 +17,9 @@ from sglang.kernels.ops.speculative.cache_locs import ( assign_extend_cache_locs as assign_extend_cache_locs, ) +from sglang.kernels.ops.speculative.cache_locs import ( + fill_front_out_cache_loc_func as fill_front_out_cache_loc_func, +) from sglang.kernels.ops.speculative.cache_locs import ( filter_finished_cache_loc_kernel as filter_finished_cache_loc_kernel, ) @@ -62,13 +65,11 @@ max_speculative_num_draft_tokens, ) from sglang.srt.utils import ( - is_cpu, is_cuda, is_hip, is_musa, is_npu, is_xpu, - next_power_of_2, ) from sglang.srt.utils.async_probe import maybe_detect_oob from sglang.srt.utils.common import fast_topk @@ -79,7 +80,6 @@ _is_npu = is_npu() _is_musa = is_musa() _is_xpu = is_xpu() -_is_cpu = is_cpu() if TYPE_CHECKING: from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject @@ -89,10 +89,6 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo -if _is_cpu: - from sgl_kernel import assign_extend_cache_locs_cpu - - logger = logging.getLogger(__name__) @@ -826,28 +822,11 @@ def move_accept_tokens_to_target_kvcache( # The targets lead each request's verify window. Read them from that window, # not req_to_token: after verify's read-done the scheduler may hand a finished # request's row to a new request. - window = batch.out_cache_loc.view(bs, -1) - window_rows = torch.arange(bs, device=device) - window_starts = torch.zeros_like(num_correct_drafts) - if _is_cpu: - assign_extend_cache_locs_cpu( - window_rows, - window.int(), - window_starts, - num_correct_drafts + 1, - tgt_cache_loc, - window.shape[1], - ) - else: - assign_extend_cache_locs[(bs,)]( - window_rows, - window, - window_starts, - num_correct_drafts + 1, - tgt_cache_loc, - window.shape[1], - next_power_of_2(bs), - ) + fill_front_out_cache_loc_func( + num_correct_drafts + 1, + batch.out_cache_loc, + tgt_cache_loc, + ) fill_accept_out_cache_loc_func( accept_index, batch.out_cache_loc, From 5160b6ebe97668a5ff96ca2c7031d1dcd50db0ff Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 20 Oct 2026 21:17:11 -0700 Subject: [PATCH 04/10] move draft cache-loc dispatch into a wrapper; drop dead forward; add kernel test --- .../kernels/ops/speculative/__init__.py | 1 + .../kernels/ops/speculative/cache_locs.py | 46 ++++++++++++++- .../srt/speculative/eagle_worker_common.py | 58 +++++++------------ python/sglang/srt/speculative/spec_utils.py | 3 - .../ops/speculative/test_next_verify_prep.py | 43 ++++++++------ 5 files changed, 90 insertions(+), 61 deletions(-) diff --git a/python/sglang/kernels/ops/speculative/__init__.py b/python/sglang/kernels/ops/speculative/__init__.py index f7b3c1079680..f92fd30ee9f0 100644 --- a/python/sglang/kernels/ops/speculative/__init__.py +++ b/python/sglang/kernels/ops/speculative/__init__.py @@ -10,6 +10,7 @@ # (module, public_fn) migrated from speculative/triton_ops. _TRITON_KERNELS = [ + ("cache_locs", "assign_draft_cache_locs_contiguous_func"), ("cache_locs", "assign_extend_cache_locs_func"), ("cache_locs", "assign_extend_cache_locs_uniform_func"), ("cache_locs", "fill_front_out_cache_loc_func"), diff --git a/python/sglang/kernels/ops/speculative/cache_locs.py b/python/sglang/kernels/ops/speculative/cache_locs.py index 2059d7fd687a..8edab660310c 100644 --- a/python/sglang/kernels/ops/speculative/cache_locs.py +++ b/python/sglang/kernels/ops/speculative/cache_locs.py @@ -25,7 +25,10 @@ _is_xpu = is_xpu() if _is_cpu: - from sgl_kernel import assign_extend_cache_locs_cpu + from sgl_kernel import ( + assign_draft_cache_locs_contiguous_cpu, + assign_extend_cache_locs_cpu, + ) def _get_oot_speculative_cache_locs_fn() -> Callable[..., torch.Tensor] | None: @@ -79,6 +82,47 @@ def assign_draft_cache_locs_contiguous( ) +def assign_draft_cache_locs_contiguous_func( + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + seq_lens: torch.Tensor, + out_cache_loc: torch.Tensor, # mutable + topk: int, + num_cols: int, + positions: Optional[torch.Tensor] = None, # mutable + mrope: Optional[torch.Tensor] = None, # mutable +): + """Copy each request's ``topk * num_cols`` slots from ``seq_lens`` on, and + write the draft positions when ``positions`` is given.""" + if _is_cpu: + assert positions is None and mrope is None, "CPU writes no positions" + assign_draft_cache_locs_contiguous_cpu( + req_pool_indices, + req_to_token, + seq_lens, + out_cache_loc, + req_to_token.shape[1], + topk, + num_cols, + ) + return + bs = seq_lens.shape[0] + assign_draft_cache_locs_contiguous[(bs,)]( + req_pool_indices, + req_to_token, + seq_lens, + out_cache_loc, + req_to_token.shape[1], + topk, + num_cols, + positions=positions, + mrope=mrope, + BS=bs, + WRITE_POSITIONS=positions is not None, + WRITE_MROPE=mrope is not None, + ) + + @triton.jit def _translate_token_ids(ids, v2p, mask, page_size: tl.constexpr): """Virtual token ids -> physical ones through the page-level v2p table, diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index 7507a695fb9b..768fc42d4e52 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/python/sglang/srt/speculative/eagle_worker_common.py @@ -5,7 +5,7 @@ import torch from sglang.kernels.ops.speculative.cache_locs import ( - assign_draft_cache_locs_contiguous, + assign_draft_cache_locs_contiguous_func, ) from sglang.kernels.ops.speculative.eagle import ( build_chain_tree, @@ -49,9 +49,6 @@ _is_cpu = is_cpu() _is_npu = is_npu() -if _is_cpu: - from sgl_kernel import assign_draft_cache_locs_contiguous_cpu - if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.tp_worker import TpModelWorker @@ -281,40 +278,25 @@ def prepare_for_draft( dtype=torch.int64, device=batch.device, ) - if _is_cpu: - assign_draft_cache_locs_contiguous_cpu( - batch.req_pool_indices, - req_to_token_pool.req_to_token, - batch.seq_lens, - batch.out_cache_loc, - req_to_token_pool.req_to_token.shape[1], - topk, - num_steps, - ) - else: - if batch.seq_lens.is_cuda and topk == 1: - prepared_positions = torch.empty_like(batch.seq_lens) - if draft_model_runner.model_config.model_is_mrope and all( - x is None for x in batch.multimodal_inputs - ): - prepared_mrope_positions = torch.empty( - (3, bs), dtype=torch.int64, device=batch.device - ) - # FIXME(lsyin): align with the default code path - assign_draft_cache_locs_contiguous[(bs,)]( - batch.req_pool_indices, - req_to_token_pool.req_to_token, - batch.seq_lens, - batch.out_cache_loc, - req_to_token_pool.req_to_token.shape[1], - topk, - window_cols, - positions=prepared_positions, - mrope=prepared_mrope_positions, - BS=bs, - WRITE_POSITIONS=prepared_positions is not None, - WRITE_MROPE=prepared_mrope_positions is not None, - ) + if batch.seq_lens.is_cuda and topk == 1: + prepared_positions = torch.empty_like(batch.seq_lens) + if draft_model_runner.model_config.model_is_mrope and all( + x is None for x in batch.multimodal_inputs + ): + prepared_mrope_positions = torch.empty( + (3, bs), dtype=torch.int64, device=batch.device + ) + # FIXME(lsyin): align with the default code path + assign_draft_cache_locs_contiguous_func( + batch.req_pool_indices, + req_to_token_pool.req_to_token, + batch.seq_lens, + batch.out_cache_loc, + topk, + window_cols, + positions=prepared_positions, + mrope=prepared_mrope_positions, + ) if plans_window: kv_loc_plan = target_translator.plan( req_pool_indices=batch.req_pool_indices, diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 0432a7fe0473..84088027990a 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -14,9 +14,6 @@ from sglang.kernels.ops.speculative.cache_locs import ( align_evict_mask_to_page_size as align_evict_mask_to_page_size, ) -from sglang.kernels.ops.speculative.cache_locs import ( - assign_extend_cache_locs as assign_extend_cache_locs, -) from sglang.kernels.ops.speculative.cache_locs import ( fill_front_out_cache_loc_func as fill_front_out_cache_loc_func, ) diff --git a/test/registered/kernels/ops/speculative/test_next_verify_prep.py b/test/registered/kernels/ops/speculative/test_next_verify_prep.py index d518d39f2c2e..c1b62e515c6d 100644 --- a/test/registered/kernels/ops/speculative/test_next_verify_prep.py +++ b/test/registered/kernels/ops/speculative/test_next_verify_prep.py @@ -250,7 +250,7 @@ def run(): @pytest.mark.parametrize("mrope", [False, True]) def test_draft_cache_positions_graph(bs, steps, dtype, mrope): from sglang.kernels.ops.speculative.cache_locs import ( - assign_draft_cache_locs_contiguous, + assign_draft_cache_locs_contiguous_func, ) table = torch.arange((bs + 3) * 64, dtype=torch.int32, device="cuda").reshape( @@ -263,19 +263,8 @@ def test_draft_cache_positions_graph(bs, steps, dtype, mrope): rope = torch.empty((3, bs), device="cuda", dtype=torch.int64) if mrope else None def run(): - assign_draft_cache_locs_contiguous[(bs,)]( - requests, - table, - seq, - cache, - 64, - 1, - steps, - positions=positions, - mrope=rope, - BS=bs, - WRITE_POSITIONS=True, - WRITE_MROPE=mrope, + assign_draft_cache_locs_contiguous_func( + requests, table, seq, cache, 1, steps, positions=positions, mrope=rope ) run() @@ -303,6 +292,26 @@ def run(): ) +@pytest.mark.parametrize("bs", [1, 3, 128]) +@pytest.mark.parametrize("width", [1, 4, 12]) +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_fill_front_out_cache_loc(bs, width, dtype): + from sglang.kernels.ops.speculative.cache_locs import ( + fill_front_out_cache_loc_func, + ) + + window = torch.randperm(bs * width * 4, device="cuda")[: bs * width] + num_front = torch.randint(0, width + 1, (bs,), device="cuda", dtype=dtype) + out = torch.full((bs * width,), -3, device="cuda", dtype=torch.int64) + fill_front_out_cache_loc_func(num_front, window, out) + + rows = window.view(bs, width) + expected = torch.cat([rows[i, :n] for i, n in enumerate(num_front.tolist())]) + torch.testing.assert_close(out[: expected.numel()], expected, rtol=0, atol=0) + # The tail past the packed slots stays as the caller filled it. + assert (out[expected.numel() :] == -3).all() + + @pytest.mark.parametrize("bs", [0, 1, 7, 128, 257]) @pytest.mark.parametrize("width", [1, 4, 7]) @pytest.mark.parametrize("interval", [0, 16, 512]) @@ -387,19 +396,15 @@ def test_dynamic_batch_reuses_kernel(operation, monkeypatch): locations = torch.empty((bs * 3,), device="cuda", dtype=torch.int64) positions = torch.empty_like(seq[:bs]) mrope = torch.empty((3, bs), device="cuda", dtype=torch.int64) - kernel[(bs,)]( + cache_locs.assign_draft_cache_locs_contiguous_func( requests[:bs], table, seq[:bs], locations, - 64, 1, 3, positions=positions, mrope=mrope, - BS=bs, - WRITE_POSITIONS=True, - WRITE_MROPE=True, ) torch.testing.assert_close(positions, seq[:bs], rtol=0, atol=0) torch.testing.assert_close( From 9206526d22be8894974e821f0cd4275a1d4c8747 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 20 Oct 2026 22:19:50 -0700 Subject: [PATCH 05/10] move dflash and uno triton kernels into kernels/ops --- .../sglang/kernels/ops/attention/__init__.py | 1 + .../ops/attention/table_qk_norm_rope.py | 83 +++++++++ .../kernels/ops/speculative/__init__.py | 2 + .../sglang/kernels/ops/speculative/dflash.py | 81 +++++++++ .../ops/speculative/dspark/dspark_accept.py | 4 +- .../ops}/speculative/uno_tree.py | 0 python/sglang/srt/models/dflash.py | 4 +- python/sglang/srt/speculative/dflash_utils.py | 161 ------------------ .../srt/speculative/dflash_worker_v2.py | 2 +- .../dspark_components/dspark_observability.py | 4 +- .../sglang/srt/speculative/uno_worker_v2.py | 2 +- 11 files changed, 178 insertions(+), 166 deletions(-) create mode 100644 python/sglang/kernels/ops/attention/table_qk_norm_rope.py rename python/sglang/{srt => kernels/ops}/speculative/uno_tree.py (100%) diff --git a/python/sglang/kernels/ops/attention/__init__.py b/python/sglang/kernels/ops/attention/__init__.py index d23cc247c56f..cab12d2b3e9e 100644 --- a/python/sglang/kernels/ops/attention/__init__.py +++ b/python/sglang/kernels/ops/attention/__init__.py @@ -47,6 +47,7 @@ ("pad", "unpad_draft_extend_output"), ("pad", "seqlens_expand_triton"), ("position", "compute_position_triton"), + ("table_qk_norm_rope", "table_qk_norm_rope_"), ("dsv4.attn_metadata_kernels", "expand_prefill_causally"), ("dsv4.attn_metadata_kernels", "build_page_table_positions"), ("dsv4.attn_metadata_kernels", "build_causal_swa_page_indices"), diff --git a/python/sglang/kernels/ops/attention/table_qk_norm_rope.py b/python/sglang/kernels/ops/attention/table_qk_norm_rope.py new file mode 100644 index 000000000000..6e59f53afb2f --- /dev/null +++ b/python/sglang/kernels/ops/attention/table_qk_norm_rope.py @@ -0,0 +1,83 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _table_qk_norm_rope_kernel( + qkv_ptr, + q_weight_ptr, + k_weight_ptr, + cos_sin_ptr, + pos_ptr, + row_stride, + q_size, + NHQ: tl.constexpr, + D: tl.constexpr, + EPS: tl.constexpr, +): + t = tl.program_id(0).to(tl.int64) + h = tl.program_id(1) + pos = tl.load(pos_ptr + t).to(tl.int64) + + HALF: tl.constexpr = D // 2 + half_ar = tl.arange(0, HALF) + d_ar = tl.arange(0, D) + cos = tl.load(cos_sin_ptr + pos * D + half_ar).to(tl.float32) + sin = tl.load(cos_sin_ptr + pos * D + HALF + half_ar).to(tl.float32) + + is_q = h < NHQ + col0 = tl.where(is_q, h * D, q_size + (h - NHQ) * D).to(tl.int64) + w_ptr = tl.where(is_q, q_weight_ptr.to(tl.int64), k_weight_ptr.to(tl.int64)).to( + tl.pointer_type(tl.bfloat16) + ) + + row = qkv_ptr + t * row_stride + col0 + x = tl.load(row + d_ar).to(tl.float32) + ms = tl.sum(x * x, 0) / D + inv = 1.0 / tl.sqrt(ms + EPS) + w1 = tl.load(w_ptr + half_ar).to(tl.float32) + w2 = tl.load(w_ptr + HALF + half_ar).to(tl.float32) + x1 = tl.load(row + half_ar).to(tl.float32) * inv * w1 + x2 = tl.load(row + HALF + half_ar).to(tl.float32) * inv * w2 + x1 = x1.to(tl.bfloat16).to(tl.float32) + x2 = x2.to(tl.bfloat16).to(tl.float32) + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + tl.store(row + half_ar, o1.to(tl.bfloat16)) + tl.store(row + HALF + half_ar, o2.to(tl.bfloat16)) + + +def table_qk_norm_rope_( + qkv: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + num_q_heads: int, + num_k_heads: int, + head_dim: int, + eps: float, +) -> None: + """In-place QK RMSNorm + table-lookup neox RoPE on the fused QKV tensor. + + Reads cos/sin from the SAME rotary table as the unfused path, so there is + no large-position angle drift (unlike theta-recompute kernels). V columns + are untouched. + """ + T = qkv.shape[0] + if T == 0: + return + grid = (T, num_q_heads + num_k_heads) + _table_qk_norm_rope_kernel[grid]( + qkv, + q_weight, + k_weight, + cos_sin_cache, + positions, + qkv.stride(0), + num_q_heads * head_dim, + NHQ=num_q_heads, + D=head_dim, + EPS=eps, + ) diff --git a/python/sglang/kernels/ops/speculative/__init__.py b/python/sglang/kernels/ops/speculative/__init__.py index f92fd30ee9f0..1017a2a43294 100644 --- a/python/sglang/kernels/ops/speculative/__init__.py +++ b/python/sglang/kernels/ops/speculative/__init__.py @@ -16,6 +16,7 @@ ("cache_locs", "fill_front_out_cache_loc_func"), ("cache_locs", "generate_draft_decode_kv_indices"), ("cache_locs", "generate_draft_decode_window_kv_indices"), + ("dflash", "compute_dflash_correct_drafts_and_bonus"), ("eagle", "fill_bonus_tokens"), ("eagle", "fill_accept_out_cache_loc"), ("gather_spec_extras", "gather_spec_extras"), @@ -23,6 +24,7 @@ ("spec_tree", "sgl_build_tree_kernel_efficient_triton"), ("spec_tree", "verify_tree_greedy_kernel_triton"), ("topk1", "draft_topk1_postprocess"), + ("uno_tree", "build_uno_tree_proposal"), ("ragged_verify_kernels", "pad_verify_lens_to_bucket"), ("ragged_verify_kernels", "build_qo_indptr"), ("lilicorr", "lilicorr_topk_lse"), diff --git a/python/sglang/kernels/ops/speculative/dflash.py b/python/sglang/kernels/ops/speculative/dflash.py index d5c7e33542f1..061e7889cb88 100644 --- a/python/sglang/kernels/ops/speculative/dflash.py +++ b/python/sglang/kernels/ops/speculative/dflash.py @@ -1,3 +1,5 @@ +from typing import Tuple + import torch import triton import triton.language as tl @@ -318,3 +320,82 @@ def selector_walk_triton( num_warps=1, ) return tokens, q_rows + + +@triton.jit +def _fused_correct_drafts_and_bonus_kernel( + candidates_ptr, + target_predict_ptr, + num_correct_drafts_ptr, + bonus_tokens_ptr, + block_size, + BLOCK: tl.constexpr, +): + b = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, BLOCK) + in_row = offs < block_size - 1 + drafts = tl.load(candidates_ptr + b * block_size + 1 + offs, mask=in_row, other=-1) + targets = tl.load(target_predict_ptr + b * block_size + offs, mask=in_row, other=-2) + eq = (drafts == targets) & in_row + # Leading-match count = index of the first mismatch lane; lanes past the + # row and all-match rows both resolve to block_size - 1 via the min. + num_correct = tl.min(tl.where(eq, BLOCK, offs), 0) + bonus_token = tl.load(target_predict_ptr + b * block_size + num_correct) + tl.store(num_correct_drafts_ptr + b, num_correct.to(tl.int32)) + tl.store(bonus_tokens_ptr + b, bonus_token.to(tl.int64)) + + +def compute_dflash_correct_drafts_and_bonus( + *, + candidates: torch.Tensor, + target_predict: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute DFlash accept lengths and bonus tokens (greedy verify rule). + + Args: + candidates: Token ids proposed by the DFlash draft, including the current token. + Shape: [bs, block_size]. candidates[:, 0] is the current token. + target_predict: Token ids predicted by the target model for each position in the block. + Shape: [bs, block_size]. target_predict[:, t] corresponds to argmax at position t. + + Returns: + correct_len: int32 tensor [bs], number of accepted *draft* tokens (excluding current token and bonus token). + bonus: int64 tensor [bs], the target-predicted token at index correct_len (the "bonus" token to append). + + Notes: + Matches the reference implementation rule: + accept while candidates[:, 1:] == target_predict[:, :-1] consecutively. + """ + if candidates.ndim != 2: + raise ValueError(f"candidates must be 2D, got shape={tuple(candidates.shape)}") + if target_predict.shape != candidates.shape: + raise ValueError( + "target_predict must have the same shape as candidates. " + f"candidates.shape={tuple(candidates.shape)}, target_predict.shape={tuple(target_predict.shape)}" + ) + + bs, block_size = candidates.shape + if bs <= 0: + raise ValueError(f"batch size must be positive, got {bs}.") + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}.") + + if candidates.is_cuda: + num_correct_drafts = torch.empty( + bs, dtype=torch.int32, device=candidates.device + ) + bonus_tokens = torch.empty(bs, dtype=torch.int64, device=candidates.device) + _fused_correct_drafts_and_bonus_kernel[(bs,)]( + candidates.contiguous(), + target_predict.contiguous(), + num_correct_drafts, + bonus_tokens, + block_size, + BLOCK=triton.next_power_of_2(max(block_size - 1, 1)), + ) + return num_correct_drafts, bonus_tokens + + matches = candidates[:, 1:] == target_predict[:, :-1] + correct_len = matches.to(torch.int32).cumprod(dim=1).sum(dim=1) + bonus = target_predict[torch.arange(bs, device=target_predict.device), correct_len] + return correct_len.to(torch.int32), bonus.to(torch.int64) diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py index 991caf2dd9ce..75c0704c9cf4 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py @@ -8,12 +8,14 @@ import triton import triton.language as tl +from sglang.kernels.ops.speculative.dflash import ( + compute_dflash_correct_drafts_and_bonus, +) from sglang.kernels.ops.speculative.dspark.dispatch import inputs_on_cuda from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 from sglang.srt.speculative.dflash_utils import ( _get_or_create_chain_verify_buffers, build_speculative_verify_target_probs, - compute_dflash_correct_drafts_and_bonus, ) from sglang.srt.utils import is_npu diff --git a/python/sglang/srt/speculative/uno_tree.py b/python/sglang/kernels/ops/speculative/uno_tree.py similarity index 100% rename from python/sglang/srt/speculative/uno_tree.py rename to python/sglang/kernels/ops/speculative/uno_tree.py diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index 03b3837779aa..cf7287bffbc3 100644 --- a/python/sglang/srt/models/dflash.py +++ b/python/sglang/srt/models/dflash.py @@ -369,7 +369,9 @@ def forward( if _is_npu: q, k, v = self.forward_prepare_npu(positions, hidden_states) elif self.use_table_qk_norm_rope and qkv.dtype == torch.bfloat16: - from sglang.srt.speculative.dflash_utils import table_qk_norm_rope_ + from sglang.kernels.ops.attention.table_qk_norm_rope import ( + table_qk_norm_rope_, + ) table_qk_norm_rope_( qkv, diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 4eab4f3df804..d288963d0d21 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -8,8 +8,6 @@ import torch import torch.nn.functional as F -import triton -import triton.language as tl from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.sampler import ( @@ -881,85 +879,6 @@ def can_dflash_use_fused_qkv_proj(qkv_proj: Any) -> Tuple[bool, str]: return True, "" -@triton.jit -def _fused_correct_drafts_and_bonus_kernel( - candidates_ptr, - target_predict_ptr, - num_correct_drafts_ptr, - bonus_tokens_ptr, - block_size, - BLOCK: tl.constexpr, -): - b = tl.program_id(0).to(tl.int64) - offs = tl.arange(0, BLOCK) - in_row = offs < block_size - 1 - drafts = tl.load(candidates_ptr + b * block_size + 1 + offs, mask=in_row, other=-1) - targets = tl.load(target_predict_ptr + b * block_size + offs, mask=in_row, other=-2) - eq = (drafts == targets) & in_row - # Leading-match count = index of the first mismatch lane; lanes past the - # row and all-match rows both resolve to block_size - 1 via the min. - num_correct = tl.min(tl.where(eq, BLOCK, offs), 0) - bonus_token = tl.load(target_predict_ptr + b * block_size + num_correct) - tl.store(num_correct_drafts_ptr + b, num_correct.to(tl.int32)) - tl.store(bonus_tokens_ptr + b, bonus_token.to(tl.int64)) - - -def compute_dflash_correct_drafts_and_bonus( - *, - candidates: torch.Tensor, - target_predict: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Compute DFlash accept lengths and bonus tokens (greedy verify rule). - - Args: - candidates: Token ids proposed by the DFlash draft, including the current token. - Shape: [bs, block_size]. candidates[:, 0] is the current token. - target_predict: Token ids predicted by the target model for each position in the block. - Shape: [bs, block_size]. target_predict[:, t] corresponds to argmax at position t. - - Returns: - correct_len: int32 tensor [bs], number of accepted *draft* tokens (excluding current token and bonus token). - bonus: int64 tensor [bs], the target-predicted token at index correct_len (the "bonus" token to append). - - Notes: - Matches the reference implementation rule: - accept while candidates[:, 1:] == target_predict[:, :-1] consecutively. - """ - if candidates.ndim != 2: - raise ValueError(f"candidates must be 2D, got shape={tuple(candidates.shape)}") - if target_predict.shape != candidates.shape: - raise ValueError( - "target_predict must have the same shape as candidates. " - f"candidates.shape={tuple(candidates.shape)}, target_predict.shape={tuple(target_predict.shape)}" - ) - - bs, block_size = candidates.shape - if bs <= 0: - raise ValueError(f"batch size must be positive, got {bs}.") - if block_size <= 0: - raise ValueError(f"block_size must be positive, got {block_size}.") - - if candidates.is_cuda: - num_correct_drafts = torch.empty( - bs, dtype=torch.int32, device=candidates.device - ) - bonus_tokens = torch.empty(bs, dtype=torch.int64, device=candidates.device) - _fused_correct_drafts_and_bonus_kernel[(bs,)]( - candidates.contiguous(), - target_predict.contiguous(), - num_correct_drafts, - bonus_tokens, - block_size, - BLOCK=triton.next_power_of_2(max(block_size - 1, 1)), - ) - return num_correct_drafts, bonus_tokens - - matches = candidates[:, 1:] == target_predict[:, :-1] - correct_len = matches.to(torch.int32).cumprod(dim=1).sum(dim=1) - bonus = target_predict[torch.arange(bs, device=target_predict.device), correct_len] - return correct_len.to(torch.int32), bonus.to(torch.int64) - - def apply_dflash_simulated_acceptance( *, candidates: torch.Tensor, @@ -1222,83 +1141,3 @@ def validate_dflash_request(req: Req, enable_overlap: bool) -> Optional[str]: return "DFLASH speculative decoding does not support return_hidden_states yet." return None - - -@triton.jit -def _table_qk_norm_rope_kernel( - qkv_ptr, - q_weight_ptr, - k_weight_ptr, - cos_sin_ptr, - pos_ptr, - row_stride, - q_size, - NHQ: tl.constexpr, - D: tl.constexpr, - EPS: tl.constexpr, -): - t = tl.program_id(0).to(tl.int64) - h = tl.program_id(1) - pos = tl.load(pos_ptr + t).to(tl.int64) - - HALF: tl.constexpr = D // 2 - half_ar = tl.arange(0, HALF) - d_ar = tl.arange(0, D) - cos = tl.load(cos_sin_ptr + pos * D + half_ar).to(tl.float32) - sin = tl.load(cos_sin_ptr + pos * D + HALF + half_ar).to(tl.float32) - - is_q = h < NHQ - col0 = tl.where(is_q, h * D, q_size + (h - NHQ) * D).to(tl.int64) - w_ptr = tl.where(is_q, q_weight_ptr.to(tl.int64), k_weight_ptr.to(tl.int64)).to( - tl.pointer_type(tl.bfloat16) - ) - - row = qkv_ptr + t * row_stride + col0 - x = tl.load(row + d_ar).to(tl.float32) - ms = tl.sum(x * x, 0) / D - inv = 1.0 / tl.sqrt(ms + EPS) - w1 = tl.load(w_ptr + half_ar).to(tl.float32) - w2 = tl.load(w_ptr + HALF + half_ar).to(tl.float32) - x1 = tl.load(row + half_ar).to(tl.float32) * inv * w1 - x2 = tl.load(row + HALF + half_ar).to(tl.float32) * inv * w2 - x1 = x1.to(tl.bfloat16).to(tl.float32) - x2 = x2.to(tl.bfloat16).to(tl.float32) - o1 = x1 * cos - x2 * sin - o2 = x2 * cos + x1 * sin - tl.store(row + half_ar, o1.to(tl.bfloat16)) - tl.store(row + HALF + half_ar, o2.to(tl.bfloat16)) - - -def table_qk_norm_rope_( - qkv: torch.Tensor, - positions: torch.Tensor, - q_weight: torch.Tensor, - k_weight: torch.Tensor, - cos_sin_cache: torch.Tensor, - num_q_heads: int, - num_k_heads: int, - head_dim: int, - eps: float, -) -> None: - """In-place QK RMSNorm + table-lookup neox RoPE on the fused QKV tensor. - - Reads cos/sin from the SAME rotary table as the unfused path, so there is - no large-position angle drift (unlike theta-recompute kernels). V columns - are untouched. - """ - T = qkv.shape[0] - if T == 0: - return - grid = (T, num_q_heads + num_k_heads) - _table_qk_norm_rope_kernel[grid]( - qkv, - q_weight, - k_weight, - cos_sin_cache, - positions, - qkv.stride(0), - num_q_heads * head_dim, - NHQ=num_q_heads, - D=head_dim, - EPS=eps, - ) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 547c7fd60343..36596e2026cb 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -13,6 +13,7 @@ from sglang.kernels.ops.speculative.dflash import ( _compute_dflash_accept_bonus_triton_unchecked, _prepare_dflash_draft_block_unchecked, + compute_dflash_correct_drafts_and_bonus, ) from sglang.kernels.ops.speculative.dspark.dspark_accept import ( accept_sampling, @@ -62,7 +63,6 @@ apply_dflash_simulated_acceptance, apply_dflash_verify_logits_adjustments, can_dflash_use_fused_qkv_proj, - compute_dflash_correct_drafts_and_bonus, compute_dflash_sampling_correct_drafts_and_bonus, is_dense_head_weight, is_dflash_sampling_verify_available, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_observability.py b/python/sglang/srt/speculative/dspark_components/dspark_observability.py index ce3395ee6181..6e30eae99f80 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_observability.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_observability.py @@ -12,11 +12,13 @@ import msgspec import torch +from sglang.kernels.ops.speculative.dflash import ( + compute_dflash_correct_drafts_and_bonus, +) from sglang.srt.environ import envs from sglang.srt.kv_canary.runner.future_tensor import FutureTensors from sglang.srt.runtime_context import get_parallel from sglang.srt.sampling.sampling_params import TOP_K_ALL -from sglang.srt.speculative.dflash_utils import compute_dflash_correct_drafts_and_bonus from sglang.srt.speculative.dspark_components.dspark_block_accept_estimator import ( create_block_accept_estimate_recorder, ) diff --git a/python/sglang/srt/speculative/uno_worker_v2.py b/python/sglang/srt/speculative/uno_worker_v2.py index b82b67740b04..85535ff08c1c 100644 --- a/python/sglang/srt/speculative/uno_worker_v2.py +++ b/python/sglang/srt/speculative/uno_worker_v2.py @@ -8,6 +8,7 @@ import torch +from sglang.kernels.ops.speculative.uno_tree import build_uno_tree_proposal from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -32,7 +33,6 @@ UnoDecodeCudaGraphRunner, ) from sglang.srt.speculative.uno_info import UnoDraftInput, UnoForwardInput -from sglang.srt.speculative.uno_tree import build_uno_tree_proposal from sglang.srt.speculative.uno_utils import ( build_uno_draft_input, pack_uno_tree_result, From fb46a739cdd2e0f4d4f6d0c1aab19f4d9ae123e8 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 20 Oct 2026 22:20:06 -0700 Subject: [PATCH 06/10] wrap reconstruct_indices_from_tree_mask dispatch --- .../kernels/ops/speculative/__init__.py | 1 + .../ops/speculative/reconstruct_tree.py | 44 ++++++++++++++++++- python/sglang/srt/speculative/ngram_worker.py | 16 +++---- 3 files changed, 49 insertions(+), 12 deletions(-) diff --git a/python/sglang/kernels/ops/speculative/__init__.py b/python/sglang/kernels/ops/speculative/__init__.py index 1017a2a43294..14c8286be653 100644 --- a/python/sglang/kernels/ops/speculative/__init__.py +++ b/python/sglang/kernels/ops/speculative/__init__.py @@ -32,6 +32,7 @@ ("reject_sampling", "chain_speculative_sampling_triton"), ("row_argmax", "row_argmax"), ("reconstruct_tree", "reconstruct_indices_from_tree_mask_triton"), + ("reconstruct_tree", "reconstruct_indices_from_tree_mask_func"), ] for _mod, _fn in _TRITON_KERNELS: register_kernel( diff --git a/python/sglang/kernels/ops/speculative/reconstruct_tree.py b/python/sglang/kernels/ops/speculative/reconstruct_tree.py index 3a9d650c7a91..56e5c3b168de 100644 --- a/python/sglang/kernels/ops/speculative/reconstruct_tree.py +++ b/python/sglang/kernels/ops/speculative/reconstruct_tree.py @@ -40,7 +40,14 @@ import triton import triton.language as tl -from sglang.srt.utils import next_power_of_2 +from sglang.srt.utils import is_cuda, is_hip, is_xpu, next_power_of_2 + +_is_xpu = is_xpu() + +if is_cuda() or is_hip(): + from sglang.kernels.ops.speculative.tree import reconstruct_indices_from_tree_mask +elif not _is_xpu: + from sgl_kernel.speculative import reconstruct_indices_from_tree_mask # Sentinel for "no ancestor / not-yet-found child": any value >= n. Reductions use # n for min-searches (first child/sibling) and -1 for max-searches (parent). @@ -151,3 +158,38 @@ def reconstruct_indices_from_tree_mask_triton( block_n, num_warps=_num_warps_for(block_n), ) + + +def reconstruct_indices_from_tree_mask_func( + tree_mask: torch.Tensor, + verified_seq_len: torch.Tensor, + positions: torch.Tensor, # mutable + retrieve_index: torch.Tensor, # mutable + retrieve_next_token: torch.Tensor, # mutable + retrieve_next_sibling: torch.Tensor, # mutable + batch_size: int, + draft_token_num: int, +) -> None: + """The Triton kernel on XPU, the compiled op everywhere else.""" + if _is_xpu: + reconstruct_indices_from_tree_mask_triton( + tree_mask, + verified_seq_len, + positions, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + batch_size, + draft_token_num, + ) + return + reconstruct_indices_from_tree_mask( + tree_mask, + verified_seq_len, + positions, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + batch_size, + draft_token_num, + ) diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index fa80e1955472..bd12d2ca5818 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -7,6 +7,9 @@ from sglang.kernels.ops.speculative.cache_locs import ( assign_extend_cache_locs_func as assign_extend_cache_locs_func, ) +from sglang.kernels.ops.speculative.reconstruct_tree import ( + reconstruct_indices_from_tree_mask_func, +) from sglang.srt.layers.logprob_processor import compute_spec_logprobs from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult @@ -31,20 +34,11 @@ prepare_mamba_track_for_verify, record_stream_for_v2_verify, ) -from sglang.srt.utils import is_cpu, is_cuda, is_hip, is_xpu +from sglang.srt.utils import is_cpu from sglang.srt.utils.async_probe import maybe_detect_inf, maybe_detect_nan _is_cpu = is_cpu() -if is_xpu(): - from sglang.kernels.ops.speculative.reconstruct_tree import ( - reconstruct_indices_from_tree_mask_triton as reconstruct_indices_from_tree_mask, - ) -elif is_cuda() or is_hip(): - from sglang.kernels.ops.speculative.tree import reconstruct_indices_from_tree_mask -else: - from sgl_kernel.speculative import reconstruct_indices_from_tree_mask - logger = logging.getLogger(__name__) @@ -345,7 +339,7 @@ def _prepare_for_speculative_decoding(self, batch: ScheduleBatch): self.grammar_tree_host = (mask, req_drafts) if batch.has_grammar else None # generate positions and some indices using tree_mask - reconstruct_indices_from_tree_mask( + reconstruct_indices_from_tree_mask_func( tree_mask, batch.seq_lens, positions, # mutable From 08ce18e09ea18bac2b64fd7856e69b4137efc0ac Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 20 Oct 2026 22:20:12 -0700 Subject: [PATCH 07/10] move tree build and greedy verify dispatch into spec_tree --- .../kernels/ops/speculative/__init__.py | 2 + .../kernels/ops/speculative/spec_tree.py | 262 ++++++++++++++++++ python/sglang/srt/speculative/eagle_utils.py | 247 +---------------- 3 files changed, 277 insertions(+), 234 deletions(-) diff --git a/python/sglang/kernels/ops/speculative/__init__.py b/python/sglang/kernels/ops/speculative/__init__.py index 14c8286be653..437c4e4c524f 100644 --- a/python/sglang/kernels/ops/speculative/__init__.py +++ b/python/sglang/kernels/ops/speculative/__init__.py @@ -23,6 +23,8 @@ ("multi_layer_eagle", "rotate_input_ids"), ("spec_tree", "sgl_build_tree_kernel_efficient_triton"), ("spec_tree", "verify_tree_greedy_kernel_triton"), + ("spec_tree", "build_tree_kernel_efficient_func"), + ("spec_tree", "verify_tree_greedy_func"), ("topk1", "draft_topk1_postprocess"), ("uno_tree", "build_uno_tree_proposal"), ("ragged_verify_kernels", "pad_verify_lens_to_bucket"), diff --git a/python/sglang/kernels/ops/speculative/spec_tree.py b/python/sglang/kernels/ops/speculative/spec_tree.py index e5c30308ad25..cf47689e553d 100644 --- a/python/sglang/kernels/ops/speculative/spec_tree.py +++ b/python/sglang/kernels/ops/speculative/spec_tree.py @@ -11,9 +11,41 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== +from enum import IntEnum + +import torch import triton import triton.language as tl +from sglang.srt.utils import is_cpu, is_cuda, is_hip, is_musa, is_npu, is_xpu + +_is_cuda = is_cuda() +_is_hip = is_hip() +_is_musa = is_musa() +_is_npu = is_npu() +_is_xpu = is_xpu() +_is_cpu = is_cpu() + +if _is_cuda or _is_hip: + from sglang.kernels.ops.speculative.tree import ( + build_tree_kernel_efficient as sgl_build_tree_kernel_efficient, + ) +elif _is_musa: + from sgl_kernel import ( + build_tree_kernel_efficient as sgl_build_tree_kernel_efficient, + ) +elif _is_cpu: + from sgl_kernel import ( + build_tree_kernel_efficient_cpu as sgl_build_tree_kernel_efficient_cpu, + ) + from sgl_kernel import verify_tree_greedy_cpu as sgl_verify_tree_greedy_cpu + + +class TreeMaskMode(IntEnum): + FULL_MASK = 0 + QLEN_ONLY = 1 + QLEN_ONLY_BITPACKING = 2 + @triton.jit def sgl_build_tree_kernel_efficient_triton( @@ -279,3 +311,233 @@ def verify_tree_greedy_kernel_triton( target_predict_ptr + target_row * num_draft_tokens + target_col ) tl.store(predicts_ptr + last_accept_retrieve_idx, final_target) + + +def sgl_build_tree_kernel_triton( + parent_list: torch.Tensor, + selected_index: torch.Tensor, + verified_seq_len: torch.Tensor, + tree_mask: torch.Tensor, + positions: torch.Tensor, + retrieve_index: torch.Tensor, + retrieve_next_token: torch.Tensor, + retrieve_next_sibling: torch.Tensor, + topk: int, + depth: int, + draft_token_num: int, + tree_mask_mode: TreeMaskMode = TreeMaskMode.FULL_MASK, +): + """Triton-based implementation.""" + # TODO: Add support for QLEN_ONLY_BITPACKING mode + if tree_mask_mode == TreeMaskMode.QLEN_ONLY_BITPACKING: + raise NotImplementedError( + "QLEN_ONLY_BITPACKING is not supported in Triton implementation" + ) + + batch_size = verified_seq_len.shape[0] + seq_len_prefix_sum = torch.cumsum(verified_seq_len, dim=0) - verified_seq_len + + # Launch kernel with one program per batch item + grid = (batch_size,) + + sgl_build_tree_kernel_efficient_triton[grid]( + parent_list, + selected_index, + verified_seq_len, + seq_len_prefix_sum, + tree_mask, + positions, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + topk=topk, + depth=depth, + draft_token_num=draft_token_num, + tree_mask_mode=int(tree_mask_mode), + batch_size=batch_size, + parent_list_stride=( + parent_list.stride(0) if parent_list.dim() > 1 else parent_list.shape[0] + ), + selected_index_stride=selected_index.stride(0), + ) + + +def build_tree_kernel_efficient_func( + parent_list: torch.Tensor, + top_scores_index: torch.Tensor, + seq_lens: torch.Tensor, + tree_mask: torch.Tensor, # mutable + positions: torch.Tensor, # mutable + retrieve_index: torch.Tensor, # mutable + retrieve_next_token: torch.Tensor, # mutable + retrieve_next_sibling: torch.Tensor, # mutable + topk: int, + spec_steps: int, + num_verify_tokens: int, + tree_mask_mode: TreeMaskMode, +) -> None: + """Build the verify tree with this platform's kernel.""" + if _is_npu: + torch.ops.npu.build_tree_kernel_efficient( + parent_list.to(dtype=torch.int64), + top_scores_index, + seq_lens, + tree_mask, + positions, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + topk, + spec_steps, + num_verify_tokens, + tree_mask_mode, + ) + elif _is_xpu: + sgl_build_tree_kernel_triton( + parent_list, + top_scores_index, + seq_lens, + tree_mask, + positions, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + topk, + spec_steps, + num_verify_tokens, + tree_mask_mode, + ) + elif _is_cpu: + sgl_build_tree_kernel_efficient_cpu( + parent_list, + top_scores_index, + seq_lens, + tree_mask, + positions, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + topk, + spec_steps, + num_verify_tokens, + tree_mask_mode, + ) + else: + sgl_build_tree_kernel_efficient( + parent_list, + top_scores_index, + seq_lens, + tree_mask, + positions, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + topk, + spec_steps, + num_verify_tokens, + tree_mask_mode, + ) + + +def verify_tree_greedy_triton( + predicts: torch.Tensor, + accept_index: torch.Tensor, + accept_token_num: torch.Tensor, + candidates: torch.Tensor, + retrieve_index: torch.Tensor, + retrieve_next_token: torch.Tensor, + retrieve_next_sibling: torch.Tensor, + target_predict: torch.Tensor, +): + """Triton-based implementation.""" + batch_size = candidates.shape[0] + num_speculative_tokens = accept_index.shape[1] + num_draft_tokens = candidates.shape[1] + + # Launch kernel with one program per batch item + grid = (batch_size,) + + verify_tree_greedy_kernel_triton[grid]( + predicts, + accept_index, + accept_token_num, + candidates, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + target_predict, + batch_size=batch_size, + num_speculative_tokens=num_speculative_tokens, + num_draft_tokens=num_draft_tokens, + ) + + +def verify_tree_greedy_func( + predicts: torch.Tensor, + accept_index: torch.Tensor, + accept_token_num: torch.Tensor, + candidates: torch.Tensor, + retrieve_index: torch.Tensor, + retrieve_next_token: torch.Tensor, + retrieve_next_sibling: torch.Tensor, + target_predict: torch.Tensor, + topk: int = -1, +): + if _is_cuda or _is_hip or _is_musa: + if _is_cuda or _is_hip: + from sglang.kernels.ops.speculative.tree import verify_tree_greedy + else: + from sgl_kernel import verify_tree_greedy + + verify_tree_greedy( + predicts=predicts, # mutable + accept_index=accept_index, # mutable + accept_token_num=accept_token_num, # mutable + candidates=candidates, + # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema. + retrive_index=retrieve_index, + retrive_next_token=retrieve_next_token, + retrive_next_sibling=retrieve_next_sibling, + target_predict=target_predict, + ) + + elif _is_cpu: + sgl_verify_tree_greedy_cpu( + predicts=predicts, # mutable + accept_index=accept_index, # mutable + accept_token_num=accept_token_num, # mutable + candidates=candidates, + # kwarg LHS retained as `retrive_*` to match the CUDA op schema, so + # the CPU/CUDA call sites stay grep-symmetric. + retrive_index=retrieve_index, + retrive_next_token=retrieve_next_token, + retrive_next_sibling=retrieve_next_sibling, + target_predict=target_predict, + ) + + elif _is_npu: + from sgl_kernel_npu.sample.verify_tree_greedy import verify_tree_greedy + + verify_tree_greedy( + predicts=predicts, + accept_index=accept_index, + accept_token_num=accept_token_num, + candidates=candidates, + # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema. + retrive_index=retrieve_index, + retrive_next_token=retrieve_next_token, + retrive_next_sibling=retrieve_next_sibling, + target_predict=target_predict, + ) + elif _is_xpu: + verify_tree_greedy_triton( + predicts=predicts, + accept_index=accept_index, + accept_token_num=accept_token_num, + candidates=candidates, + retrieve_index=retrieve_index, + retrieve_next_token=retrieve_next_token, + retrieve_next_sibling=retrieve_next_sibling, + target_predict=target_predict, + ) + return predicts, accept_index, accept_token_num diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 149b7e56c174..90ef29ba59e0 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -2,7 +2,6 @@ import logging import math -from enum import IntEnum from functools import partial from typing import TYPE_CHECKING, List, Optional, Tuple @@ -10,8 +9,9 @@ from sglang.kernels.ops.sampling import softmax as sampling_softmax from sglang.kernels.ops.speculative.spec_tree import ( - sgl_build_tree_kernel_efficient_triton, - verify_tree_greedy_kernel_triton, + TreeMaskMode, + build_tree_kernel_efficient_func, + verify_tree_greedy_func, ) from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( maybe_build_dsv4_verify_bundle, @@ -26,7 +26,6 @@ is_cpu, is_cuda, is_hip, - is_musa, is_npu, is_xpu, ) @@ -46,26 +45,11 @@ _is_cuda = is_cuda() _is_hip = is_hip() _is_npu = is_npu() -_is_musa = is_musa() _is_xpu = is_xpu() _is_cpu = is_cpu() logger = logging.getLogger(__name__) -if _is_cuda or _is_hip: - from sglang.kernels.ops.speculative.tree import ( - build_tree_kernel_efficient as sgl_build_tree_kernel_efficient, - ) -elif _is_musa: - from sgl_kernel import ( - build_tree_kernel_efficient as sgl_build_tree_kernel_efficient, - ) -elif _is_cpu: - from sgl_kernel import ( - build_tree_kernel_efficient_cpu as sgl_build_tree_kernel_efficient_cpu, - ) - from sgl_kernel import verify_tree_greedy_cpu as sgl_verify_tree_greedy_cpu - def per_step_draft_out_cache_loc( out_cache_loc: torch.Tensor, @@ -142,12 +126,6 @@ def organize_draft_results( return parent_list, top_scores_index, draft_tokens -class TreeMaskMode(IntEnum): - FULL_MASK = 0 - QLEN_ONLY = 1 - QLEN_ONLY_BITPACKING = 2 - - def default_tree_mask_mode() -> TreeMaskMode: # The CPU verify attention kernel (intel_amx) consumes the qlen x qlen # QLEN_ONLY tree mask directly; FULL_MASK is for the GPU kernels. @@ -229,229 +207,30 @@ def build_tree_kernel_efficient( # then, positions = [7, 8, 8, 9] positions = torch.empty((bs * num_verify_tokens,), device=device, dtype=torch.long) - if _is_npu: - torch.ops.npu.build_tree_kernel_efficient( - parent_list.to(dtype=torch.int64), - top_scores_index, - seq_lens, - tree_mask, - positions, - retrieve_index, - retrieve_next_token, - retrieve_next_sibling, - topk, - spec_steps, - num_verify_tokens, - tree_mask_mode, - ) - elif _is_xpu: - sgl_build_tree_kernel_triton( - parent_list, - top_scores_index, - seq_lens, - tree_mask, - positions, - retrieve_index, - retrieve_next_token, - retrieve_next_sibling, - topk, - spec_steps, - num_verify_tokens, - tree_mask_mode, - ) - elif _is_cpu: - sgl_build_tree_kernel_efficient_cpu( - parent_list, - top_scores_index, - seq_lens, - tree_mask, - positions, - retrieve_index, - retrieve_next_token, - retrieve_next_sibling, - topk, - spec_steps, - num_verify_tokens, - tree_mask_mode, - ) - else: - sgl_build_tree_kernel_efficient( - parent_list, - top_scores_index, - seq_lens, - tree_mask, - positions, - retrieve_index, - retrieve_next_token, - retrieve_next_sibling, - topk, - spec_steps, - num_verify_tokens, - tree_mask_mode, - ) - return ( + build_tree_kernel_efficient_func( + parent_list, + top_scores_index, + seq_lens, tree_mask, positions, retrieve_index, retrieve_next_token, retrieve_next_sibling, - draft_tokens, + topk, + spec_steps, + num_verify_tokens, + tree_mask_mode, ) - - -def sgl_build_tree_kernel_triton( - parent_list: torch.Tensor, - selected_index: torch.Tensor, - verified_seq_len: torch.Tensor, - tree_mask: torch.Tensor, - positions: torch.Tensor, - retrieve_index: torch.Tensor, - retrieve_next_token: torch.Tensor, - retrieve_next_sibling: torch.Tensor, - topk: int, - depth: int, - draft_token_num: int, - tree_mask_mode: TreeMaskMode = TreeMaskMode.FULL_MASK, -): - """Triton-based implementation.""" - # TODO: Add support for QLEN_ONLY_BITPACKING mode - if tree_mask_mode == TreeMaskMode.QLEN_ONLY_BITPACKING: - raise NotImplementedError( - "QLEN_ONLY_BITPACKING is not supported in Triton implementation" - ) - - batch_size = verified_seq_len.shape[0] - seq_len_prefix_sum = torch.cumsum(verified_seq_len, dim=0) - verified_seq_len - - # Launch kernel with one program per batch item - grid = (batch_size,) - - sgl_build_tree_kernel_efficient_triton[grid]( - parent_list, - selected_index, - verified_seq_len, - seq_len_prefix_sum, + return ( tree_mask, positions, retrieve_index, retrieve_next_token, retrieve_next_sibling, - topk=topk, - depth=depth, - draft_token_num=draft_token_num, - tree_mask_mode=int(tree_mask_mode), - batch_size=batch_size, - parent_list_stride=( - parent_list.stride(0) if parent_list.dim() > 1 else parent_list.shape[0] - ), - selected_index_stride=selected_index.stride(0), - ) - - -def verify_tree_greedy_triton( - predicts: torch.Tensor, - accept_index: torch.Tensor, - accept_token_num: torch.Tensor, - candidates: torch.Tensor, - retrieve_index: torch.Tensor, - retrieve_next_token: torch.Tensor, - retrieve_next_sibling: torch.Tensor, - target_predict: torch.Tensor, -): - """Triton-based implementation.""" - batch_size = candidates.shape[0] - num_speculative_tokens = accept_index.shape[1] - num_draft_tokens = candidates.shape[1] - - # Launch kernel with one program per batch item - grid = (batch_size,) - - verify_tree_greedy_kernel_triton[grid]( - predicts, - accept_index, - accept_token_num, - candidates, - retrieve_index, - retrieve_next_token, - retrieve_next_sibling, - target_predict, - batch_size=batch_size, - num_speculative_tokens=num_speculative_tokens, - num_draft_tokens=num_draft_tokens, + draft_tokens, ) -def verify_tree_greedy_func( - predicts: torch.Tensor, - accept_index: torch.Tensor, - accept_token_num: torch.Tensor, - candidates: torch.Tensor, - retrieve_index: torch.Tensor, - retrieve_next_token: torch.Tensor, - retrieve_next_sibling: torch.Tensor, - target_predict: torch.Tensor, - topk: int = -1, -): - if _is_cuda or _is_hip or _is_musa: - if _is_cuda or _is_hip: - from sglang.kernels.ops.speculative.tree import verify_tree_greedy - else: - from sgl_kernel import verify_tree_greedy - - verify_tree_greedy( - predicts=predicts, # mutable - accept_index=accept_index, # mutable - accept_token_num=accept_token_num, # mutable - candidates=candidates, - # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema. - retrive_index=retrieve_index, - retrive_next_token=retrieve_next_token, - retrive_next_sibling=retrieve_next_sibling, - target_predict=target_predict, - ) - - elif _is_cpu: - sgl_verify_tree_greedy_cpu( - predicts=predicts, # mutable - accept_index=accept_index, # mutable - accept_token_num=accept_token_num, # mutable - candidates=candidates, - # kwarg LHS retained as `retrive_*` to match the CUDA op schema, so - # the CPU/CUDA call sites stay grep-symmetric. - retrive_index=retrieve_index, - retrive_next_token=retrieve_next_token, - retrive_next_sibling=retrieve_next_sibling, - target_predict=target_predict, - ) - - elif _is_npu: - from sgl_kernel_npu.sample.verify_tree_greedy import verify_tree_greedy - - verify_tree_greedy( - predicts=predicts, - accept_index=accept_index, - accept_token_num=accept_token_num, - candidates=candidates, - # kwarg LHS retained as `retrive_*` to match sgl_kernel op schema. - retrive_index=retrieve_index, - retrive_next_token=retrieve_next_token, - retrive_next_sibling=retrieve_next_sibling, - target_predict=target_predict, - ) - elif _is_xpu: - verify_tree_greedy_triton( - predicts=predicts, - accept_index=accept_index, - accept_token_num=accept_token_num, - candidates=candidates, - retrieve_index=retrieve_index, - retrieve_next_token=retrieve_next_token, - retrieve_next_sibling=retrieve_next_sibling, - target_predict=target_predict, - ) - return predicts, accept_index, accept_token_num - - def get_draft_input_from_target_hidden_dim(model_runner: ModelRunner) -> int: """Width of the target hidden states fed into the draft model. From 736fb6fed94ad03f70ff4998349259dfc3032677 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 20 Oct 2026 22:20:32 -0700 Subject: [PATCH 08/10] resolve spec sampling ops in one per-platform table --- python/sglang/kernels/ops/sampling/renorm.py | 37 ++++++++++++++ .../ops/speculative/verify_sampling.py | 34 +++++++++++++ python/sglang/srt/speculative/dflash_utils.py | 49 ++++--------------- python/sglang/srt/speculative/eagle_utils.py | 42 ++-------------- .../runner_utils/test_graph_pool_borrow.py | 5 +- 5 files changed, 88 insertions(+), 79 deletions(-) create mode 100644 python/sglang/kernels/ops/sampling/renorm.py create mode 100644 python/sglang/kernels/ops/speculative/verify_sampling.py diff --git a/python/sglang/kernels/ops/sampling/renorm.py b/python/sglang/kernels/ops/sampling/renorm.py new file mode 100644 index 000000000000..d93921380316 --- /dev/null +++ b/python/sglang/kernels/ops/sampling/renorm.py @@ -0,0 +1,37 @@ +"""Top-k / top-p probability renormalization, resolved once per platform. + +A name is ``None`` where the platform has no kernel for it. +""" + +from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu +from sglang.utils import LazyImport + +if is_cuda(): + # Resolved on first call: CUDA builds without FlashInfer must still import + # this module. + top_k_renorm_prob = LazyImport("flashinfer.sampling", "top_k_renorm_probs") + top_p_renorm_prob = LazyImport("flashinfer.sampling", "top_p_renorm_probs") +elif is_hip(): + from sglang.kernels.ops.sampling.renorm_triton import ( + top_k_renorm_probs_triton as top_k_renorm_prob, + ) + from sglang.kernels.ops.sampling.renorm_triton import ( + top_p_renorm_probs_triton as top_p_renorm_prob, + ) +elif is_npu(): + try: + from sgl_kernel_npu.sample import top_k_renorm_prob, top_p_renorm_prob + except ImportError: + top_k_renorm_prob = None + top_p_renorm_prob = None +elif is_musa(): + try: + from sgl_kernel import top_k_renorm_prob, top_p_renorm_prob + except Exception: + top_k_renorm_prob = None + top_p_renorm_prob = None +else: + top_k_renorm_prob = None + top_p_renorm_prob = None + +__all__ = ["top_k_renorm_prob", "top_p_renorm_prob"] diff --git a/python/sglang/kernels/ops/speculative/verify_sampling.py b/python/sglang/kernels/ops/speculative/verify_sampling.py new file mode 100644 index 000000000000..655d21cd3ca6 --- /dev/null +++ b/python/sglang/kernels/ops/speculative/verify_sampling.py @@ -0,0 +1,34 @@ +"""Speculative verify sampling ops, resolved once per platform. + +A name is ``None`` where the platform has no kernel for it. +""" + +from sglang.srt.utils import is_cuda, is_musa, is_npu + +if is_npu(): + try: + from sgl_kernel_npu.sample import ( + chain_speculative_sampling_triton, + tree_speculative_sampling_target_only, + ) + except ImportError: + chain_speculative_sampling_triton = None + tree_speculative_sampling_target_only = None +else: + from sglang.kernels.ops.speculative.reject_sampling import ( + chain_speculative_sampling_triton, + ) + + if is_cuda(): + from sglang.kernels.ops.speculative.sampling import ( + tree_speculative_sampling_target_only, + ) + elif is_musa(): + try: + from sgl_kernel import tree_speculative_sampling_target_only + except Exception: + tree_speculative_sampling_target_only = None + else: + tree_speculative_sampling_target_only = None + +__all__ = ["chain_speculative_sampling_triton", "tree_speculative_sampling_target_only"] diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index d288963d0d21..85e7bcadb3b9 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -9,6 +9,10 @@ import torch import torch.nn.functional as F +from sglang.kernels.ops.sampling.renorm import top_k_renorm_prob, top_p_renorm_prob +from sglang.kernels.ops.speculative.verify_sampling import ( + tree_speculative_sampling_target_only, +) from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.sampler import ( apply_custom_logit_processor, @@ -18,14 +22,12 @@ from sglang.srt.model_executor.runner_utils.pool import borrow_graph_pool from sglang.srt.runtime_context import get_spec from sglang.srt.speculative.spec_utils import sample_simulated_acc_len -from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu -from sglang.utils import LazyImport +from sglang.srt.utils import is_npu DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>" logger = logging.getLogger(__name__) -_DFLASH_SAMPLING_VERIFY_AVAILABLE = False _DFLASH_CHAIN_VERIFY_BUFFERS: dict[tuple[Optional[int], int], dict[str, Any]] = {} _DFLASH_VERIFY_SKIP_CUSTOM_MASK_BACKENDS = frozenset( { @@ -40,43 +42,10 @@ ) -if is_cuda(): - from sglang.kernels.ops.speculative.sampling import ( - tree_speculative_sampling_target_only, - ) - - # Resolved on first call: CUDA builds without FlashInfer must still import - # this module. - top_k_renorm_prob = LazyImport("flashinfer.sampling", "top_k_renorm_probs") - top_p_renorm_prob = LazyImport("flashinfer.sampling", "top_p_renorm_probs") - - _DFLASH_SAMPLING_VERIFY_AVAILABLE = True -elif is_musa(): - try: - from sgl_kernel import ( - top_k_renorm_prob, - top_p_renorm_prob, - tree_speculative_sampling_target_only, - ) - - _DFLASH_SAMPLING_VERIFY_AVAILABLE = True - except Exception: - top_k_renorm_prob = None - top_p_renorm_prob = None - tree_speculative_sampling_target_only = None -elif is_hip(): - from sglang.kernels.ops.sampling.renorm_triton import ( - top_k_renorm_probs_triton as top_k_renorm_prob, - ) - from sglang.kernels.ops.sampling.renorm_triton import ( - top_p_renorm_probs_triton as top_p_renorm_prob, - ) - - tree_speculative_sampling_target_only = None -else: - top_k_renorm_prob = None - top_p_renorm_prob = None - tree_speculative_sampling_target_only = None +# DFlash has never run sampling verify on NPU; enabling it is its own change. +_DFLASH_SAMPLING_VERIFY_AVAILABLE = ( + tree_speculative_sampling_target_only is not None and not is_npu() +) def is_dflash_sampling_verify_available() -> bool: diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 90ef29ba59e0..1c1350d40ef4 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -8,11 +8,16 @@ import torch from sglang.kernels.ops.sampling import softmax as sampling_softmax +from sglang.kernels.ops.sampling.renorm import top_k_renorm_prob, top_p_renorm_prob from sglang.kernels.ops.speculative.spec_tree import ( TreeMaskMode, build_tree_kernel_efficient_func, verify_tree_greedy_func, ) +from sglang.kernels.ops.speculative.verify_sampling import ( + chain_speculative_sampling_triton, + tree_speculative_sampling_target_only, +) from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( maybe_build_dsv4_verify_bundle, ) @@ -26,7 +31,6 @@ is_cpu, is_cuda, is_hip, - is_npu, is_xpu, ) from sglang.srt.utils.async_probe import maybe_detect_oob @@ -44,7 +48,6 @@ _is_cuda = is_cuda() _is_hip = is_hip() -_is_npu = is_npu() _is_xpu = is_xpu() _is_cpu = is_cpu() @@ -709,48 +712,13 @@ def eagle_sample( tp_group.broadcast(accept_index, src=0) tp_group.broadcast(num_correct_drafts, src=0) else: - if _is_npu: - from sgl_kernel_npu.sample import ( - chain_speculative_sampling_triton, - top_k_renorm_prob, - top_p_renorm_prob, - tree_speculative_sampling_target_only, - ) - else: - from sglang.kernels.ops.speculative.reject_sampling import ( - chain_speculative_sampling_triton, - ) - - # if/else, not a ternary: the CUDA-only name still has to resolve in the - # branch not taken, and HIP only reaches here with rejection sampling on. if use_rejection_sampling: sampling_fn = chain_speculative_sampling_triton if get_spec().speculative_use_block_verification: sampling_fn = partial(sampling_fn, block_verification=True) else: - if _is_cuda: - from sglang.kernels.ops.speculative.sampling import ( - tree_speculative_sampling_target_only, - ) - elif not _is_npu: - from sgl_kernel import tree_speculative_sampling_target_only - sampling_fn = tree_speculative_sampling_target_only - if _is_hip: - # Same names, same contract: dflash_utils.py aliases these too. - from sglang.kernels.ops.sampling.renorm_triton import ( - top_k_renorm_probs_triton as top_k_renorm_prob, - ) - from sglang.kernels.ops.sampling.renorm_triton import ( - top_p_renorm_probs_triton as top_p_renorm_prob, - ) - elif _is_cuda: - from flashinfer.sampling import top_k_renorm_probs as top_k_renorm_prob - from flashinfer.sampling import top_p_renorm_probs as top_p_renorm_prob - elif not _is_npu: - from sgl_kernel import top_k_renorm_prob, top_p_renorm_prob - expanded_temperature = torch.repeat_interleave( sampling_info.temperatures, verify_input.draft_token_num, dim=0 ) # (bs * num_draft_tokens, 1) diff --git a/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py index 784c91b0b8f2..c215fd4d194c 100644 --- a/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py +++ b/test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py @@ -224,8 +224,9 @@ def fake_sampling(**kwargs): return_value=False, ), get_parallel().override(tp_group=tp_group), - patch( - "sglang.kernels.ops.speculative.sampling.tree_speculative_sampling_target_only", + patch.object( + eagle_utils, + "tree_speculative_sampling_target_only", side_effect=fake_sampling, ), ): From f594cd2a9354577e3cf6ca7747ba8a52b70d5c11 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 20 Oct 2026 23:17:04 -0700 Subject: [PATCH 09/10] remove dead spec kernels, helpers, and forwards; drop the unused decoupled-spec rid codec and inbox --- .../kernels/ops/speculative/cache_locs.py | 125 --------------- .../srt/speculative/decoupled_spec_io.py | 127 +-------------- .../dspark_components/dspark_config.py | 6 - .../sglang/srt/speculative/ragged_verify.py | 22 --- python/sglang/srt/speculative/spec_utils.py | 29 ---- .../unit/spec/test_decoupled_spec_io.py | 147 +----------------- 6 files changed, 4 insertions(+), 452 deletions(-) diff --git a/python/sglang/kernels/ops/speculative/cache_locs.py b/python/sglang/kernels/ops/speculative/cache_locs.py index 8edab660310c..147a58371f88 100644 --- a/python/sglang/kernels/ops/speculative/cache_locs.py +++ b/python/sglang/kernels/ops/speculative/cache_locs.py @@ -428,131 +428,6 @@ def generate_draft_decode_window_kv_indices( tl.store(window_kv_indptr + zid, base) -@triton.jit -def align_evict_mask_to_page_size( - seq_lens, - evict_mask, - page_size: tl.constexpr, - num_draft_tokens: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - t_range = tl.arange(0, BLOCK_SIZE) - - bid = tl.program_id(axis=0) - seq_len = tl.load(seq_lens + bid) - io_mask = t_range < num_draft_tokens - mask_row = tl.load( - evict_mask + bid * num_draft_tokens + t_range, mask=io_mask, other=0 - ) - - num_trues = tl.sum(mask_row) - num_false = num_draft_tokens - num_trues - - start = (seq_len + num_false - 1) // page_size * page_size - seq_len - for i in range(max(start, 0), min(start + page_size, num_draft_tokens)): - tl.store(evict_mask + bid * num_draft_tokens + i, False) - - -@torch.compile(dynamic=True, disable=_is_npu) -def get_src_tgt_cache_loc( - seq_lens: torch.Tensor, - out_cache_loc: torch.Tensor, - accept_index: torch.Tensor, - num_correct_drafts: torch.Tensor, - draft_token_num: int, - page_size: int, -): - src_cache_loc = out_cache_loc[accept_index] - # zeros_like, not empty_like: any uncovered tail stays at slot 0 (padding) - # instead of caching-allocator garbage. - tgt_cache_loc = torch.zeros_like(src_cache_loc) - extended_len = seq_lens + draft_token_num - keep_len = torch.minimum( - (seq_lens + num_correct_drafts + 1 + page_size - 1) // page_size * page_size, - extended_len, - ) - to_free_num_slots = extended_len - keep_len - return src_cache_loc, tgt_cache_loc, to_free_num_slots - - -@triton.jit -def get_target_cache_loc( - tgt_cache_loc, - to_free_slots, - num_correct_drafts, - to_free_num_slots, - out_cache_loc, - num_verify_tokens: tl.constexpr, - num_verify_tokens_upper: tl.constexpr, - bs_upper: tl.constexpr, -): - bid = tl.program_id(axis=0) - offset = tl.arange(0, num_verify_tokens_upper) - bs_offset = tl.arange(0, bs_upper) - - # write the first part to tgt_cache_loc - accept_len_all = tl.load(num_correct_drafts + bs_offset, mask=bs_offset < bid) - tgt_cache_loc_start = tl.sum(accept_len_all) + bid - copy_len = tl.load(num_correct_drafts + bid) + 1 - out_cache_loc_row = tl.load( - out_cache_loc + bid * num_verify_tokens + offset, mask=offset < copy_len - ) - tl.store( - tgt_cache_loc + tgt_cache_loc_start + offset, - out_cache_loc_row, - mask=offset < copy_len, - ) - - # write the second part to to_free_num_pages - to_free_num_slots_all = tl.load(to_free_num_slots + bs_offset, mask=bs_offset < bid) - to_free_num_slots_cur = tl.load(to_free_num_slots + bid) - out_cache_loc_start = num_verify_tokens - to_free_num_slots_cur - to_free_slots_start = tl.sum(to_free_num_slots_all) - - copy_len = to_free_num_slots_cur - out_cache_loc_row = tl.load( - out_cache_loc + bid * num_verify_tokens + out_cache_loc_start + offset, - mask=offset < copy_len, - ) - tl.store( - to_free_slots + to_free_slots_start + offset, - out_cache_loc_row, - mask=offset < copy_len, - ) - - -@triton.jit -def filter_finished_cache_loc_kernel( - out_cache_loc, - tgt_cache_loc, - num_correct_drafts, - num_accept_tokens_filter, - bs_upper: tl.constexpr, - num_verify_tokens_upper: tl.constexpr, -): - bid = tl.program_id(0) - bs_offset = tl.arange(0, bs_upper) - - num_correct_drafts_all = tl.load( - num_correct_drafts + bs_offset, mask=bs_offset < bid - ) - old_start = tl.sum(num_correct_drafts_all) + bid - - num_accept_tokens_filter_all = tl.load( - num_accept_tokens_filter + bs_offset, mask=bs_offset < bid - ) - new_start = tl.sum(num_accept_tokens_filter_all) - - copy_len = tl.load(num_accept_tokens_filter + bid) - copy_offset = tl.arange(0, num_verify_tokens_upper) - value = tl.load( - tgt_cache_loc + old_start + copy_offset, mask=copy_offset < copy_len - ) - tl.store( - out_cache_loc + new_start + copy_offset, value, mask=copy_offset < copy_len - ) - - @triton.jit def rebuild_compact_draft_req_to_token( draft_req_to_token, diff --git a/python/sglang/srt/speculative/decoupled_spec_io.py b/python/sglang/srt/speculative/decoupled_spec_io.py index 65744bbb5e62..ec22088518c3 100644 --- a/python/sglang/srt/speculative/decoupled_spec_io.py +++ b/python/sglang/srt/speculative/decoupled_spec_io.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, field from enum import Enum -from typing import Callable, Optional +from typing import Optional class DraftMeshMessageType(str, Enum): @@ -23,23 +23,6 @@ class DraftReqKey: request_id: str -def build_draft_scheduler_rid(draft_key: DraftReqKey) -> str: - return f"draft:{int(draft_key.src_verifier_rank)}:{draft_key.request_id}" - - -def parse_draft_scheduler_rid(rid: str) -> DraftReqKey: - if rid.startswith("draft:"): - encoded = rid[len("draft:") :] - rank_text, sep, request_id = encoded.partition(":") - if sep and request_id: - return DraftReqKey( - src_verifier_rank=int(rank_text), - request_id=request_id, - ) - - raise ValueError(f"Invalid decoupled draft scheduler rid: {rid}") - - @dataclass class DraftSync: """Open or re-open a drafter request from a verifier-owned prefix. @@ -246,114 +229,6 @@ def extract_prefix(self, num_tokens: int) -> VerifierCommitSegment: return prefix_segment -@dataclass -class DraftControlInbox: - """Drafter-side inbox for verifier control messages. - - The TokenSync thread temporarily stores incoming control messages here. - The drafter scheduler extracts and consumes them each time it finishes a decoding step. - """ - - sync_messages: list[DraftSync] = field(default_factory=list) - verifier_commit_segments: dict[DraftReqKey, VerifierCommitSegment] = field( - default_factory=dict - ) - close_keys: set[DraftReqKey] = field(default_factory=set) - - def is_empty(self) -> bool: - return ( - not self.sync_messages - and not self.verifier_commit_segments - and not self.close_keys - ) - - def pending_control_count(self) -> int: - return ( - len(self.sync_messages) - + len(self.verifier_commit_segments) - + len(self.close_keys) - ) - - def add_control_batch_locked(self, batch: DraftControlBatch) -> None: - for message in batch.close_messages: - self.add_close_key_locked(message.draft_key) - for message in batch.sync_messages: - if message.draft_key not in self.close_keys: - self.sync_messages.append(message) - for message in batch.verify_commit_messages: - self.add_verify_commit_locked(message) - - def add_close_key_locked(self, draft_key: DraftReqKey) -> None: - self.close_keys.add(draft_key) - self.verifier_commit_segments.pop(draft_key, None) - self.sync_messages = [ - message for message in self.sync_messages if message.draft_key != draft_key - ] - - def add_verify_commit_locked(self, message: VerifyCommit) -> None: - if message.draft_key in self.close_keys: - return - segment = self.verifier_commit_segments.get(message.draft_key) - if segment is None: - segment = VerifierCommitSegment( - draft_key=message.draft_key, - dst_drafter_rank=int(message.dst_drafter_rank), - pre_verify_committed_len=int(message.pre_verify_committed_len), - ) - segment.append_message(message) - self.verifier_commit_segments[message.draft_key] = segment - return - segment.append_message(message) - - def extract_ready_controls_locked( - self, - consumable_commit_len: Callable[[VerifierCommitSegment], int], - ) -> ReadyDraftControls: - ready_controls = ReadyDraftControls() - - if self.close_keys: - ready_controls.close_keys = self.close_keys - self.close_keys = set() - - if self.sync_messages: - ready_controls.sync_messages = self.sync_messages - self.sync_messages = [] - - for draft_key, segment in list(self.verifier_commit_segments.items()): - consumable_len = consumable_commit_len(segment) - if consumable_len <= 0: - continue - - ready_controls.ready_commit_segments.append( - segment.extract_prefix(consumable_len) - ) - if not segment.committed_tokens: - self.verifier_commit_segments.pop(draft_key, None) - - return ready_controls - - -@dataclass -class ReadyDraftControls: - sync_messages: list[DraftSync] = field(default_factory=list) - close_keys: set[DraftReqKey] = field(default_factory=set) - ready_commit_segments: list[VerifierCommitSegment] = field(default_factory=list) - - def is_empty(self) -> bool: - return ( - not self.sync_messages - and not self.close_keys - and not self.ready_commit_segments - ) - - def extracted_control_count(self) -> int: - return ( - len(self.sync_messages) - + len(self.close_keys) - + len(self.ready_commit_segments) - ) - - @dataclass class DraftMeshMessage: message_type: DraftMeshMessageType diff --git a/python/sglang/srt/speculative/dspark_components/dspark_config.py b/python/sglang/srt/speculative/dspark_components/dspark_config.py index 48143d3fc7c5..841b211d4ecf 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_config.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_config.py @@ -155,12 +155,6 @@ def read_draft_checkpoint_config(*, server_args: ServerArgs) -> DSparkDraftConfi return parse_dspark_draft_config(draft_hf_config=draft_hf_config) -def read_draft_checkpoint_gamma(*, server_args: ServerArgs) -> Optional[int]: - return read_draft_checkpoint_config(server_args=server_args).resolve_gamma( - default=None - ) - - def checkpoint_bundles_dspark_draft(hf_config: Any) -> bool: """The checkpoint carries a bundled DSpark draft head, marked by the prefixed dspark_* keys on the target hf config. Single source of truth diff --git a/python/sglang/srt/speculative/ragged_verify.py b/python/sglang/srt/speculative/ragged_verify.py index 92542e553040..f82a82aca554 100644 --- a/python/sglang/srt/speculative/ragged_verify.py +++ b/python/sglang/srt/speculative/ragged_verify.py @@ -302,25 +302,3 @@ def compute_uniform_extend_lengths( num_tokens=num_tokens, extend_start_loc=None, ) - - -def compute_ragged_extend_lengths( - *, - seq_lens: torch.Tensor, - seq_lens_cpu: List[int], - ragged_layout: RaggedVerifyLayout, -) -> VerifyExtendLengths: - extend_seq_lens_cpu = list(ragged_layout.verify_lens_cpu) - seq_lens_extended = seq_lens + ragged_layout.verify_lens - seq_lens_cpu_extended = [ - raw + length for raw, length in zip(seq_lens_cpu, extend_seq_lens_cpu) - ] - num_tokens = ragged_layout.total_verify_tokens - extend_start_loc = ragged_layout.extend_start_loc - return VerifyExtendLengths( - seq_lens_extended=seq_lens_extended, - seq_lens_cpu_extended=seq_lens_cpu_extended, - extend_seq_lens_cpu=extend_seq_lens_cpu, - num_tokens=num_tokens, - extend_start_loc=extend_start_loc, - ) diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 84088027990a..40e595591df7 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -11,27 +11,15 @@ from huggingface_hub import snapshot_download from sglang.kernels.ops.sampling import softmax as sampling_softmax -from sglang.kernels.ops.speculative.cache_locs import ( - align_evict_mask_to_page_size as align_evict_mask_to_page_size, -) from sglang.kernels.ops.speculative.cache_locs import ( fill_front_out_cache_loc_func as fill_front_out_cache_loc_func, ) -from sglang.kernels.ops.speculative.cache_locs import ( - filter_finished_cache_loc_kernel as filter_finished_cache_loc_kernel, -) from sglang.kernels.ops.speculative.cache_locs import ( generate_draft_decode_kv_indices as generate_draft_decode_kv_indices, ) from sglang.kernels.ops.speculative.cache_locs import ( generate_draft_decode_window_kv_indices as generate_draft_decode_window_kv_indices, ) -from sglang.kernels.ops.speculative.cache_locs import ( - get_src_tgt_cache_loc as get_src_tgt_cache_loc, -) -from sglang.kernels.ops.speculative.cache_locs import ( - get_target_cache_loc as get_target_cache_loc, -) from sglang.kernels.ops.speculative.eagle import ( fill_accept_out_cache_loc_func as fill_accept_out_cache_loc_func, ) @@ -48,9 +36,6 @@ from sglang.srt.layers.attention.linear.utils import pp_spec_stable_rows_enabled from sglang.srt.managers.schedule_batch import set_mamba_track_indices_from_reqs from sglang.srt.managers.utils import _async_d2h -from sglang.srt.mem_cache.allocation import ( - assign_req_to_token_pool as assign_req_to_token_pool, -) from sglang.srt.mem_cache.allocation import ( assign_req_to_token_pool_func as assign_req_to_token_pool_func, ) @@ -344,20 +329,6 @@ def spec_need_hidden_states() -> bool: return not spec.enable_multi_layer_eagle -@torch.compile(dynamic=True, disable=_is_npu or _is_xpu) -def create_num_accept_tokens_filter( - num_correct_drafts: torch.Tensor, - unfinished_index_device: torch.Tensor, - seq_lens: torch.Tensor, -): - num_accept_tokens_filter = torch.zeros_like(num_correct_drafts) - num_accept_tokens_filter[unfinished_index_device] = ( - num_correct_drafts[unfinished_index_device] + 1 - ) - seq_lens.add_(num_correct_drafts + 1) - return num_accept_tokens_filter - - def _select_top_k_tokens_first( topk_p: torch.Tensor, topk_index: torch.Tensor, diff --git a/test/registered/unit/spec/test_decoupled_spec_io.py b/test/registered/unit/spec/test_decoupled_spec_io.py index 388a76a4da97..561310f61230 100644 --- a/test/registered/unit/spec/test_decoupled_spec_io.py +++ b/test/registered/unit/spec/test_decoupled_spec_io.py @@ -1,25 +1,20 @@ """Unit tests for srt/speculative/decoupled_spec_io. decoupled_spec_io is the schema-only IPC layer for decoupled speculative -decoding: protocol message dataclasses, the cross-process request id codec, and -the drafter-side reconciliation helpers. These tests drive the real logic (id -round-trip + parse errors, commit validation, segment coalescing / contiguity / -prefix extraction, and inbox routing) on CPU; there is no GPU or transport here. +decoding. These tests drive its logic (commit validation, segment coalescing / +contiguity / prefix extraction, and the message envelope) on CPU; there is no +GPU or transport here. """ import unittest from sglang.srt.speculative.decoupled_spec_io import ( - DraftClose, DraftControlBatch, DraftMeshMessage, DraftMeshMessageType, DraftReqKey, - DraftSync, VerifierCommitSegment, VerifyCommit, - build_draft_scheduler_rid, - parse_draft_scheduler_rid, ) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -47,25 +42,6 @@ def _segment( ) -class TestDraftSchedulerRid(CustomTestCase): - def test_round_trip(self): - key = DraftReqKey(src_verifier_rank=3, request_id="req-1") - rid = build_draft_scheduler_rid(key) - self.assertEqual(rid, "draft:3:req-1") - self.assertEqual(parse_draft_scheduler_rid(rid), key) - - def test_request_id_containing_colon_round_trips(self): - # request_id may contain ':' — parse splits on the first ':' only. - key = DraftReqKey(src_verifier_rank=1, request_id="a:b:c") - self.assertEqual(parse_draft_scheduler_rid(build_draft_scheduler_rid(key)), key) - - def test_parse_invalid_rid_raises(self): - for bad in ["no-prefix", "draft:", "draft:0:", "draft:notint:r"]: - with self.subTest(rid=bad): - with self.assertRaises(ValueError): - parse_draft_scheduler_rid(bad) - - class TestVerifyCommitValidation(CustomTestCase): def test_empty_tokens_raises(self): with self.assertRaises(ValueError): @@ -131,123 +107,6 @@ def test_extract_prefix_bounds(self): seg.extract_prefix(3) # exceeds segment length -class TestDraftControlInbox(CustomTestCase): - def _inbox(self): - from sglang.srt.speculative.decoupled_spec_io import DraftControlInbox - - return DraftControlInbox() - - def _sync(self, rid, drafter_rank=0): - return DraftSync( - request_id=rid, src_verifier_rank=0, dst_drafter_rank=drafter_rank - ) - - def _close(self, rid, drafter_rank=0): - return DraftClose( - request_id=rid, - src_verifier_rank=0, - dst_drafter_rank=drafter_rank, - reason="x", - ) - - def test_add_control_batch_routes_each_message_type(self): - inbox = self._inbox() - inbox.add_control_batch_locked( - DraftControlBatch( - dst_drafter_rank=0, - sync_messages=[self._sync("s")], - verify_commit_messages=[_commit("c", pre=0, tokens=[1])], - close_messages=[self._close("x")], - ) - ) - self.assertEqual([m.request_id for m in inbox.sync_messages], ["s"]) - self.assertIn(DraftReqKey(0, "c"), inbox.verifier_commit_segments) - self.assertIn(DraftReqKey(0, "x"), inbox.close_keys) - - def test_close_drops_pending_segment_and_sync(self): - inbox = self._inbox() - inbox.add_control_batch_locked( - DraftControlBatch( - dst_drafter_rank=0, - sync_messages=[self._sync("r")], - verify_commit_messages=[_commit("r", pre=0, tokens=[1])], - ) - ) - inbox.add_close_key_locked(DraftReqKey(0, "r")) - self.assertEqual(inbox.sync_messages, []) - self.assertNotIn(DraftReqKey(0, "r"), inbox.verifier_commit_segments) - self.assertIn(DraftReqKey(0, "r"), inbox.close_keys) - - def test_verify_commit_for_closed_key_is_ignored(self): - inbox = self._inbox() - inbox.add_close_key_locked(DraftReqKey(0, "r")) - inbox.add_verify_commit_locked(_commit("r", pre=0, tokens=[1])) - self.assertNotIn(DraftReqKey(0, "r"), inbox.verifier_commit_segments) - - def test_extract_ready_controls_full_consume(self): - inbox = self._inbox() - inbox.add_control_batch_locked( - DraftControlBatch( - dst_drafter_rank=0, - sync_messages=[self._sync("s")], - verify_commit_messages=[_commit("c", pre=0, tokens=[1, 2])], - close_messages=[self._close("x")], - ) - ) - ready = inbox.extract_ready_controls_locked( - lambda seg: len(seg.committed_tokens) - ) - self.assertEqual([m.request_id for m in ready.sync_messages], ["s"]) - self.assertEqual({k.request_id for k in ready.close_keys}, {"x"}) - self.assertEqual(len(ready.ready_commit_segments), 1) - self.assertEqual(ready.ready_commit_segments[0].committed_tokens, [1, 2]) - # Fully consumed -> the segment is gone; inbox drained. - self.assertTrue(inbox.is_empty()) - - def test_extract_ready_controls_zero_consumable_keeps_segment(self): - inbox = self._inbox() - inbox.add_verify_commit_locked(_commit("c", pre=0, tokens=[1, 2])) - ready = inbox.extract_ready_controls_locked(lambda seg: 0) - self.assertEqual(ready.ready_commit_segments, []) - # Segment is left buffered for a later step. - self.assertIn(DraftReqKey(0, "c"), inbox.verifier_commit_segments) - - def test_extract_ready_controls_partial_consume_buffers_remainder(self): - inbox = self._inbox() - inbox.add_verify_commit_locked(_commit("c", pre=0, tokens=[1, 2, 3])) - ready = inbox.extract_ready_controls_locked(lambda seg: 1) - self.assertEqual(ready.ready_commit_segments[0].committed_tokens, [1]) - # Remainder [2, 3] stays buffered with pre advanced to 1. - seg = inbox.verifier_commit_segments[DraftReqKey(0, "c")] - self.assertEqual(seg.committed_tokens, [2, 3]) - self.assertEqual(seg.pre_verify_committed_len, 1) - - def test_close_in_same_batch_drops_sync_and_commit(self): - # add_control_batch applies close first, so a same-key sync/commit in the - # same batch is dropped/ignored: close wins within one batch. - inbox = self._inbox() - inbox.add_control_batch_locked( - DraftControlBatch( - dst_drafter_rank=0, - sync_messages=[self._sync("r")], - verify_commit_messages=[_commit("r", pre=0, tokens=[1])], - close_messages=[self._close("r")], - ) - ) - self.assertEqual(inbox.sync_messages, []) - self.assertNotIn(DraftReqKey(0, "r"), inbox.verifier_commit_segments) - self.assertEqual({k.request_id for k in inbox.close_keys}, {"r"}) - - def test_two_commits_same_key_coalesce_in_inbox(self): - # A second commit for an existing key appends to the buffered segment. - inbox = self._inbox() - inbox.add_verify_commit_locked(_commit("r", pre=0, tokens=[1, 2])) - inbox.add_verify_commit_locked(_commit("r", pre=2, tokens=[3])) - seg = inbox.verifier_commit_segments[DraftReqKey(0, "r")] - self.assertEqual(seg.committed_tokens, [1, 2, 3]) - self.assertEqual(seg.end_committed_len, 3) - - class TestDraftMeshMessageEnvelope(CustomTestCase): def test_from_control_batch_sets_discriminant_and_slot(self): batch = DraftControlBatch(dst_drafter_rank=0) From dfb93aee7e1f353a518c45b153959e52dd371fba Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Tue, 20 Oct 2026 23:38:28 -0700 Subject: [PATCH 10/10] trim comments; trim spec tests --- .../ops/speculative/reconstruct_tree.py | 1 - .../kernels/ops/speculative/spec_tree.py | 1 - python/sglang/srt/speculative/dflash_utils.py | 2 +- .../srt/speculative/eagle_worker_common.py | 2 +- .../ops/speculative/test_next_verify_prep.py | 29 +--- .../unit/spec/test_decoupled_spec_io.py | 131 ------------------ 6 files changed, 6 insertions(+), 160 deletions(-) delete mode 100644 test/registered/unit/spec/test_decoupled_spec_io.py diff --git a/python/sglang/kernels/ops/speculative/reconstruct_tree.py b/python/sglang/kernels/ops/speculative/reconstruct_tree.py index 56e5c3b168de..eb84644a1fa1 100644 --- a/python/sglang/kernels/ops/speculative/reconstruct_tree.py +++ b/python/sglang/kernels/ops/speculative/reconstruct_tree.py @@ -170,7 +170,6 @@ def reconstruct_indices_from_tree_mask_func( batch_size: int, draft_token_num: int, ) -> None: - """The Triton kernel on XPU, the compiled op everywhere else.""" if _is_xpu: reconstruct_indices_from_tree_mask_triton( tree_mask, diff --git a/python/sglang/kernels/ops/speculative/spec_tree.py b/python/sglang/kernels/ops/speculative/spec_tree.py index cf47689e553d..138145d44d6e 100644 --- a/python/sglang/kernels/ops/speculative/spec_tree.py +++ b/python/sglang/kernels/ops/speculative/spec_tree.py @@ -376,7 +376,6 @@ def build_tree_kernel_efficient_func( num_verify_tokens: int, tree_mask_mode: TreeMaskMode, ) -> None: - """Build the verify tree with this platform's kernel.""" if _is_npu: torch.ops.npu.build_tree_kernel_efficient( parent_list.to(dtype=torch.int64), diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 85e7bcadb3b9..6deb4facf50d 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -42,7 +42,7 @@ ) -# DFlash has never run sampling verify on NPU; enabling it is its own change. +# DFlash does not run sampling verify on NPU, even where the op exists. _DFLASH_SAMPLING_VERIFY_AVAILABLE = ( tree_speculative_sampling_target_only is not None and not is_npu() ) diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index 768fc42d4e52..df4ae734bc92 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/python/sglang/srt/speculative/eagle_worker_common.py @@ -286,7 +286,7 @@ def prepare_for_draft( prepared_mrope_positions = torch.empty( (3, bs), dtype=torch.int64, device=batch.device ) - # FIXME(lsyin): align with the default code path + # TODO(lsyin): align with the default code path assign_draft_cache_locs_contiguous_func( batch.req_pool_indices, req_to_token_pool.req_to_token, diff --git a/test/registered/kernels/ops/speculative/test_next_verify_prep.py b/test/registered/kernels/ops/speculative/test_next_verify_prep.py index c1b62e515c6d..1e4cae58a2b2 100644 --- a/test/registered/kernels/ops/speculative/test_next_verify_prep.py +++ b/test/registered/kernels/ops/speculative/test_next_verify_prep.py @@ -244,11 +244,10 @@ def run(): assert out[7] is None -@pytest.mark.parametrize("bs", [1, 3, 128]) -@pytest.mark.parametrize("steps", [1, 3, 7]) -@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("bs", [1, 128]) +@pytest.mark.parametrize("steps", [1, 7]) @pytest.mark.parametrize("mrope", [False, True]) -def test_draft_cache_positions_graph(bs, steps, dtype, mrope): +def test_draft_cache_positions_graph(bs, steps, mrope): from sglang.kernels.ops.speculative.cache_locs import ( assign_draft_cache_locs_contiguous_func, ) @@ -257,7 +256,7 @@ def test_draft_cache_positions_graph(bs, steps, dtype, mrope): bs + 3, 64 ) requests = torch.arange(bs, device="cuda") - seq = torch.arange(bs, device="cuda", dtype=dtype) % 48 + seq = torch.arange(bs, device="cuda", dtype=torch.int64) % 48 cache = torch.empty((bs * steps,), device="cuda", dtype=torch.int64) positions = torch.empty_like(seq) rope = torch.empty((3, bs), device="cuda", dtype=torch.int64) if mrope else None @@ -292,26 +291,6 @@ def run(): ) -@pytest.mark.parametrize("bs", [1, 3, 128]) -@pytest.mark.parametrize("width", [1, 4, 12]) -@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) -def test_fill_front_out_cache_loc(bs, width, dtype): - from sglang.kernels.ops.speculative.cache_locs import ( - fill_front_out_cache_loc_func, - ) - - window = torch.randperm(bs * width * 4, device="cuda")[: bs * width] - num_front = torch.randint(0, width + 1, (bs,), device="cuda", dtype=dtype) - out = torch.full((bs * width,), -3, device="cuda", dtype=torch.int64) - fill_front_out_cache_loc_func(num_front, window, out) - - rows = window.view(bs, width) - expected = torch.cat([rows[i, :n] for i, n in enumerate(num_front.tolist())]) - torch.testing.assert_close(out[: expected.numel()], expected, rtol=0, atol=0) - # The tail past the packed slots stays as the caller filled it. - assert (out[expected.numel() :] == -3).all() - - @pytest.mark.parametrize("bs", [0, 1, 7, 128, 257]) @pytest.mark.parametrize("width", [1, 4, 7]) @pytest.mark.parametrize("interval", [0, 16, 512]) diff --git a/test/registered/unit/spec/test_decoupled_spec_io.py b/test/registered/unit/spec/test_decoupled_spec_io.py deleted file mode 100644 index 561310f61230..000000000000 --- a/test/registered/unit/spec/test_decoupled_spec_io.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Unit tests for srt/speculative/decoupled_spec_io. - -decoupled_spec_io is the schema-only IPC layer for decoupled speculative -decoding. These tests drive its logic (commit validation, segment coalescing / -contiguity / prefix extraction, and the message envelope) on CPU; there is no -GPU or transport here. -""" - -import unittest - -from sglang.srt.speculative.decoupled_spec_io import ( - DraftControlBatch, - DraftMeshMessage, - DraftMeshMessageType, - DraftReqKey, - VerifierCommitSegment, - VerifyCommit, -) -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=10, suite="base-a-test-cpu") - - -def _commit(rid, *, pre, tokens, src_verifier_rank=0, drafter_rank=0) -> VerifyCommit: - return VerifyCommit( - request_id=rid, - src_verifier_rank=src_verifier_rank, - dst_drafter_rank=drafter_rank, - pre_verify_committed_len=pre, - committed_tokens=list(tokens), - ) - - -def _segment( - rid, *, pre=0, drafter_rank=0, src_verifier_rank=0 -) -> VerifierCommitSegment: - return VerifierCommitSegment( - draft_key=DraftReqKey(src_verifier_rank=src_verifier_rank, request_id=rid), - dst_drafter_rank=drafter_rank, - pre_verify_committed_len=pre, - ) - - -class TestVerifyCommitValidation(CustomTestCase): - def test_empty_tokens_raises(self): - with self.assertRaises(ValueError): - _commit("r", pre=0, tokens=[]).validate_committed_tokens() - - def test_negative_pre_len_raises(self): - with self.assertRaises(ValueError): - _commit("r", pre=-1, tokens=[1]).validate_committed_tokens() - - def test_valid_commit_passes(self): - # Should not raise. - _commit("r", pre=0, tokens=[1, 2]).validate_committed_tokens() - - -class TestVerifierCommitSegment(CustomTestCase): - def test_append_coalesces_contiguous_commits(self): - seg = _segment("r", pre=0) - seg.append_message(_commit("r", pre=0, tokens=[10, 11])) - self.assertEqual(seg.committed_tokens, [10, 11]) - self.assertEqual(seg.end_committed_len, 2) - seg.append_message(_commit("r", pre=2, tokens=[12])) - self.assertEqual(seg.committed_tokens, [10, 11, 12]) - self.assertEqual(seg.end_committed_len, 3) - - def test_append_wrong_request_raises(self): - seg = _segment("r", pre=0) - with self.assertRaises(RuntimeError): - seg.append_message(_commit("other", pre=0, tokens=[1])) - - def test_append_wrong_drafter_rank_raises(self): - seg = _segment("r", pre=0, drafter_rank=0) - with self.assertRaises(RuntimeError): - seg.append_message(_commit("r", pre=0, tokens=[1], drafter_rank=7)) - - def test_append_non_contiguous_raises(self): - seg = _segment("r", pre=0) - seg.append_message(_commit("r", pre=0, tokens=[10])) # end -> 1 - with self.assertRaises(RuntimeError): - seg.append_message(_commit("r", pre=5, tokens=[11])) # gap - - def test_append_runs_message_validation(self): - seg = _segment("r", pre=0) - with self.assertRaises(ValueError): - seg.append_message(_commit("r", pre=0, tokens=[])) # empty -> validate - - def test_extract_prefix_splits_segment(self): - seg = _segment("r", pre=0) - seg.append_message(_commit("r", pre=0, tokens=[10, 11, 12, 13])) - prefix = seg.extract_prefix(2) - self.assertEqual(prefix.committed_tokens, [10, 11]) - self.assertEqual(prefix.pre_verify_committed_len, 0) - # Remainder stays in the original segment, with pre advanced by 2. - self.assertEqual(seg.committed_tokens, [12, 13]) - self.assertEqual(seg.pre_verify_committed_len, 2) - self.assertEqual(seg.end_committed_len, 4) - - def test_extract_prefix_bounds(self): - seg = _segment("r", pre=0) - seg.append_message(_commit("r", pre=0, tokens=[10, 11])) - with self.assertRaises(ValueError): - seg.extract_prefix(0) - with self.assertRaises(ValueError): - seg.extract_prefix(3) # exceeds segment length - - -class TestDraftMeshMessageEnvelope(CustomTestCase): - def test_from_control_batch_sets_discriminant_and_slot(self): - batch = DraftControlBatch(dst_drafter_rank=0) - msg = DraftMeshMessage.from_control_batch(batch) - self.assertEqual(msg.message_type, DraftMeshMessageType.CONTROL_BATCH) - self.assertIs(msg.control_batch, batch) - self.assertIsNone(msg.tail_stream_output_batch) - - def test_from_tail_stream_output_batch_sets_discriminant_and_slot(self): - from sglang.srt.speculative.decoupled_spec_io import DraftTailStreamOutputBatch - - batch = DraftTailStreamOutputBatch() - msg = DraftMeshMessage.from_tail_stream_output_batch(batch) - self.assertEqual( - msg.message_type, DraftMeshMessageType.TAIL_STREAM_OUTPUT_BATCH - ) - self.assertIs(msg.tail_stream_output_batch, batch) - self.assertIsNone(msg.control_batch) - - -if __name__ == "__main__": - unittest.main(verbosity=3)