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

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(

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),
+        indices.min() >= low,
-        f"OOB indices not in [{low}, {high}): {msg}",
+        f"index < {low} (negative / unmasked sentinel?): {msg}",
+    )
+    torch._assert_async(
+        indices.max() < high,
+        f"index >= {high} (out of range): {msg}",
     )
 
 

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,

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))