执行摘要
- 一句话:修复KV校验测试因Radix缓存去重导致的flaky问题
- 推荐动作:值得精读PR body中的根因分析,它揭示了
cache_unfinished_req与send_kv_chunk之间的时序竞态如何导致去重后的槽位被错误释放,是理解PD架构中KV传输、Radix缓存和canary验证三者交互的绝佳案例。
功能与动机
test_self_e2e_pd_perturb.py在CI中flaky:P侧扰动触发成功,但D侧从未报告预期的verify_real_kv_hash违例。根因是并行请求共享相同prompt,Prefill端RadixCache.cache_unfinished_req在send_kv_chunk之前去重了KV槽位,导致被扰动的非规范副本在传输前就被释放。PR body详细描述了根因分析过程。
实现拆解
- 在
pd_fixture.py的send_parallel_short_requests方法新增distinct_prompts参数:当为True时,为每个请求拼接"{i} {nonce} {_SHORT_PROMPT_BODY}",其中nonce为uuid.uuid4().hex[:8];否则保持原有行为。
- 在
test_self_e2e_pd_perturb.py的测试方法中,调用self.send_parallel_short_requests(n=4, distinct_prompts=True)代替原调用,并添加详细注释说明原因。
- 回滚了#27425中禁用该测试的提交,重新启用CI中的测试。
关键文件:
python/sglang/test/kv_canary/pd_fixture.py(模块 测试夹具;类别 test;类型 test-coverage;符号 send_parallel_short_requests): 新增distinct_prompts参数,允许测试用例生成差异化prompt以避免Radix去重
test/registered/kv_canary/test_self_e2e_pd_perturb.py(模块 端到端测试;类别 test;类型 test-coverage;符号 _PDPerturbBase.test_p_side_perturb_surfaces_real_kv_hash_violation_on_decode_side): 使用distinct_prompts=True调用send_parallel_short_requests并添加注释说明根本原因
关键符号:send_parallel_short_requests, test_p_side_perturb_surfaces_real_kv_hash_violation_on_decode_side
关键源码片段
python/sglang/test/kv_canary/pd_fixture.py
新增distinct_prompts参数,允许测试用例生成差异化prompt以避免Radix去重
def send_parallel_short_requests(
self,
n: int,
*,
assert_all_success: bool = True,
max_new_tokens: int = 100,
timeout: float = 60.0,
distinct_prompts: bool = False, # new parameter
) -> list[dict]:
if distinct_prompts:
# 每个请求使用不同的前缀(请求索引 + 随机 nonce),
# 使得请求间除 BOS token 外没有共享前缀,
# 从而避免 RadixCache.cache_unfinished_req 进行去重。
nonce = uuid.uuid4().hex[:8]
prompts = [f"{i} {nonce} {_SHORT_PROMPT_BODY}" for i in range(n)]
else:
prompts = [_SHORT_PROMPT_BODY] * n
results = post_parallel_generate(
url=self.lb_url + "/generate",
prompts=prompts,
max_new_tokens=max_new_tokens,
timeout=timeout,
)
if assert_all_success:
for result in results:
self.assertEqual(result.get("status_code"), 200, result)
return results
test/registered/kv_canary/test_self_e2e_pd_perturb.py
使用distinct_prompts=True调用send_parallel_short_requests并添加注释说明根本原因
def test_p_side_perturb_surfaces_real_kv_hash_violation_on_decode_side(
self,
) -> None:
# distinct_prompts 对于可靠触发违例是必须的:
# 当请求共享同一个 prompt 时,P 端 radix 缓存会在 cache_unfinished_req 中
# 将每个请求的 req_to_token 行改写为第一个插入的(规范)副本的槽位,
# 这发生在 send_kv_chunk 快照索引之前。因此,如果扰动落在被去重的副本槽位上,
# 该槽位会被释放而不会被传输或重新验证,导致两侧均无法报告违例。
self.send_parallel_short_requests(n=4, distinct_prompts=True)
# D 端:第一个 decode forward 重新验证已传输的前缀槽位,
# 因此扰动必须作为 real_kv_hash 违例出现。
self.assert_per_forward_violation_reported(
fail_reason="verify_real_kv_hash",
target_group=self.target_group,
side="decode",
flush_wait_seconds=4.0,
)
# P 端:扰动发生在 prefill forward 的 TAIL 之后,
# 而 PD prefill 不会在 P 上运行另一个 forward 来验证被扰动的槽位,
# 因此 P 必须保持静默(无假阳性违例)。
self.assert_no_violation(side="prefill", wait_seconds=0.5)
评论区精华
无人工review评论,仅有gemini-code-assist的自动评论,但未提供实质性反馈。PR作者fzyzcjy通过详尽的PR body和commit history独立完成了根因分析、方案演进和验证。
风险与影响
- 风险:该变更仅影响测试代码,不涉及生产逻辑。使用distinct_prompts后,Radix缓存依然启用(与生产一致),但请求之间不再共享前缀,因此Radix去重路径仅在
test_self_e2e_pd_baseline.py中覆盖(该测试仍使用相同prompt)。如果未来生产代码的Radix去重行为发生变化,该测试可能无法检测到相关回归;但已有基线测试覆盖。
- 影响:影响范围仅限于
test_self_e2e_pd_perturb.py及其依赖的pd_fixture.py。修复后测试稳定性显著提升:在2-GPU H200上15/15轮全文件通过、30/30轮子类专项通过;在原始flaky环境2-gpu-h100上4/5轮通过(失败为不相关的基础设施问题)。
- 风险标记:测试覆盖调整
关联脉络
- PR #27425 Temporarily disable test_self_e2e_pd_perturb in CI: 前序PR临时禁用了该测试,本PR重新启用并修复了flaky问题
参与讨论