From 346a4131cfbac5fe2f42ba01d7c4062853e2ded5 Mon Sep 17 00:00:00 2001 From: kpham-sgl Date: Thu, 5 Mar 2026 13:51:05 -0800 Subject: [PATCH] [Spec] Refactor NaN/OOB checks to async `maybe_detect_*` with env-var control (#19899) Co-authored-by: hnyls2002 --- python/sglang/srt/environ.py | 2 ++ python/sglang/srt/server_args.py | 10 +++++- python/sglang/srt/speculative/eagle_worker.py | 28 ++++++++++------ .../sglang/srt/speculative/eagle_worker_v2.py | 31 ++++++++++++++---- .../speculative/multi_layer_eagle_worker.py | 20 +++++++----- .../multi_layer_eagle_worker_v2.py | 19 ++++++++--- python/sglang/srt/speculative/spec_utils.py | 23 +++++++++---- .../srt/speculative/standalone_worker.py | 1 - .../srt/speculative/standalone_worker_v2.py | 1 - .../test/server_fixtures/eagle_fixture.py | 32 +++++++++++-------- .../eagle/test_deepseek_v3_fp4_mtp_small.py | 8 ++++- .../eagle/test_eagle_constrained_decoding.py | 8 ++++- .../spec/eagle/test_eagle_dp_attention.py | 16 ++++++---- .../spec/eagle/test_eagle_infer_beta.py | 8 ++++- .../test_eagle_infer_beta_dp_attention.py | 8 ++++- ...est_eagle_infer_beta_dp_attention_large.py | 8 ++++- ...est_constrained_decoding_spec_reasoning.py | 16 ++++++---- 17 files changed, 171 insertions(+), 68 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 882da78d2..90a73c00c 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -422,6 +422,8 @@ class Envs: # Spec Config SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True) + SGLANG_SPEC_NAN_DETECTION = EnvBool(False) + SGLANG_SPEC_OOB_DETECTION = EnvBool(False) # VLM SGLANG_VLM_CACHE_SIZE_MB = EnvInt(100) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 03fe80e03..2885a7562 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -889,6 +889,14 @@ class ServerArgs: ) self.tool_call_parser = deprecated_tool_call_parsers[self.tool_call_parser] + if self.enable_nan_detection: + logger.warning( + "--enable-nan-detection is deprecated. " + "Use SGLANG_SPEC_NAN_DETECTION=1 and SGLANG_SPEC_OOB_DETECTION=1 instead." + ) + envs.SGLANG_SPEC_NAN_DETECTION.set(True) + envs.SGLANG_SPEC_OOB_DETECTION.set(True) + def _handle_prefill_delayer_env_compat(self): if envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.get(): self.enable_prefill_delayer = True @@ -5013,7 +5021,7 @@ class ServerArgs: parser.add_argument( "--enable-nan-detection", action="store_true", - help="Enable the NaN detection for debugging purposes.", + help="[Deprecated] Use SGLANG_SPEC_NAN_DETECTION=1 and SGLANG_SPEC_OOB_DETECTION=1 instead.", ) parser.add_argument( "--enable-p2p-check", diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 32b3a520a..b5277fb88 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -49,12 +49,13 @@ from sglang.srt.speculative.eagle_utils import ( from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( assign_draft_cache_locs, - detect_nan, draft_tp_context, fast_topk, generate_token_bitmask, get_last_loc_large_page_size_large_top_k, load_token_map, + maybe_detect_nan, + maybe_detect_oob, select_top_k_tokens, ) from sglang.srt.utils import ( @@ -94,7 +95,6 @@ class EAGLEWorker(TpModelWorker): 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 @@ -622,6 +622,9 @@ class EAGLEWorker(TpModelWorker): spec_info.topk_index, spec_info.hidden_states, ) + + maybe_detect_nan(topk_p, "draft_forward: NaN in initial topk_p from spec_info") + if self.hot_token_id is not None: topk_index = self.hot_token_id[topk_index] # TODO: We only need self.speculative_num_steps - 1 cache loc @@ -670,10 +673,15 @@ class EAGLEWorker(TpModelWorker): logits_output = self.draft_model_runner.forward( forward_batch, skip_attn_backend_init=True ).logits_output - if self.server_args.enable_nan_detection: - detect_nan(logits_output) + maybe_detect_nan(logits_output.next_token_logits, f"draft_forward step {i}") probs = torch.softmax(logits_output.next_token_logits, dim=-1) topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) + maybe_detect_oob( + topk_index, + 0, + logits_output.next_token_logits.shape[-1], + f"draft_forward step {i}: topk_index OOB vs vocab_size={logits_output.next_token_logits.shape[-1]}", + ) if self.hot_token_id is not None: topk_index = self.hot_token_id[topk_index] hidden_states = logits_output.hidden_states @@ -741,8 +749,7 @@ class EAGLEWorker(TpModelWorker): # and will be applied to produce wrong results batch.sampling_info.vocab_mask = None - if self.enable_nan_detection: - detect_nan(logits_output) + maybe_detect_nan(logits_output.next_token_logits, "verify: target model logits") spec_info.hidden_states = logits_output.hidden_states res: EagleVerifyOutput = spec_info.verify( @@ -893,8 +900,7 @@ class EAGLEWorker(TpModelWorker): if mm_input_embeds is not None: forward_batch.mm_input_embeds = mm_input_embeds logits_output = self.draft_model_runner.forward(forward_batch).logits_output - if self.enable_nan_detection: - detect_nan(logits_output) + maybe_detect_nan(logits_output.next_token_logits, "draft_extend_for_prefill") assert isinstance(forward_batch.spec_info, EagleDraftInput) assert forward_batch.spec_info is batch.spec_info self.capture_for_decode(logits_output, forward_batch.spec_info) @@ -975,8 +981,10 @@ class EAGLEWorker(TpModelWorker): ).logits_output self.capture_for_decode(logits_output, forward_batch.spec_info) - if self.enable_nan_detection: - detect_nan(logits_output) + maybe_detect_nan( + logits_output.next_token_logits, + f"draft_extend_after_decode (cuda_graph={can_cuda_graph})", + ) # Restore backup. # This is because `seq_lens` can be modified in `prepare_extend_after_decode` diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index b6bf6d5ef..5aa4d2c63 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -44,10 +44,11 @@ from sglang.srt.speculative.eagle_info_v2 import ( from sglang.srt.speculative.eagle_utils import TreeMaskMode, build_tree_kernel_efficient from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( - detect_nan, draft_tp_context, generate_token_bitmask, load_token_map, + maybe_detect_nan, + maybe_detect_oob, select_top_k_tokens, ) from sglang.srt.utils.common import ( @@ -387,6 +388,9 @@ class EagleDraftWorker(BaseDraftWorker): spec_info.topk_index, spec_info.hidden_states, ) + + maybe_detect_nan(topk_p, "draft_forward: NaN in initial topk_p from spec_info") + if self.hot_token_id is not None: topk_index = self.hot_token_id[topk_index] @@ -427,10 +431,15 @@ class EagleDraftWorker(BaseDraftWorker): logits_output = self.draft_runner.forward( forward_batch, skip_attn_backend_init=True ).logits_output - if self.server_args.enable_nan_detection: - detect_nan(logits_output) + maybe_detect_nan(logits_output.next_token_logits, f"draft_forward step {i}") probs = torch.softmax(logits_output.next_token_logits, dim=-1) topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) + maybe_detect_oob( + topk_index, + 0, + logits_output.next_token_logits.shape[-1], + f"draft_forward step {i}: topk_index OOB vs vocab_size={logits_output.next_token_logits.shape[-1]}", + ) if self.hot_token_id is not None: topk_index = self.hot_token_id[topk_index] hidden_states = logits_output.hidden_states @@ -447,6 +456,12 @@ class EagleDraftWorker(BaseDraftWorker): ) top_scores_index = top_scores.indices top_scores_index = torch.sort(top_scores_index).values + maybe_detect_oob( + top_scores_index, + 0, + ss_token_list.shape[1], + "draft_forward: top_scores_index OOB for gather on ss_token_list", + ) draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1) if len(parents_list) > 1: @@ -502,6 +517,7 @@ class EagleDraftWorker(BaseDraftWorker): if mm_input_embeds is not None: forward_batch.mm_input_embeds = mm_input_embeds logits_output = self.draft_runner.forward(forward_batch).logits_output + maybe_detect_nan(logits_output.next_token_logits, "draft_extend_for_prefill") # Update spec_info for the next draft step probs = torch.softmax(logits_output.next_token_logits, dim=-1) @@ -559,6 +575,11 @@ class EagleDraftWorker(BaseDraftWorker): forward_batch, skip_attn_backend_init=True ).logits_output + maybe_detect_nan( + draft_logits_output.next_token_logits, + f"draft_extend_for_decode (cuda_graph={can_cuda_graph})", + ) + # Reorganize the spec info for the next batch draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[ select_index @@ -601,7 +622,6 @@ class EAGLEWorkerV2(BaseSpecWorker): 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.tp_rank = tp_rank self.gpu_id = gpu_id self.device = server_args.device @@ -781,8 +801,7 @@ class EAGLEWorkerV2(BaseSpecWorker): batch.sampling_info.vocab_mask = None # Sample - if self.enable_nan_detection: - detect_nan(logits_output) + maybe_detect_nan(logits_output.next_token_logits, "verify: target model logits") ( predict, accept_length, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index dfd98b943..c5ea19981 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker.py @@ -47,11 +47,11 @@ from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner imp ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( - detect_nan, draft_tp_context, fast_topk, generate_token_bitmask, load_token_map, + maybe_detect_nan, select_top_k_tokens, ) from sglang.srt.utils import empty_context, get_available_gpu_memory, is_cuda, is_npu @@ -86,7 +86,6 @@ class MultiLayerEagleWorker(TpModelWorker): 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 @@ -382,6 +381,8 @@ class MultiLayerEagleWorker(TpModelWorker): spec_info.hidden_states, ) + maybe_detect_nan(topk_p, "draft: NaN in initial topk_p from spec_info") + # Return values score_list: List[torch.Tensor] = [] token_list: List[torch.Tensor] = [] @@ -515,8 +516,7 @@ class MultiLayerEagleWorker(TpModelWorker): # and will be applied to produce wrong results batch.sampling_info.vocab_mask = None - if self.enable_nan_detection: - detect_nan(logits_output) + maybe_detect_nan(logits_output.next_token_logits, "verify: target model logits") spec_info.hidden_states = logits_output.hidden_states res: EagleVerifyOutput = spec_info.verify( @@ -623,8 +623,10 @@ class MultiLayerEagleWorker(TpModelWorker): logits_output = ( self.mtp_model_runner(step).forward(forward_batch).logits_output ) - if self.enable_nan_detection: - detect_nan(logits_output) + maybe_detect_nan( + logits_output.next_token_logits, + f"draft_extend_for_prefill step {step}", + ) probs = torch.softmax(logits_output.next_token_logits, dim=-1) topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) topk_p_list.append(topk_p) @@ -718,8 +720,10 @@ class MultiLayerEagleWorker(TpModelWorker): .logits_output ) - if self.enable_nan_detection: - detect_nan(logits_output) + maybe_detect_nan( + logits_output.next_token_logits, + f"draft_extend_after_decode step {step} (cuda_graph={can_cuda_graph})", + ) probs = torch.softmax(logits_output.next_token_logits, dim=-1) topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) topk_p_list.append(topk_p) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 0d5f08556..32f264361 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -38,8 +38,9 @@ from sglang.srt.speculative.multi_layer_eagle_utils import ( ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( - detect_nan, draft_tp_context, + maybe_detect_nan, + maybe_detect_oob, select_top_k_tokens, ) from sglang.srt.utils.common import empty_context, fast_topk @@ -286,6 +287,8 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): spec_info.hidden_states, ) + maybe_detect_nan(topk_p, "draft_forward: NaN in initial topk_p from spec_info") + # Return values score_list: List[torch.Tensor] = [] token_list: List[torch.Tensor] = [] @@ -329,6 +332,12 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): ) top_scores_index = top_scores.indices top_scores_index = torch.sort(top_scores_index).values + maybe_detect_oob( + top_scores_index, + 0, + ss_token_list.shape[1], + "draft_forward: top_scores_index OOB for gather on ss_token_list", + ) draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1) if len(parents_list) > 1: @@ -387,6 +396,10 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): output: ModelRunnerOutput = self.draft_runner_list[step].forward( forward_batch ) + maybe_detect_nan( + output.logits_output.next_token_logits, + f"draft_extend_for_prefill step {step}", + ) probs = torch.softmax(output.logits_output.next_token_logits, dim=-1) topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) topk_p_list.append(topk_p) @@ -560,7 +573,6 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): 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 @@ -702,8 +714,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): logits_output = forward_batch_output.logits_output # Sample - if self.enable_nan_detection: - detect_nan(logits_output) + maybe_detect_nan(logits_output.next_token_logits, "verify: target model logits") ( predict, accept_length, diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 39e2806f5..ca6e36ca3 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -17,7 +17,6 @@ from sglang.srt.distributed.parallel_state import ( patch_tensor_parallel_group, ) from sglang.srt.environ import envs -from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.common import get_last_loc from sglang.srt.server_args import ServerArgs, get_global_server_args @@ -706,11 +705,23 @@ def draft_tp_context(tp_group: GroupCoordinator): yield -def detect_nan(logits_output: LogitsProcessorOutput): - logits = logits_output.next_token_logits - if torch.any(torch.isnan(logits)): - logger.error("Detected errors during sampling! NaN in the logits.") - raise ValueError("Detected errors during sampling! NaN in the logits.") +def maybe_detect_nan(tensor: torch.Tensor, msg: str = ""): + """Async NaN check — no GPU-CPU sync, error surfaces at next sync point.""" + if not envs.SGLANG_SPEC_NAN_DETECTION.get(): + return + torch._assert_async(~torch.any(torch.isnan(tensor)), f"NaN detected! {msg}") + + +def maybe_detect_oob(indices: torch.Tensor, low: int, high: int, msg: str): + """Async OOB check — no GPU-CPU sync, error surfaces at next sync point.""" + if not envs.SGLANG_SPEC_OOB_DETECTION.get(): + return + if indices.numel() == 0: + return + torch._assert_async( + (indices.min() >= low) & (indices.max() < high), + f"OOB indices not in [{low}, {high}): {msg}", + ) # Disable torch.compile for this function because it will be diff --git a/python/sglang/srt/speculative/standalone_worker.py b/python/sglang/srt/speculative/standalone_worker.py index 4d7ca30e3..da21b989b 100644 --- a/python/sglang/srt/speculative/standalone_worker.py +++ b/python/sglang/srt/speculative/standalone_worker.py @@ -40,7 +40,6 @@ class StandaloneWorker(EAGLEWorker): 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 diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index 26ef88548..f5ee84508 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -147,7 +147,6 @@ class StandaloneWorkerV2(EAGLEWorkerV2): 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 diff --git a/python/sglang/test/server_fixtures/eagle_fixture.py b/python/sglang/test/server_fixtures/eagle_fixture.py index 9c5b28ed0..3330d2868 100644 --- a/python/sglang/test/server_fixtures/eagle_fixture.py +++ b/python/sglang/test/server_fixtures/eagle_fixture.py @@ -4,6 +4,7 @@ import time import requests +from sglang.srt.environ import envs from sglang.srt.utils.common import kill_process_tree from sglang.test.test_utils import ( DEFAULT_DRAFT_MODEL_EAGLE, @@ -36,20 +37,23 @@ class EagleServerBase(CustomTestCase): @classmethod def setUpClass(cls): cls.base_url = DEFAULT_URL_FOR_TEST - cls.process = popen_launch_server( - cls.target_model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - f"--speculative-algorithm={cls.spec_algo}", - f"--speculative-draft-model-path={cls.draft_model}", - f"--speculative-num-steps={cls.spec_steps}", - f"--speculative-eagle-topk={cls.spec_topk}", - f"--speculative-num-draft-tokens={cls.spec_tokens}", - f"--mem-fraction-static={cls.mem_fraction_static}", - ] - + cls.extra_args, - ) + with envs.SGLANG_SPEC_NAN_DETECTION.override( + True + ), envs.SGLANG_SPEC_OOB_DETECTION.override(True): + cls.process = popen_launch_server( + cls.target_model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + f"--speculative-algorithm={cls.spec_algo}", + f"--speculative-draft-model-path={cls.draft_model}", + f"--speculative-num-steps={cls.spec_steps}", + f"--speculative-eagle-topk={cls.spec_topk}", + f"--speculative-num-draft-tokens={cls.spec_tokens}", + f"--mem-fraction-static={cls.mem_fraction_static}", + ] + + cls.extra_args, + ) @classmethod def tearDownClass(cls): diff --git a/test/registered/spec/eagle/test_deepseek_v3_fp4_mtp_small.py b/test/registered/spec/eagle/test_deepseek_v3_fp4_mtp_small.py index 6abdd19cc..58cba7abb 100644 --- a/test/registered/spec/eagle/test_deepseek_v3_fp4_mtp_small.py +++ b/test/registered/spec/eagle/test_deepseek_v3_fp4_mtp_small.py @@ -50,7 +50,13 @@ class TestDeepseekV3FP4MTP(CustomTestCase): "--model-loader-extra-config", '{"enable_multithread_load": true,"num_threads": 64}', ] - with envs.SGLANG_ENABLE_SPEC_V2.override(True): + with envs.SGLANG_ENABLE_SPEC_V2.override( + True + ), envs.SGLANG_SPEC_NAN_DETECTION.override( + True + ), envs.SGLANG_SPEC_OOB_DETECTION.override( + True + ): cls.process = popen_launch_server( cls.model, cls.base_url, diff --git a/test/registered/spec/eagle/test_eagle_constrained_decoding.py b/test/registered/spec/eagle/test_eagle_constrained_decoding.py index 2c010e894..e6d1271ef 100644 --- a/test/registered/spec/eagle/test_eagle_constrained_decoding.py +++ b/test/registered/spec/eagle/test_eagle_constrained_decoding.py @@ -59,7 +59,13 @@ class TestEagleConstrainedDecoding( cls.grammar_backend, ] launch_args.extend(cls.other_launch_args) - with envs.SGLANG_ENABLE_SPEC_V2.override(cls.spec_v2): + with envs.SGLANG_ENABLE_SPEC_V2.override( + cls.spec_v2 + ), envs.SGLANG_SPEC_NAN_DETECTION.override( + True + ), envs.SGLANG_SPEC_OOB_DETECTION.override( + True + ): cls.process = popen_launch_server( cls.model, cls.base_url, diff --git a/test/registered/spec/eagle/test_eagle_dp_attention.py b/test/registered/spec/eagle/test_eagle_dp_attention.py index e25893c7c..bba8ecf28 100644 --- a/test/registered/spec/eagle/test_eagle_dp_attention.py +++ b/test/registered/spec/eagle/test_eagle_dp_attention.py @@ -3,6 +3,7 @@ from types import SimpleNamespace import requests +from sglang.srt.environ import envs from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k from sglang.test.send_one import BenchArgs, send_one_prompt @@ -55,12 +56,15 @@ class TestEAGLE3EngineDPAttention(CustomTestCase): "--cuda-graph-max-bs", "64", ] - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=other_args, - ) + with envs.SGLANG_SPEC_NAN_DETECTION.override( + True + ), envs.SGLANG_SPEC_OOB_DETECTION.override(True): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + ) @classmethod def tearDownClass(cls): diff --git a/test/registered/spec/eagle/test_eagle_infer_beta.py b/test/registered/spec/eagle/test_eagle_infer_beta.py index c67ff334a..313fd5512 100644 --- a/test/registered/spec/eagle/test_eagle_infer_beta.py +++ b/test/registered/spec/eagle/test_eagle_infer_beta.py @@ -59,7 +59,13 @@ class TestEagleServerBase(CustomTestCase, MatchedStopMixin): launch_args.extend(cls.other_launch_args) with envs.SGLANG_ENABLE_SPEC_V2.override( True - ), envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1): + ), envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override( + 1 + ), envs.SGLANG_SPEC_NAN_DETECTION.override( + True + ), envs.SGLANG_SPEC_OOB_DETECTION.override( + True + ): cls.process = popen_launch_server( cls.model, cls.base_url, diff --git a/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention.py b/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention.py index b55dd964c..b990e5990 100644 --- a/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention.py +++ b/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention.py @@ -65,7 +65,13 @@ class TestEagleDPAttnServerSmall(CustomTestCase): "--speculative-num-draft-tokens", "4", ] - with envs.SGLANG_ENABLE_SPEC_V2.override(True): + with envs.SGLANG_ENABLE_SPEC_V2.override( + True + ), envs.SGLANG_SPEC_NAN_DETECTION.override( + True + ), envs.SGLANG_SPEC_OOB_DETECTION.override( + True + ): cls.process = popen_launch_server( cls.model, cls.base_url, diff --git a/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention_large.py b/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention_large.py index fc8c4b12d..8a7fcd00d 100644 --- a/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention_large.py +++ b/test/registered/spec/eagle/test_eagle_infer_beta_dp_attention_large.py @@ -73,7 +73,13 @@ class TestEagleDPAttnServerLarge(CustomTestCase): "--model-loader-extra-config", '{"enable_multithread_load": true,"num_threads": 64}', ] - with envs.SGLANG_ENABLE_SPEC_V2.override(True): + with envs.SGLANG_ENABLE_SPEC_V2.override( + True + ), envs.SGLANG_SPEC_NAN_DETECTION.override( + True + ), envs.SGLANG_SPEC_OOB_DETECTION.override( + True + ): cls.process = popen_launch_server( cls.model, cls.base_url, diff --git a/test/registered/spec/test_constrained_decoding_spec_reasoning.py b/test/registered/spec/test_constrained_decoding_spec_reasoning.py index 89765d952..18931ec4d 100644 --- a/test/registered/spec/test_constrained_decoding_spec_reasoning.py +++ b/test/registered/spec/test_constrained_decoding_spec_reasoning.py @@ -3,6 +3,7 @@ import unittest import openai +from sglang.srt.environ import envs from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import ( @@ -50,12 +51,15 @@ class ServerWithGrammar(CustomTestCase): "--speculative-num-draft-tokens=8", ] - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=launch_args, - ) + with envs.SGLANG_SPEC_NAN_DETECTION.override( + True + ), envs.SGLANG_SPEC_OOB_DETECTION.override(True): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=launch_args, + ) @classmethod def tearDownClass(cls):