执行摘要

• 一句话：新增 ngram decode 专用快速更新 kernel
• 推荐动作：值得精读，展示如何通过简化 kernel 假设实现数十倍性能提升。尤其关注 review 中对 int64 溢出的讨论——这是一个在长上下文场景中容易被忽略的缺陷。

功能与动机

PR body 说明 decode 阶段每次更新仅处理一个 token（req_lens==1），且无忽略 token，因此可以简化 kernel 实现。原始通用 kernel 需要为每个 token 计算 req_lens 偏移，造成额外开销。

实现拆解

1. 在 CUDA kernel 文件 ngram_embedding.cuh 中新增 UpdateTokenTableDecodeKernel，移除 req_lens 和 ignore_tokens 参数，无需偏移计算，使用 int64_t 指针避免长上下文表偏移溢出。
2. 在 ngram_embedding.py 中注册 JIT 模块并导出 update_token_table_decode Python 包装函数，类型标注明确为 decode 快速路径。
3. 在 model_runner.py 中修改 maybe_update_ngram_token_table，将调用从 update_token_table 切换为 update_token_table_decode，移除 req_lens 和 ignore_tokens 参数。
4. 新增基准文件 bench_ngram_update_token_table.py，使用 triton 测试框架对比 general 与 decode 路径，支持 CI 运行。
5. 在 test_ngram_embedding.py 中增加 test_update_token_table_decode_matches_general 参数化测试，验证 decode 快速路径在 req_lens==1 时与通用路径输出一致。

关键文件：

• python/sglang/jit_kernel/ngram_embedding.py（模块 JIT内核；类别 source；类型 core-logic；符号 update_token_table_decode）: 新增 update_token_table_decode 函数，导出 JIT kernel 的快速路径，是客户端调用的直接入口。
• python/sglang/srt/model_executor/model_runner.py（模块 模型执行器；类别 source；类型 data-contract）: 修改 maybe_update_ngram_token_table 方法，路由到 decode 快速路径，移除不再需要的参数。
• python/sglang/jit_kernel/csrc/ngram_embedding.cuh（模块 JIT内核；类别 other；类型 core-logic）: 新增 UpdateTokenTableDecodeKernel CUDA kernel，专为 decode 优化，移除不必要参数和计算，使用 int64_t 防止溢出。
• python/sglang/jit_kernel/tests/test_ngram_embedding.py（模块 测试；类别 test；类型 test-coverage；符号 test_update_token_table_decode_matches_general）: 新增 test_update_token_table_decode_matches_general 测试，验证 decode 快速路径正确性。
• python/sglang/jit_kernel/benchmark/bench_ngram_update_token_table.py（模块 基准测试；类别 source；类型 core-logic；符号 benchmark, fn）: 新增基准测试文件，对比 general 与 decode 路径性能，提供量化收益证据。

关键符号：update_token_table_decode, UpdateTokenTableDecodeKernel, maybe_update_ngram_token_table, benchmark, test_update_token_table_decode_matches_general

评论区精华

Review 中 @yuan-luo 指出通用路径中 row_indices * max_context_len 可能溢出 32 位整数（当 max_context_len 较大时），导致表偏移错误和越界写入。@BBuf 确认已在合并 origin/main 时解决：将 general kernel 和 decode kernel 中的偏移量统一使用 int64_t 计算，并在 H200 上重新验证测试和基准均通过。

• int64 溢出问题：row_indices * max_context_len 可能溢出 (correctness): BBuf 在合并冲突解决时统一改为 int64_t，并验证 H200 上测试和基准通过。

风险与影响

• 风险：原有通用路径的 int64 溢出风险已修复，但 decode 快速路径假设 req_lens==1 且无忽略 token，若未来调用场景变化（如 prefill 误用此路径）可能引入隐蔽错误。不过调用点 maybe_update_ngram_token_table 在 decode 阶段始终设置 req_lens=1，当前没有风险。另外，新 kernel 仅在 H200 上验证，其他 GPU 架构上可能因寄存器压力或 warp 调度不同而表现不一致，但功能应正确。
• 影响：仅影响启用 ngram embedding 的模型（如某些 speculative decoding 配置）。在 H200 上 decode 阶段 token 表更新吞吐提升显著（batch 4096 时延迟降低 98%）。代码变更量小，回归风险低。
• 风险标记：溢出风险（已修复）, kernel 假设 req_lens==1, 其他 GPU 架构未测试

关联脉络

• 暂无明显关联 PR