Enhance sglang engine dumping tests in dump comparator (#19681)

test/registered/debug_utils/comparator/conftest.py

@@ -1,7 +1,30 @@
+import warnings
+
+warnings.filterwarnings(
+    "ignore", message="builtin type Swig.*", category=DeprecationWarning
+)
+
 import pytest
 
 from sglang.srt.debug_utils.comparator.output_types import report_sink
 
+collect_ignore_glob: list[str] = []
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    config.addinivalue_line(
+        "filterwarnings",
+        "ignore:Unknown config option. asyncio_mode:pytest.PytestConfigWarning",
+    )
+    config.addinivalue_line(
+        "filterwarnings",
+        "ignore:builtin type Swig.*:DeprecationWarning",
+    )
+    config.addinivalue_line(
+        "filterwarnings",
+        "ignore:Named tensors and all their associated APIs:UserWarning",
+    )
+
 
 @pytest.fixture(autouse=True)
 def _reset_report_sink() -> None:

test/registered/debug_utils/test_engine_dumper_comparator_e2e.py

@@ -1,9 +1,13 @@
 """E2E test: source patcher + dumper + comparator on SGLang server.
 
-Patches Qwen3DecoderLayer.forward to insert dumper.dump() calls,
+Patches Qwen3MoeDecoderLayer.forward (and related methods) to insert
-launches 1-GPU baseline and 2-GPU TP=2 target servers, runs inference,
+dumper.dump() calls at 7 points, launches servers with Qwen3-30B-A3B
-verifies patched dump fields exist, then runs comparator to verify
+(MOE model), runs inference, verifies patched dump fields exist, then
-numerical consistency.
+runs comparator to verify numerical consistency.
+
+Test cases:
+- test_patch_dump_and_compare: TP=2 baseline vs TP=4 target
+- test_dp_attention: TP=2 baseline vs TP=2+DP=2+dp-attention target
 
 The dumper.apply_source_patches() auto-injects ``from ... import dumper``
 so the YAML only needs ``dumper.dump(...)`` calls.
@@ -13,15 +17,15 @@ import os
 import subprocess
 import tempfile
 from pathlib import Path
+from typing import Optional
 
 import pytest
 import requests
 
-from sglang.srt.debug_utils.comparator.output_types import (
+pytestmark = pytest.mark.filterwarnings(
-    AnyRecord,
+    "ignore:Unknown config option. asyncio_mode:pytest.PytestConfigWarning",
-    SummaryRecord,
-    parse_record_json,
 )
+
 from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import (
@@ -30,115 +34,233 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=120, suite="nightly-2-gpu", nightly=True)
+register_cuda_ci(est_time=300, suite="nightly-4-gpu", nightly=True)
 
-MODEL = "Qwen/Qwen3-0.6B"
+MODEL = "Qwen/Qwen3-30B-A3B"
+BASELINE_TP = 2
+TARGET_TP = 4
 EXP_NAME = "e2e_source_patcher"
 DUMPER_FILTER = "layer_id in [0, 1, 2]"
 
+_FIELDS_TO_VERIFY: list[str] = [
+    # decoder layer level (aligned with miles)
+    "layer_input",
+    "attn_output",
+    "pre_mlp_residual",
+    "mlp_output",
+    # attention internals
+    "attn_pre_o_proj",
+    # moe internals
+    "moe_router_logits",
+    "moe_expert_output",
+]
+
 PATCH_CONFIG_YAML: str = """\
 patches:
-  - target: sglang.srt.models.qwen3.Qwen3DecoderLayer.forward
+  # --- decoder layer level (aligned with miles test) ---
+  - target: sglang.srt.models.qwen3_moe.Qwen3MoeDecoderLayer.forward
     edits:
-      - match: "hidden_states = self.mlp(hidden_states)"
+      - match: |
-        prepend: "dumper.dump('patched_attn_output', hidden_states, dims='t h')"
+          hidden_states, residual = (
-      - match: "return hidden_states, residual"
+              self.layer_communicator.prepare_attn_and_capture_last_layer_outputs(
-        prepend: "dumper.dump('patched_mlp_output', hidden_states, dims='t h')"
+                  hidden_states,
+                  residual,
+                  forward_batch,
+                  captured_last_layer_outputs=captured_last_layer_outputs,
+                  **kwargs,
+              )
+          )
+        append: "dumper.dump('layer_input', hidden_states, dims='t h # tp:replicated')"
+      - match: |
+          hidden_states = self.self_attn(
+              positions=positions,
+              hidden_states=hidden_states,
+              forward_batch=forward_batch,
+          )
+        append: "dumper.dump('attn_output', hidden_states, dims='t h[tp:partial]')"
+      - match: |
+          hidden_states, residual = self.layer_communicator.prepare_mlp(
+              hidden_states, residual, forward_batch
+          )
+        append: "dumper.dump('pre_mlp_residual', hidden_states, dims='t h # tp:replicated')"
+      - match: |
+          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:partial]')"
+
+  # --- attention internals ---
+  - target: sglang.srt.models.qwen3_moe.Qwen3MoeAttention.forward_core
+    edits:
+      - match: "output, _ = self.o_proj(attn_output)"
+        prepend: "dumper.dump('attn_pre_o_proj', attn_output, dims='t attn_h[tp]')"
+
+  # --- moe internals ---
+  - target: sglang.srt.models.qwen3_moe.Qwen3MoeSparseMoeBlock.forward_normal
+    edits:
+      - match: "router_logits, _ = self.gate(hidden_states)"
+        append: "dumper.dump('moe_router_logits', router_logits, dims='t num_experts # tp:replicated')"
+      - match: "final_hidden_states = self.experts(hidden_states, topk_output)"
+        append: "dumper.dump('moe_expert_output', final_hidden_states, dims='t h[tp:partial]')"
+"""
+
+PATCH_CONFIG_DP_ATTENTION_YAML: str = """\
+patches:
+  # --- decoder layer level (aligned with miles test) ---
+  # In dp-attention mode: attn tensors are NOT TP-sharded (attn_tp_size=1),
+  # and mlp_output is already all-reduced inside forward_normal().
+  # 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
+    edits:
+      - match: |
+          hidden_states, residual = (
+              self.layer_communicator.prepare_attn_and_capture_last_layer_outputs(
+                  hidden_states,
+                  residual,
+                  forward_batch,
+                  captured_last_layer_outputs=captured_last_layer_outputs,
+                  **kwargs,
+              )
+          )
+        append: "dumper.dump('layer_input', hidden_states, dims='t h # tp:replicated dp:=attn_dp')"
+      - match: |
+          hidden_states = self.self_attn(
+              positions=positions,
+              hidden_states=hidden_states,
+              forward_batch=forward_batch,
+          )
+        append: "dumper.dump('attn_output', hidden_states, dims='t h # tp:replicated')"
+      - match: |
+          hidden_states, residual = self.layer_communicator.prepare_mlp(
+              hidden_states, residual, forward_batch
+          )
+        append: "dumper.dump('pre_mlp_residual', hidden_states, dims='t h # tp:replicated')"
+      - match: |
+          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')"
+
+  # --- attention internals ---
+  - target: sglang.srt.models.qwen3_moe.Qwen3MoeAttention.forward_core
+    edits:
+      - match: "output, _ = self.o_proj(attn_output)"
+        prepend: "dumper.dump('attn_pre_o_proj', attn_output, dims='t attn_h # tp:replicated')"
+
+  # --- moe internals ---
+  - target: sglang.srt.models.qwen3_moe.Qwen3MoeSparseMoeBlock.forward_normal
+    edits:
+      - match: "router_logits, _ = self.gate(hidden_states)"
+        append: "dumper.dump('moe_router_logits', router_logits, dims='t num_experts # tp:replicated')"
+      - match: "final_hidden_states = self.experts(hidden_states, topk_output)"
+        append: "dumper.dump('moe_expert_output', final_hidden_states, dims='t h[tp:partial]')"
 """
 
 
 class TestSourcePatcherE2ESGLang:
-    """E2E: patch Qwen3 forward -> dump -> compare 1gpu vs 2gpu-tp2."""
+    """E2E: patch Qwen3Moe forward -> dump -> compare."""
 
-    @pytest.mark.timeout(300)
     def test_patch_dump_and_compare(self, tmp_path: Path) -> None:
-        patched_fields: list[str] = ["patched_attn_output", "patched_mlp_output"]
+        """TP=2 baseline vs TP=4 target."""
-        base_url: str = DEFAULT_URL_FOR_TEST
+        _run_e2e_scenario(
-
+            tmp_path=tmp_path,
-        config_path: Path = tmp_path / "patch_config.yaml"
+            target_tp=TARGET_TP,
-        config_path.write_text(PATCH_CONFIG_YAML)
-
-        # Run 1: baseline (1 GPU)
-        baseline_dir: Path = tmp_path / "baseline"
-        _run_server_and_generate(
-            dump_dir=baseline_dir,
-            config_path=config_path,
-            tp=1,
-            base_url=base_url,
-        )
-        _verify_patched_fields(dump_dir=baseline_dir, field_names=patched_fields)
-
-        # Run 2: target (2 GPU TP=2)
-        target_dir: Path = tmp_path / "target"
-        _run_server_and_generate(
-            dump_dir=target_dir,
-            config_path=config_path,
-            tp=2,
-            base_url=base_url,
-        )
-        _verify_patched_fields(dump_dir=target_dir, field_names=patched_fields)
-
-        # Compare baseline vs target
-        baseline_exp: Path = baseline_dir / EXP_NAME
-        target_exp: Path = target_dir / EXP_NAME
-
-        result: subprocess.CompletedProcess[str] = subprocess.run(
-            [
-                "python",
-                "-m",
-                "sglang.srt.debug_utils.comparator",
-                "--baseline-path",
-                str(baseline_exp),
-                "--target-path",
-                str(target_exp),
-                "--output-format",
-                "json",
-                "--grouping",
-                "logical",
-            ],
-            capture_output=True,
-            text=True,
         )
 
-        debug_file: Path = _save_comparator_output(
+    def test_dp_attention(self, tmp_path: Path) -> None:
-            stdout=result.stdout, stderr=result.stderr
+        """TP=2 baseline vs TP=2+DP=2+dp-attention target.
-        )
-        print(f"Comparator debug output: {debug_file}")
 
-        assert result.returncode == 0, (
+        In dp-attention mode (attn_tp_size=1, attn_dp_size=2), attention
-            f"Comparator failed (rc={result.returncode}). "
+        tensors are NOT TP-sharded and mlp_output is already all-reduced.
-            f"Debug output: {debug_file}"
+        A separate patch config with corrected dims is used for the target.
-        )
+        """
-
+        _run_e2e_scenario(
-        records: list[AnyRecord] = [
+            tmp_path=tmp_path,
-            parse_record_json(line)
+            target_tp=BASELINE_TP,
-            for line in result.stdout.strip().splitlines()
+            extra_target_server_args=["--dp", "2", "--enable-dp-attention"],
-            if line.strip()
+            target_patch_config_yaml=PATCH_CONFIG_DP_ATTENTION_YAML,
-        ]
-        assert (
-            len(records) > 0
-        ), f"Comparator produced no output records. Debug: {debug_file}"
-
-        summary: SummaryRecord = _find_summary(records=records, debug_file=debug_file)
-        assert (
-            summary.passed > 0
-        ), f"No comparisons passed (total={summary.total}). Debug: {debug_file}"
-        assert summary.failed == 0, (
-            f"{summary.failed} comparisons failed "
-            f"(passed={summary.passed}, skipped={summary.skipped}). "
-            f"Debug: {debug_file}"
         )
 
 
 # --------------------------------- helpers ---------------------------------
 
 
+def _run_e2e_scenario(
+    *,
+    tmp_path: Path,
+    target_tp: int,
+    extra_target_server_args: Optional[list[str]] = None,
+    target_patch_config_yaml: Optional[str] = None,
+) -> None:
+    """Full e2e: write patch config -> baseline run -> target run -> compare."""
+    base_url: str = DEFAULT_URL_FOR_TEST
+
+    baseline_config_path: Path = tmp_path / "patch_config.yaml"
+    baseline_config_path.write_text(PATCH_CONFIG_YAML)
+
+    target_config_path: Path = tmp_path / "patch_config_target.yaml"
+    target_config_path.write_text(target_patch_config_yaml or PATCH_CONFIG_YAML)
+
+    baseline_dir: Path = tmp_path / "baseline"
+    _run_server_and_generate(
+        dump_dir=baseline_dir,
+        config_path=baseline_config_path,
+        tp=BASELINE_TP,
+        base_url=base_url,
+    )
+    _verify_patched_fields(dump_dir=baseline_dir, field_names=_FIELDS_TO_VERIFY)
+
+    target_dir: Path = tmp_path / "target"
+    _run_server_and_generate(
+        dump_dir=target_dir,
+        config_path=target_config_path,
+        tp=target_tp,
+        base_url=base_url,
+        extra_server_args=extra_target_server_args,
+    )
+    _verify_patched_fields(dump_dir=target_dir, field_names=_FIELDS_TO_VERIFY)
+
+    baseline_exp: Path = baseline_dir / EXP_NAME
+    target_exp: Path = target_dir / EXP_NAME
+
+    cmd: list[str] = [
+        "python",
+        "-m",
+        "sglang.srt.debug_utils.comparator",
+        "--baseline-path",
+        str(baseline_exp),
+        "--target-path",
+        str(target_exp),
+        "--output-format",
+        "json",
+        "--allow-skipped-pattern",
+        "input_ids|positions",
+    ]
+
+    result: subprocess.CompletedProcess[str] = subprocess.run(
+        cmd,
+        capture_output=True,
+        text=True,
+    )
+
+    debug_file: Path = _save_comparator_output(
+        stdout=result.stdout, stderr=result.stderr
+    )
+    print(f"Comparator debug output: {debug_file}")
+
+    assert result.returncode == 0, (
+        f"Comparator failed (rc={result.returncode}). " f"Debug output: {debug_file}"
+    )
+
+
 def _run_server_and_generate(
     *,
     dump_dir: Path,
     config_path: Path,
     tp: int,
     base_url: str,
+    extra_server_args: Optional[list[str]] = None,
 ) -> None:
     """Launch SGLang server with source patcher + dumper, send a generate request."""
     env: dict[str, str] = {
@@ -149,11 +271,24 @@ def _run_server_and_generate(
         "DUMPER_SERVER_PORT": "reuse",
     }
 
+    server_args: list[str] = [
+        "--tp",
+        str(tp),
+        "--max-total-tokens",
+        "128",
+        "--mem-fraction-static",
+        "0.5",
+        "--disable-cuda-graph",
+        "--disable-radix-cache",
+    ]
+    if extra_server_args:
+        server_args.extend(extra_server_args)
+
     proc = popen_launch_server(
         MODEL,
         base_url,
         timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-        other_args=["--tp", str(tp), "--max-total-tokens", "128"],
+        other_args=server_args,
         env=env,
     )
     try:
@@ -170,7 +305,7 @@ def _run_server_and_generate(
             f"{base_url}/generate",
             json={
                 "text": "The capital of France is",
-                "sampling_params": {"max_new_tokens": 8},
+                "sampling_params": {"max_new_tokens": 1, "temperature": 0},
             },
         )
         assert resp.status_code == 200, f"Generate failed: {resp.text}"
@@ -188,19 +323,6 @@ def _verify_patched_fields(*, dump_dir: Path, field_names: list[str]) -> None:
         )
 
 
-def _find_summary(*, records: list[AnyRecord], debug_file: Path) -> SummaryRecord:
-    """Extract the SummaryRecord from comparator output."""
-    summaries: list[SummaryRecord] = [
-        r for r in records if isinstance(r, SummaryRecord)
-    ]
-    assert len(summaries) == 1, (
-        f"Expected 1 summary record, got {len(summaries)}. "
-        f"Record types: {[type(r).__name__ for r in records]}. "
-        f"Debug: {debug_file}"
-    )
-    return summaries[0]
-
-
 def _save_comparator_output(*, stdout: str, stderr: str) -> Path:
     """Save comparator stdout+stderr to a temp file that persists for debugging."""
     fd, path_str = tempfile.mkstemp(prefix="comparator_e2e_", suffix=".log", dir="/tmp")