[Spec] Async-assert probes across EAGLE/MTP; zero tgt_cache_loc (#26335)
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user