From ec6f8d61f707952a87fedd2a3491d1e28130c86a Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 26 May 2026 11:34:03 -0700 Subject: [PATCH] [Spec] Async-assert probes across EAGLE/MTP; zero `tgt_cache_loc` (#26335) --- python/sglang/srt/environ.py | 6 ++- python/sglang/srt/layers/sampler.py | 16 +------ python/sglang/srt/mem_cache/memory_pool.py | 6 +++ python/sglang/srt/server_args.py | 14 ------ .../eagle_draft_cuda_graph_runner.py | 5 +- python/sglang/srt/speculative/eagle_info.py | 33 +++++++++++++ .../sglang/srt/speculative/eagle_info_v2.py | 16 +++++++ python/sglang/srt/speculative/eagle_worker.py | 11 ++++- .../sglang/srt/speculative/eagle_worker_v2.py | 22 ++++++++- .../srt/speculative/frozen_kv_mtp_worker.py | 12 ++++- .../speculative/multi_layer_eagle_worker.py | 2 +- .../multi_layer_eagle_worker_v2.py | 12 ++++- python/sglang/srt/speculative/spec_utils.py | 23 ++------- python/sglang/srt/utils/async_probe.py | 48 +++++++++++++++++++ .../test/server_fixtures/eagle_fixture.py | 5 +- test/registered/spec/dflash/test_dflash.py | 3 +- .../eagle/test_deepseek_v3_fp4_mtp_small.py | 5 +- .../eagle/test_eagle_constrained_decoding.py | 3 +- .../spec/eagle/test_eagle_dp_attention.py | 5 +- .../spec/eagle/test_eagle_infer_beta.py | 3 +- .../test_eagle_infer_beta_dp_attention.py | 5 +- ...est_eagle_infer_beta_dp_attention_large.py | 5 +- ...est_constrained_decoding_spec_reasoning.py | 5 +- 23 files changed, 171 insertions(+), 94 deletions(-) create mode 100644 python/sglang/srt/utils/async_probe.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 6ca70d5ea..a91ab5b37 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -526,8 +526,10 @@ class Envs: # Spec Config SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True) - SGLANG_SPEC_NAN_DETECTION = EnvBool(False) - SGLANG_SPEC_OOB_DETECTION = EnvBool(False) + # Master switch for all async-asserted invariant probes (NaN, Inf, OOB, + # page alignment). Off in prod; tests turn it on to fail-fast on + # numerical / index violations instead of getting silent NaN cascades. + SGLANG_ENABLE_ASYNC_ASSERT = EnvBool(False) # VLM SGLANG_VLM_CACHE_SIZE_MB = EnvInt(100) diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index 9181fbac5..816702f3c 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -17,7 +17,6 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.sampling.sampling_params import TOP_K_ALL from sglang.srt.server_args import get_global_server_args from sglang.srt.utils.common import ( - crash_on_warnings, get_bool_env_var, is_cuda, is_musa, @@ -57,7 +56,6 @@ _BUILT_IN_SAMPLING_BACKENDS = {"flashinfer", "pytorch", "ascend"} class Sampler(nn.Module): def __init__(self): super().__init__() - self.use_nan_detection = get_global_server_args().enable_nan_detection self.tp_sync_group = get_tp_group().device_group if is_dp_attention_enabled(): self.tp_sync_group = get_attention_tp_group().device_group @@ -74,20 +72,9 @@ class Sampler(nn.Module): def _preprocess_logits( self, logits: torch.Tensor, sampling_info: SamplingBatchInfo ) -> torch.Tensor: - """Apply custom logit processors and handle NaN detection.""" - # Apply the custom logit processors if registered in the sampling info + """Apply custom logit processors.""" if sampling_info.has_custom_logit_processor: apply_custom_logit_processor(logits, sampling_info) - - # Detect and handle NaN values in logits - if self.use_nan_detection and torch.any(torch.isnan(logits)): - logger.warning("Detected errors during sampling! NaN in the logits.") - logits = torch.where( - torch.isnan(logits), torch.full_like(logits, -1e5), logits - ) - if crash_on_warnings(): - raise ValueError("Detected errors during sampling! NaN in the logits.") - return logits def forward( @@ -238,7 +225,6 @@ class Sampler(nn.Module): sampling_info.top_ks, sampling_info.top_ps, filter_apply_order="joint", - check_nan=self.use_nan_detection, ) elif backend == "pytorch": # A slower fallback implementation with torch native operations. diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 0d520a5ba..8bd91fad9 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -64,6 +64,7 @@ from sglang.srt.utils import ( is_npu, next_power_of_2, ) +from sglang.srt.utils.async_probe import maybe_detect_oob from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter if TYPE_CHECKING: @@ -1091,6 +1092,11 @@ class MHATokenToKVPool(KVCache): ) def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + # Catch stale indices here instead of as illegal-addr or silent KV corruption. + size_limit = self.size + self.page_size + maybe_detect_oob(tgt_loc, 0, size_limit, "move_kv_cache tgt_loc") + maybe_detect_oob(src_loc, 0, size_limit, "move_kv_cache src_loc") + if envs.SGLANG_NATIVE_MOVE_KV_CACHE.get(): move_kv_cache_native(self.k_buffer, self.v_buffer, tgt_loc, src_loc) return diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index eb4041cef..c2d455a9f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -732,7 +732,6 @@ class ServerArgs: piecewise_cuda_graph_tokens: Optional[List[int]] = None piecewise_cuda_graph_compiler: str = "eager" torchao_config: str = "" - enable_nan_detection: bool = False enable_p2p_check: bool = False triton_attention_reduce_in_fp32: bool = False triton_attention_num_kv_splits: int = 8 @@ -1110,14 +1109,6 @@ 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) - # Deprecated attention-backend alias: "compressed" -> "dsv4". for attr in ( "attention_backend", @@ -6387,11 +6378,6 @@ class ServerArgs: default=ServerArgs.torchao_config, help="Optimize the model with torchao. Experimental feature. Current choices are: int8dq, int8wo, int4wo-, fp8wo, fp8dq-per_tensor, fp8dq-per_row", ) - parser.add_argument( - "--enable-nan-detection", - action="store_true", - help="[Deprecated] Use SGLANG_SPEC_NAN_DETECTION=1 and SGLANG_SPEC_OOB_DETECTION=1 instead.", - ) parser.add_argument( "--enable-p2p-check", action="store_true", 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 c0d180231..79927e989 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -27,16 +27,13 @@ from sglang.srt.model_executor.forward_batch_info import ( from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.input_buffers import ForwardInputBuffers from sglang.srt.speculative.eagle_info import EagleDraftInput -from sglang.srt.speculative.spec_utils import ( - maybe_detect_nan, - maybe_detect_oob, -) from sglang.srt.utils import ( require_attn_tp_gather, require_gathered_buffer, require_mlp_sync, require_mlp_tp_gather, ) +from sglang.srt.utils.async_probe import maybe_detect_nan, maybe_detect_oob if TYPE_CHECKING: from sglang.srt.speculative.eagle_worker import EAGLEWorker diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 79f797b1e..106423b0a 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -44,6 +44,7 @@ from sglang.srt.speculative.spec_utils import ( get_target_cache_loc, ) from sglang.srt.utils import is_cuda, is_musa, next_power_of_2 +from sglang.srt.utils.async_probe import maybe_detect_nan, maybe_detect_oob if is_cuda() or is_musa(): from sgl_kernel import ( @@ -125,6 +126,12 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin): return batch.input_ids = self.draft_token + maybe_detect_oob( + batch.input_ids, + 0, + batch.model_config.vocab_size, + "eagle prepare_for_verify input_ids", + ) if page_size == 1: batch.out_cache_loc = alloc_token_slots( @@ -349,12 +356,14 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin): target_probs = F.softmax( logits_output.next_token_logits / expanded_temperature, dim=-1 ) # (bs * draft_token_num, vocab_size) + maybe_detect_nan(target_probs, "verify: target_probs after softmax") target_probs = top_k_renorm_prob( target_probs, torch.repeat_interleave( sampling_info.top_ks, self.draft_token_num, dim=0 ), ) # (bs * draft_token_num, vocab_size) + maybe_detect_nan(target_probs, "verify: target_probs after top_k_renorm") if sampling_info.need_top_p_sampling: target_probs = top_p_renorm_prob( target_probs, @@ -362,6 +371,9 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin): sampling_info.top_ps, self.draft_token_num, dim=0 ), ) + maybe_detect_nan( + target_probs, "verify: target_probs after top_p_renorm" + ) target_probs = target_probs.reshape(bs, self.draft_token_num, -1) draft_probs = torch.zeros( @@ -418,6 +430,21 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin): spec_steps=self.spec_steps, ) + # accept_index values index batch.out_cache_loc (size = bs * draft_token_num); + # -1 is the reject sentinel. + maybe_detect_oob( + accept_index, + -1, + bs * self.draft_token_num, + "eagle verify accept_index post-sampling", + ) + maybe_detect_oob( + num_correct_drafts, + 0, + self.draft_token_num + 1, + "eagle verify num_correct_drafts post-sampling", + ) + unfinished_index = [] unfinished_accept_index = [] accept_index_cpu = accept_index.tolist() @@ -475,6 +502,12 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin): # TODO: fuse them accept_index = accept_index[accept_index != -1] accept_tokens = predict[accept_index] + maybe_detect_oob( + accept_tokens, + 0, + batch.model_config.vocab_size, + "eagle verify accept_tokens", + ) evict_mask = torch.full_like(self.draft_token, True, dtype=torch.bool) evict_mask[accept_index] = False num_correct_drafts_cpu = num_correct_drafts.cpu() diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index dd426f2e8..390e2e78d 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -38,6 +38,7 @@ from sglang.srt.speculative.spec_utils import ( SIMULATE_ACC_LEN, generate_simulated_accept_index, ) +from sglang.srt.utils.async_probe import maybe_detect_nan, maybe_detect_oob from sglang.srt.utils.common import is_cuda, is_hip, is_musa, is_npu, next_power_of_2 _is_cuda = is_cuda() @@ -226,6 +227,12 @@ class EagleDraftInputV2Mixin: batch.spec_info = self batch.input_ids = predict + maybe_detect_oob( + batch.input_ids, + 0, + batch.model_config.vocab_size, + "v2 prepare_for_extend_to_fill_draft_kvcache input_ids", + ) batch.extend_lens = [num_draft_tokens for _ in range(len(batch.seq_lens))] batch.prefix_lens = seq_lens_cpu_.tolist() batch.extend_num_tokens = extend_num_tokens @@ -264,6 +271,12 @@ class EagleVerifyInputV2Mixin: # Assign cache locations bs = len(batch.req_pool_indices) batch.input_ids = self.draft_token + maybe_detect_oob( + batch.input_ids, + 0, + batch.model_config.vocab_size, + "v2 prepare_for_verify input_ids", + ) device = batch.input_ids.device batch.out_cache_loc = assign_extend_cache_locs_func( req_pool_indices=batch.req_pool_indices, @@ -398,18 +411,21 @@ class EagleVerifyInputV2Mixin: target_probs = F.softmax( next_token_logits / expanded_temperature, dim=-1 ) # (bs * num_draft_tokens, vocab_size) + maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax") target_probs = top_k_renorm_prob( target_probs, torch.repeat_interleave( sampling_info.top_ks, self.draft_token_num, dim=0 ), ) # (bs * num_draft_tokens, vocab_size) + maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm") target_probs = top_p_renorm_prob( target_probs, torch.repeat_interleave( sampling_info.top_ps, self.draft_token_num, dim=0 ), ) + maybe_detect_nan(target_probs, "v2 verify: target_probs after top_p_renorm") target_probs = target_probs.reshape(bs, self.draft_token_num, -1) draft_probs = torch.zeros_like(target_probs) diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 3f3964d57..1bf696791 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -65,8 +65,6 @@ from sglang.srt.speculative.spec_utils import ( 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 ( @@ -79,6 +77,11 @@ from sglang.srt.utils import ( log_info_on_rank0, next_power_of_2, ) +from sglang.srt.utils.async_probe import ( + maybe_detect_inf, + maybe_detect_nan, + maybe_detect_oob, +) from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions _is_npu = is_npu() @@ -896,6 +899,7 @@ class EAGLEWorker(TpModelWorker): forward_batch, skip_attn_backend_init=True ).logits_output maybe_detect_nan(logits_output.next_token_logits, f"draft_forward step {i}") + maybe_detect_inf(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( @@ -907,6 +911,8 @@ class EAGLEWorker(TpModelWorker): if self.hot_token_id is not None: topk_index = self.hot_token_id[topk_index] hidden_states = logits_output.hidden_states + maybe_detect_nan(hidden_states, f"draft_forward step {i}: hidden_states") + maybe_detect_inf(hidden_states, f"draft_forward step {i}: hidden_states") forward_batch.positions.add_(1) parent_list, top_scores_index, draft_tokens = organize_draft_results( @@ -969,6 +975,7 @@ class EAGLEWorker(TpModelWorker): batch.sampling_info.vocab_mask = None maybe_detect_nan(logits_output.next_token_logits, "verify: target model logits") + maybe_detect_inf(logits_output.next_token_logits, "verify: target model logits") spec_info.hidden_states = logits_output.hidden_states res: EagleVerifyOutput = spec_info.verify( diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index eb5c40a77..bb1b51686 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -63,12 +63,15 @@ from sglang.srt.speculative.spec_utils import ( draft_tp_context, generate_token_bitmask, load_token_map, - maybe_detect_nan, - maybe_detect_oob, record_stream_each, record_stream_for_v2_verify, select_top_k_tokens, ) +from sglang.srt.utils.async_probe import ( + maybe_detect_inf, + maybe_detect_nan, + maybe_detect_oob, +) from sglang.srt.utils.common import ( MultiprocessingSerializer, empty_context, @@ -483,6 +486,7 @@ class EagleDraftWorker(BaseDraftWorker): forward_batch, skip_attn_backend_init=True ).logits_output maybe_detect_nan(logits_output.next_token_logits, f"draft_forward step {i}") + maybe_detect_inf(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( @@ -578,6 +582,7 @@ class EagleDraftWorker(BaseDraftWorker): 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") + maybe_detect_inf(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) @@ -642,6 +647,10 @@ class EagleDraftWorker(BaseDraftWorker): draft_logits_output.next_token_logits, f"draft_extend_for_decode (cuda_graph={can_cuda_graph})", ) + maybe_detect_inf( + 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[ @@ -1072,6 +1081,7 @@ class EAGLEWorkerV2(BaseSpecWorker): # Sample maybe_detect_nan(logits_output.next_token_logits, "verify: target model logits") + maybe_detect_inf(logits_output.next_token_logits, "verify: target model logits") ( predict, accept_lens, @@ -1204,6 +1214,14 @@ class EAGLEWorkerV2(BaseSpecWorker): bs = len(batch.seq_lens) size = bs * self.speculative_num_draft_tokens + # fill_accepted_out_cache_loc reads out_cache_loc[accept_index]; -1 sentinel ok. + maybe_detect_oob( + accept_index, + -1, + batch.out_cache_loc.size(0), + "eagle v2 move_accepted_tokens accept_index", + ) + tgt_cache_loc = torch.zeros( size, dtype=torch.int64, diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index d1e8daf93..ea8bcec72 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py @@ -70,11 +70,14 @@ from sglang.srt.speculative.spec_utils import ( draft_tp_context, fast_topk, generate_token_bitmask, - maybe_detect_nan, - maybe_detect_oob, select_top_k_tokens, ) from sglang.srt.utils import empty_context +from sglang.srt.utils.async_probe import ( + maybe_detect_inf, + maybe_detect_nan, + maybe_detect_oob, +) logger = logging.getLogger(__name__) @@ -406,6 +409,7 @@ class FrozenKVMTPWorker(TpModelWorker): forward_batch, skip_attn_backend_init=True ).logits_output maybe_detect_nan(logits_output.next_token_logits, "frozen_kv_mtp_seed") + maybe_detect_inf(logits_output.next_token_logits, "frozen_kv_mtp_seed") self._capture_for_decode(logits_output, draft_input) finally: batch.forward_mode = forward_mode_backup @@ -694,6 +698,9 @@ class FrozenKVMTPWorker(TpModelWorker): maybe_detect_nan( logits_output.next_token_logits, f"frozen_kv_mtp_draft step {i}" ) + maybe_detect_inf( + logits_output.next_token_logits, f"frozen_kv_mtp_draft 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( @@ -752,6 +759,7 @@ class FrozenKVMTPWorker(TpModelWorker): batch.sampling_info.vocab_mask = None maybe_detect_nan(logits_output.next_token_logits, "frozen_kv_mtp_verify") + maybe_detect_inf(logits_output.next_token_logits, "frozen_kv_mtp_verify") spec_info.hidden_states = logits_output.hidden_states res: FrozenKVMTPVerifyOutput = spec_info.verify( diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index 974d2a669..4c055535f 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker.py @@ -54,10 +54,10 @@ from sglang.srt.speculative.spec_utils import ( 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 +from sglang.srt.utils.async_probe import maybe_detect_nan if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner 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 f5f592aa8..3e18e81b4 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -48,12 +48,15 @@ from sglang.srt.speculative.multi_layer_eagle_utils import ( from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( draft_tp_context, - maybe_detect_nan, - maybe_detect_oob, record_stream_each, record_stream_for_v2_verify, select_top_k_tokens, ) +from sglang.srt.utils.async_probe import ( + maybe_detect_inf, + maybe_detect_nan, + maybe_detect_oob, +) from sglang.srt.utils.common import empty_context, fast_topk if TYPE_CHECKING: @@ -425,6 +428,10 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): output.logits_output.next_token_logits, f"draft_extend_for_prefill step {step}", ) + maybe_detect_inf( + 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) @@ -764,6 +771,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): # Sample maybe_detect_nan(logits_output.next_token_logits, "verify: target model logits") + maybe_detect_inf(logits_output.next_token_logits, "verify: target model logits") ( predict, accept_lens, diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 3a39dcd49..e1b0f9fe8 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -464,7 +464,9 @@ def get_src_tgt_cache_loc( page_size: int, ): src_cache_loc = out_cache_loc[accept_index] - tgt_cache_loc = torch.empty_like(src_cache_loc) + # 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, @@ -803,25 +805,6 @@ def draft_tp_context(tp_group: GroupCoordinator): yield -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 # even slower. # @torch.compile(dynamic=True) diff --git a/python/sglang/srt/utils/async_probe.py b/python/sglang/srt/utils/async_probe.py new file mode 100644 index 000000000..aefc63bcc --- /dev/null +++ b/python/sglang/srt/utils/async_probe.py @@ -0,0 +1,48 @@ +"""Async invariant probes — fire torch._assert_async without CPU sync. + +All probes are gated on SGLANG_ENABLE_ASYNC_ASSERT (default off in prod). +When the gate is on, a violation surfaces as an assertion at the next CUDA +sync point instead of as a silent NaN cascade or illegal-address crash. +""" + +import torch + +from sglang.srt.environ import envs + + +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_ENABLE_ASYNC_ASSERT.get(): + return + torch._assert_async(~torch.any(torch.isnan(tensor)), f"NaN detected! {msg}") + + +def maybe_detect_inf(tensor: torch.Tensor, msg: str = ""): + """Async Inf check — fp16 overflow surfaces as Inf before NaN.""" + if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): + return + torch._assert_async(~torch.any(torch.isinf(tensor)), f"Inf 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_ENABLE_ASYNC_ASSERT.get(): + return + if indices.numel() == 0: + return + torch._assert_async( + (indices.min() >= low) & (indices.max() < high), + f"OOB indices not in [{low}, {high}): {msg}", + ) + + +def maybe_detect_page_aligned(indices: torch.Tensor, page_size: int, msg: str): + """Async page-alignment check on slot ids.""" + if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): + return + if indices.numel() == 0 or page_size <= 1: + return + torch._assert_async( + (indices % page_size == 0).all(), + f"page-misaligned indices (page_size={page_size}): {msg}", + ) diff --git a/python/sglang/test/server_fixtures/eagle_fixture.py b/python/sglang/test/server_fixtures/eagle_fixture.py index 512f6d976..b628831de 100644 --- a/python/sglang/test/server_fixtures/eagle_fixture.py +++ b/python/sglang/test/server_fixtures/eagle_fixture.py @@ -37,10 +37,7 @@ class EagleServerBase(CustomTestCase): @classmethod def setUpClass(cls): cls.base_url = DEFAULT_URL_FOR_TEST - with ( - envs.SGLANG_SPEC_NAN_DETECTION.override(True), - envs.SGLANG_SPEC_OOB_DETECTION.override(True), - ): + with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): cls.process = popen_launch_server( cls.target_model, cls.base_url, diff --git a/test/registered/spec/dflash/test_dflash.py b/test/registered/spec/dflash/test_dflash.py index d139cbf54..221957861 100644 --- a/test/registered/spec/dflash/test_dflash.py +++ b/test/registered/spec/dflash/test_dflash.py @@ -55,8 +55,7 @@ class TestDFlashServerBase(CustomTestCase, MatchedStopMixin, GSM8KMixin): try: with ( envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1), - envs.SGLANG_SPEC_NAN_DETECTION.override(True), - envs.SGLANG_SPEC_OOB_DETECTION.override(True), + envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True), ): cls.process = popen_launch_server( cls.model, 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 99b3a2e88..b6fd25cb0 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 @@ -49,10 +49,7 @@ class TestDeepseekV3FP4MTP(CustomTestCase): "--model-loader-extra-config", '{"enable_multithread_load": true,"num_threads": 64}', ] - with ( - envs.SGLANG_SPEC_NAN_DETECTION.override(True), - envs.SGLANG_SPEC_OOB_DETECTION.override(True), - ): + with envs.SGLANG_ENABLE_ASYNC_ASSERT.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 85e7d75e7..267897c5c 100644 --- a/test/registered/spec/eagle/test_eagle_constrained_decoding.py +++ b/test/registered/spec/eagle/test_eagle_constrained_decoding.py @@ -62,8 +62,7 @@ class TestEagleConstrainedDecoding( launch_args.extend(cls.other_launch_args) with ( envs.SGLANG_ENABLE_SPEC_V2.override(cls.spec_v2), - envs.SGLANG_SPEC_NAN_DETECTION.override(True), - envs.SGLANG_SPEC_OOB_DETECTION.override(True), + envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True), ): cls.process = popen_launch_server( cls.model, diff --git a/test/registered/spec/eagle/test_eagle_dp_attention.py b/test/registered/spec/eagle/test_eagle_dp_attention.py index b31b38ebf..8fc48cb1b 100644 --- a/test/registered/spec/eagle/test_eagle_dp_attention.py +++ b/test/registered/spec/eagle/test_eagle_dp_attention.py @@ -60,10 +60,7 @@ class TestEAGLE3EngineDPAttention(CustomTestCase): "--cuda-graph-max-bs", "64", ] - with ( - envs.SGLANG_SPEC_NAN_DETECTION.override(True), - envs.SGLANG_SPEC_OOB_DETECTION.override(True), - ): + with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): cls.process = popen_launch_server( cls.model, cls.base_url, diff --git a/test/registered/spec/eagle/test_eagle_infer_beta.py b/test/registered/spec/eagle/test_eagle_infer_beta.py index 92577ca0f..e0b50b0a8 100644 --- a/test/registered/spec/eagle/test_eagle_infer_beta.py +++ b/test/registered/spec/eagle/test_eagle_infer_beta.py @@ -65,8 +65,7 @@ class TestEagle3ServerBase(CustomTestCase, MatchedStopMixin): launch_args.extend(cls.other_launch_args) with ( envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1), - envs.SGLANG_SPEC_NAN_DETECTION.override(True), - envs.SGLANG_SPEC_OOB_DETECTION.override(True), + envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True), envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.override(True), ): cls.process = popen_launch_server( 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 9e4bf5676..ebd6cc1bf 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,10 +65,7 @@ class TestEagleDPAttnServerSmall(CustomTestCase): "--speculative-num-draft-tokens", "4", ] - with ( - envs.SGLANG_SPEC_NAN_DETECTION.override(True), - envs.SGLANG_SPEC_OOB_DETECTION.override(True), - ): + with envs.SGLANG_ENABLE_ASYNC_ASSERT.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 0c1d63ec9..8e6c9c1b2 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,10 +73,7 @@ class TestEagleDPAttnServerLarge(CustomTestCase): "--model-loader-extra-config", '{"enable_multithread_load": true,"num_threads": 64}', ] - with ( - envs.SGLANG_SPEC_NAN_DETECTION.override(True), - envs.SGLANG_SPEC_OOB_DETECTION.override(True), - ): + with envs.SGLANG_ENABLE_ASYNC_ASSERT.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 60e695bc0..add3849cf 100644 --- a/test/registered/spec/test_constrained_decoding_spec_reasoning.py +++ b/test/registered/spec/test_constrained_decoding_spec_reasoning.py @@ -51,10 +51,7 @@ class ServerWithGrammar(CustomTestCase): "--speculative-num-draft-tokens=8", ] - with ( - envs.SGLANG_SPEC_NAN_DETECTION.override(True), - envs.SGLANG_SPEC_OOB_DETECTION.override(True), - ): + with envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True): cls.process = popen_launch_server( cls.model, cls.base_url,