[Spec] Async-assert probes across EAGLE/MTP; zero tgt_cache_loc (#26335)

This commit is contained in:
Liangsheng Yin
2026-05-26 11:34:03 -07:00
committed by GitHub
parent 6c8128650e
commit ec6f8d61f7
23 changed files with 171 additions and 94 deletions
+4 -2
View File
@@ -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)
+1 -15
View File
@@ -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.
@@ -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
-14
View File
@@ -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-<group_size>, 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",
@@ -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
@@ -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()
@@ -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)
@@ -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(
@@ -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,
@@ -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(
@@ -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
@@ -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,
+3 -20
View File
@@ -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)
+48
View File
@@ -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}",
)
@@ -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,