# PR #27482 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[Bug] Fix out-of-range token id crashing tp=1 `VocabParallelEmbedding`
- 合并时间：2026-06-07 14:00
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/27482

---

# 执行摘要

- 一句话：修复 token id 越界导致的 GPU crash
- 推荐动作：值得合并：修复了实际 CI 干扰的 flaky bug，且通过添加防御性断言改善了调试体验。

# 功能与动机

由于 `random.randint(0, 32000)` 是 inclusive 的，可能生成 32000 这个非法 token id（合法范围为 0-31999），在 `tp=1` 时 `VocabParallelEmbedding.forward` 不做输入掩码，导致 `F.embedding` 触发 `vectorized_gather_kernel` device-side assert 并造成 GPU coredump。此问题表现为 `TestEagleLlama2Retract.test_radix_attention` 的 flaky 失败。

# 实现拆解

1. **修复测试数据生成**：在 `python/sglang/test/kits/radix_cache_server_kit.py` 的 `gen_radix_tree` 中，将 `random.randint(0, 32000)` 改为 `random.randint(0, 31999)`，确保 token id 严格在合法范围内。
2. **添加嵌入层越界检测**：在 `python/sglang/srt/layers/vocab_parallel_embedding.py` 的 `VocabParallelEmbedding.forward` 中，在 TP=1 路径（无掩码）开始处调用 `maybe_detect_oob(input_, 0, self.num_embeddings, ...)`，使非法 token id 被捕获为 positioned async assert。
3. **改进断言消息粒度**：在 `python/sglang/srt/utils/async_probe.py` 的 `maybe_detect_oob` 中，将原来合并的 lower/upper 断言拆为两条独立的 `torch._assert_async` 调用，分别报告“index < low”（可能为负值或未掩码 sentinel）和“index >= high”（越界），便于快速定位问题类型。
4. **修复单元测试中的边界**：在 `test/registered/unit/mem_cache/test_radix_cache_unit.py` 的 `test_memory_allocated` 中，将 `random.randint(1, vocab_size)` 改为 `random.randint(1, vocab_size - 1)`，避免生成非法 token id（`vocab_size` 为 1000 时合法 id 最大为 999）。

关键文件：
- `python/sglang/srt/utils/async_probe.py`（模块 调试工具；类别 source；类型 core-logic；符号 maybe_detect_oob）: 核心改进：拆分 OOB 断言为两条独立 assert，分别报告 low 和 high 违规，提升定位精度。
- `python/sglang/srt/layers/vocab_parallel_embedding.py`（模块 嵌入层；类别 source；类型 dependency-wiring；符号 VocabParallelEmbedding.forward）: 关键防御：在 forward 中调用 maybe_detect_oob，使得 tp=1 时非法 token id 被捕获为可定位的 async assert。
- `python/sglang/test/kits/radix_cache_server_kit.py`（模块 测试夹具；类别 test；类型 test-coverage；符号 gen_radix_tree）: 直接修复 flaky 测试的根因：gen_radix_tree 中的 token id 上限修正。
- `test/registered/unit/mem_cache/test_radix_cache_unit.py`（模块 单元测试；类别 test；类型 test-coverage；符号 test_memory_allocated）: 单元测试 token id 范围修正，防止测试中使用非法 token id。

关键符号：maybe_detect_oob, VocabParallelEmbedding.forward, gen_radix_tree

## 关键源码片段

### `python/sglang/srt/utils/async_probe.py`

核心改进：拆分 OOB 断言为两条独立 assert，分别报告 low 和 high 违规，提升定位精度。

```python
# python/sglang/srt/utils/async_probe.py

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.

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

```

### `python/sglang/srt/layers/vocab_parallel_embedding.py`

关键防御：在 forward 中调用 maybe_detect_oob，使得 tp=1 时非法 token id 被捕获为可定位的 async assert。

```python
# python/sglang/srt/layers/vocab_parallel_embedding.py

from sglang.srt.utils.async_probe import maybe_detect_oob

# ... (class definition) ...

def forward(self, input_):
    # 在 tp=1 时，不进行输入掩码，直接在嵌入前检测越界 token id
    maybe_detect_oob(
        input_, 0, self.num_embeddings, "VocabParallelEmbedding input id"
    )
    if self.tp_size > 1:
        masked_input, input_mask = get_masked_input_and_mask(
            input_,
            self.shard_indices.org_vocab_start_index,
            self.shard_indices.org_vocab_end_index,
            self.shard_indices.num_org_vocab_padding,
            self.shard_indices.added_vocab_start_index,
            self.shard_indices.added_vocab_end_index,
        )
    else:
        masked_input = input_

    with use_symmetric_memory(
        get_tp_group(), disabled=not is_allocation_symmetric()
    ):
        output_parallel = self.quant_method.embedding(self, masked_input.long())

    if self.tp_size > 1:
        output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
        if not get_attn_tp_context().input_scattered:
            if self.use_attn_tp_group:
                output_parallel = attn_tp_all_reduce(output_parallel)
            else:
                output_parallel = tensor_model_parallel_all_reduce(output_parallel)
    return output_parallel

```

### `python/sglang/test/kits/radix_cache_server_kit.py`

直接修复 flaky 测试的根因：gen_radix_tree 中的 token id 上限修正。

```python
# python/sglang/test/kits/radix_cache_server_kit.py

def gen_radix_tree(num_nodes=400, chunk_len=256):
    num0 = num_nodes // 2
    num1 = num_nodes - num0
    nodes = [{"input_ids": [37] * 117, "decode_len": 217}]
    for _ in range(num0):
        parent = random.choice(nodes)
        unique_len = random.randint(0, chunk_len)
        decode_len = random.randint(0, chunk_len)
        token_id = random.randint(
            0, 31999
        )  # randint 是 inclusive 的；vocab_size-1 = 31999
        child = {
            "input_ids": parent["input_ids"] + [token_id] * unique_len,
            "decode_len": decode_len,
        }
        nodes.append(child)
    # ... ( 后续循环类似修改 )
    return nodes

```

# 评论区精华

暂无 reviewer 评论，PR body 和 commit 消息已清晰说明问题与解决方案。

- 暂无高价值评论线程

# 风险与影响

- 风险：风险较低：测试数据生成修复确保不会产生越界 token id；`maybe_detect_oob` 仅在 `SGLANG_ENABLE_ASYNC_ASSERT` 启用时才会执行（CI 默认开启），且包含 `None` / empty 张量的 early return，backward 兼容性好；异步断言不引入 GPU-CPU 同步，性能开销极小。
- 影响：影响范围：直接修复 `TestEagleLlama2Retract.test_radix_attention` 的 flaky 失败，提升 CI 稳定性；同时为 `VocabParallelEmbedding` 的所有调用方（tp=1 场景）提供越界检测防御，避免类似问题难以诊断。对用户无功能影响。
- 风险标记：暂无

# 关联脉络

- PR #27461 Enable async-assert invariant probes by default in CI: 此前 PR 在 CI 中默认启用了 async assert，使得此 PR 添加的 maybe_detect_oob 调用能真正生效
- PR #27478 [Spec] Guard async-assert probes against `None` tensor: 同 async probe 相关，此前对 None 张量的防护也适用于此 PR 的调用上下文