[Bug] Fix out-of-range token id crashing tp=1 VocabParallelEmbedding (#27482)

This commit is contained in:
Liangsheng Yin
2026-06-06 23:00:17 -07:00
committed by GitHub
parent 80eee2d92c
commit 0ce3db3c0a
4 changed files with 27 additions and 7 deletions
@@ -43,6 +43,7 @@ from sglang.srt.utils import (
is_npu,
set_weight_attrs,
)
from sglang.srt.utils.async_probe import maybe_detect_oob
DEFAULT_VOCAB_PADDING_SIZE = 64
@@ -495,6 +496,11 @@ class VocabParallelEmbedding(torch.nn.Module):
param[loaded_weight.shape[0] :].data.fill_(0)
def forward(self, input_):
# Surface a bad token id (>= vocab_size, or a negative / unmasked sentinel) as a
# located async assert instead of a silent OOB embedding gather (tp=1 does not mask).
maybe_detect_oob(
input_, 0, self.num_embeddings, "VocabParallelEmbedding input id"
)
if self.tp_size > 1:
# Build the mask.
masked_input, input_mask = get_masked_input_and_mask(
+11 -3
View File
@@ -33,14 +33,22 @@ def maybe_detect_inf(tensor: Optional[torch.Tensor], msg: str = ""):
def maybe_detect_oob(indices: Optional[torch.Tensor], low: int, high: int, msg: str):
"""Async OOB check — no GPU-CPU sync, error surfaces at next sync point."""
"""Async OOB check — no GPU-CPU sync, error surfaces at next sync point.
Low/high asserted separately so the message names which failed (low =
negative/sentinel, high = out of range).
"""
if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
return
if indices is None or indices.numel() == 0:
return
torch._assert_async(
(indices.min() >= low) & (indices.max() < high),
f"OOB indices not in [{low}, {high}): {msg}",
indices.min() >= low,
f"index < {low} (negative / unmasked sentinel?): {msg}",
)
torch._assert_async(
indices.max() < high,
f"index >= {high} (out of range): {msg}",
)
@@ -11,7 +11,9 @@ def gen_radix_tree(num_nodes=400, chunk_len=256):
parent = random.choice(nodes)
unique_len = random.randint(0, chunk_len)
decode_len = random.randint(0, chunk_len)
token_id = random.randint(0, 32000)
token_id = random.randint(
0, 31999
) # randint is inclusive; vocab_size-1 = 31999
child = {
"input_ids": parent["input_ids"] + [token_id] * unique_len,
"decode_len": decode_len,
@@ -24,7 +26,9 @@ def gen_radix_tree(num_nodes=400, chunk_len=256):
for _ in range(num_branch):
unique_len = random.randint(0, chunk_len)
decode_len = random.randint(0, chunk_len)
token_id = random.randint(0, 32000)
token_id = random.randint(
0, 31999
) # randint is inclusive; vocab_size-1 = 31999
child = {
"input_ids": parent["input_ids"] + [token_id] * unique_len,
"decode_len": decode_len,