Support data parallel attention in dump comparator (#19602)

This commit is contained in:
fzyzcjy
2026-03-01 10:51:21 +08:00
committed by GitHub
parent ea6ff7b01f
commit e64095c3c7
19 changed files with 783 additions and 325 deletions

View File

@@ -2074,6 +2074,33 @@ class TestDumperE2E:
assert "rank" in loaded["meta"]
assert "step" in loaded["meta"]
par = loaded["meta"].get("sglang_parallel_info", {})
expected_keys = [
"tp_rank",
"tp_size",
"pp_rank",
"pp_size",
"moe_ep_rank",
"moe_ep_size",
"moe_tp_rank",
"moe_tp_size",
"moe_dp_rank",
"moe_dp_size",
"enable_dp_attention",
"attn_tp_rank",
"attn_tp_size",
"attn_dp_rank",
"attn_dp_size",
"local_attn_dp_rank",
"local_attn_dp_size",
"attn_cp_rank",
"attn_cp_size",
]
for key in expected_keys:
assert (
key in par
), f"Missing {key} in sglang_parallel_info, got: {sorted(par)}"
rids_files = [f for f in dump_files if "name=rids" in f.name]
rids_loaded = torch.load(
rids_files[0], map_location="cpu", weights_only=False