Add a debug toggle for selectively reverting PR fixes (#26802)

python/sglang/srt/debug_utils/pr_fix_toggle.py

@@ -0,0 +1,86 @@
+"""Reverse-apply historical PR fixes for regression-style tests."""
+
+from __future__ import annotations
+
+from typing import Dict
+
+from sglang.srt.debug_utils.source_patcher import apply_patches_from_config
+from sglang.srt.environ import envs
+
+_PR_REVERT_YAML_25015 = """
+patches:
+  - target: sglang.srt.speculative.eagle_worker.EAGLEWorker.draft_forward
+    edits:
+      - match: |
+          forward_batch.out_cache_loc = out_cache_loc[i]
+          spec_info.hidden_states = hidden_states
+        replacement: |
+          forward_batch.out_cache_loc = out_cache_loc[i]
+          forward_batch.positions.add_(1)
+          spec_info.hidden_states = hidden_states
+      - match: |
+          hidden_states = logits_output.hidden_states
+          maybe_detect_nan(hidden_states, f"draft_forward step {i}: hidden_states")
+          maybe_detect_inf(hidden_states, f"draft_forward step {i}: hidden_states")
+          forward_batch.positions.add_(1)
+        replacement: |
+          hidden_states = logits_output.hidden_states
+          maybe_detect_nan(hidden_states, f"draft_forward step {i}: hidden_states")
+          maybe_detect_inf(hidden_states, f"draft_forward step {i}: hidden_states")
+
+  - target: sglang.srt.speculative.eagle_worker_v2.EagleDraftWorker.draft_forward
+    edits:
+      - match: |
+          forward_batch.out_cache_loc = out_cache_loc[i]
+          spec_info.hidden_states = hidden_states
+        replacement: |
+          forward_batch.out_cache_loc = out_cache_loc[i]
+          forward_batch.positions.add_(1)
+          spec_info.hidden_states = hidden_states
+      - match: |
+          hidden_states = logits_output.hidden_states
+          forward_batch.positions.add_(1)
+        replacement: |
+          hidden_states = logits_output.hidden_states
+
+  - target: sglang.srt.speculative.eagle_draft_cuda_graph_runner.EAGLEDraftCudaGraphRunner.capture_one_batch_size
+    edits:
+      - match: |
+          forward_batch.spec_info.hidden_states = hidden_states_backup
+          forward_batch.positions.sub_(self.eagle_worker.speculative_num_steps - 1)
+          return ret
+        replacement: |
+          forward_batch.spec_info.hidden_states = hidden_states_backup
+          return ret
+"""
+
+
+_PR_REVERT_YAML_26329 = """
+patches:
+  - target: sglang.srt.speculative.eagle_utils._eagle_prefill_tail_tokens
+    edits:
+      - match: |
+          tail_tokens = next_token_ids.to(batch.input_ids.dtype)
+        prepend: |
+          return next_token_ids.to(batch.input_ids.dtype)
+"""
+
+
+_PR_FIX_REVERT_YAML: Dict[int, str] = {
+    25015: _PR_REVERT_YAML_25015,
+    26329: _PR_REVERT_YAML_26329,
+}
+
+
+def maybe_revert_pr_fix() -> None:
+    if pr_num := envs.SGLANG_DEBUG_REVERT_PR.get():
+        _revert_pr_fix(pr_num)
+
+
+def _revert_pr_fix(pr_num: int) -> None:
+    if pr_num not in _PR_FIX_REVERT_YAML:
+        raise NotImplementedError(
+            f"PR #{pr_num} revert is not registered; "
+            f"available: {sorted(_PR_FIX_REVERT_YAML.keys())}"
+        )
+    apply_patches_from_config(_PR_FIX_REVERT_YAML[pr_num])

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

@@ -67,8 +67,9 @@ def _find_match(*, source_lines: list[str], match_lines: list[str]) -> int:
     ]
 
     if len(found_indices) == 0:
-        preview: str = "\n".join(match_lines)
-        raise PatchApplicationError(f"match text not found in source:\n{preview}")
+        raise PatchApplicationError(
+            _not_found_diagnostic(stripped_source, stripped_match)
+        )
     if len(found_indices) > 1:
         preview = "\n".join(match_lines)
         raise PatchApplicationError(
@@ -78,6 +79,41 @@ def _find_match(*, source_lines: list[str], match_lines: list[str]) -> int:
     return found_indices[0]
 
 
+def _not_found_diagnostic(stripped_source: list[str], stripped_match: list[str]) -> str:
+    preview = "\n".join(stripped_match)
+    lines = [
+        f"match text not found in source:\n{preview}",
+        "",
+        f"source_len={len(stripped_source)} lines",
+    ]
+
+    if not stripped_match:
+        return "\n".join(lines)
+    first_match_line = stripped_match[0]
+    hits = [i for i, line in enumerate(stripped_source) if line == first_match_line]
+    if not hits:
+        lines.append(
+            f"first match line {first_match_line!r} does NOT appear anywhere in source"
+        )
+        return "\n".join(lines)
+
+    lines.append(
+        f"first match line {first_match_line!r} appears {len(hits)} time(s); showing up to 8 windows with context:"
+    )
+    for i in hits[:8]:
+        lo = max(0, i - 2)
+        hi = min(len(stripped_source), i + len(stripped_match) + 2)
+        block: list[str] = []
+        for j in range(lo, hi):
+            marker = (
+                ">" if lo + (j - lo) >= i and (j - i) < len(stripped_match) else " "
+            )
+            block.append(f"{marker} {j:4d}: {stripped_source[j]}")
+        lines.append("--")
+        lines.extend(block)
+    return "\n".join(lines)
+
+
 def _realign_replacement(
     *, replacement_lines: list[str], original_indent: int
 ) -> list[str]: