执行摘要

• 一句话：Ascend 采样改主机侧判定，消除每 token 设备同步
• 推荐动作：建议精读，但只需要 10 分钟左右：改动量极小，却是理解 Ascend 采样热路径与「设备标量判定 → 隐式同步」这一常见性能陷阱的好样本。重点关注三处设计决策：一是 npu_top_k_top_p_eligible 的计算位置选在 from_schedule_batch（而非每次采样），依赖「top_k 在采样期间不变」这一前提；二是 merge_batch 用 AND、默认值取 False 的保守取向，使漏传与合并只带来性能退化而不带来正确性风险；三是新增测试用 torch.all 抛 AssertionError 做护栏的写法，可作为同类优化的测试模板。需要在真实 NPU 上验证收益的团队，建议补充带同步计数的 profiler 证据，而不是依赖吞吐数字。

功能与动机

PR body 明确指出根因：调用 torch_npu.npu_top_k_top_p 前需要校验 top_k 是否落在 [1, 1024]，而 top_ks 此时已存放在 NPU 上，「Although torch.all() produces a device scalar, using that scalar in a Python condition requires converting it to a host boolean」——这就在 per-token 采样路径引入隐式 NPU 到 CPU 同步，且因该条件每个 decode step 都求值而反复打断异步执行、抬高采样延迟。与此同时 top_k 完全来自请求参数、采样期间不变，因此「It can therefore be evaluated once on the CPU when constructing the batch metadata」，这构成本次改造的全部动机：只前移可用性判定，不改采样数学。

实现拆解

1. 元数据层新增主机侧标志（python/sglang/srt/sampling/sampling_batch_info.py）：dataclass 新增 npu_top_k_top_p_eligible: bool = False，注释写明「Keeping this off the device avoids a scalar synchronization in the per-token sampling path」；在 from_schedule_batch 中与 is_all_greedy、need_top_p_sampling 等既有布尔标志并列计算 all(1 <= r.sampling_params.top_k <= 1024 for r in reqs)，全程只读 CPU 侧请求参数，不触碰设备张量。TOP_K_ALL（-1）天然不满足下界，继续走通用路径，因此行为等价性由区间表达式保证。
2. 批次生命周期语义（同文件 merge_batch）：新增 self.npu_top_k_top_p_eligible &= other.npu_top_k_top_p_eligible，与 is_all_greedy 的 AND 语义一致、与 need_*_sampling 的 OR 语义相反；filter_batch 不做任何「重新提升」，即过滤掉非法请求后仍保持保守。PR body 明确这是有意为之：「This cannot incorrectly dispatch unsupported values to the fused operator」。
3. 采样分派点改造（python/sglang/srt/layers/sampler.py）：Sampler._sample_from_logits 在调用 top_k_top_p_min_p_sampling_from_logits_ascend 时新增关键字参数 npu_top_k_top_p_eligible=sampling_info.npu_top_k_top_p_eligible；被调函数签名新增同名参数（默认 False），融合分支条件由 hasattr(torch_npu, "npu_top_k_top_p") and torch.all(...) 变为 hasattr(...) and npu_top_k_top_p_eligible，运行时算子可用性检查保持不变，两个条件仍是 AND。参数默认值取 False 属于「安全方向」：漏传只会导致性能退化，不会把非法 top_k 送进融合算子。
4. 测试配套：新增 test/registered/unit/sampling/test_ascend_sampler.py（register_cpu_ci(est_time=1, suite="base-a-test-cpu")），类 TestAscendSamplerDispatch 通过 patch.object(sampler_module.torch, "all", side_effect=AssertionError(...)) 建立回归护栏，并断言 eligible 为 True/False 时 npu_top_k_top_p 的调用与否；test/registered/unit/sampling/test_sampling_batch_info.py 扩展 test_merge_combines_flags 增加 npu_top_k_top_p_eligible 的 AND 断言，新增 test_npu_top_k_top_p_eligibility_uses_request_params 覆盖 (1, 1024)、(1, 1025)、(4, TOP_K_ALL) 三组边界。
5. 无配置与部署改动：没有新增用户参数、环境变量或 CLI 开关，优化对 Ascend 采样后端自动生效；非 Ascend 路径（CUDA/ROCm/MUSA 等）因默认值为 False 成为 no-op，行为与改动前完全一致。

关键文件：

• python/sglang/srt/sampling/sampling_batch_info.py（模块 采样批次信息；类别 source；类型 core-logic；符号 SamplingBatchInfo, from_schedule_batch, merge_batch）: 新增主机侧判定字段并在 batch 构造阶段一次性计算，同时定义合并期的 AND 保守语义，是本次优化的语义核心。
• python/sglang/srt/layers/sampler.py（模块 采样器；类别 source；类型 core-logic；符号 _sample_from_logits, top_k_top_p_min_p_sampling_from_logits_ascend）: 采样热路径的实际改动点：把设备侧 torch.all 判定替换为主机标志透传，并新增带保守默认值的函数参数。
• test/registered/unit/sampling/test_ascend_sampler.py（模块 采样器测试；类别 test；类型 test-coverage；符号 TestAscendSamplerDispatch, test_top_k_dispatch_does_not_read_device_values_on_cpu）: 新增的分派路径单测，用让 torch.all 抛 AssertionError 的方式建立「不得再读设备标量」的回归护栏，是本 PR 最具借鉴价值的测试写法。
• test/registered/unit/sampling/test_sampling_batch_info.py（模块 批次信息测试；类别 test；类型 test-coverage；符号 test_merge_combines_flags, test_npu_top_k_top_p_eligibility_uses_request_params）: 覆盖新增标志的取值边界与合并语义，保证 (1,1024)、(1,1025)、TOP_K_ALL 三组边界以及 merge 的 AND 行为有回归保护。

关键符号：Sampler._sample_from_logits, top_k_top_p_min_p_sampling_from_logits_ascend, SamplingBatchInfo.from_schedule_batch, SamplingBatchInfo.merge_batch, TestAscendSamplerDispatch.test_top_k_dispatch_does_not_read_device_values_on_cpu, test_npu_top_k_top_p_eligibility_uses_request_params

关键源码片段

python/sglang/srt/sampling/sampling_batch_info.py

新增主机侧判定字段并在 batch 构造阶段一次性计算，同时定义合并期的 AND 保守语义，是本次优化的语义核心。

# python/sglang/srt/sampling/sampling_batch_info.py
​
@dataclasses.dataclass
class SamplingBatchInfo:
    # ... 其余字段略 ...
    device: str = "cuda"
​
    # Handle logit bias
    logit_bias: Optional[torch.Tensor] = None
​
    # 主机侧（host）判定：本 batch 是否可以使用 torch_npu.npu_top_k_top_p。
    # 判定留在 CPU 上，是为了避免在每 token 的采样路径里做隐式标量同步。
    npu_top_k_top_p_eligible: bool = False
​
    @classmethod
    def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int):
        reqs = batch.reqs
        device = batch.device
        _pin = is_pin_memory_available(device)
        # ... temperatures / top_ps / top_ks / min_ps 等张量构造略 ...
​
        ret = cls(
            temperatures=temperatures,
            top_ps=top_ps,
            top_ks=top_ks,
            min_ps=min_ps,
            sampling_seed=sampling_seed,
            is_all_greedy=all(r.sampling_params.top_k <= 1 for r in reqs),
            is_any_greedy=any(r.sampling_params.top_k <= 1 for r in reqs),
            need_top_p_sampling=any(r.sampling_params.top_p != 1.0 for r in reqs),
            need_top_k_sampling=any(r.sampling_params.top_k != TOP_K_ALL for r in reqs),
            need_min_p_sampling=any(r.sampling_params.min_p > 0 for r in reqs),
            # 只读请求参数：采样过程中 top_k 不会变化，因此在这里算一次就够。
            # TOP_K_ALL（-1）不满足下界，会继续走通用采样路径。
            npu_top_k_top_p_eligible=all(
                1 <= r.sampling_params.top_k <= 1024 for r in reqs
            ),
            vocab_size=vocab_size,
            penalizer_orchestrator=penalizer_orchestrator,
            has_custom_logit_processor=has_custom_logit_processor,
            custom_params=custom_params,
            custom_logit_processor=merged_custom_logit_processor,
            device=device,
            logit_bias=logit_bias,
            return_sampling_masks=return_sampling_masks,
            sampling_mask_batch_indices=sampling_mask_batch_indices,
        )
        ret.adjusted_from_schedule_batch(batch, vocab_size)
        return ret
​
    def merge_batch(self, other: "SamplingBatchInfo"):
        # ... 张量拼接与 sampling_mask 合并略 ...
​
        self.is_all_greedy &= other.is_all_greedy
        self.is_any_greedy |= other.is_any_greedy
        self.need_top_p_sampling |= other.need_top_p_sampling
        self.need_top_k_sampling |= other.need_top_k_sampling
        self.need_min_p_sampling |= other.need_min_p_sampling
        # AND 语义：合并后只要有一个请求的 top_k 落在 [1, 1024] 之外，
        # 整个 batch 就退回通用采样路径（与 need_*_sampling 的 OR 语义相反）。
        self.npu_top_k_top_p_eligible &= other.npu_top_k_top_p_eligible
​
        self.adjusted_merge_batch(other)

python/sglang/srt/layers/sampler.py

采样热路径的实际改动点：把设备侧 torch.all 判定替换为主机标志透传，并新增带保守默认值的函数参数。

# python/sglang/srt/layers/sampler.py
​
# Sampler._sample_from_logits 的 Ascend 分支：把 batch 构造阶段算好的标志透传下去，
# 采样器本身不再接触 top_ks 的设备侧取值范围判定。
batch_next_token_ids = top_k_top_p_min_p_sampling_from_logits_ascend(
    logits,
    sampling_info.top_ks,
    sampling_info.top_ps,
    sampling_info.min_ps,
    sampling_info.need_min_p_sampling,
    sampling_info.sampling_seed,
    positions,
    npu_top_k_top_p_eligible=sampling_info.npu_top_k_top_p_eligible,
)
return batch_next_token_ids.to(torch.int32)
​
​
def top_k_top_p_min_p_sampling_from_logits_ascend(
    logits: torch.Tensor,
    top_ks: torch.Tensor,
    top_ps: torch.Tensor,
    min_ps: torch.Tensor,
    need_min_p_sampling: bool,
    sampling_seed: Optional[torch.Tensor],
    positions: torch.Tensor,
    # 默认 False：任何未显式传入标志的调用点都会退回通用采样路径，
    # 保证「漏传」只带来性能退化，不会把非法 top_k 送进融合算子。
    npu_top_k_top_p_eligible: bool = False,
):
    """A top-k, top-p and min-p sampling implementation for ascend npu with torch_npu interface."""
    # torch_npu.npu_top_k_top_p 要求 top_k 取值范围为 [1, 1024]。
    # 旧实现是 torch.all((top_ks <= 1024) & (top_ks >= 1))：top_ks 已在 NPU 上，
    # torch.all 产生设备标量，在 Python 条件里求值会触发隐式 NPU 到 CPU 同步，
    # 而该分支每个 decode step 都要走一次，等于反复打断异步执行。
    # 现在条件只读主机侧布尔值，同步彻底消失；算子可用性检查保持不变。
    if hasattr(torch_npu, "npu_top_k_top_p") and npu_top_k_top_p_eligible:
        logits_top_k_top_p = torch_npu.npu_top_k_top_p(logits, top_ps, top_ks)
        probs_top_k_top_p = logits_top_k_top_p.softmax(dim=-1)
​
        if need_min_p_sampling:
            min_p_thresholds = probs_top_k_top_p.max(dim=-1) * min_ps
            min_p_mask = probs_top_k_top_p < min_p_thresholds.view(-1, 1)
            probs_top_k_top_p.masked_fill_(min_p_mask, 0.0)
​
        if sampling_seed is None:
            batch_next_token_ids = torch.multinomial(probs_top_k_top_p, num_samples=1)
        else:
            logprobs_top_k_top_p = probs_top_k_top_p.to(torch.float64)
            del probs_top_k_top_p
            logprobs_top_k_top_p.log_()
            batch_next_token_ids = multinomial_with_seed(
                logprobs_top_k_top_p, sampling_seed, positions
            )
    else:
        # 通用路径保持不变。注意这里会就地改写 top_ks（TOP_K_ALL 改写为 vocab size），
        # 主机侧标志在 batch 构造时已定型，因此不会被这种就地改写反向影响。
        probs = torch.softmax(logits, dim=-1)
        probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
        # ... top-k / top-p / min-p 掩码与多项式采样略 ...

test/registered/unit/sampling/test_ascend_sampler.py

新增的分派路径单测，用让 torch.all 抛 AssertionError 的方式建立「不得再读设备标量」的回归护栏，是本 PR 最具借鉴价值的测试写法。

# test/registered/unit/sampling/test_ascend_sampler.py
​
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
​
​
class TestAscendSamplerDispatch(unittest.TestCase):
    def test_top_k_dispatch_does_not_read_device_values_on_cpu(self):
        logits = torch.tensor([[0.1, 0.2, 0.3, 0.4]])
        top_ks = torch.tensor([2], dtype=torch.int32)
        top_ps = torch.ones(1)
        min_ps = torch.zeros(1)
        positions = torch.zeros(1, dtype=torch.int64)
​
        # 分别验证 eligible=True / False 两个分支，避免只覆盖融合路径
        for eligible in (True, False):
            with self.subTest(eligible=eligible):
                npu_top_k_top_p = MagicMock(return_value=logits)
                with (
                    patch.object(
                        sampler_module,
                        "torch_npu",
                        SimpleNamespace(npu_top_k_top_p=npu_top_k_top_p),
                        create=True,
                    ),
                    # 关键回归护栏：如果 Ascend 分派路径再次触碰设备张量判定
                    # （例如重新引入 torch.all(top_ks <= 1024)），这里会直接抛
                    # AssertionError 让测试失败，从机制上阻止同步被重新引入。
                    patch.object(
                        sampler_module.torch,
                        "all",
                        side_effect=AssertionError("device predicate was inspected"),
                    ),
                ):
                    result = (
                        sampler_module.top_k_top_p_min_p_sampling_from_logits_ascend(
                            logits.clone(),
                            top_ks.clone(),
                            top_ps,
                            min_ps,
                            False,
                            None,
                            positions,
                            npu_top_k_top_p_eligible=eligible,
                        )
                    )
​
                self.assertEqual(tuple(result.shape), (1,))
                # 融合算子是否被调用，完全由主机标志决定
                self.assertEqual(npu_top_k_top_p.called, eligible)

评论区精华

本次 PR 的 review comments 为 0，评审意见只有 ping1jing2 的两次空 body APPROVED，实质讨论集中在 Issue 评论区：

作者 xiedeyantu："Two requests have already been verified on the 910B4 machine, and the gen throughput of both requests is basically the same."（附 DeepSeek-V4-Flash 上 top_k=20 与默认参数两条 curl 请求）

作者 xiedeyantu："This CI [PR Test (NPU) / base-b-test-4-npu-a3 / run (0)] failure is unrelated to the change. It occurs during Python module import while initializing the test environment, before any Ascend sampling code or test logic is executed."

可见两点：其一，作者给出的收益证据是「吞吐基本一致」+ 消除同步的理论论证，属于不引入回退的验证而非量化的性能提升证明；其二，CI 多次 rerun-failed-ci 后的失败被归因于测试环境 Python 模块导入阶段，与采样代码无逻辑关联。批次合并的 AND 语义、filter 不提升的保守策略均只在 PR body 中自述，没有 reviewer 交叉质询。

• 优化收益的实测验证方式 (performance): 作者以「吞吐基本一致」作为不回退证据，评审未进一步要求同步计数或延迟量化数据；真实收益仍缺少系统性 benchmark 支撑。
• CI 失败归因与重跑 (testing): 归因为测试环境导入问题，与本变更无逻辑关联；NPU 相关用例未提供真实设备上的采样分派验证结论。
• 批次生命周期中的保守资格判定 (design): 作者明确选择保守方向，牺牲部分性能换取「绝不把不支持的 top_k 送入融合算子」的正确性保证；该设计被评审接受但未经交叉验证。

风险与影响

• 风险：
  1. 就地改写与固定标志的交互：通用分支中 top_ks.masked_fill_(topk_all_mask, probs.shape[1]) 会把 TOP_K_ALL 就地改成 vocab size，而主机标志在 batch 构造时已定型、之后不再重算。方向是安全的（不会被「改写成合法区间」后误判为可用），但也没有任何路径能把标志从 False 提升为 True。
  2. filter_batch 不重算导致长期保守：若批次中仅个别请求的 top_k 越界，只要它们还在 batch 内融合路径就不可用；被过滤后依然不恢复，长生命周期 batch 可能长期停留在通用采样路径。影响限于 Ascend 侧性能，不影响输出正确性。
  3. 收益缺少端到端验证：新增单测注册在 CPU 套件（base-a-test-cpu），只能验证分派分支走向，无法度量真实同步消除；作者仅在 910B4 上用两条请求做手工验证，样本量小，且「吞吐基本一致」意味着收益更偏延迟与异步流水特性，而非吞吐提升，后续若要宣称收益需要更系统的 benchmark。
  4. 回归护栏覆盖面有限：patch.object(sampler_module.torch, "all", side_effect=AssertionError(...)) 只能拦住重新引入 torch.all 的写法；若未来改用 .item()、torch.min、torch.any 等方式读回设备标量，测试不会失败。
  5. 共享结构体的影响面：SamplingBatchInfo 被所有后端共用，新字段会参与 dataclasses.replace/copy_for_forward 与 merge 语义；默认 False 使非 Ascend 路径成为 no-op，但后续任何新增构造或复制路径都需记得维护该标志，否则 Ascend 侧会静默退化。
     - 影响：影响范围集中在 Ascend/NPU 部署（如 910B 系列），且无需任何配置即可生效。对用户而言，采样语义（top-k、top-p、min-p、概率归一化、多项式采样、sampling_seed 行为）完全不变，不存在采样结果回归或参数兼容性问题。对系统而言，收益是移除了每 decode step 一次的设备与主机同步点，有利于 NPU 异步执行流水与采样延迟的稳定性；但由于实测吞吐持平，短期不宜期待明显吞吐提升。对团队而言，本 PR 沉淀了一个可复用范式：把「算子可用性/取值范围」这类静态判定前移到 batch 构造阶段、危险默认值取保守方向、合并用 AND 语义，并配一条以 mock 抛异常实现的防回归护栏；该模式可推广到其他非 CUDA 后端的分派门槛校验。评审强度偏低（两次空 body APPROVED、0 条 review comments）意味着这些语义主要靠作者自证，读者需要自行确认。
     - 风险标记：热路径同步语义变更, 保守标志可能长期停留通用路径, 缺少真实 NPU 端到端性能验证, 回归护栏仅覆盖 torch.all 写法, 共享结构体新增字段需维护

关联脉络

• PR #37134 [ROCm] Fix EAGLE spec-decode verify silently sampling greedy on HIP: 同属非 CUDA 加速器后端的采样分派治理：一个修复 HIP 上 verify 静默贪心（走错分支），一个消除 Ascend 采样判定引入的设备同步（分支判定方式本身的成本），两者都围绕采样后端分派路径的正确性与效率。
• PR #39561 [NPU] [bugfix] Fix undefined require_attn_tp_allgather: 同为 NPU/Ascend 后端路径的持续加固，反映该后端从「能跑通」向「热路径稳定可控」推进的同一演进线。
• PR #38402 fix(npu): fix hybrid KV transfer with PP prefill in PD disaggregation: 同为 Ascend 侧改动并附带注册单测（NPU/PD 场景），与本 PR 一样强调在 NPU 上保持行为等价与测试可回归。
• PR #39000 [Router] Fleet-wide sampling contract 1/3: the config surface: 概念上有交集：那条 PR 把采样参数约束下沉为路由侧配置面，本 PR 则是把同一个 top_k 范围约束在引擎侧前移到 CPU 计算，属于「采样契约」在不同层次上的实现。