diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index 3deb1ad0e..cbba8cd06 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/python/sglang/srt/layers/vocab_parallel_embedding.py @@ -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( diff --git a/python/sglang/srt/utils/async_probe.py b/python/sglang/srt/utils/async_probe.py index f937fd85c..84f510e45 100644 --- a/python/sglang/srt/utils/async_probe.py +++ b/python/sglang/srt/utils/async_probe.py @@ -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}", ) diff --git a/python/sglang/test/kits/radix_cache_server_kit.py b/python/sglang/test/kits/radix_cache_server_kit.py index 12caa2b8c..9c4dd872f 100644 --- a/python/sglang/test/kits/radix_cache_server_kit.py +++ b/python/sglang/test/kits/radix_cache_server_kit.py @@ -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, diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index c09997436..b72476217 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -863,9 +863,11 @@ class TestRadixCache(unittest.TestCase): torch_allocated_before = torch.cuda.memory_allocated() # build dataset with common prefix - common_prefix = [random.randint(1, vocab_size) for _ in range(base_prefix_len)] + common_prefix = [ + random.randint(1, vocab_size - 1) for _ in range(base_prefix_len) + ] for _ in range(num_seqs): - suffix = [random.randint(1, vocab_size) for _ in range(suffix_len)] + suffix = [random.randint(1, vocab_size - 1) for _ in range(suffix_len)] seq = common_prefix + suffix keys.append(seq) values.append(torch.zeros(len(seq), device="cuda", dtype=torch.int32))