[CI] Wait for a killed test server's GPU memory before the next launch (#39545)

.claude/skills/write-sglang-test/SKILL.md

@@ -5,20 +5,36 @@ description: Guide for writing SGLang CI/UT tests. Covers CustomTestCase, CI reg
 
 # Writing SGLang CI / UT Tests
 
-This skill covers **how to write and register tests**. For CI pipeline internals (stage ordering, fast-fail, gating, partitioning, debugging CI failures), see the [CI workflow guide](../ci-workflow-guide/SKILL.md).
+This skill covers **how to write and register tests**. For CI pipeline internals (stage ordering, fast-fail, gating, partitioning, debugging CI failures), see the [CI workflow guide](../ci-workflow-guide/SKILL.md). Whether a case is worth adding at all is decided by [`unit-test-admission`](../../rules/unit-test-admission.md) — read it before writing the case, not after.
 
 ## Core Rules
 
 1. **Always use `CustomTestCase`** — never raw `unittest.TestCase`. It ensures `tearDownClass` runs even when `setUpClass` fails, preventing resource leaks in CI.
-2. **`tearDownClass` must be defensive** — use `hasattr`/null checks before accessing resources (e.g. `cls.process`) that `setUpClass` may not have finished allocating.
+2. **`tearDownClass` must shut the server down gracefully** — call `terminate_and_kill_process_tree(cls.process)`, never a bare `kill_process_tree`. SIGKILL alone skips the server's userspace cleanup and leaves its GPU memory charged to the dead process; the next class then OOMs while loading weights. Keep it defensive too: `hasattr`/null checks before accessing resources (e.g. `cls.process`) that `setUpClass` may not have finished allocating.
 3. **Place tests in `test/registered/<kind>/<subsystem>/`** — `<kind>` is `unit`, `kernel`, `e2e`, `accuracy`, `perf`, or `stress`; hardware belongs in registrations, not directory names
 4. **Reuse server fixtures** — inherit from `DefaultServerBase` or write `setUpClass`/`tearDownClass` with `popen_launch_server`
 5. **Mock boundaries, not SGLang behavior** — mock slow or external dependencies only when the assertion still checks an observable result, state transition, or error. A test whose evidence is only `assert_called*` mirrors its mock and is not admissible. Launch a real server only when inference results or lifecycle behavior are the contract under test.
 
+> **Existing files are not the reference.** About 290 test files still call the bare
+> `kill_process_tree(cls.process.pid)`, against 43 on the current helper. They predate
+> rule 2 and are being migrated, so grepping the repo for a teardown pattern finds the
+> wrong one roughly seven times out of eight. The same goes for `register_cuda_ci(suite=...)`:
+> four files still pass it, all of them under `test/registered/stress/`.
+
+```python
+# Bad:  kill_process_tree(cls.process.pid)           # SIGKILL only; GPU memory lingers
+# Good: terminate_and_kill_process_tree(cls.process)
+
+# Bad:  register_cuda_ci(est_time=80, suite="base-b-test-1-gpu-small")
+# Good: register_cuda_ci(est_time=80, stage="base-b", runner_config="1-gpu-small")
+```
+
 JIT kernel notes:
 - If the task is adding or updating code under `python/sglang/kernels/jit/`, prefer the `add-jit-kernel` skill first.
 - New JIT kernel correctness tests use `test/registered/kernel/jit/**/test_*.py`.
 - New JIT kernel benchmarks use `test/registered/kernel/jit/benchmark/**/bench_*.py`.
+- `test/registered/jit/` also exists and still runs. It is a leftover from the kernel
+  reclassification (RFC #29630) that was never finished; do not add files there.
 - Those files are executed by `test/run_suite.py` through dedicated kernel suites (`base-b-kernel-*`); a `register_*_ci(...)` call placed under `python/sglang/` is rejected by the `check-no-registered-tests-in-package` pre-commit hook.
 
 ---
@@ -31,7 +47,7 @@ JIT kernel notes:
 | **Common / backend-independent** (middleware, abort, routing, config, arg parsing) | `DEFAULT_SMALL_MODEL_NAME_FOR_TEST` (1B) | `register_cuda_ci` only | `base-b-test-1-gpu-small` |
 | **Model-agnostic functionality** (sampling, session, OpenAI API features) | `DEFAULT_SMALL_MODEL_NAME_FOR_TEST` (1B) | `register_cuda_ci` (+ AMD if relevant) | `base-b-test-1-gpu-small` |
 | **General performance** (single node, no spec/DP/parallelism) | `DEFAULT_MODEL_NAME_FOR_TEST` (8B) | `register_cuda_ci` | `base-b-test-1-gpu-large` |
-| **Bigger features** (spec, DP, TP, disaggregation) | Case by case | Case by case | See suite table below |
+| **Bigger features** (spec, DP, TP, disaggregation) | Case by case | Case by case | See **Choosing a Suite** below |
 
 **Key principle for E2E tests**: Do NOT add `register_amd_ci` unless the test specifically exercises AMD/ROCm code paths. Common E2E tests just need any GPU to run — duplicating across backends wastes CI time with no extra coverage.
 
@@ -60,74 +76,23 @@ A per-commit suite name is **generated** from registration metadata as `{stage}-
 
 ### All CI Suites
 
-#### Per-commit (CUDA)
+Do not work from a list copied into this file; it goes stale silently. Read the
+current one:
 
-| Suite | Runner (label) | Description |
-|-------|----------------|-------------|
-| `base-a-test-1-gpu-small` | `1-gpu-5090` | Quick checks on a small NVIDIA GPU before heavier stages |
-| `base-a-test-cpu` | `ubuntu-latest` | CPU-only unit tests |
-| `base-b-test-1-gpu-small` | `1-gpu-5090` | Core engine tests that fit a 5090-class card |
-| `base-b-test-1-gpu-large` | `1-gpu-h100` | Tests that need H100-class memory or kernels (e.g. FA3) |
-| `base-b-test-2-gpu-large` | `2-gpu-h100` | Two-GPU correctness and parallelism (TP/PP) on H100 |
-| `base-b-test-4-gpu-b200` | `4-gpu-b200` | Early Blackwell coverage (SM100+ paths) on four GPUs |
-| `base-b-kernel-unit-test-1-gpu-large` | `1-gpu-h100` | JIT kernel correctness tests under `test/registered/kernel/jit/` |
-| `base-b-kernel-unit-test-4-gpu-b200` | `4-gpu-b200` | JIT kernel correctness tests for Blackwell / SM100-specific paths |
-| `base-b-kernel-unit-test-8-gpu-h200` | `8-gpu-h200` | Multi-GPU JIT kernel correctness tests under `test/registered/kernel/jit/` |
-| `base-b-kernel-benchmark-test-1-gpu-large` | `1-gpu-h100` | JIT kernel benchmark files under `test/registered/kernel/jit/benchmark/` |
-| `base-c-test-4-gpu-h100` | `4-gpu-h100` | Large 4-GPU H100 integration and scaling tests |
-| `base-c-test-8-gpu-h200` | `8-gpu-h200` | Large 8-GPU H200 runs for big models and parallelism |
-| `base-c-test-8-gpu-h20` | `8-gpu-h20` | Large 8-GPU H20 runs for big models |
-| `base-c-test-deepep-4-gpu-h100` | `4-gpu-h100` | DeepEP expert-parallel and networking on four H100s |
-| `base-c-test-8-gpu-b200` | `8-gpu-b200` | 8-GPU B200 suite (registered but not yet wired to a workflow) |
-| `base-c-test-4-gpu-b200` | `4-gpu-b200` | 4-GPU B200 suite for large models on Blackwell |
-| `base-c-test-4-gpu-b200-small` | `4-gpu-b200` | Smaller 4-GPU B200 suite split onto low-disk B200 runners |
-| `base-c-test-4-gpu-gb200` | `4-gpu-gb200` | 4-GPU GB200 suite for Grace Blackwell; registered in `run_suite.py`, but the PR workflow is currently disabled until a runner is provisioned |
+```bash
+grep -n "_SUITES = {" test/run_suite.py   # PER_COMMIT_SUITES, NIGHTLY_SUITES, OTHER_SUITES
+cat scripts/ci/runner_configs.yml         # runner_config -> physical runner label
+```
 
-#### Per-commit (AMD)
+`scripts/ci/runner_configs.yml` calls itself the single source of truth for the
+`runner_config` field, and `run_suite.py` is what actually dispatches, so those two
+files settle any disagreement with prose anywhere else.
 
-| Suite | Runner (label) | Description |
-|-------|----------------|-------------|
-| `stage-a-test-1-gpu-small-amd` | `linux-mi325-1gpu-sglang` | Quick checks on one MI325-class GPU |
-| `stage-b-test-1-gpu-small-amd` | `linux-mi325-1gpu-sglang` | Core 1-GPU AMD tests (14 partitions) |
-| `stage-b-test-1-gpu-small-amd-nondeterministic` | `linux-mi325-1gpu-sglang` | Non-deterministic 1-GPU AMD tests |
-| `stage-b-test-1-gpu-small-amd-mi35x` | `linux-mi35x-gpu-1` | 1-GPU tests on MI35x hardware |
-| `stage-b-test-1-gpu-large-amd` | `linux-mi325-1gpu-sglang` | Large 1-GPU AMD tests (2 partitions) |
-| `stage-b-test-2-gpu-large-amd` | `linux-mi325-2gpu-sglang` | 2-GPU ROCm correctness and parallel setups |
-| `stage-b-test-large-8-gpu-mi35x-disaggregation-amd` | `linux-mi35x-gpu-8.fabric` | PD disaggregation and RDMA on 8×MI35x fabric |
-| `stage-c-test-4-gpu-amd` | `linux-mi325-4gpu-sglang` | 4-GPU AMD integration (2 partitions) |
-| `stage-c-test-large-8-gpu-amd` | `linux-mi325-8gpu-sglang` | 8-GPU MI325 scaling and integration |
-| `stage-c-test-large-8-gpu-amd-mi35x` | `linux-mi35x-gpu-8` | 8-GPU MI35x scaling (2 partitions) |
-
-
-### Per-commit (Ascend NPU)
-
-| Suite | Runner (label) | Description |
-| --- | --- | --- |
-| `per-commit-1-npu-a2` | `linux-aarch64-a2-1` | 1-NPU LLM CI machine |
-| `per-commit-2-npu-a2` | `linux-aarch64-a2-2` | 2-NPU LLM CI machine |
-| `per-commit-4-npu-a3` | `linux-aarch64-a3-4` | 4-NPU LLM CI machine |
-| `per-commit-16-npu-a3` | `linux-aarch64-a3-16` | 16-NPU LLM CI machine  |
-| `multimodal-gen-test-1-npu-a3` | `linux-aarch64-a3-2` | 1-NPU multimodal CI machine |
-| `multimodal-gen-test-2-npu-a3` | `linux-aarch64-a3-16` | 2-NPU multimodal CI machine |
-| `multimodal-gen-test-8-npu-a3` | `linux-aarch64-a3-16` | 8-NPU multimodal CI machine |
-
-#### Nightly
-
-Nightly suites are listed in `NIGHTLY_SUITES` in [`test/run_suite.py`](../../../test/run_suite.py). They run via `nightly-test-nvidia.yml`, `nightly-test-amd.yml`, and `nightly-test-npu.yml`, not `pr-test.yml`.
-
-CUDA nightly suites are named `nightly-test-{runner_config}` — one per machine type, holding everything that runs nightly on it. There is no per-purpose split (kernel / eval / perf / precision all share their machine's suite); `auto_partition` splits the work. Examples:
-
-- `nightly-test-1-gpu-large` (CUDA)
-- `nightly-test-2-gpu-large` (CUDA)
-- `nightly-test-8-gpu-h200` (CUDA)
-- `nightly-test-4-gpu-gb300` (CUDA)
-- `nightly-amd` (AMD)
-- `nightly-amd-8-gpu-mi35x` (AMD)
-- `nightly-1-npu-a3` (NPU)
-- `nightly-2-npu-a3` (NPU)
-- `nightly-4-npu-a3` (NPU)
-- `nightly-8-npu-a3` (NPU)
-- `nightly-16-npu-a3` (NPU)
+Nightly suites live in `NIGHTLY_SUITES` and run via `nightly-test-nvidia.yml`,
+`nightly-test-amd.yml`, and `nightly-test-npu.yml`, not `pr-test.yml`. CUDA nightly is
+named `nightly-test-{runner_config}` — one suite per machine type, holding everything
+that runs nightly on it, with `auto_partition` splitting the work. There is no
+per-purpose split; kernel, eval, perf, and precision all share their machine's suite.
 
 > **Note**: Multimodal diffusion uses `python/sglang/multimodal_gen/test/run_suite.py`, not `test/run_suite.py`.
 
@@ -197,7 +162,6 @@ import unittest
 
 import requests
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import (
     DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
@@ -205,9 +169,10 @@ from sglang.test.test_utils import (
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
     popen_launch_server,
+    terminate_and_kill_process_tree,
 )
 
-register_cuda_ci(est_time=60, suite="base-b-test-1-gpu-small")
+register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-small")
 
 
 class TestMyFeature(CustomTestCase):
@@ -225,7 +190,7 @@ class TestMyFeature(CustomTestCase):
     @classmethod
     def tearDownClass(cls):
         if hasattr(cls, "process") and cls.process:
-            kill_process_tree(cls.process.pid)
+            terminate_and_kill_process_tree(cls.process)
 
     def test_basic_functionality(self):
         response = requests.post(
@@ -239,6 +204,9 @@ if __name__ == "__main__":
     unittest.main(verbosity=3)
 ```
 
+Copy the `tearDownClass` above verbatim. Most existing E2E files still show the bare
+`kill_process_tree`; that form is being migrated out and must not be reproduced.
+
 ### E2E test (8B model, server needed, performance)
 
 ```python
@@ -247,7 +215,6 @@ import unittest
 
 import requests
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import (
     DEFAULT_MODEL_NAME_FOR_TEST,
@@ -255,9 +222,10 @@ from sglang.test.test_utils import (
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
     popen_launch_server,
+    terminate_and_kill_process_tree,
 )
 
-register_cuda_ci(est_time=300, suite="base-b-test-1-gpu-large")
+register_cuda_ci(est_time=300, stage="base-b", runner_config="1-gpu-large")
 
 
 class TestMyFeaturePerf(CustomTestCase):
@@ -274,7 +242,7 @@ class TestMyFeaturePerf(CustomTestCase):
     @classmethod
     def tearDownClass(cls):
         if hasattr(cls, "process") and cls.process:
-            kill_process_tree(cls.process.pid)
+            terminate_and_kill_process_tree(cls.process)
 
     def test_latency(self):
         start = time.perf_counter()
@@ -332,26 +300,29 @@ from sglang.test.ci.ci_register import (
 )
 
 # Per-commit test (small 1-gpu, runs on 5090)
-register_cuda_ci(est_time=80, suite="base-b-test-1-gpu-small")
+register_cuda_ci(est_time=80, stage="base-b", runner_config="1-gpu-small")
 
 # Per-commit test (large 1-gpu, runs on H100)
-register_cuda_ci(est_time=120, suite="base-b-test-1-gpu-large")
+register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large")
 
 # Nightly-only test (same shape as per-commit, stage is just "nightly")
 register_cuda_ci(est_time=200, stage="nightly", runner_config="1-gpu-large")
 
 # Multi-backend test (only when testing backend-specific code paths)
-register_cuda_ci(est_time=80, suite="base-a-test-1-gpu-small")
+register_cuda_ci(est_time=80, stage="base-a", runner_config="1-gpu-small")
 register_amd_ci(est_time=120, suite="stage-a-test-1-gpu-small-amd")
 register_npu_ci(est_time=400, suite="nightly-8-npu-a3", nightly=True)
 
 # Temporarily disabled test
-register_cuda_ci(est_time=80, suite="base-b-test-1-gpu-small", disabled="flaky - see #12345")
+register_cuda_ci(
+    est_time=80, stage="base-b", runner_config="1-gpu-small", disabled="flaky - see #12345"
+)
 ```
 
 Parameters:
 - `est_time`: estimated runtime in seconds (used for CI partitioning)
-- `suite`: which CI suite to run in (see suite tables above)
+- `stage` + `runner_config`: the canonical pair for CUDA; the suite name is generated from them (see Naming Conventions)
+- `suite`: legacy single-string form. Only `stress` and some AMD/CPU/NPU pools still take it; `register_cpu_ci(suite="base-a-test-cpu")` is correct and is not being migrated
 - `nightly=True`: legacy cadence flag, for non-CUDA nightly suites only. CUDA nightly uses `stage="nightly"` and must leave this unset
 - `disabled="reason"`: temporarily disable with explanation
 
@@ -359,17 +330,18 @@ Parameters:
 
 ### JIT Kernel Registration
 
-JIT kernel files live outside `test/registered/` but still use registration:
+`run_suite.py` discovers every `test/registered/**/*.py`, JIT kernel files included.
+They are ordinary registered tests; only their stage differs:
 
 ```python
 from sglang.test.ci.ci_register import register_cuda_ci
 
-# Correctness tests in test/registered/jit/
+# Correctness tests in test/registered/kernel/jit/
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
 register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="8-gpu-h200")
 
-# Benchmarks in test/registered/jit/benchmark/
+# Benchmarks in test/registered/kernel/jit/benchmark/
 register_cuda_ci(est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large")
 
 # Optional nightly registration — same form, stage is just "nightly"
@@ -393,13 +365,14 @@ test/
 │   ├── perf/<family>/         # scheduled latency/throughput contracts
 │   └── stress/<subsystem>/    # stress/weekly coverage
 ├── manual/              # Non-CI: debugging, one-off, manual verification
-└── run_suite.py         # CI runner (scans registered/ plus jit_kernel test/benchmark files)
+└── run_suite.py         # CI runner (globs test/registered/**/*.py; nothing outside it)
 
-python/sglang/kernels/jit/
-├── tests/               # JIT kernel correctness tests (CI-discovered by test/run_suite.py)
-└── benchmark/           # JIT kernel benchmarks (CI-discovered by test/run_suite.py)
+python/sglang/kernels/jit/   # implementation + test-only helpers, never registered tests
 ```
 
+A `register_*_ci(...)` under `python/sglang/` is rejected by the
+`check-no-registered-tests-in-package` pre-commit hook.
+
 **Decision rule** (see also `test/registered/README.md`):
 - CPU component logic, no server → `registered/unit/<subsystem>/`
 - JIT kernel correctness / benchmarks → `registered/kernel/jit/`
@@ -432,11 +405,12 @@ class TestMyFeature(CustomTestCase, MMLUMixin):
 from sglang.test.test_utils import (
     CustomTestCase,              # base class with retry logic
     popen_launch_server,         # launch server subprocess
+    terminate_and_kill_process_tree,    # SIGTERM, then SIGKILL, then wait for
+                                        # the GPU memory to come back
     DEFAULT_URL_FOR_TEST,        # auto-configured base URL
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,  # 600s default
     run_bench_serving,           # benchmark helper (launch + bench)
 )
-from sglang.srt.utils import kill_process_tree  # cleanup server
 ```
 
 ---
@@ -451,7 +425,15 @@ Before submitting a test:
 - [ ] JIT kernel work: test files live in `test/registered/kernel/jit/`; only test-only helpers stay under `python/sglang/kernels/jit/`
 - [ ] Backend-independent tests: `register_cuda_ci` only + smallest model
 - [ ] Logic that doesn't need a server / engine launch → unit test in `registered/unit/` (see Unit Tests section)
-- [ ] `setUpClass` launches server, `tearDownClass` kills it (if server-based)
 - [ ] `tearDownClass` is defensive — uses `hasattr`/null checks before accessing resources that may not have been allocated
-- [ ] Has `if __name__ == "__main__": unittest.main()`
+- [ ] Every case answers "what future diff turns this red?" — see [`unit-test-admission`](../../rules/unit-test-admission.md)
 - [ ] `est_time` is reasonable (measure locally)
+
+Run these against the new file and paste the output rather than self-attesting:
+
+```bash
+f=<your new test file>
+grep -n "kill_process_tree" $f          # every hit must be terminate_and_kill_process_tree
+grep -n "register_.*_ci(" $f            # CUDA: stage= + runner_config=, never suite=
+grep -n "CustomTestCase\|unittest.main" $f   # both must appear
+```

python/sglang/test/runners.py

@@ -35,7 +35,12 @@ from sglang.srt.entrypoints.engine import Engine
 from sglang.srt.model_loader.ci_weight_validation import ci_validate_and_clean_hf_cache
 from sglang.srt.utils import get_device, is_npu, load_image
 from sglang.srt.utils.hf_transformers_utils import get_tokenizer
-from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, calculate_rouge_l
+from sglang.test.test_utils import (
+    DEFAULT_PORT_FOR_SRT_TEST_RUNNER,
+    calculate_rouge_l,
+    collect_process_tree_pids,
+    wait_for_gpu_release,
+)
 
 if is_npu():
     from sglang.srt.hardware_backend.npu.utils import init_npu_backend
@@ -428,12 +433,14 @@ class HFRunner:
         # Fire-and-forget terminate() leaves the child holding the accelerator
         # during teardown; a follow-on SRTRunner on the same device can then
         # deadlock in driver init (observed on Intel XPU B580).
+        pid = self.model_proc.pid
         self.model_proc.terminate()
         self.model_proc.join(timeout=30)
         if self.model_proc.is_alive():
             self.model_proc.kill()
             self.model_proc.join()
         self.in_queue = self.out_queue = None
+        wait_for_gpu_release([pid])
 
     def terminate(self):
         self._stop_model_proc()
@@ -751,8 +758,13 @@ class SRTRunner:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
+        # Wait only on the pids this shutdown actually killed;
+        # a nested HFRunner or SRTRunner is deliberately still alive.
+        before = collect_process_tree_pids(os.getpid(), include_parent=False)
         self.engine.shutdown()
         del self.engine
+        alive = set(collect_process_tree_pids(os.getpid(), include_parent=False))
+        wait_for_gpu_release([pid for pid in before if pid not in alive])
 
     @staticmethod
     def forward_generation_raw(

python/sglang/test/test_utils.py

@@ -28,6 +28,7 @@ from typing import Any, Awaitable, Callable, List, Optional, Tuple
 import aiohttp
 import msgspec
 import numpy as np
+import psutil
 import requests
 import torch
 import torch.nn.functional as F
@@ -847,14 +848,17 @@ def terminate_and_kill_process_tree(
     and unpin the host memory during process reclaim, which can hold GPU memory
     for minutes on a busy host -- long enough to trip the per-class GPU-idle
     gate in the next ``setUpClass``. SIGTERM first so the server releases those
-    resources in userspace.
+    resources in userspace, then wait for the memory to come back:
+    a reaped tree does not mean the driver is done with it.
     """
+    pids = collect_process_tree_pids(process.pid)
     process.terminate()
     try:
         process.wait(timeout=terminate_timeout)
     except subprocess.TimeoutExpired:
         pass
     kill_process_tree(process.pid, **kill_kwargs)
+    wait_for_gpu_release(pids)
 
 
 def popen_launch_pd_server(
@@ -2001,6 +2005,8 @@ def maybe_stub_sgl_kernel():
 _GPU_IDLE_TIMEOUT_SECS = 30.0
 _GPU_IDLE_POLL_INTERVAL_SECS = 2.0
 _GPU_IDLE_USED_MEMORY_THRESHOLD = 2 << 30  # 2 GiB
+_GPU_RELEASE_TIMEOUT_SECS = 60.0
+_GPU_RELEASE_POLL_INTERVAL_SECS = 0.5
 
 
 def _format_gib(num_bytes: Optional[int]) -> str:
@@ -2102,6 +2108,86 @@ def _wait_for_gpu_idle_in_ci(
             pass
 
 
+def collect_process_tree_pids(pid: int, include_parent: bool = True) -> List[int]:
+    """Snapshot a process tree's pids, for a later ``wait_for_gpu_release``.
+
+    Call it BEFORE the kill; afterwards the tree cannot be walked.
+    """
+    try:
+        pids = [child.pid for child in psutil.Process(pid).children(recursive=True)]
+    except psutil.Error:
+        pids = []
+    if include_parent:
+        pids.append(pid)
+    return pids
+
+
+def _gpu_memory_holders(pynvml, gpu_indices: List[int], pids: set) -> List[str]:
+    reports = []
+    for index in gpu_indices:
+        handle = pynvml.nvmlDeviceGetHandleByIndex(index)
+        try:
+            procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
+        except pynvml.NVMLError:
+            # No per-pid enumeration in this container; nothing to wait on.
+            continue
+        reports.extend(
+            f"GPU {index} pid={proc.pid} {_format_gib(proc.usedGpuMemory)}"
+            for proc in procs
+            if proc.pid in pids
+        )
+    return reports
+
+
+def wait_for_gpu_release(
+    pids: List[int],
+    timeout: float = _GPU_RELEASE_TIMEOUT_SECS,
+    poll_interval: float = _GPU_RELEASE_POLL_INTERVAL_SECS,
+) -> None:
+    """Block until none of ``pids`` is still charged device memory.
+
+    Killing a server only queues the driver-side teardown,
+    so the next launch can OOM against memory charged to a reaped process.
+    Waiting on these pids, rather than on an idle GPU,
+    keeps this usable while other servers of the same test still run.
+    Best effort: a timeout or a dead NVML warns, never raises.
+    """
+    if not pids:
+        return
+    try:
+        import pynvml
+
+        pynvml.nvmlInit()
+    except Exception:
+        # Non-NVIDIA runner (CPU/AMD) or NVML unavailable; nothing to check.
+        return
+    try:
+        gpu_indices = _visible_gpu_indices(pynvml)
+        pending = set(pids)
+        deadline = time.monotonic() + timeout
+        while True:
+            holders = _gpu_memory_holders(pynvml, gpu_indices, pending)
+            if not holders:
+                return
+            if time.monotonic() >= deadline:
+                print(
+                    f"[CI GPU Release] Still charged after {timeout:.0f}s:"
+                    f" {'; '.join(holders)}",
+                    flush=True,
+                )
+                return
+            time.sleep(poll_interval)
+    except Exception as e:
+        # NVML can go away after a successful init (GPU lost, driver reset).
+        # Raising here would fail a teardown whose test already passed.
+        print(f"[CI GPU Release] Giving up, {type(e).__name__}: {e}", flush=True)
+    finally:
+        try:
+            pynvml.nvmlShutdown()
+        except Exception:
+            pass
+
+
 # Names the runner kits stamp onto a record that are not members of it.
 # `ModelRunner` computes `use_mla_backend` on itself; the kits copy that bool
 # onto the record they hand the runner, and `hasattr` cannot see it.

test/registered/ep/test_deepep_large.py

@@ -3,7 +3,6 @@ from types import SimpleNamespace
 
 import requests
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.send_one import BenchArgs, send_one_prompt
@@ -13,6 +12,7 @@ from sglang.test.test_utils import (
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
     popen_launch_server,
+    terminate_and_kill_process_tree,
 )
 
 register_cuda_ci(est_time=569, stage="extra-b", runner_config="8-gpu-h200")
@@ -63,7 +63,7 @@ class TestDeepseek(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(
@@ -132,7 +132,7 @@ class TestDeepseekMTP(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(
@@ -190,7 +190,7 @@ class TestDeepseekV32TBO(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_a_gsm8k(
         self,

test/registered/ep/test_deepep_small.py

@@ -4,7 +4,6 @@ from types import SimpleNamespace
 
 import requests
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
@@ -14,6 +13,7 @@ from sglang.test.test_utils import (
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
     popen_launch_server,
+    terminate_and_kill_process_tree,
 )
 
 register_cuda_ci(est_time=407, stage="base-c", runner_config="4-gpu-h100")
@@ -30,10 +30,10 @@ class TestPureDP(CustomTestCase):
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=[
                 "--trust-remote-code",
-                "--tp",
+                "--tp-size",
                 "4",
                 "--enable-dp-attention",
-                "--dp",
+                "--dp-size",
                 "4",
                 "--moe-a2a-backend",
                 "deepep",
@@ -48,7 +48,7 @@ class TestPureDP(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(
@@ -77,7 +77,7 @@ class TestTP(CustomTestCase):
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=[
                 "--trust-remote-code",
-                "--tp",
+                "--tp-size",
                 "4",
                 "--moe-a2a-backend",
                 "deepep",
@@ -90,7 +90,7 @@ class TestTP(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(
@@ -119,10 +119,10 @@ class TestTBO(CustomTestCase):
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=[
                 "--trust-remote-code",
-                "--tp",
+                "--tp-size",
                 "4",
                 "--enable-dp-attention",
-                "--dp",
+                "--dp-size",
                 "4",
                 "--moe-dense-tp-size",
                 "1",
@@ -142,7 +142,7 @@ class TestTBO(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(
@@ -205,83 +205,7 @@ class TestMTPWithTBO(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_gsm8k(self):
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="gsm8k",
-            api="completion",
-            max_tokens=512,
-            num_examples=200,
-            num_threads=128,
-        )
-        metrics = run_eval(args)
-        print(metrics)
-
-        self.assertGreater(metrics["score"], 0.60)
-
-        server_info = requests.get(self.base_url + "/server_info")
-        avg_spec_accept_length = server_info.json()["internal_states"][0][
-            "avg_spec_accept_length"
-        ]
-        print(
-            f"###test_gsm8k (deepseek-v3 mtp + dp + tbo):\n"
-            f"accuracy={metrics['score']=:.3f}\n"
-            f"{avg_spec_accept_length=:.3f}\n"
-        )
-        self.assertGreater(avg_spec_accept_length, 2.1)
-
-
-@unittest.skip("skipped due to bug when using MTP & TBO & attn_tp_size > 1")
-class TestMTPWithTPAttnAndTBO(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--tp-size",
-                "4",
-                "--moe-dense-tp-size",
-                "1",
-                "--enable-two-batch-overlap",
-                "--moe-a2a-backend",
-                "deepep",
-                "--trust-remote-code",
-                "--speculative-algorithm",
-                "EAGLE",
-                "--speculative-num-steps",
-                "2",
-                "--speculative-eagle-topk",
-                "3",
-                "--speculative-num-draft-tokens",
-                "3",
-                "--speculative-draft-model-path",
-                DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
-                "--chunked-prefill-size",
-                "256",
-                "--cuda-graph-max-bs-decode",
-                "32",
-                "--max-running-requests",
-                "128",
-                "--mem-fraction-static",  # temp fix as DeepEP buffer is too large.
-                "0.7",
-            ],
-            env={
-                **os.environ,
-                "SGLANG_TBO_DEBUG": "1",
-            },
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(

test/registered/ep/test_deepep_small_extra.py

@@ -2,7 +2,6 @@ import os
 import unittest
 from types import SimpleNamespace
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
@@ -11,6 +10,7 @@ from sglang.test.test_utils import (
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
     popen_launch_server,
+    terminate_and_kill_process_tree,
 )
 
 register_cuda_ci(est_time=283, stage="extra-b", runner_config="4-gpu-h100")
@@ -43,7 +43,7 @@ class TestHybridDPTP(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(
@@ -92,7 +92,7 @@ class TestTBOWithTPAttn(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(
@@ -145,7 +145,7 @@ class TestTBOWithTPAttnAndDenseDP(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(

test/registered/ep/test_eplb_no_a2a.py

@@ -11,7 +11,6 @@ degrades output instead of failing.
 import unittest
 from types import SimpleNamespace
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
@@ -20,6 +19,7 @@ from sglang.test.test_utils import (
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
     popen_launch_server,
+    terminate_and_kill_process_tree,
     try_cached_model,
 )
 
@@ -72,7 +72,7 @@ class TestEPLBNoA2A(CustomTestCase):
     @classmethod
     def tearDownClass(cls):
         if hasattr(cls, "process") and cls.process:
-            kill_process_tree(cls.process.pid)
+            terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(

test/registered/ep/test_flashinfer_a2a.py

@@ -3,7 +3,6 @@ from types import SimpleNamespace
 
 import requests
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
@@ -11,6 +10,7 @@ from sglang.test.test_utils import (
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
     popen_launch_server,
+    terminate_and_kill_process_tree,
 )
 
 register_cuda_ci(est_time=561, stage="base-c", runner_config="4-gpu-gb300")
@@ -63,7 +63,7 @@ class TestFlashinferA2ATrtllmRoutedFP4(CustomTestCase):
     @classmethod
     def tearDownClass(cls):
         if hasattr(cls, "process") and cls.process:
-            kill_process_tree(cls.process.pid)
+            terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(
@@ -126,7 +126,7 @@ class TestFlashinferA2ACutedslStaticFP4(CustomTestCase):
     @classmethod
     def tearDownClass(cls):
         if hasattr(cls, "process") and cls.process:
-            kill_process_tree(cls.process.pid)
+            terminate_and_kill_process_tree(cls.process)
 
     def test_generate(self):
         response = requests.post(
@@ -179,7 +179,7 @@ class TestFlashinferA2ATrtllmRoutedFP8(CustomTestCase):
     @classmethod
     def tearDownClass(cls):
         if hasattr(cls, "process") and cls.process:
-            kill_process_tree(cls.process.pid)
+            terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(

test/registered/ep/test_mooncake_ep_small.py

@@ -4,7 +4,6 @@ from types import SimpleNamespace
 
 import requests
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.server_fixtures.disaggregation_fixture import get_rdma_devices_args
@@ -15,6 +14,7 @@ from sglang.test.test_utils import (
     CustomTestCase,
     is_in_ci,
     popen_launch_server,
+    terminate_and_kill_process_tree,
     try_cached_model,
 )
 
@@ -73,7 +73,7 @@ class TestTP(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(

test/registered/ep/test_routed_experts_dp_readback.py

@@ -10,13 +10,13 @@ import pybase64
 import requests
 import torch
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
     popen_launch_server,
+    terminate_and_kill_process_tree,
 )
 
 register_cuda_ci(est_time=72, stage="base-c", runner_config="4-gpu-h100")
@@ -102,7 +102,7 @@ class _ReadbackMixin:
     @classmethod
     def tearDownClass(cls):
         if getattr(cls, "process", None):
-            kill_process_tree(cls.process.pid)
+            terminate_and_kill_process_tree(cls.process)
 
     def _one_request(self, i: int):
         resp = requests.post(

test/registered/ep/test_tbo_shared_experts_fusion.py

@@ -2,7 +2,6 @@ import os
 import unittest
 from types import SimpleNamespace
 
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
@@ -10,6 +9,7 @@ from sglang.test.test_utils import (
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
     popen_launch_server,
+    terminate_and_kill_process_tree,
 )
 
 register_cuda_ci(est_time=253, stage="extra-b", runner_config="8-gpu-h200")
@@ -51,7 +51,7 @@ class TestTBOWithSharedExpertsFusion(CustomTestCase):
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        terminate_and_kill_process_tree(cls.process)
 
     def test_gsm8k(self):
         args = SimpleNamespace(