执行摘要

• 一句话：修复NPU上推测解码与XGrammar冲突，确保语法掩码正确应用。
• 推荐动作：此PR值得精读，尤其是xgrammar_backend.py中设备分支的设计决策和torch_ops/bitmask_ops.py的回退实现，展示了如何在多硬件环境下优雅处理特定设备问题。关注review讨论中关于函数位置和命名的权衡，这对代码组织有借鉴意义。

功能与动机

根据Issue #20148报告，在NPU上启用eagle/eagle3推测解码并发送带有json_schema的请求时，服务可能间歇性崩溃。原因是apply_token_bitmask_inplace_triton在NPU上无法正确强制执行语法掩码，导致被掩码的token仍保持高logits并被选中，引发下游'Tokens not accepted'错误。此PR旨在修复此问题，确保NPU设备上的约束解码正常工作。

实现拆解

1. 修改XGrammar后端以支持NPU回退：在python/sglang/srt/constrained/xgrammar_backend.py中，调整XGrammarGrammar.apply_vocab_mask方法，将NPU设备从通用CUDA/XPU/MUSA分支分离，引入专用回退路径，调用新增的apply_token_bitmask_inplace_torch函数。
2. 新增通用torch回退函数：创建文件python/sglang/srt/constrained/torch_ops/bitmask_ops.py，定义apply_token_bitmask_inplace_torch函数，提供CPU-based bitmask解包和GPU掩码填充的后端无关回退，用于NPU设备。
3. 添加NPU专用测试覆盖：新增测试文件python/sglang/test/ascend/test_ascend_vocab_mask.py，包含test_mask_blocks_disallowed_token_on_npu和test_npu_path_matches_reference_random两个测试函数，验证NPU路径的掩码效果和与参考实现的一致性。
4. 调整导入依赖：在xgrammar_backend.py中导入新增的torch_ops模块，确保NPU分支能正确调用回退函数。

关键文件：

• python/sglang/srt/constrained/xgrammar_backend.py（模块 语法后端；类别 source；类型 core-logic；符号 apply_vocab_mask）: 核心源码文件，修改了apply_vocab_mask方法，为NPU设备引入专用回退路径，直接影响语法掩码应用逻辑。
• python/sglang/srt/constrained/torch_ops/bitmask_ops.py（模块 掩码操作；类别 infra；类型 infrastructure；符号 apply_token_bitmask_inplace_torch）: 新增基础设施文件，提供后端无关的torch回退函数，用于NPU设备的掩码应用，是修复的关键组件。
• python/sglang/test/ascend/test_ascend_vocab_mask.py（模块 NPU测试；类别 test；类型 test-coverage；符号 _pack_mask, _apply_ref_cpu, test_mask_blocks_disallowed_token_on_npu, test_npu_path_matches_reference_random）: 新增测试文件，验证NPU路径的掩码应用正确性，确保修复效果和与参考实现的一致性。

关键符号：apply_vocab_mask, apply_token_bitmask_inplace_torch, _pack_mask, _apply_ref_cpu

关键源码片段

python/sglang/srt/constrained/xgrammar_backend.py

核心源码文件，修改了apply_vocab_mask方法，为NPU设备引入专用回退路径，直接影响语法掩码应用逻辑。

def apply_vocab_mask(self, logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
    if logits.device.type in {"cuda", "xpu", "musa"}: # 保持原有CUDA、XPU、MUSA设备路径不变
        if _is_hip:
            apply_token_bitmask_inplace_cuda(logits, vocab_mask)
        else:
            apply_token_bitmask_inplace_triton(logits, vocab_mask)
    elif logits.device.type == "npu": # 新增NPU专用分支，使用torch回退路径
        apply_token_bitmask_inplace_torch(logits, vocab_mask) # 调用通用回退函数，确保掩码正确应用
    else:
        raise RuntimeError(f"Unsupported device: {logits.device.type}") # 其他设备抛出错误

python/sglang/srt/constrained/torch_ops/bitmask_ops.py

新增基础设施文件，提供后端无关的torch回退函数，用于NPU设备的掩码应用，是修复的关键组件。

def apply_token_bitmask_inplace_torch(
    logits: torch.Tensor,
    bitmask: torch.Tensor,
) -> None:
    """Backend-agnostic torch fallback for packed-bitmask application.
​
    This path is currently used as a fallback on NPU in xgrammar backend.
    """
    vocab_size = logits.shape[-1] # 获取词汇表大小
    bitmask_cpu = bitmask.detach().cpu() # 将掩码移至CPU处理
    token_ids = torch.arange(vocab_size, device="cpu", dtype=torch.int32) # 生成token ID数组
    word_idx = token_ids // 32 # 计算掩码中的字索引
    bit_idx = token_ids % 32 # 计算位索引
    words = bitmask_cpu[:, word_idx].to(torch.int32) # 提取对应字
    allowed = ((words >> bit_idx) & 1).to(torch.bool) # 解包位掩码，得到允许的token布尔数组
    allowed = allowed.to(logits.device, non_blocking=True) # 移回原设备
    logits.masked_fill_(~allowed, float("-inf")) # 将不允许的token logits设为负无穷，确保不被选中

评论区精华

reviewer kpham-sgl指出，初始实现中在triton_ops/bitmask_ops.py添加的NPU专用函数不特定于Triton或NPU，建议重命名并移至更通用的位置。ChefWu551采纳建议，将函数移至constrained/torch_ops/bitmask_ops.py，重命名为apply_token_bitmask_inplace_torch，并移除NPU-only断言，使其成为后端无关的回退选项。最终决策是保持代码清晰和可维护性，避免硬件特定耦合。

• 函数位置和命名优化 (design): ChefWu551采纳建议，将函数移至constrained/torch_ops/bitmask_ops.py，重命名为apply_token_bitmask_inplace_torch，并移除NPU-only断言。

风险与影响

• 风险：回归风险：修改了xgrammar_backend.py中的设备分支逻辑，将NPU从通用分支移除，可能影响其他设备（如CUDA、XPU、MUSA）的现有行为，但review确认非NPU路径保持不变。
  性能风险：NPU回退路径涉及CPU解包和GPU数据传输，可能引入额外开销，但PR body中的基准测试显示性能影响可接受。
  兼容性风险：新增的torch回退函数可能不适用于所有设备类型，但错误处理已保留，对于不支持设备会抛出RuntimeError。
  测试覆盖风险：新增测试仅覆盖NPU设备，如果NPU不可用则跳过，但测试设计合理，确保功能正确性。

• 影响：对用户的影响：使用NPU设备并启用eagle/eagle3推测解码与XGrammar的用户将不再遇到间歇性崩溃，解码结果保持一致性。
  对系统的影响：核心解码路径增加了NPU专用处理逻辑，但非NPU设备不受影响，系统整体稳定性提升。
  对团队的影响：修复了特定硬件下的关键bug，减少了维护负担，并为未来NPU相关功能提供了可扩展的回退机制。

• 风险标记：核心路径变更, 硬件特定回退, 测试覆盖有限

关联脉络

• PR #20168 [Fix] eagle/eagle3 speculative decoding conflicts with xgrammar in NPU: 此PR的直接前身，因强制推送而关闭，PR #20989在其基础上重新实现并修复冲突。