From a7e4f78870d8efdbd507010e3ea4032557092b7c Mon Sep 17 00:00:00 2001 From: Xianjie <5410381+qiaoxj07@users.noreply.github.com> Date: Sat, 5 Sep 2026 17:22:21 +0800 Subject: [PATCH] [None][feat] Add opt-in GPU keepalive to the benchmark fill gate In benchmark disagg mode a generation worker blocks at the fill gate (TLLM_BENCHMARK_REQ_QUEUES_SIZE) until the context tier has filled it; the wait scales with the concurrency and can take tens of minutes, during which its GPU does nothing and GPU-activity metrics read 0. With TRTLLM_GPU_KEEPALIVE=1 (default off) every closed gate retry queues a short chunk of GPU work (a resident warp on every SM, ~100 ms, two deep) instead of only sleeping, and the chunks are drained when the gate opens. The Triton spin kernel is compile-tested once in a subprocess that is polled from the gate (a compile failure can SIGABRT; the executor thread never blocks on it) with a low-duty, wall-clock paced torch.mm fallback; chunk length is calibrated at runtime. The stream and buffers are allocated at the first tick, outside the creator's executor_extra memory scope that sleep() releases, and freed when the gate opens or the executor loop exits. Warmup never reaches the gate. With the variable unset the gate is unchanged. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> --- docs/source/features/disagg-serving.md | 2 + .../_torch/pyexecutor/gpu_keepalive.py | 425 ++++++++++++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 23 +- .../_torch/executor/test_gpu_keepalive.py | 342 ++++++++++++++ .../_torch/executor/test_gpu_keepalive_gpu.py | 100 +++++ 5 files changed, 889 insertions(+), 3 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/gpu_keepalive.py create mode 100644 tests/unittest/_torch/executor/test_gpu_keepalive.py create mode 100644 tests/unittest/_torch/executor/test_gpu_keepalive_gpu.py diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index 9b9efeae9dd8..65ec65ae0a8b 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -337,6 +337,8 @@ TRT-LLM uses some environment variables to control the behavior of disaggregated * `TRTLLM_NIXL_KVCACHE_BACKEND`: Selects the transport NIXL itself uses. Valid values are `UCX` (default) and `LIBFABRIC`; an unsupported value logs a warning and falls back to `UCX`. `LIBFABRIC` additionally requires a NIXL build carrying the libfabric plugin — see the [disaggregated serving examples](source:examples/disaggregated/README.md). +* `TRTLLM_GPU_KEEPALIVE`: If set to `1`, a generation worker that is waiting at the benchmark fill gate (`TLLM_BENCHMARK_REQ_QUEUES_SIZE`) keeps a resident warp on every SM in ~100 ms chunks instead of idling through the wait, so GPU-activity metrics do not read idle while the context tier fills it. The work is drained when the gate opens and never overlaps a forward pass. The default value is `0`. + There are some other useful environment variables that may help when encountering failures or performance issues. * `NCCL_GRAPH_MIXING_SUPPORT`: TensorRT-LLM now initializes common NCCL communicators with graph diff --git a/tensorrt_llm/_torch/pyexecutor/gpu_keepalive.py b/tensorrt_llm/_torch/pyexecutor/gpu_keepalive.py new file mode 100644 index 000000000000..a92755093e9f --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/gpu_keepalive.py @@ -0,0 +1,425 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Opt-in GPU keepalive for the benchmark fill gate. + +With ``TRTLLM_GPU_KEEPALIVE=1`` a generation worker that is blocked on the +benchmark fill gate (``TLLM_BENCHMARK_REQ_QUEUES_SIZE``) keeps its GPU visibly +busy instead of idling through the wait, so GPU-activity metrics do not read +0 while the context tier fills it. Off by default; does no useful work. + +The primitive is a Triton kernel with one warp per CTA and two CTAs per SM +spinning on an FMA chain in ~100 ms chunks. SM-activity metrics count SMs +with at least one resident warp, so this reads like a loaded GPU at ~3% warp +occupancy and slows co-running kernels by only a few percent. Chunk length is +calibrated at runtime in ns per iteration, so nothing is GPU-specific. + +Design points: + +* The kernel is compiled in a subprocess first: a Triton compile failure can + abort the process (SIGABRT), which no in-process ``try`` can catch. The + child is polled from :meth:`tick` and never blocks the executor thread; if + it does not pass, a low-duty ``torch.mm`` fallback is used. +* Everything device-side is allocated at the first tick, on the executor + thread and outside the creator's ``executor_extra`` memory scope (which an + engine sleep releases), and freed again when the gate opens. +* Work is launched only while the gate is closed (never during warmup) and + drained when the gate opens, so it never overlaps CUDA-graph capture or a + forward. +""" + +from __future__ import annotations + +import collections +import os +import signal +import subprocess # nosec B404 - runs this very file with sys.executable to isolate a Triton compiler abort +import sys +import tempfile +import time +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from logging import Logger + +KEEPALIVE_ENV_VAR_NAME = "TRTLLM_GPU_KEEPALIVE" + +_PERIOD_SEC = 0.1 # chunk length; keep under NVML's ~1/6 s utilization sample +_GRID_MULT = 2 # CTAs per SM, so every SM gets a warp even if another kernel holds some +_QUEUE_DEPTH = 2 # chunks queued ahead of the GPU; bounds the drain when the gate opens +_MM_DUTY = 0.05 # duty of the torch.mm fallback; a saturating GEMM slows other kernels ~8x +_MM_DIM = 4096 +_SPIN_CALIBRATION_ITERS = 2_000_000 +_SPIN_MIN_ITERS = 1_000 +_SPIN_MAX_ITERS = (1 << 31) - 1 # n_iters is an unspecialized i32 kernel argument +_SELFTEST_TIMEOUT_SEC = 120.0 +_SELFTEST_OK_MARKER = "KEEPALIVE_SELFTEST_OK" +_SELFTEST_RC_NO_CONTEXT = 3 # child could not create a CUDA context (e.g. EXCLUSIVE_PROCESS) +# Verdict per device for this process: PyExecutor may be built twice per rank +# (KV-cache estimation) and the ~15 s compile check cannot change in between. +_SELFTEST_VERDICT: dict[int, bool] = {} + +try: + import triton + import triton.language as tl + + _HAVE_TRITON = True +except ImportError: # pragma: no cover - depends on the environment + triton = None + tl = None + _HAVE_TRITON = False + + +if _HAVE_TRITON: + + @triton.jit(do_not_specialize=["n_iters"]) + def _spin_kernel(n_iters, sink_ptr): + # The store keeps the chain alive so the loop is not optimised away. + # do_not_specialize: one compiled variant for every iteration count, + # so no JIT compile ever happens on the executor thread. + x = tl.zeros([32], dtype=tl.float32) + tl.program_id(0) + for _ in range(n_iters): + x = x * 0.999 + 1.0 + tl.store(sink_ptr + tl.arange(0, 32), x) + + +def _log() -> "Logger": + # Lazy: the subprocess self-test runs this file by path and must not + # import the tensorrt_llm package. + from tensorrt_llm.logger import logger + + return logger + + +class _SelftestChild: + """The subprocess that compiles and launches the spin kernel once. + + Output goes to temporary files, not pipes, so the child can never block + on a full pipe while nobody reads it. + """ + + def __init__(self, device_index: int): + self._out = tempfile.TemporaryFile() + self._err = tempfile.TemporaryFile() + self._proc = subprocess.Popen( # nosec B603 - fixed argv, no shell, interpreter is sys.executable + [sys.executable, os.path.abspath(__file__), "--selftest", str(device_index)], + stdout=self._out, + stderr=self._err, + env={**os.environ, "PYTHONSAFEPATH": "1"}, # keep this directory off sys.path + start_new_session=True, # own process group: kill() also reaches ptxas et al. + ) + self.started = time.monotonic() + + def poll(self) -> int | None: + return self._proc.poll() + + def kill(self) -> None: + """Stop the whole compiler process tree, not just the Python wrapper.""" + try: + try: + os.killpg(self._proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + self._proc.wait() + finally: + self._close() + + def output(self) -> tuple[bytes, str]: + """(stdout, last stderr lines); call once, after the child exited.""" + self._out.seek(0) + self._err.seek(0) + out = self._out.read() + tail = " | ".join(self._err.read().decode(errors="replace").strip().splitlines()[-3:]) + self._close() + return out, tail + + def _close(self) -> None: + self._out.close() + self._err.close() + + +def _run_selftest(device_index: int) -> int: + """Subprocess entry: compile, launch once, print the marker. Exit code is the verdict.""" + if not _HAVE_TRITON: + return 2 + try: + torch.cuda.set_device(device_index) + sink = torch.zeros(32, dtype=torch.float32, device=f"cuda:{device_index}") + except RuntimeError: + return _SELFTEST_RC_NO_CONTEXT + _spin_kernel[(4,)](1000, sink, num_warps=1) + torch.cuda.synchronize() + print(_SELFTEST_OK_MARKER, flush=True) + return 0 + + +class GpuKeepalive: + """Keeps the GPU visibly non-idle while the executor waits at the fill gate. + + Call :meth:`tick` on every iteration that finds the gate closed, + :meth:`drain` when it opens and :meth:`close` when the executor loop + exits. All are rank-local, issue no collectives and never raise: any + error disables the keepalive. + """ + + def __init__(self, device): + # No CUDA work here: PyExecutor is constructed inside a memory scope + # that sleep() releases. Everything device-side happens in _initialize. + self.device = torch.device(device if not isinstance(device, int) else f"cuda:{device}") + self.period_s = _PERIOD_SEC + self.queue_depth = _QUEUE_DEPTH + self.mode: str | None = None + self._initialized = False + self._disabled = False + self._stream: torch.cuda.Stream | None = None + self._selftest: _SelftestChild | None = None + self._inflight: collections.deque[torch.cuda.Event] = collections.deque() + self._launches = 0 + self._last_launch = 0.0 + self._active = False + + # ------------------------------------------------------------------ setup + def _initialize(self) -> bool: + """Device-side setup, driven from tick(); False while the self-test child runs.""" + if self._stream is None: + self._stream = torch.cuda.Stream(device=self.device) + num_sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self._grid = _GRID_MULT * num_sms + device_index = self.device.index if self.device.index is not None else 0 + if device_index not in _SELFTEST_VERDICT: + verdict = self._poll_selftest(device_index) + if verdict is None: + return False + _SELFTEST_VERDICT[device_index] = verdict + if _SELFTEST_VERDICT[device_index]: + self._init_spin() + else: + self._init_mm() + self._initialized = True + _log().info( + f"[gpu-keepalive] enabled mode={self.mode} grid={self._grid} " + f"period={self.period_s}s depth={self.queue_depth}" + ) + return True + + def _poll_selftest(self, device_index: int) -> bool | None: + """Drive the subprocess compile check without blocking; None while pending.""" + if not _HAVE_TRITON: + return False + if self._selftest is None: + self._selftest = _SelftestChild(device_index) + return None + rc = self._selftest.poll() + if rc is None: + if time.monotonic() - self._selftest.started < _SELFTEST_TIMEOUT_SEC: + return None + self._selftest.kill() + _log().warning( + f"[gpu-keepalive] spin self-test did not finish within {_SELFTEST_TIMEOUT_SEC:.0f}s; " + "using torch.mm" + ) + verdict = False + else: + out, tail = self._selftest.output() + verdict = rc == 0 and _SELFTEST_OK_MARKER.encode() in out + if not verdict and rc == _SELFTEST_RC_NO_CONTEXT: + _log().warning( + "[gpu-keepalive] spin self-test could not create a CUDA context in a child " + f"process (compute mode EXCLUSIVE_PROCESS?): {tail}; using torch.mm" + ) + elif not verdict: + _log().warning( + f"[gpu-keepalive] spin self-test failed (rc={rc}): {tail}; using torch.mm" + ) + self._selftest = None + return verdict + + def _use_on_private_stream(self, *tensors: torch.Tensor) -> None: + """Hand buffers created on the current stream to the private stream. + + The private stream first waits for their initialisation, and the + allocator is told they are in use there, so their blocks are not + recycled under in-flight chunks. They stay in the current stream's + pool, so releasing them gives the memory back to the model's pool. + """ + self._stream.wait_stream(torch.cuda.current_stream(self.device)) + for t in tensors: + t.record_stream(self._stream) + + def _init_spin(self) -> None: + self._sink = torch.zeros(32, dtype=torch.float32, device=self.device) + self._use_on_private_stream(self._sink) + self._ns_per_iter = self._calibrate_spin() + if not (self._ns_per_iter > 0.0): + raise RuntimeError(f"spin calibration returned {self._ns_per_iter} ns/iter") + self._recalibrate_when_hot = True + self.mode = "spin" + + def _calibrate_spin(self) -> float: + """Measure ns per spin iteration on the private stream.""" + beg = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + with torch.cuda.stream(self._stream): + _spin_kernel[(self._grid,)](10_000, self._sink, num_warps=1) + beg.record(self._stream) + _spin_kernel[(self._grid,)](_SPIN_CALIBRATION_ITERS, self._sink, num_warps=1) + end.record(self._stream) + end.synchronize() + return beg.elapsed_time(end) * 1e6 / _SPIN_CALIBRATION_ITERS + + def _init_mm(self) -> None: + dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + kw = dict(dtype=dtype, device=self.device) + # torch.ones, not randn: do not consume the default CUDA generator. + self._a = torch.ones(_MM_DIM, _MM_DIM, **kw) + self._b = torch.ones(_MM_DIM, _MM_DIM, **kw) + self._c = torch.empty(_MM_DIM, _MM_DIM, **kw) + self._use_on_private_stream(self._a, self._b, self._c) + beg = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + with torch.cuda.stream(self._stream): + for _ in range(5): + torch.mm(self._a, self._b, out=self._c) + beg.record(self._stream) + for _ in range(20): + torch.mm(self._a, self._b, out=self._c) + end.record(self._stream) + end.synchronize() + self._mm_ms = beg.elapsed_time(end) / 20.0 + self.mode = "mm" + + def _release(self) -> None: + """Free everything device-side; a later closed gate re-initialises lazily.""" + if self._selftest is not None: + try: + self._selftest.kill() + except Exception: # noqa: BLE001 - best-effort + pass + self._selftest = None + self._inflight.clear() + for name in ("_sink", "_a", "_b", "_c"): + if hasattr(self, name): + setattr(self, name, None) + self._stream = None + self._initialized = False + self.mode = None + + # ---------------------------------------------------------------- launch + def _launch_chunk(self) -> None: + """Queue one chunk of ~period_s GPU work on the private stream.""" + done = torch.cuda.Event() # created first, so a failure here submits nothing + with torch.cuda.stream(self._stream): + if self.mode == "spin": + iters = int(self.period_s * 1e9 / self._ns_per_iter) + iters = max(_SPIN_MIN_ITERS, min(_SPIN_MAX_ITERS, iters)) + _spin_kernel[(self._grid,)](iters, self._sink, num_warps=1) + else: + n = max(1, int(self.period_s * _MM_DUTY * 1e3 / self._mm_ms)) + for _ in range(n): + torch.mm(self._a, self._b, out=self._c) + done.record(self._stream) + self._inflight.append(done) + self._launches += 1 + + def _reap_inflight(self) -> int: + """Drop completed chunks from the front; return how many remain.""" + while self._inflight and self._inflight[0].query(): + self._inflight.popleft() + return len(self._inflight) + + def _disable(self, reason: str) -> None: + if self._disabled: + return + self._disabled = True + self._active = False + _log().warning(f"[gpu-keepalive] disabled: {reason}") + try: + if self._stream is not None: + self._stream.synchronize() # anything already submitted must finish + except Exception: # noqa: BLE001 - best-effort + pass + self._release() + + # ---------------------------------------------------------------- public + @property + def launches(self) -> int: + return self._launches + + def tick(self) -> bool: + """Queue a chunk while the gate is closed; True if one was launched.""" + if self._disabled: + return False + try: + if not self._initialized and not self._initialize(): + return False + if self._reap_inflight() >= self.queue_depth: + return False + now = time.monotonic() + if self.mode == "mm" and now - self._last_launch < self.period_s: + return False # one ~5 ms GEMM burst per period keeps the fallback at its duty + if not self._active: + self._active = True + _log().info("[gpu-keepalive] fill gate closed: keeping the GPU busy") + self._launch_chunk() + self._last_launch = now + if self.mode == "spin" and self._recalibrate_when_hot and self._launches >= 2: + # Clocks have boosted by now; re-measure once so chunks match period_s. + self._recalibrate_when_hot = False + self._ns_per_iter = self._calibrate_spin() + return True + except Exception as exc: # noqa: BLE001 - best-effort: never fail the executor loop + self._disable(f"error: {exc!r}") + return False + + def drain(self) -> None: + """Wait for queued chunks (at most queue_depth * period_s) and free the device side.""" + self._finish(f"fill gate opened after {self._launches} chunks") + + def close(self) -> None: + """Executor shutdown: like :meth:`drain`, also while the gate is still closed.""" + self._finish(f"closed with the fill gate still closed after {self._launches} chunks") + + def _finish(self, message: str) -> None: + if self._inflight: + try: + self._inflight[-1].synchronize() + except Exception as exc: # noqa: BLE001 - best-effort: never fail the executor loop + self._disable(f"error: {exc!r}") + if self._active: + self._active = False + _log().info(f"[gpu-keepalive] {message}") + self._release() # also stops a still-running self-test child + + @classmethod + def create_from_env(cls, device) -> GpuKeepalive | None: + """Build a keepalive if ``TRTLLM_GPU_KEEPALIVE=1``, else None. Never raises. + + Construction touches no CUDA state; the device side is set up at the + first :meth:`tick`, where any failure disables the keepalive. + """ + if os.environ.get(KEEPALIVE_ENV_VAR_NAME, "0") != "1": + return None + try: + return cls(device) + except Exception as exc: # noqa: BLE001 - opt-in best-effort feature must not fail executor init + _log().warning(f"[gpu-keepalive] not active: {exc!r}") + return None + + +if __name__ == "__main__": # pragma: no cover - subprocess self-test entry + if len(sys.argv) >= 3 and sys.argv[1] == "--selftest": + sys.exit(_run_selftest(int(sys.argv[2]))) + sys.exit(2) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index cebdbc942c27..81734bde1253 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -76,6 +76,7 @@ from .error_classification import ErrorBudget from .executor_request_queue import (ExecutorRequestQueue, RequestAdmissionState, RequestQueueItem) +from .gpu_keepalive import GpuKeepalive from .guided_decoder import GuidedDecoder from .handle_additional_outputs import HandleAdditionalOutputs from .handle_logits import HandleLogits @@ -947,6 +948,9 @@ def on_detected(): # _pop_from_waiting_queue). 0 = uninitialised; first throttled iter # seeds it to tp_size and each subsequent iter doubles it. self._fill_admit_cap: int = 0 + # Optional GPU keepalive for the fill gate (TRTLLM_GPU_KEEPALIVE=1). + # Allocates nothing here; its device side is set up at its first tick. + self._gpu_keepalive = GpuKeepalive.create_from_env(self.device_id) # Initialize disagg PP termination handler if needed self._disagg_pp_termination_handler = None @@ -2511,6 +2515,11 @@ def _executor_loop_cleanup(self): self.response_cv.notify_all() self.shutdown_event.set() + # The loop may exit while the benchmark fill gate is still closed. + keepalive = getattr(self, "_gpu_keepalive", None) + if keepalive is not None: + keepalive.close() + for i in range(self.num_micro_batches): try: self.wait_on_pp_send_handles(self.send_handles, i) @@ -4126,7 +4135,9 @@ def _check_benchmark_disagg_gate(self, scheduled_batch: ScheduledRequests, A short sleep (0.1s) yields the CPU between retries that made no transfer progress while keeping the polling interval short enough to avoid KV transfer backpressure on the CTX server. Retries that complete - a transfer do not sleep. + a transfer do not sleep. With ``TRTLLM_GPU_KEEPALIVE=1`` every closed + retry also queues a short GPU keepalive chunk, so the GPU does not read + idle for the whole wait; the chunks are drained when the gate opens. Args: scheduled_batch: The current scheduled batch. @@ -4142,13 +4153,19 @@ def _check_benchmark_disagg_gate(self, scheduled_batch: ScheduledRequests, can_forward = self._is_benchmark_disagg_fill_complete( scheduled_batch, transfer_made_progress) transfer_made_progress = self._benchmark_transfer_progress_global + keepalive = getattr(self, "_gpu_keepalive", None) if can_forward: self._benchmark_fill_phase_active = False self._fill_admit_cap = 0 self._benchmark_fill_stall_since = None self._benchmark_completed_gen_transfer_ids.clear() - elif not transfer_made_progress: - time.sleep(0.1) + if keepalive is not None: + keepalive.drain() + else: + if keepalive is not None: + keepalive.tick() + if not transfer_made_progress: + time.sleep(0.1) if not can_forward: self._fail_if_fill_gate_stalled(transfer_made_progress) return can_forward, True diff --git a/tests/unittest/_torch/executor/test_gpu_keepalive.py b/tests/unittest/_torch/executor/test_gpu_keepalive.py new file mode 100644 index 000000000000..1adb3d9f0955 --- /dev/null +++ b/tests/unittest/_torch/executor/test_gpu_keepalive.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU-only tests for the opt-in fill-gate GPU keepalive. + +The keepalive's CUDA surface is replaced by a fake. The gate wiring is tested +against the production method bound onto a bare stub; the no-keepalive path +is covered by the existing gate tests in test_benchmark_disagg.py. +""" + +import ast +import collections +import inspect +import threading +from unittest.mock import Mock, patch + +import pytest + +from tensorrt_llm._torch.pyexecutor import gpu_keepalive as ka + +pytestmark = pytest.mark.cpu_only + + +class _FakeEvent: + """Event on one in-order fake stream: synchronize() completes this chunk + and every earlier one, as CUDA stream ordering guarantees.""" + + def __init__(self, done_ref, all_chunks): + self._done_ref = done_ref + self._all_chunks = all_chunks + + def query(self): + return self._done_ref["done"] + + def synchronize(self): + for ref in self._all_chunks: + ref["done"] = True + if ref is self._done_ref: + break + + +class FakeKeepalive(ka.GpuKeepalive): + """GpuKeepalive with every CUDA touch replaced by bookkeeping.""" + + def __init__(self, queue_depth=ka._QUEUE_DEPTH, mode="spin"): + # Bypass GpuKeepalive.__init__ (device probe); pretend the first tick already ran. + self.device = "fake" + self.period_s = ka._PERIOD_SEC + self.queue_depth = queue_depth + self.mode = mode + self._initialized = True + self._disabled = False + self._stream = Mock() + self._selftest = None + self._inflight = collections.deque() + self._launches = 0 + self._last_launch = 0.0 + self._active = False + self._grid = 304 + self._recalibrate_when_hot = False + self._sink = self._a = self._b = self._c = Mock() # the buffers _release must drop + self.chunks = [] + self.raise_on_launch = None + + def _launch_chunk(self): + if self.raise_on_launch is not None: + raise self.raise_on_launch + ref = {"done": False} + self.chunks.append(ref) + self._inflight.append(_FakeEvent(ref, self.chunks)) + self._launches += 1 + + +@pytest.fixture +def clock(): + now = {"t": 1000.0} + with patch.object(ka.time, "monotonic", side_effect=lambda: now["t"]): + yield now + + +@pytest.fixture(autouse=True) +def quiet_logger(): + log = Mock() + with patch.object(ka, "_log", return_value=log): + yield log + + +def _released(k): + return ( + k._sink is None + and k._a is None + and k._b is None + and k._c is None + and k._stream is None + and not k._inflight + and not k._initialized + ) + + +# ---------------------------------------------------------------- keepalive + + +def test_tick_drain_release_and_errors(clock): + k = FakeKeepalive(queue_depth=2) + assert k.tick() is True + assert k.tick() is True # spin: two ~period chunks queued back to back + assert k.tick() is False # two resident, none done + k.chunks[0]["done"] = True # oldest finishes -> one slot opens + assert k.tick() is True + assert k.launches == 3 + k.drain() + assert all(ref["done"] for ref in k.chunks) + assert _released(k) and not k._active # the gate is one-shot: memory goes back + + mm = FakeKeepalive(queue_depth=2, mode="mm") + assert mm.tick() is True + assert mm.tick() is False # one GEMM burst per period, however fast the loop retries + clock["t"] += ka._PERIOD_SEC + assert mm.tick() is True + + # Any error disables the keepalive without raising; submitted work is waited for. + class _CompileError(Exception): # Triton raises its own classes, not RuntimeError + pass + + k = FakeKeepalive() + stream = k._stream + k.raise_on_launch = _CompileError("ptxas failed") + assert k.tick() is False + stream.synchronize.assert_called_once() + assert k._disabled and _released(k) + assert k.tick() is False # inert afterwards + + k = FakeKeepalive() + assert k.tick() is True + k._inflight[-1].synchronize = Mock(side_effect=RuntimeError("event sync failed")) + k.drain() # drain error: disabled, released, no exception + assert k._disabled and _released(k) + + k = FakeKeepalive() + k._initialized = False + k._initialize = Mock(side_effect=RuntimeError("no CUDA device")) + assert k.tick() is False # first-tick device setup fails: disabled, nothing launched + assert k._disabled and k.launches == 0 + + +def test_selftest_runs_asynchronously(monkeypatch, clock, quiet_logger): + """The compile check is a polled child: ticks launch nothing until it reports, + the verdict is cached per device, and a hung child is killed into the fallback.""" + children = [] + + class FakeChild: + def __init__(self, device_index): + self.rc, self.out, self.killed, self.started = None, b"", False, clock["t"] + children.append(self) + + def poll(self): + return self.rc + + def kill(self): + self.killed = True + + def output(self): + return self.out, "tail" + + monkeypatch.setattr(ka, "_SelftestChild", FakeChild) + monkeypatch.setattr(ka, "_HAVE_TRITON", True) + monkeypatch.setattr(ka, "_SELFTEST_VERDICT", {}) + monkeypatch.setattr(ka.torch.cuda, "Stream", Mock()) + monkeypatch.setattr( + ka.torch.cuda, "get_device_properties", Mock(return_value=Mock(multi_processor_count=152)) + ) + + def fake_init_spin(self): + self.mode, self._ns_per_iter, self._recalibrate_when_hot = "spin", 3.4, False + + monkeypatch.setattr(ka.GpuKeepalive, "_init_spin", fake_init_spin) + monkeypatch.setattr(ka.GpuKeepalive, "_init_mm", lambda self: setattr(self, "mode", "mm")) + monkeypatch.setattr( + ka.GpuKeepalive, + "_launch_chunk", + lambda self: setattr(self, "_launches", self._launches + 1), + ) + + k = ka.GpuKeepalive(0) + assert ( + k.tick() is False and len(children) == 1 and k.launches == 0 + ) # child started, nothing launched + assert k.tick() is False and len(children) == 1 # still running: keep polling, no new child + children[0].rc, children[0].out = 0, ka._SELFTEST_OK_MARKER.encode() + assert k.tick() is True and k.mode == "spin" and k._grid == 304 + assert ka.GpuKeepalive(0).tick() is True and len(children) == 1 # verdict cached per device + + monkeypatch.setattr(ka, "_SELFTEST_VERDICT", {}) + k = ka.GpuKeepalive(1) + k.tick() + children[-1].rc = ka._SELFTEST_RC_NO_CONTEXT + assert k.tick() is True and k.mode == "mm" + assert "CUDA context" in quiet_logger.warning.call_args.args[0] + + monkeypatch.setattr(ka, "_SELFTEST_VERDICT", {}) + k = ka.GpuKeepalive(2) + k.tick() + clock["t"] += ka._SELFTEST_TIMEOUT_SEC + 1 + assert k.tick() is True and k.mode == "mm" and children[-1].killed # hung compiler: fallback + assert ( + ka.GpuKeepalive(2).tick() is True and len(children) == 3 + ) # a failed verdict is cached too + + monkeypatch.setattr(ka, "_SELFTEST_VERDICT", {}) + k = ka.GpuKeepalive(3) + k.tick() + children[-1].rc = 0 # exited 0 but never printed the marker: not trusted + assert k.tick() is True and k.mode == "mm" + + monkeypatch.setattr(ka, "_SELFTEST_VERDICT", {}) + k = ka.GpuKeepalive(4) + k.tick() # child pending ... + k.close() # ... and the executor loop exits: child stopped, nothing leaks, no verdict recorded + assert ( + children[-1].killed + and k._selftest is None + and k._stream is None + and 4 not in ka._SELFTEST_VERDICT + ) + + +def test_create_from_env_gating(monkeypatch): + with patch.object(ka.GpuKeepalive, "__init__", return_value=None) as init: + for raw in (None, "0", "true"): + if raw is None: + monkeypatch.delenv(ka.KEEPALIVE_ENV_VAR_NAME, raising=False) + else: + monkeypatch.setenv(ka.KEEPALIVE_ENV_VAR_NAME, raw) + assert ka.GpuKeepalive.create_from_env(0) is None + init.assert_not_called() # disabled: the constructor is never reached + monkeypatch.setenv(ka.KEEPALIVE_ENV_VAR_NAME, "1") + assert ka.GpuKeepalive.create_from_env(3) is not None + assert init.call_count == 1 and init.call_args.args[-1] == 3 + with patch.object(ka.GpuKeepalive, "__init__", side_effect=RuntimeError("bad device")): + assert ka.GpuKeepalive.create_from_env(0) is None # never fails executor init + # Construction touches no CUDA state: the device side waits for the first tick. + with patch.object(ka.torch.cuda, "Stream", Mock(side_effect=AssertionError("CUDA touched"))): + k = ka.GpuKeepalive.create_from_env(0) + assert k is not None and not k._initialized and k.mode is None + + +# ------------------------------------------------------------ gate wiring + + +def test_fill_gate_wiring(): + """A closed gate ticks the keepalive and still sleeps as before; an opening gate + drains it; an already-open gate and warmup touch neither. The constructor must + build it from the environment and the loop cleanup must close it.""" + from tensorrt_llm._torch.pyexecutor import py_executor as pe + + def gate(*, can_forward=False, complete=False, progress=False, warmup=False): + class Stub: + _check_benchmark_disagg_gate = pe.PyExecutor._check_benchmark_disagg_gate + + stub = Stub() + stub.is_warmup = warmup + stub._gpu_keepalive = Mock() + stub._disagg_gen_transfer_made_progress = False + stub._benchmark_transfer_progress_global = progress + stub._is_benchmark_disagg_fill_complete = lambda batch, made_progress: complete + stub._fail_if_fill_gate_stalled = Mock() + stub._benchmark_fill_phase_active = True + stub._fill_admit_cap = 4 + stub._benchmark_fill_stall_since = 1.0 + stub._benchmark_completed_gen_transfer_ids = {7} + order = [] + stub._gpu_keepalive.tick.side_effect = lambda: order.append("tick") + with patch.object(pe.time, "sleep", side_effect=lambda s: order.append("sleep")): + result = stub._check_benchmark_disagg_gate(Mock(), can_forward) + return stub._gpu_keepalive, result, order + + k, result, order = gate() # closed, no transfer progress: tick first, then the usual sleep + assert result == (False, True) and order == ["tick", "sleep"] + k.drain.assert_not_called() + k, result, order = gate(progress=True) # closed, progress: tick, no sleep (as before) + assert result == (False, True) and order == ["tick"] + k, result, order = gate(complete=True) # opens: drain before the first forward + assert result == (True, False) and order == [] + k.drain.assert_called_once() + for kw in ({"can_forward": True}, {"warmup": True}): # gate bypassed: nothing at all + k, result, order = gate(**kw) + assert result[1] is False and order == [] + k.drain.assert_not_called() + + # Exactly one `self._gpu_keepalive = GpuKeepalive.create_from_env(self.device_id)` in __init__. + tree = ast.parse(inspect.getsource(pe)) + cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "PyExecutor") + init = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "__init__") + constructions = [ + n + for n in ast.walk(init) + if isinstance(n, ast.Assign) + and ast.unparse(n.targets[0]) == "self._gpu_keepalive" + and ast.unparse(n.value) == "GpuKeepalive.create_from_env(self.device_id)" + ] + assert len(constructions) == 1 + + # Loop cleanup (the loop may exit while the gate is still closed): waiters are + # notified first, then the keepalive is closed, then the PP handles are awaited. + class CleanupStub: + _executor_loop_cleanup = pe.PyExecutor._executor_loop_cleanup + + def __init__(self, keepalive): + self.events = [] + self.response_cv = threading.Condition(threading.Lock()) + self.response_cv.notify_all = lambda: self.events.append("notify_all") + self.is_shutdown = False + self.shutdown_event = threading.Event() + self.num_micro_batches = 1 + self.send_handles = self.send_schedule_handles = ( + self.send_expected_batch_num_handles + ) = {} + self._gpu_keepalive = keepalive + if keepalive is not None: + keepalive.close.side_effect = lambda: self.events.append("close") + + def wait_on_pp_send_handles(self, handles, idx): + self.events.append("wait_pp") + + stub = CleanupStub(Mock()) + stub._executor_loop_cleanup() + assert stub.is_shutdown and stub.shutdown_event.is_set() + assert stub.events == ["notify_all", "close", "wait_pp", "wait_pp", "wait_pp"] + plain = CleanupStub(None) + plain._executor_loop_cleanup() # without a keepalive: unchanged + assert plain.events == ["notify_all", "wait_pp", "wait_pp", "wait_pp"] diff --git a/tests/unittest/_torch/executor/test_gpu_keepalive_gpu.py b/tests/unittest/_torch/executor/test_gpu_keepalive_gpu.py new file mode 100644 index 000000000000..e95dbed878ed --- /dev/null +++ b/tests/unittest/_torch/executor/test_gpu_keepalive_gpu.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GPU smoke tests for the fill-gate GPU keepalive: the production path on a device.""" + +import subprocess +import sys +import time + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import gpu_keepalive as ka + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + + +def _run_production_path(k, timeout_s=300.0): + """tick() until initialised, launch, then drain; returns (mode, evidence the chunks ran). + + Initialisation and the one-time hot recalibration both write the output buffer + themselves, so the helper first gets past them, then clears the buffer and requires + at least one further chunk. The buffer is read back only after drain(), without any + other synchronisation, so the check proves both that a runtime chunk computed and + that drain() waited for it. + """ + deadline = time.monotonic() + timeout_s + try: + while not k._initialized and not k._disabled and time.monotonic() < deadline: + k.tick() # polls the self-test child; launches nothing until it reports + time.sleep(0.2) + assert k._initialized and not k._disabled, "keepalive did not initialise" + mode = k.mode + while ( + getattr(k, "_recalibrate_when_hot", False) + and not k._disabled + and time.monotonic() < deadline + ): + k.tick() # spin: the second launch triggers the recalibration kernels + time.sleep(0.05) + assert not k._disabled, "keepalive disabled during recalibration" + k._stream.synchronize() + out = k._sink if mode == "spin" else k._c # keep a reference: drain() drops the keepalive's + with torch.cuda.stream(k._stream): + out.zero_() + k._stream.synchronize() + launches_before = k.launches + while k.launches <= launches_before and not k._disabled and time.monotonic() < deadline: + k.tick() # a free queue slot and (mm) the period cadence gate permitting + time.sleep(0.05) + assert not k._disabled, "keepalive disabled while launching" + assert k.launches > launches_before, "no runtime chunk was launched" + finally: + k.drain() + assert not k._inflight and k._stream is None and k.mode is None # released when the gate opens + # No synchronize() here on purpose: only drain() may have waited for the chunks. + computed = bool(torch.all(out > 0)) if mode == "spin" else out[0, 0].item() == ka._MM_DIM + return mode, computed + + +def test_mm_fallback_on_device(monkeypatch): + """The torch.mm path taken when the Triton self-test does not pass.""" + monkeypatch.setattr(ka, "_SELFTEST_VERDICT", {torch.cuda.current_device(): False}) + k = ka.GpuKeepalive(torch.cuda.current_device()) + assert k.mode is None # nothing on the device until the first tick + assert _run_production_path(k) == ("mm", True) + assert k._a is None # the 96 MiB of operands are gone + + +def test_spin_kernel_on_device(monkeypatch): + """The real spin kernel through the production path: async self-test, then launches.""" + if not ka._HAVE_TRITON: + pytest.skip("triton not installed") + device = torch.cuda.current_device() + monkeypatch.setattr(ka, "_SELFTEST_VERDICT", {}) + k = ka.GpuKeepalive(device) + mode, computed = _run_production_path(k) + if mode != "spin": # diagnose before failing: a compute mode may forbid the child's context + child = subprocess.run( + [sys.executable, ka.__file__, "--selftest", str(device)], + capture_output=True, + timeout=300, + ) + if child.returncode == ka._SELFTEST_RC_NO_CONTEXT: + pytest.skip("compute mode forbids a second CUDA context on this device") + pytest.fail( + f"self-test rc={child.returncode}: {child.stderr.decode(errors='replace')[-2000:]}" + ) + assert computed