From 7518dc3532b1d1188183a05d44569f1010dc7ac7 Mon Sep 17 00:00:00 2001 From: Gaoji Liu <34803073+attack204@users.noreply.github.com> Date: Tue, 30 Dec 2025 17:55:04 +0800 Subject: [PATCH] feat(SpecEagleV2): add standalone_worker_v2 (#12625) Co-authored-by: liugaoji.lgj --- .../srt/layers/attention/triton_backend.py | 4 +- python/sglang/srt/managers/overlap_utils.py | 14 +- python/sglang/srt/managers/schedule_batch.py | 4 +- python/sglang/srt/managers/scheduler.py | 6 +- .../srt/model_executor/cpu_graph_runner.py | 5 +- .../srt/model_executor/cuda_graph_runner.py | 2 + .../model_runner_kv_cache_mixin.py | 2 +- python/sglang/srt/server_args.py | 2 +- .../eagle_draft_cuda_graph_runner.py | 2 + .../eagle_draft_extend_cuda_graph_runner.py | 1 + .../sglang/srt/speculative/eagle_worker_v2.py | 16 +- python/sglang/srt/speculative/spec_info.py | 9 + .../srt/speculative/standalone_worker_v2.py | 169 ++++++++++++++++++ .../test_standalone_speculative_decoding.py | 98 ++++++++++ 14 files changed, 317 insertions(+), 17 deletions(-) create mode 100644 python/sglang/srt/speculative/standalone_worker_v2.py diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 344a479cd..7550ecbfd 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -602,7 +602,7 @@ class TritonAttnBackend(AttentionBackend): num_kv_splits = None attn_logits = None attn_lse = None - elif forward_mode.is_draft_extend(): + elif forward_mode.is_draft_extend(include_v2=True): num_tokens_per_bs = self.speculative_num_steps + 1 qo_indptr = self.qo_indptr[: bs + 1] qo_indptr[: bs + 1] = torch.arange( @@ -749,7 +749,7 @@ class TritonAttnBackend(AttentionBackend): seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens) mask_indptr = self.mask_indptr[: bs + 1] mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0) - elif forward_mode.is_draft_extend(): + elif forward_mode.is_draft_extend(include_v2=True): seq_lens = seq_lens[:bs] accept_lens = spec_info.accept_length[:bs] qo_indptr = self.qo_indptr[: bs + 1] diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 3056cdc54..4d60f5180 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -116,7 +116,9 @@ class FutureMap: return FutureIndices(indices=indices, interval=slice(start, end)) def resolve_future(self, model_worker_batch: ModelWorkerBatch): - if self.spec_algo.is_eagle(): + if self.spec_algo.is_none(): + _resolve_future_token_ids(model_worker_batch.input_ids, self.token_ids_buf) + else: # TODO(lsyin): write future indices into spec_info.future_indices draft_input: EagleDraftInput = model_worker_batch.spec_info if draft_input is None: @@ -129,8 +131,6 @@ class FutureMap: draft_input.new_seq_lens = self.new_seq_lens_buf[indices] if spec_need_hidden_states(): draft_input.hidden_states = self.hidden_states_buf[indices] - else: - _resolve_future_token_ids(model_worker_batch.input_ids, self.token_ids_buf) def is_empty_slice(self, s: slice) -> bool: start, stop, step = s.indices(self.future_buffer_len) @@ -142,12 +142,12 @@ class FutureMap: def store_to_map( self, future_indices: FutureIndices, batch_result: GenerationBatchResult ): - if self.spec_algo.is_eagle(): - draft_input: EagleDraftInput = batch_result.next_draft_input - self.store_to_map_for_new_batch(future_indices, draft_input) - else: + if self.spec_algo.is_none(): intv = future_indices.interval self.token_ids_buf[intv] = batch_result.next_token_ids + else: + draft_input: EagleDraftInput = batch_result.next_draft_input + self.store_to_map_for_new_batch(future_indices, draft_input) def store_to_map_for_new_batch( self, future_indices: FutureIndices, draft_input: EagleDraftInput diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 373ffbcc3..8b807613e 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1849,7 +1849,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): @property def is_spec_v2(self): # FIXME: finally deprecate is_spec_v2 - return self.enable_overlap and self.spec_algorithm.is_eagle() + ret = self.enable_overlap and not self.spec_algorithm.is_none() + assert not ret or self.spec_algorithm.supports_spec_v2() + return ret def prepare_for_decode(self): self.forward_mode = ForwardMode.DECODE diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 14cbf0bb1..241787fde 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -499,7 +499,7 @@ class Scheduler( # Draft workers are looked up via `SpeculativeAlgorithm` registry; new # algorithms should register their factory instead of patching this code. - if self.spec_algorithm.is_eagle(): + if self.spec_algorithm.supports_spec_v2(): draft_worker_kwargs["enable_overlap"] = self.enable_overlap # FIXME: refactor the draft worker registration logic @@ -852,7 +852,7 @@ class Scheduler( if self.draft_worker is None or self.spec_algorithm.is_ngram(): draft_token_to_kv_pool = None - elif self.spec_algorithm.is_eagle() and self.enable_overlap: + elif self.spec_algorithm.supports_spec_v2() and self.enable_overlap: if self.server_args.enable_multi_layer_eagle: draft_runner = self.draft_worker.draft_worker.draft_runner_list[0] else: @@ -930,11 +930,13 @@ class Scheduler( hidden_size=( model_config.hidden_size if self.spec_algorithm.is_eagle() + or self.spec_algorithm.is_standalone() else 16 # minimal padding size for RDMA ), hidden_states_dtype=( model_config.dtype if self.spec_algorithm.is_eagle() + or self.spec_algorithm.is_standalone() else torch.float32 ), custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(), diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 9eda46722..20579a3c2 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -606,7 +606,10 @@ class CPUGraphRunner: def get_spec_info(self, num_tokens: int): spec_info = None - if self.model_runner.spec_algorithm.is_eagle(): + if ( + self.model_runner.spec_algorithm.is_eagle() + or self.model_runner.spec_algorithm.is_standalone() + ): from sglang.srt.speculative.eagle_info import EagleVerifyInput if self.model_runner.is_draft_worker: diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 2d96431a3..6ef74ca80 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -375,6 +375,7 @@ class CudaGraphRunner: cuda_graph_bs = ( max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs if self.model_runner.spec_algorithm.is_eagle() + or self.model_runner.spec_algorithm.is_standalone() else max(forward_batch.global_num_tokens_cpu) ) else: @@ -777,6 +778,7 @@ class CudaGraphRunner: max_batch_size = ( max_num_tokens / self.num_tokens_per_bs if self.model_runner.spec_algorithm.is_eagle() + or self.model_runner.spec_algorithm.is_standalone() else max_num_tokens ) index = bisect.bisect_left(self.capture_bs, max_batch_size) diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index fc753436a..128ee2f90 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -303,7 +303,7 @@ class ModelRunnerKVCacheMixin: max_num_reqs, self.server_args.max_mamba_cache_size // ratio ) - if self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone(): + if not self.spec_algorithm.is_none(): if self.is_draft_worker: self.max_total_num_tokens = self.server_args.draft_runner_cache_size max_num_reqs = self.server_args.max_num_reqs diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 3022bae40..2927b04d1 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1998,7 +1998,7 @@ class ServerArgs: ) if ( - self.speculative_algorithm in ["EAGLE", "EAGLE3"] + self.speculative_algorithm in ["EAGLE", "EAGLE3", "STANDALONE"] and envs.SGLANG_ENABLE_SPEC_V2.get() ): self.disable_overlap_schedule = False diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 3f5070aec..5fe45086c 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -141,6 +141,7 @@ class EAGLEDraftCudaGraphRunner: cuda_graph_bs = ( max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs if self.model_runner.spec_algorithm.is_eagle() + or self.model_runner.spec_algorithm.is_standalone() else max(forward_batch.global_num_tokens_cpu) ) else: @@ -328,6 +329,7 @@ class EAGLEDraftCudaGraphRunner: max_batch_size = ( max_num_tokens // self.num_tokens_per_bs if self.model_runner.spec_algorithm.is_eagle() + or self.model_runner.spec_algorithm.is_standalone() else max_num_tokens ) index = bisect.bisect_left(self.capture_bs, max_batch_size) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 916136cec..231531171 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -189,6 +189,7 @@ class EAGLEDraftExtendCudaGraphRunner: cuda_graph_bs = ( max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs if self.model_runner.spec_algorithm.is_eagle() + or self.model_runner.spec_algorithm.is_standalone() else max(forward_batch.global_num_tokens_cpu) ) else: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index ab0b2d6a3..7ba55f714 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -12,6 +12,7 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_r from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import ( EAGLEDraftNpuGraphRunner, ) +from sglang.srt.layers.attention.triton_backend import TritonMultiStepDraftBackend from sglang.srt.layers.moe.utils import ( speculative_moe_a2a_backend_context, speculative_moe_backend_context, @@ -50,12 +51,14 @@ from sglang.srt.utils.common import ( empty_context, fast_topk, get_available_gpu_memory, + is_cuda, is_npu, next_power_of_2, ) from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions _is_npu = is_npu() +_is_cuda = is_cuda() logger = logging.getLogger(__name__) @@ -251,8 +254,14 @@ class EagleDraftWorker(BaseDraftWorker): "cuda": EAGLEDraftExtendCudaGraphRunner, } # Capture extend - # FIXME cuda not support draft_extend capture - if self.draft_extend_attn_backend and _is_npu: + # TODO: support draft extend cuda graph for more attention backends + if self.draft_extend_attn_backend and ( + _is_npu + or ( + _is_cuda + and isinstance(self.draft_attn_backend, TritonMultiStepDraftBackend) + ) + ): tic = time.perf_counter() before_mem = get_available_gpu_memory(self.device, self.gpu_id) logger.info( @@ -506,6 +515,9 @@ class EagleDraftWorker(BaseDraftWorker): self.plan_stream ) + if forward_batch.spec_info.accept_length is None: + forward_batch.spec_info.accept_length = batch_result.accept_lens + # Run draft extend batch in the main compute stream can_cuda_graph = ( self.cuda_graph_runner_for_draft_extend diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 2b0d48aaa..824e905dd 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -161,6 +161,9 @@ class SpeculativeAlgorithm(metaclass=_SpeculativeAlgorithmMeta): def is_none(self) -> bool: return self is SpeculativeAlgorithm.NONE + def supports_spec_v2(self) -> bool: + return self.is_eagle() or self.is_eagle3() or self.is_standalone() + def is_eagle(self) -> bool: return self._has_flag("EAGLE") @@ -274,6 +277,12 @@ def _create_eagle_worker(**kwargs: Any) -> Any: def _create_standalone_worker(**kwargs: Any) -> Any: + enable_overlap = kwargs.pop("enable_overlap", False) + if enable_overlap: + from sglang.srt.speculative.standalone_worker_v2 import StandaloneWorkerV2 + + return StandaloneWorkerV2(**kwargs) + from sglang.srt.speculative.standalone_worker import StandaloneWorker return StandaloneWorker(**kwargs) diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py new file mode 100644 index 000000000..da6e3523d --- /dev/null +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -0,0 +1,169 @@ +import contextlib +import logging +from typing import Optional, Tuple + +import torch + +from sglang.srt.environ import envs +from sglang.srt.layers.moe.utils import speculative_moe_backend_context +from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.server_args import ServerArgs +from sglang.srt.speculative.eagle_utils import TreeMaskMode +from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2 +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.speculative.spec_utils import draft_tp_context +from sglang.srt.utils import empty_context, get_bool_env_var, is_cuda + +if is_cuda(): + from sgl_kernel import segment_packbits # noqa: F401 + +logger = logging.getLogger(__name__) +SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB") + + +def _get_plan_stream( + device: str, +) -> Tuple[any, contextlib.AbstractContextManager]: + if envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get(): + plan_stream = torch.get_device_module(device).Stream() + plan_stream_ctx = torch.get_device_module(device).stream(plan_stream) + return plan_stream, plan_stream_ctx + else: + return None, contextlib.nullcontext() + + +class StandaloneDraftWorker(EagleDraftWorker): + """Custom EagleDraftWorker that doesn't share embeddings/lm_head with target model.""" + + def __init__( + self, + server_args: ServerArgs, + gpu_id: int, + tp_rank: int, + dp_rank: int, + moe_ep_rank: int, + nccl_port: int, + target_worker: TpModelWorker, + ): + # copy args + self.server_args = server_args + self.gpu_id = gpu_id + self.tp_rank = tp_rank + self.dp_rank = dp_rank + self.moe_ep_rank = moe_ep_rank + self.nccl_port = nccl_port + self.target_worker = target_worker + + # Args for easy access + self.device = server_args.device + self.topk = server_args.speculative_eagle_topk + self.speculative_num_steps = server_args.speculative_num_steps + self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens + self.speculative_algorithm = SpeculativeAlgorithm.from_string( + server_args.speculative_algorithm + ) + + # Set constant + from sglang.srt.speculative.eagle_info import EagleDraftInput + + EagleDraftInput.ALLOC_LEN_PER_DECODE = max( + self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens + ) + + # Do not capture cuda graph in `TpModelWorker` init, + # will capture later with init_cuda_graphs() + backup_disable_cuda_graph = server_args.disable_cuda_graph + server_args.disable_cuda_graph = True + + # Share the allocator with a target worker. + # Draft and target worker own their own KV cache pools. + self.req_to_token_pool, self.token_to_kv_pool_allocator = ( + target_worker.get_memory_pool() + ) + with empty_context(): + # Init draft worker + self.draft_worker = TpModelWorker( + server_args=server_args, + gpu_id=gpu_id, + tp_rank=tp_rank, + pp_rank=0, # FIXME + dp_rank=dp_rank, + moe_ep_rank=moe_ep_rank, + nccl_port=nccl_port, + is_draft_worker=True, + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + ) + + # Alias for better readability + self.draft_runner = self.draft_worker.model_runner + + self.init_token_map() + self.init_lm_head() + + # Init attention backend and cuda graphs + self.draft_runner.server_args.disable_cuda_graph = backup_disable_cuda_graph + self.draft_tp_context = ( + draft_tp_context if server_args.enable_dp_attention else empty_context + ) + with self.draft_tp_context( + self.draft_runner.tp_group + ), speculative_moe_backend_context(): + self.init_attention_backend() + self.init_cuda_graphs() + self.tree_mask_mode = TreeMaskMode.FULL_MASK + + self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device) + + def init_lm_head(self): + """Override to prevent sharing embeddings and lm_head with target model.""" + # For standalone worker, we don't share embeddings and lm_head + # The draft model uses its own embeddings and lm_head + pass + + +class StandaloneWorkerV2(EAGLEWorkerV2): + + def __init__( + self, + server_args: ServerArgs, + gpu_id: int, + tp_rank: int, + dp_rank: Optional[int], + moe_ep_rank: int, + nccl_port: int, + target_worker: TpModelWorker, + ): + # Parse arguments + self.server_args = server_args + self.topk = server_args.speculative_eagle_topk + self.speculative_num_steps = server_args.speculative_num_steps + self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens + self.enable_nan_detection = server_args.enable_nan_detection + self.gpu_id = gpu_id + self.device = server_args.device + self._target_worker = target_worker + self.page_size = server_args.page_size + self.speculative_algorithm = SpeculativeAlgorithm.from_string( + server_args.speculative_algorithm + ) + + self.req_to_token_pool, self.token_to_kv_pool_allocator = ( + target_worker.get_memory_pool() + ) + + # Override the context length of the draft model to be the same as the target model. + server_args.context_length = target_worker.model_runner.model_config.context_len + + # Create our custom draft worker that doesn't share embeddings/lm_head + self._draft_worker = StandaloneDraftWorker( + server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, nccl_port, target_worker + ) + + # Some dummy tensors + self.num_new_pages_per_topk = torch.empty( + (), dtype=torch.int64, device=self.device + ) + self.extend_lens = torch.empty((), dtype=torch.int64, device=self.device) + + self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device) diff --git a/test/registered/spec/test_standalone_speculative_decoding.py b/test/registered/spec/test_standalone_speculative_decoding.py index 67475f341..fb125c6c9 100644 --- a/test/registered/spec/test_standalone_speculative_decoding.py +++ b/test/registered/spec/test_standalone_speculative_decoding.py @@ -1,3 +1,4 @@ +import os import unittest from types import SimpleNamespace @@ -41,6 +42,25 @@ DEFAULT_SERVER_ARGS = [ 0.7, ] +# Default server arguments for V2 tests +DEFAULT_SERVER_ARGS_V2 = [ + "--trust-remote-code", + "--cuda-graph-max-bs", + "8", + "--speculative-algorithm", + "STANDALONE", + "--speculative-draft-model-path", + DEFAULT_DRAFT_MODEL_STANDALONE, + "--speculative-num-steps", + "4", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "5", + "--mem-fraction-static", + 0.7, +] + class TestStandaloneSpeculativeDecodingBase(CustomTestCase): @@ -100,6 +120,67 @@ class TestStandaloneSpeculativeDecodingBase(CustomTestCase): self.assertGreater(avg_spec_accept_length, self.spec_decode_threshold) +class TestStandaloneV2SpeculativeDecodingBase(CustomTestCase): + + model = DEFAULT_TARGET_MODEL_STANDALONE + draft_model = DEFAULT_DRAFT_MODEL_STANDALONE + base_url = DEFAULT_URL_FOR_TEST + accuracy_threshold = 0.7 # derived tests need to override this + spec_decode_threshold = 3.6 # derived spec decoding tests need to override this + + @classmethod + def get_server_args(cls): + """Return the arguments for the server launch. Override in subclasses.""" + return DEFAULT_SERVER_ARGS_V2 + ["--attention-backend", "fa3"] + + @classmethod + def setUpClass(cls): + # disable deep gemm precompile to make launch server faster + # please don't do this if you want to make your inference workload faster + envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.set(False) + envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False) + envs.SGLANG_ENABLE_SPEC_V2.set(True) # Enable Speculative Decoding V2 + model = cls.model + cls.process = popen_launch_server( + model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=cls.get_server_args(), + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + if "SGLANG_ENABLE_SPEC_V2" in os.environ: + envs.SGLANG_ENABLE_SPEC_V2.set(False) + + def test_gsm8k(self): + requests.get(self.base_url + "/flush_cache") + + args = SimpleNamespace( + num_shots=4, + num_questions=100, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + data_path=GSM_DATASET_PATH, + ) + metrics = run_eval_few_shot_gsm8k(args) + print(f"{metrics=}") + + # Use the appropriate metric key based on the test class + metric_key = "accuracy" + self.assertGreater(metrics[metric_key], self.accuracy_threshold) + + server_info = requests.get(self.base_url + "/get_server_info") + avg_spec_accept_length = server_info.json()["internal_states"][0][ + "avg_spec_accept_length" + ] + print(f"{avg_spec_accept_length=}") + self.assertGreater(avg_spec_accept_length, self.spec_decode_threshold) + + class TestStandaloneSpeculativeDecodingTriton(TestStandaloneSpeculativeDecodingBase): @classmethod @@ -115,5 +196,22 @@ class TestStandaloneSpeculativeDecodingFlashinfer( return DEFAULT_SERVER_ARGS + ["--attention-backend", "flashinfer"] +class TestStandaloneV2SpeculativeDecodingTriton( + TestStandaloneV2SpeculativeDecodingBase +): + + @classmethod + def get_server_args(cls): + return DEFAULT_SERVER_ARGS_V2 + ["--attention-backend", "triton"] + + +class TestStandaloneV2SpeculativeDecodingFlashinfer( + TestStandaloneV2SpeculativeDecodingBase +): + @classmethod + def get_server_args(cls): + return DEFAULT_SERVER_ARGS_V2 + ["--attention-backend", "flashinfer"] + + if __name__ == "__main__": unittest.main()