执行摘要

• 一句话：修复采样掩码测试的断言逻辑，允许 top-k 截止概率并列，消除 CI 中的 flaky test。
• 推荐动作：这是一个值得精读的测试修复示例。它清晰地展示了如何基于对底层系统（此处为 FlashInfer 采样器）行为的正确理解来修复测试断言，而不是简单地放宽或删除断言。这种“让断言匹配正确行为”的方法值得借鉴。对于关注测试稳定性和采样逻辑的工程师，此 PR 具有参考价值。

功能与动机

该测试自 PR #36630 合并后在 main 分支上不稳定（例如：11 not less than or equal to 10）。根本原因是采样器（flashinfer_joint_cutoff_ties_match_capture）的行为：它会保留所有概率等于第 k 大概率的 token（即处理并列情况），因此 len(mask) <= top_k 这个不变量并不总是成立。需要调整断言以匹配采样器的实际语义。

实现拆解

1. 变更范围：仅修改测试文件 test/registered/sampling/test_sampling_mask.py 中的 test_sampling_mask_matches_topk_logprobs 方法，不涉及任何生产代码。
2. 断言逻辑重构：移除了 self.assertLessEqual(len(mask_set), top_k) 这一断言。新增逻辑：首先计算出 top-k 截止概率 top_k_cutoff（排序后的第 k 大概率值），然后遍历掩码中的每个 token，使用带 0.1% 容差的断言 self.assertGreaterEqual(probs[token_id], top_k_cutoff * (1 - 1e-3))，确保所有 token 都处于或高于截止概率。此容差是为了处理内核计算概率与日志概率转换之间的微小数值差异。
3. 文档注释更新：同步更新了方法文档字符串，明确说明掩码可能因并列情况而超过 top_k，并将检查项描述从“top-k-bounded mask”修正为“returned mask”和“at or above the top-k cutoff”。
4. 无配套改动：未涉及配置、schema 或部署变更。修复验证通过作者手动重跑 10 次测试全部通过。

关键文件：

• test/registered/sampling/test_sampling_mask.py（模块 采样测试；类别 test；类型 test-coverage；符号 test_sampling_mask_matches_topk_logprobs）: 本次 PR 唯一修改的文件，包含了需要修复的 flaky test 方法。

关键符号：test_sampling_mask_matches_topk_logprobs

关键源码片段

test/registered/sampling/test_sampling_mask.py

本次 PR 唯一修改的文件，包含了需要修复的 flaky test 方法。

# 从修改后的 test_sampling_mask_matches_topk_logprobs 方法中提取关键断言逻辑。
# 该逻辑确保采样掩码中的每个 token 的概率都接近或超过 top-k 截止概率。
​
mask_set = set(mask)
​
self.assertIn(output_id, mask_set) # 采样出的 token 必须在掩码中
self.assertTrue(mask_set.issubset(probs)) # 掩码 token 必须在返回的对数概率集中
​
# 计算 top-k 截止概率：排序后第 k 大的概率值。
top_k_cutoff = sorted(probs.values(), reverse=True)[top_k - 1]
for token_id in mask_set:
    # 1e-3 的容差：内核使用自己的概率计算，可能与从对数概率转换来的值有微小差异。
    # 断言每个掩码 token 的概率 >= 截止概率 * (1 - 0.001)。
    self.assertGreaterEqual(probs[token_id], top_k_cutoff * (1 - 1e-3))

评论区精华

本 PR 没有正式的 review 评论。讨论主要体现在 PR body 中：

• 作者详细解释了 flaky test 的根因，引用了仓库中另一个测试用例 test_flashinfer_joint_cutoff_ties_match_capture 作为采样器行为的证据。

• 提供了 10 个成功重跑的 CI 链接作为修复有效性的证据。这种主动的稳定性验证是处理 flaky test 的良好实践。

• 暂无高价值评论线程

风险与影响

• 风险：风险极低。此变更仅限于测试逻辑，不影响任何生产代码、调度逻辑或采样行为。唯一风险是新的断言逻辑可能因引入 0.1% 的容差而掩盖极小概率的、真正的逻辑错误，但鉴于内核与 Python 概率计算可能的浮点精度差异，此容差是合理且必要的。不会引入回归、性能、安全或兼容性问题。
• 影响：直接影响是消除了 CI 流水线中的一个不稳定测试源，提升了开发体验和 CI 可靠性。间接影响是确保了采样掩码相关功能的测试覆盖率更准确、更健壮。对用户和系统无直接影响。
• 风险标记：测试断言不准确, CI 间歇性失败

关联脉络

• PR #36630 Unknown (根据上下文推断): PR body 明确指出本 PR 修复的 flaky test 自 PR #36630 合并后开始出现，表明 #36630 的变更引入了导致该测试失败的条件。