执行摘要

• 一句话：修复 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/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/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/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 的调用上下文