[Bug] Fix out-of-range token id crashing tp=1 VocabParallelEmbedding (#27482)
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user