chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Spec] Refactor NaN/OOB checks to async maybe_detect_* with env-var control (#19899)

Browse Source 
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
This commit is contained in:
kpham-sgl
2026-03-05 13:51:05 -08:00
committed by GitHub
parent 1c1712d8e5
commit 346a4131cf
17 changed files with 171 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/environ.py
View File

@@ -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)

10
python/sglang/srt/server_args.py
View File

@@ -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",

28
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -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`

31
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -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,

20
python/sglang/srt/speculative/multi_layer_eagle_worker.py
View File

@@ -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)

19
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
View File

@@ -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,

23
python/sglang/srt/speculative/spec_utils.py
View File

@@ -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

1
python/sglang/srt/speculative/standalone_worker.py
View File

@@ -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

1
python/sglang/srt/speculative/standalone_worker_v2.py
View File

@@ -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

32
python/sglang/test/server_fixtures/eagle_fixture.py
View File

@@ -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):