Support partial tensors waiting for reduction and pipeline parallel in dump comparator (#19595)

python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py

@@ -6,6 +6,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     ConcatParams,
     CpThdConcatParams,
     PickParams,
+    ReduceSumParams,
     UnsharderParams,
     UnsharderPlan,
 )
@@ -63,6 +64,14 @@ def _apply_unshard(
             seq_lens_per_rank=params.seq_lens_per_rank,
         )
 
+    if isinstance(params, ReduceSumParams):
+        stripped: list[torch.Tensor] = [t.rename(None) for t in ordered_tensors]
+        result: torch.Tensor = torch.stack(stripped).sum(dim=0)
+        names: tuple[Optional[str], ...] = ordered_tensors[0].names
+        if names[0] is not None:
+            result = result.refine_names(*names)
+        return result
+
     raise ValueError(f"Unsupported unshard operation: {type(params).__name__}")
 
 

python/sglang/srt/debug_utils/comparator/aligner/unsharder/planner.py

@@ -6,6 +6,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
     ConcatParams,
     CpThdConcatParams,
     PickParams,
+    ReduceSumParams,
     UnsharderParams,
     UnsharderPlan,
 )
@@ -155,9 +156,7 @@ def _resolve_unshard_params(
     thd_global_seq_lens: Optional[list[int]] = None,
 ) -> UnsharderParams:
     if spec.reduction is not None:
-        raise NotImplementedError(
-            f"Unshard for reduction={spec.reduction} not yet implemented (Phase 2)"
-        )
+        return ReduceSumParams()
 
     if (
         spec.name == TOKEN_DIM_NAME

python/sglang/srt/debug_utils/comparator/aligner/unsharder/types.py

@@ -38,8 +38,12 @@ class PickParams(_FrozenBase):
     op: Literal["pick"] = "pick"
 
 
+class ReduceSumParams(_FrozenBase):
+    op: Literal["reduce_sum"] = "reduce_sum"
+
+
 UnsharderParams = Annotated[
-    Union[ConcatParams, CpThdConcatParams, PickParams],
+    Union[ConcatParams, CpThdConcatParams, PickParams, ReduceSumParams],
     Field(discriminator="op"),
 ]
 

python/sglang/srt/debug_utils/source_patcher/code_patcher.py

@@ -118,7 +118,7 @@ def _inject_preamble(*, config: PatchConfig, extra_imports: list[str]) -> PatchC
     for spec in config.patches:
         existing: str = spec.preamble
         combined: str = (
-            existing + "\n" + import_block if existing.strip() else import_block
+            import_block + "\n" + existing if existing.strip() else import_block
         )
         new_patches.append(
             PatchSpec(target=spec.target, edits=spec.edits, preamble=combined)