diff --git a/python/pyproject.toml b/python/pyproject.toml index cbf19703a..2341abb28 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -139,6 +139,7 @@ test = [ "pandas", "parameterized", "peft>=0.18.0", + "polars", "pytest", "pytest-cov", "diff-cover", diff --git a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py index b16c13982..98ef9d7a0 100644 --- a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py +++ b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py @@ -120,7 +120,9 @@ patches: # All sub-axes (attn_tp, moe_tp, attn_dp) are uniquely determined by tp_rank, # so only tp:replicated is needed — sub-axes are auto-resolved as implicitly replicated. # - # Attn tensors are NOT TP-sharded, mlp_output is already all-reduced. + # Attn tensors are NOT TP-sharded (attn_tp_size=1). + # mlp_output is still moe_tp:partial — the reduce-scatter happens in + # postprocess_layer(), after the dump point. # layer_input is dumped after prepare_attn which DP-distributes tokens, # so it needs dp:=attn_dp to filter to the non-empty DP rank. - target: sglang.srt.models.qwen3_moe.Qwen3MoeDecoderLayer.forward @@ -152,7 +154,7 @@ patches: hidden_states = self.mlp( hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter ) - append: "dumper.dump('mlp_output', hidden_states, dims='t h # tp:replicated')" + append: "dumper.dump('mlp_output', hidden_states, dims='t h[moe_tp:partial] # tp:replicated')" # --- attention internals --- - target: sglang.srt.models.qwen3_moe.Qwen3MoeAttention.forward_core @@ -184,14 +186,33 @@ class TestSourcePatcherE2ESGLang: """TP=2 baseline vs TP=2+DP=2+dp-attention target. In dp-attention mode (attn_tp_size=1, attn_dp_size=2), attention - tensors are NOT TP-sharded and mlp_output is already all-reduced. - A separate patch config with corrected dims is used for the target. + tensors are NOT TP-sharded and mlp_output is still moe_tp:partial + (the reduce-scatter happens in postprocess_layer, after the dump + point). A separate patch config with corrected dims is used for + the target. + + Comparison is limited to step 0 (prefill) because the decode + step has tokens on both DP ranks, which breaks the dp:=attn_dp + single-rank assumption and causes comparator errors. + + mlp_output is allowed to fail because the FusedMoE dispatcher + combine path may include an implicit all-reduce that makes the + dumped value differ from the raw partial expert output. All + other tensors (layer_input, attn_output, attn_pre_o_proj, + pre_mlp_residual, moe_router_logits, moe_expert_output) must + pass at step 0. """ _run_e2e_scenario( tmp_path=tmp_path, target_tp=BASELINE_TP, extra_target_server_args=["--dp", "2", "--enable-dp-attention"], target_patch_config_yaml=PATCH_CONFIG_DP_ATTENTION_YAML, + extra_comparator_args=[ + "--end-step", + "0", + "--allow-failed-pattern", + "mlp_output", + ], ) @@ -204,6 +225,7 @@ def _run_e2e_scenario( target_tp: int, extra_target_server_args: Optional[list[str]] = None, target_patch_config_yaml: Optional[str] = None, + extra_comparator_args: Optional[list[str]] = None, ) -> None: """Full e2e: write patch config -> baseline run -> target run -> compare.""" base_url: str = DEFAULT_URL_FOR_TEST @@ -249,6 +271,8 @@ def _run_e2e_scenario( "--allow-skipped-pattern", "input_ids|positions", ] + if extra_comparator_args: + cmd.extend(extra_comparator_args) result: subprocess.CompletedProcess[str] = subprocess.run( cmd,