diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index 5f86c100f8c7..ccf0055b0d7e 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -308,7 +308,10 @@ def __init__(self, attn: TrtllmAttention) -> None: # normal forwarding never pays for an environment lookup. self._cache_sanity_check_enabled = _is_fmha_cache_sanity_check_enabled() self.fmha_libs: list[Fmha] = [] - for fmha_cls in get_enabled_fmha_lib_classes(): + for fmha_cls in ( + *getattr(attn, "prepended_fmha_lib_classes", ()), + *get_enabled_fmha_lib_classes(), + ): if fmha_cls.is_available(attn): self.fmha_libs.append(fmha_cls(attn)) self._layer_idx = getattr(attn, "layer_idx", None) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index a4d3c25b7bce..d69ea5be0128 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py @@ -1987,9 +1987,24 @@ def sparse_attn_indexer( ) if has_decode and not metadata.skip_indexer_for_gen_reqs and reuse_topk: - topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = ( - metadata.shared_topk_indices[:num_generations, :] + self._copy_reused_topk( + metadata, + topk_indices_buffer, + token_offset, + num_gen_tokens, + num_generations, + has_prefill, + gvr_prior_indices, + ) + elif ( + has_decode + and not metadata.skip_indexer_for_gen_reqs + and self._fill_decode_topk( + metadata, q_fp8, weights, topk_indices_buffer, token_offset, num_gen_tokens ) + ): + # A subclass filled the generation rows. + pass elif has_decode and not metadata.skip_indexer_for_gen_reqs: # Get decode lengths per request (from seq_lens) for validation gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations] @@ -2268,30 +2283,83 @@ def sparse_attn_indexer( prev_topk_dst.copy_(last_mtp_topk) if self.mtp_index_share and metadata.in_mtp_draft_loop and not reuse_topk: - rows = None - if has_decode: - next_n = num_gen_tokens // num_generations - gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens] - rows = self._mtp_last_accepted_rows( - gen_topk, metadata, num_contexts, num_generations, next_n - ) - if has_prefill: - ctx_last = ( - torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1 - ) - ctx_rows = topk_indices_buffer[ctx_last] - rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows]) - if rows is not None: - shared_topk_indices = metadata.shared_topk_indices - if shared_topk_indices is None: - metadata.shared_topk_indices = rows.contiguous() - else: - row_start = num_contexts if is_generation else 0 - shared_topk_indices[ - row_start : row_start + rows.shape[0], : rows.shape[1] - ].copy_(rows) + self._store_mtp_shared_rows( + topk_indices_buffer, + metadata, + num_contexts, + num_generations, + num_gen_tokens, + token_offset, + has_prefill, + has_decode, + is_generation, + ) return topk_indices_buffer + def _copy_reused_topk( + self, + metadata: DSAtrtllmAttentionMetadata, + topk_indices_buffer: torch.Tensor, + token_offset: int, + num_gen_tokens: int, + num_generations: int, + has_prefill: bool, + gvr_prior_indices: Optional[torch.Tensor], + ) -> None: + """Copy the shared TopK rows an MTP draft step reuses into its generation rows.""" + topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = ( + metadata.shared_topk_indices[:num_generations, :] + ) + + def _fill_decode_topk( + self, + metadata: DSAtrtllmAttentionMetadata, + q_fp8: torch.Tensor, + weights: torch.Tensor, + topk_indices_buffer: torch.Tensor, + token_offset: int, + num_gen_tokens: int, + ) -> bool: + """Fill the generation TopK rows ahead of the decode branch; True if it filled them. + + This class leaves them to the decode branch of ``sparse_attn_indexer``. + """ + return False + + def _store_mtp_shared_rows( + self, + topk_indices_buffer: torch.Tensor, + metadata: DSAtrtllmAttentionMetadata, + num_contexts: int, + num_generations: int, + num_gen_tokens: int, + token_offset: int, + has_prefill: bool, + has_decode: bool, + is_generation: Optional[bool], + ) -> None: + """Store each request's last accepted TopK row for the next MTP draft step to reuse.""" + rows = None + if has_decode: + next_n = num_gen_tokens // num_generations + gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens] + rows = self._mtp_last_accepted_rows( + gen_topk, metadata, num_contexts, num_generations, next_n + ) + if has_prefill: + ctx_last = torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1 + ctx_rows = topk_indices_buffer[ctx_last] + rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows]) + if rows is not None: + shared_topk_indices = metadata.shared_topk_indices + if shared_topk_indices is None: + metadata.shared_topk_indices = rows.contiguous() + else: + row_start = num_contexts if is_generation else 0 + shared_topk_indices[row_start : row_start + rows.shape[0], : rows.shape[1]].copy_( + rows + ) + def _mtp_last_accepted_rows( self, gen_topk: torch.Tensor, diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 9b5ac32120db..bc020324acc2 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -1704,6 +1704,8 @@ def is_sm_version_trtllm_gen_kernel(self, sm): class TrtllmAttention(AttentionBackend[TrtllmAttentionMetadata]): Metadata = TrtllmAttentionMetadata + # FMHA library classes this instance tries before the registry's libraries. + prepended_fmha_lib_classes: tuple = () @staticmethod def is_sm_version_trtllm_gen_kernel(sm): diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index 2ec483ecba15..7bb868bdf4ed 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -30,6 +30,7 @@ ResourceManager, ResourceManagerType, ) +from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import RemoteTailPhase @@ -324,6 +325,9 @@ class DecoderInputPreparer: path overrides ``_new_extra_inputs`` and ``_prepare_inputs_fast_path``. """ + # The model's decode prologue (RuntimeHooks.decode_prologue), set by the runner. + decode_prologue: DecodePrologue | None = None + def __init__( self, *, @@ -2561,6 +2565,16 @@ def preprocess_inputs( Make some changes to the device inputs and avoid blocking the async data transfer """ attn_meta = inputs.get("attn_metadata") + # A model's decode prologue may replace the device work below with one + # launch; it decides on the host before any device work and returns + # False to keep the code below. + prologue = self.decode_prologue + if prologue is not None and prologue( + self, inputs, enable_spec_decode=enable_spec_decode, runtime_draft_len=runtime_draft_len + ): + if self.guided_decoder is not None: + self.guided_decoder.token_event.record() + return inputs # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. if attn_meta is not None: attn_meta.on_update_kv_lens() diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 8f3c7e8a48ec..03242760a2a8 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -36,6 +36,11 @@ from tensorrt_llm._torch.pyexecutor.guided_decoder import CapturableGuidedDecoder from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( + DEFAULT_RUNTIME_HOOKS, + RuntimeHooks, + runtime_hooks_of, +) from tensorrt_llm._torch.pyexecutor.sampler.beam_search import PendingBeamHistories from tensorrt_llm._torch.pyexecutor.sampler.sampler_common import SampleType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -127,6 +132,9 @@ def __getattr__(self, name: str) -> Any: class DecoderRunner(ScheduledModelRunner): """Run scheduled decoder generation with its own buffers, graphs and metadata.""" + # Step-loop behaviour the model opted into (RuntimeHooks); all off unless it declares otherwise. + runtime_hooks: RuntimeHooks = DEFAULT_RUNTIME_HOOKS + # Optional execution paths that a subclass can turn off. _context_warmups_supported = True _context_graph_promotion_supported = True @@ -204,6 +212,8 @@ def __init__( self._dummy_requests = self._make_dummy_requests() self._init_decoder_state() self._inputs = self._make_input_preparer(mm_item_scheduler) + self.runtime_hooks = runtime_hooks_of(self._model) + self._inputs.decode_prologue = self.runtime_hooks.decode_prologue self._graphs = self._make_cuda_graphs() self._inputs.cuda_graph_runner = self.cuda_graph_runner self._initialize_moe_workspace_reclaimer() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a0e323196a53..02378b6334f9 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -120,6 +120,7 @@ get_from_waiting_queue, merge_requests) from .resource_manager import (NoFreeSlotsError, ResourceManager, ResourceManagerType, request_context) +from .runtime_hooks import ExecutorEvents, runtime_hooks_of from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, SampleStateTensors) from .scheduler import (RequestScheduler, ScheduledRequests, @@ -451,6 +452,16 @@ def _single_rank_adp_skips_padding(executor: "PyExecutor") -> bool: class PyExecutor: + # True while the overlap loop runs with the execution stream current + # (_overlap_loop_stream_scope). + _loop_on_execution_stream = False + # Whether the model opted into running the overlap loop on the execution + # stream (RuntimeHooks.loop_on_execution_stream); read once at construction. + _execution_stream_loop_armed = False + # The model's executor events (RuntimeHooks.executor_events); read once at + # construction. + _executor_events: Optional[ExecutorEvents] = None + # Minimum number of async micro batches for async PP execution. # This is a trade-off between memory usage and performance. # If the number of micro batches is too small, the executor will spend too much time in synchronization. @@ -571,6 +582,9 @@ def __init__( # related modules self.resource_manager = resource_manager self.model_engine = model_engine + hooks = runtime_hooks_of(model_engine.model) + self._execution_stream_loop_armed = hooks.loop_on_execution_stream + self._executor_events = hooks.executor_events self._enable_adp_dummy_fixes = enable_adp_dummy_fixes self._enable_scheduler_aware_adp_dummy = enable_scheduler_aware_adp_dummy self._enable_non_overlap_adp_forward_intent = ( @@ -980,6 +994,12 @@ def __init__( # Ensure the default stream waits for execution_stream to complete # before subsequent operations. torch.cuda.current_stream().wait_stream(self.execution_stream) + if self._executor_events is not None: + # The model's actions after warmup and graph capture, on every + # rank (some are collective, so every rank reaches this point, + # and a rank-local failure arms the rank-crash watchdog). + with _distributed_warmup_guard(self.dist, self.dist.mapping): + self._executor_events.after_warmup() self.encoder_launch_executor = (ThreadPoolExecutor( max_workers=1, thread_name_prefix="encoder-launch") if self.is_encoder_decoder else None) @@ -3144,6 +3164,9 @@ def _executor_loop_cleanup(self): f"Error waiting on PP send handles during cleanup: " f"{traceback.format_exc()}") + if self._executor_events is not None and self._event_loop_completed: + self._executor_events.at_shutdown() + def _pp_rank_schedules(self) -> bool: """Whether this rank computes the schedule that the other PP ranks receive. @@ -6197,11 +6220,44 @@ def control_action(self, def _wait_for_model_engine_input_copy(self): self.model_engine.wait_for_input_copy() + @contextmanager + def _overlap_loop_stream_scope(self): + """Run the overlap loop with the execution stream current, when the model opted in. + + _forward_step makes the execution stream wait for the current stream + before every forward and the current stream wait for the execution + stream after it. On the thread's default stream, which in steady decode + carries only these waits and timing events, step N+1's inputs wait for + decode graph N through two cross-stream event hops. With the execution + stream current the loop's stream work is queued in the same order on one + stream and _forward_step skips both waits. + + Not taken with a two-model drafter, a guided decoder, a KV connector, + an encoder-decoder model or multimodal encoder item scheduling, whose + side streams synchronize with the current stream. + """ + if not (self._execution_stream_loop_armed and self.drafter is None and + self.draft_model_engine is None and self.guided_decoder is None + and self.kv_connector_manager is None + and not self.is_encoder_decoder + and not self._mm_encoder_item_scheduling_enabled): + yield + return + with torch.cuda.stream(self.execution_stream): + self._loop_on_execution_stream = True + try: + yield + finally: + self._loop_on_execution_stream = False + def _executor_loop_overlap(self) -> None: torch.cuda.set_device(self.device_id) # ensure the context is created, otherwise, some MPI calls will fail. CUASSERT(cudart.cudaSetDevice(self.device_id)) - with self._profiler() as profile_step, self.hang_detector: + if self._executor_events is not None: + self._executor_events.before_overlap_loop() + with self._overlap_loop_stream_scope(), self._profiler() as profile_step, \ + self.hang_detector: iter_start_time = time.time() iter_stats = None previous_tensors_device = None @@ -8946,8 +9002,10 @@ def forward(scheduled_requests, resource_manager, new_tensors_device, cache_indirection_buffer = self.sampler.get_cache_indirection() # Run model forward on the execution stream for proper synchronization - # with KVCacheTransferManager's onboard/offload operations. - self.execution_stream.wait_stream(torch.cuda.current_stream()) + # with KVCacheTransferManager's onboard/offload operations. With the + # loop on the execution stream both waits are program order. + if not self._loop_on_execution_stream: + self.execution_stream.wait_stream(torch.cuda.current_stream()) self._attach_encoder_output_to_execution_stream(scheduled_requests) with torch.cuda.stream(self.execution_stream): outputs = forward(scheduled_requests, self.resource_manager, @@ -8964,7 +9022,8 @@ def forward(scheduled_requests, resource_manager, new_tensors_device, # Ensure the default stream waits for execution_stream to complete # before downstream operations use the outputs. - torch.cuda.current_stream().wait_stream(self.execution_stream) + if not self._loop_on_execution_stream: + torch.cuda.current_stream().wait_stream(self.execution_stream) self._kv_connector_wait_for_save() if route_capture is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py new file mode 100644 index 000000000000..772845919949 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Step-loop behaviour a model may opt into, read through ``runtime_hooks_of``. + +A model declares its hooks as ``runtime_hooks`` (an instance attribute, a class +attribute or a property). Every field defaults to off, and a model that declares +nothing gets ``RuntimeHooks()``: the step loop then runs exactly as it does +without this module. Every engine call site +of a hook runs its original statements when the hook is off. + +The hooks change *where* and *how* the step loop moves data, never what a +forward computes. The protocols below state only the methods an engine call +site needs; the objects that implement them are owned by the model. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional, Protocol + + +class DecodePrologue(Protocol): + """Replaces the decode step's input prologue with one launch. + + Called first in the step's input preprocessing, inside CUDA graph capture + as well. It decides on the host before any device work; returning False + keeps the native prologue, and returning True means every engine-visible + tensor the native prologue writes holds the value it would have written. + """ + + def __call__( + self, preparer: Any, inputs: Any, *, enable_spec_decode: bool, runtime_draft_len: int + ) -> bool: + """``preparer`` is the ``DecoderInputPreparer`` running the step's preprocessing.""" + + +class ExecutorEvents(Protocol): + """Points in the executor's life a model may act at.""" + + def after_warmup(self) -> None: + """After warmup and CUDA graph capture, before the first request is served.""" + + def before_overlap_loop(self) -> None: + """Right before the overlap executor loop starts, on every rank; must not raise.""" + + def at_shutdown(self) -> None: + """When the executor loop has completed, during cleanup; must not raise.""" + + +@dataclass(frozen=True) +class RuntimeHooks: + """The step-loop hooks a model opts into. All off by default.""" + + # Run the overlap executor loop with the execution stream current, so the + # loop's own stream work and the forward are ordered on one stream. + loop_on_execution_stream: bool = False + # The decode step's input prologue in one launch (see DecodePrologue). + decode_prologue: Optional[DecodePrologue] = None + # Actions after warmup, before the overlap loop and at shutdown (see ExecutorEvents). + executor_events: Optional[ExecutorEvents] = None + + +DEFAULT_RUNTIME_HOOKS = RuntimeHooks() + + +def runtime_hooks_of(model: object) -> RuntimeHooks: + """The hooks ``model`` declares, or the all-off default when it declares none. + + Declaring hooks is optional for every model class, so absence is the common + case and means the default. + """ + hooks = getattr(model, "runtime_hooks", None) + if hooks is None: + return DEFAULT_RUNTIME_HOOKS + if not isinstance(hooks, RuntimeHooks): + raise TypeError( + f"{type(model).__name__}.runtime_hooks must be a RuntimeHooks, got {type(hooks).__name__}" + ) + return hooks diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index d16299ede1b9..5d09cdc3d680 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -5327,3 +5327,103 @@ def _run_indexer(): assert "indexer_topk_out_buffer" in metadata.cuda_graph_buffers.buffers, ( "indexer topk-output buffer must be drawn from the cuda_graph_buffers arena" ) + + +# --- The subclass seams of Indexer.sparse_attn_indexer ---------------------------------------- + + +def _seam_indexer(): + indexer = Indexer.__new__(Indexer) + torch.nn.Module.__init__(indexer) + return indexer + + +def test_reused_topk_copy_is_the_row_copy(): + gen = torch.Generator().manual_seed(0) + shared = torch.randint(0, 1000, (4, 16), generator=gen, dtype=torch.int32) + buffer = torch.full((12, 16), -1, dtype=torch.int32) + expected = buffer.clone() + expected[2:5, :] = shared[:3, :] + + _seam_indexer()._copy_reused_topk( + SimpleNamespace(shared_topk_indices=shared), buffer, 2, 3, 3, False, None + ) + + assert torch.equal(buffer, expected) + + +def test_fill_decode_topk_leaves_the_rows_to_the_decode_branch(): + buffer = torch.full((6, 8), -1, dtype=torch.int32) + + filled = _seam_indexer()._fill_decode_topk(SimpleNamespace(), None, None, buffer, 0, 6) + + assert filled is False + assert bool((buffer == -1).all()) + + +def _stored_rows_inline( + indexer, + topk_indices_buffer, + metadata, + num_contexts, + num_generations, + num_gen_tokens, + token_offset, + has_prefill, + has_decode, + is_generation, +): + """The statements ``sparse_attn_indexer`` ran in place of ``_store_mtp_shared_rows``.""" + rows = None + if has_decode: + next_n = num_gen_tokens // num_generations + gen_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens] + rows = indexer._mtp_last_accepted_rows( + gen_topk, metadata, num_contexts, num_generations, next_n + ) + if has_prefill: + ctx_last = torch.cumsum(metadata.seq_lens_cuda[:num_contexts].to(torch.long), dim=0) - 1 + ctx_rows = topk_indices_buffer[ctx_last] + rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows]) + if rows is not None: + shared_topk_indices = metadata.shared_topk_indices + if shared_topk_indices is None: + metadata.shared_topk_indices = rows.contiguous() + else: + row_start = num_contexts if is_generation else 0 + shared_topk_indices[row_start : row_start + rows.shape[0], : rows.shape[1]].copy_(rows) + + +@pytest.mark.parametrize("has_prefill, has_decode", [(False, True), (True, True), (True, False)]) +@pytest.mark.parametrize("is_generation", [True, False]) +@pytest.mark.parametrize("preset", [False, True], ids=["unset", "preset"]) +@pytest.mark.parametrize("accepted", [None, (0, 0, 1, 4, 2)], ids=["no_counts", "counts"]) +def test_mtp_shared_rows_are_the_inline_statements( + has_prefill, has_decode, is_generation, preset, accepted +): + num_contexts, num_generations, next_n = 2, 3, 4 + seq_lens = torch.tensor([5, 3] + [next_n] * num_generations, dtype=torch.int32) + token_offset = int(seq_lens[:num_contexts].sum()) + num_gen_tokens = num_generations * next_n + gen = torch.Generator().manual_seed(1) + buffer = torch.randint( + 0, 1000, (token_offset + num_gen_tokens, 6), generator=gen, dtype=torch.int32 + ) + + def metadata(): + return SimpleNamespace( + seq_lens_cuda=seq_lens.clone(), + mtp_num_accepted=None if accepted is None else torch.tensor(accepted), + shared_topk_indices=(torch.full((8, 8), -7, dtype=torch.int32) if preset else None), + ) + + args = (num_contexts, num_generations, num_gen_tokens, token_offset) + flags = (has_prefill, has_decode, is_generation) + expected, got = metadata(), metadata() + indexer = _seam_indexer() + _stored_rows_inline(indexer, buffer.clone(), expected, *args, *flags) + seam_buffer = buffer.clone() + indexer._store_mtp_shared_rows(seam_buffer, got, *args, *flags) + + assert torch.equal(seam_buffer, buffer) + assert torch.equal(got.shared_topk_indices, expected.shared_topk_indices) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py index 623a33a5e536..04aa0c8da075 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py @@ -177,6 +177,9 @@ def prefill( aux_stream=None, ) indexer.forward_prepared = MethodType(Indexer.forward_prepared, indexer) + indexer._copy_reused_topk = MethodType(Indexer._copy_reused_topk, indexer) + indexer._fill_decode_topk = MethodType(Indexer._fill_decode_topk, indexer) + indexer._store_mtp_shared_rows = MethodType(Indexer._store_mtp_shared_rows, indexer) with ( patch.object( top_k, "update_gvr_prior_from_prefill", wraps=top_k.update_gvr_prior_from_prefill diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py index 7cfb59eee093..397bfd93a4b5 100644 --- a/tests/unittest/_torch/attention/test_fmha_manager.py +++ b/tests/unittest/_torch/attention/test_fmha_manager.py @@ -974,3 +974,60 @@ def __init__(self, attn: TrtllmAttention) -> None: assert isinstance(attn._fmha_manager.fmha_libs[0], _NewFmha) assert quant_states_during_construction == [True] assert events == [("support", "old", None)] + + +def _listed_fmha(name: str, available: bool, checked: list) -> type: + """An FMHA library class that records its availability check.""" + + def is_available(cls, attn) -> bool: + checked.append((name, attn)) + return available + + return type( + name, + (), + {"is_available": classmethod(is_available), "__init__": lambda self, attn: None}, + ) + + +@pytest.mark.parametrize("declared", [False, True], ids=["undeclared", "empty"]) +def test_manager_without_prepended_libraries_builds_the_registry_list(declared: bool) -> None: + checked: list = [] + registry = [_listed_fmha("First", True, checked), _listed_fmha("Second", True, checked)] + attn = FakeAttention() + if declared: + attn.prepended_fmha_lib_classes = () + + with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry): + manager = FmhaManager(attn) + + assert [type(lib) for lib in manager.fmha_libs] == registry + assert TrtllmAttention.prepended_fmha_lib_classes == () + + +def test_manager_tries_a_prepended_library_first() -> None: + checked: list = [] + registry = [_listed_fmha("First", True, checked), _listed_fmha("Second", True, checked)] + prepended = _listed_fmha("Prepended", True, checked) + attn = FakeAttention() + attn.prepended_fmha_lib_classes = (prepended,) + + with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry): + manager = FmhaManager(attn) + + assert [type(lib) for lib in manager.fmha_libs] == [prepended, *registry] + assert [name for name, _ in checked] == ["Prepended", "First", "Second"] + assert all(checked_attn is attn for _, checked_attn in checked) + + +def test_manager_skips_an_unavailable_prepended_library() -> None: + checked: list = [] + registry = [_listed_fmha("First", True, checked)] + attn = FakeAttention() + attn.prepended_fmha_lib_classes = (_listed_fmha("Prepended", False, checked),) + + with patch.object(fmha_manager, "get_enabled_fmha_lib_classes", return_value=registry): + manager = FmhaManager(attn) + + assert [type(lib) for lib in manager.fmha_libs] == registry + assert [name for name, _ in checked] == ["Prepended", "First"] diff --git a/tests/unittest/_torch/executor/test_gpu_keepalive.py b/tests/unittest/_torch/executor/test_gpu_keepalive.py index 5c3318c00690..51d9b63b277e 100644 --- a/tests/unittest/_torch/executor/test_gpu_keepalive.py +++ b/tests/unittest/_torch/executor/test_gpu_keepalive.py @@ -315,6 +315,7 @@ class Stub: # notified first, then the keepalive is closed, then the PP handles are awaited. class CleanupStub: _executor_loop_cleanup = pe.PyExecutor._executor_loop_cleanup + _executor_events = None def __init__(self, keepalive): self.events = [] diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index b12ca615f74d..2d58a0534378 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -83,6 +83,7 @@ def _forward_step_executor( runtime_draft_len=runtime_draft_len, sampler=Mock(get_cache_indirection=Mock(return_value=object())), execution_stream=Mock(), + _loop_on_execution_stream=False, _iter_adp_dummy_ctx_tokens=0, _iter_adp_dummy_gen_tokens=0, _compute_adp_dummy_tokens=PyExecutor._compute_adp_dummy_tokens, @@ -2585,6 +2586,8 @@ class _CleanupStub: """Stub for _executor_loop_cleanup: records the order in which the shutdown notification and PP-handle wait happen.""" + _executor_events = None + def __init__(self, pp_handles_raise=False): self.response_lock = threading.Lock() self.response_cv = threading.Condition(self.response_lock) diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py new file mode 100644 index 000000000000..d2459e6804de --- /dev/null +++ b/tests/unittest/_torch/executor/test_runtime_hooks.py @@ -0,0 +1,685 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""How the engine reads the step-loop hooks a model declares (``runtime_hooks_of``). + +The engine's call sites run their original statements when a hook is off: every +object that holds a hook declares it off at class level, every call a call site +makes on a hook is a method of the hook's protocol, every read of a hook in the +engine files is decided by its all-off value toward the original statements +(checked on the source), and call sites of each hook run on stand-ins with the +hook off and on. +""" + +import ast +import contextlib +import dataclasses +import inspect +import types +from pathlib import Path + +import pytest +import torch + +import tensorrt_llm +from tensorrt_llm._torch.pyexecutor.runtime_hooks import ( + DEFAULT_RUNTIME_HOOKS, + ExecutorEvents, + RuntimeHooks, + runtime_hooks_of, +) + +pytestmark = pytest.mark.cpu_only + + +class _PlainModel(torch.nn.Module): + """A model that declares no hooks, like every model outside the opt-in.""" + + +class _ClassHooksModel(torch.nn.Module): + runtime_hooks = RuntimeHooks(loop_on_execution_stream=True) + + +class _Events: + def after_warmup(self) -> None: + pass + + def before_overlap_loop(self) -> None: + pass + + def at_shutdown(self) -> None: + pass + + +@pytest.mark.parametrize("model", [_PlainModel(), object()], ids=["module", "object"]) +def test_a_model_without_hooks_gets_the_all_off_default(model): + hooks = runtime_hooks_of(model) + + assert hooks is DEFAULT_RUNTIME_HOOKS + assert hooks == RuntimeHooks() + assert not hooks.loop_on_execution_stream + assert hooks.decode_prologue is None + assert hooks.executor_events is None + + +def test_every_field_defaults_to_off(): + for field in dataclasses.fields(RuntimeHooks): + assert getattr(DEFAULT_RUNTIME_HOOKS, field.name) in (False, None), field.name + + +def test_hooks_declared_as_none_mean_the_default(): + model = _PlainModel() + model.runtime_hooks = None + + assert runtime_hooks_of(model) is DEFAULT_RUNTIME_HOOKS + + +def test_class_attribute_hooks_are_returned_as_declared(): + assert runtime_hooks_of(_ClassHooksModel()) is _ClassHooksModel.runtime_hooks + + +def test_instance_attribute_hooks_are_returned_as_declared(): + hooks = RuntimeHooks(executor_events=_Events()) + model = _PlainModel() + model.runtime_hooks = hooks + + assert runtime_hooks_of(model) is hooks + + +def test_instance_attribute_overrides_class_attribute(): + hooks = RuntimeHooks() + model = _ClassHooksModel() + model.runtime_hooks = hooks + + assert runtime_hooks_of(model) is hooks + + +@pytest.mark.parametrize( + "declared", + [True, {"loop_on_execution_stream": True}, "RuntimeHooks()"], + ids=["bool", "dict", "str"], +) +def test_hooks_of_the_wrong_type_raise(declared): + model = _PlainModel() + model.runtime_hooks = declared + + with pytest.raises(TypeError, match=r"_PlainModel\.runtime_hooks must be a RuntimeHooks"): + runtime_hooks_of(model) + + +def test_hooks_are_immutable_once_declared(): + hooks = RuntimeHooks() + + with pytest.raises(dataclasses.FrozenInstanceError): + hooks.loop_on_execution_stream = True + + +def test_engine_classes_declare_the_hooks_off(): + """Every object that holds a hook declares it off at class level.""" + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + assert DecoderInputPreparer.decode_prologue is None + assert DecoderRunner.runtime_hooks is DEFAULT_RUNTIME_HOOKS + assert PyExecutor._execution_stream_loop_armed is False + assert PyExecutor._loop_on_execution_stream is False + assert PyExecutor._executor_events is None + + +# Per hook protocol: the local names and attribute names a call site holds the hook in, and +# the engine files with call sites. +_HOOK_CONSUMERS = { + ExecutorEvents: ((), ("_executor_events",), ("pyexecutor/py_executor.py",)), +} + + +def _protocol_methods(protocol) -> set: + return { + name + for name, value in vars(protocol).items() + if inspect.isfunction(value) and not name.startswith("_") + } + + +def _hook_calls(tree: ast.AST, local_names, attr_names): + """``(method, line)`` of every call on a hook held in one of the locals or attributes.""" + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + receiver = node.func.value + if (isinstance(receiver, ast.Name) and receiver.id in local_names) or ( + isinstance(receiver, ast.Attribute) and receiver.attr in attr_names + ): + yield node.func.attr, node.lineno + + +@pytest.mark.parametrize("protocol", list(_HOOK_CONSUMERS), ids=lambda p: p.__name__) +def test_every_hook_call_site_uses_its_protocol(protocol): + """A call site calls only its hook's protocol methods, and the protocol has no unused method.""" + local_names, attr_names, consumers = _HOOK_CONSUMERS[protocol] + root = Path(tensorrt_llm.__file__).resolve().parent / "_torch" + methods = _protocol_methods(protocol) + used, unknown = set(), [] + for rel in consumers: + tree = ast.parse((root / rel).read_text()) + calls = list(_hook_calls(tree, local_names, attr_names)) + assert calls, f"{rel} has no {protocol.__name__} call site" + for method, line in calls: + used.add(method) + if method not in methods: + unknown.append(f"{rel}:{line}: {method}") + assert not unknown, f"calls outside the {protocol.__name__} protocol:\n" + "\n".join(unknown) + assert used == methods + + +def test_executor_events_protocol(): + assert _protocol_methods(ExecutorEvents) == { + "after_warmup", + "before_overlap_loop", + "at_shutdown", + } + events = _Events() + assert RuntimeHooks(executor_events=events).executor_events is events + + +class _StreamRecorder: + """Stands in for ``torch.cuda.stream``: records the streams it is entered with.""" + + def __init__(self) -> None: + self.entered = [] + + @contextlib.contextmanager + def __call__(self, stream): + self.entered.append(stream) + yield + + +def _scope_executor(armed: bool, **overrides): + """A ``PyExecutor`` with only the attributes ``_overlap_loop_stream_scope`` reads.""" + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + executor = PyExecutor.__new__(PyExecutor) + executor._execution_stream_loop_armed = armed + executor.execution_stream = object() + attrs = dict( + drafter=None, + draft_model_engine=None, + guided_decoder=None, + kv_connector_manager=None, + is_encoder_decoder=False, + _mm_encoder_item_scheduling_enabled=False, + ) + attrs.update(overrides) + for name, value in attrs.items(): + setattr(executor, name, value) + return executor + + +def test_overlap_loop_scope_is_a_no_op_when_not_armed(monkeypatch): + executor = _scope_executor(armed=False) + recorder = _StreamRecorder() + monkeypatch.setattr(torch.cuda, "stream", recorder) + + with executor._overlap_loop_stream_scope(): + assert executor._loop_on_execution_stream is False + assert recorder.entered == [] + assert "_loop_on_execution_stream" not in vars(executor) + + +def test_overlap_loop_scope_engages_when_armed(monkeypatch): + executor = _scope_executor(armed=True) + recorder = _StreamRecorder() + monkeypatch.setattr(torch.cuda, "stream", recorder) + + with executor._overlap_loop_stream_scope(): + assert executor._loop_on_execution_stream is True + assert recorder.entered == [executor.execution_stream] + assert executor._loop_on_execution_stream is False + + +def test_overlap_loop_scope_resets_the_flag_when_the_loop_raises(monkeypatch): + executor = _scope_executor(armed=True) + monkeypatch.setattr(torch.cuda, "stream", _StreamRecorder()) + + with pytest.raises(RuntimeError, match="loop failed"): + with executor._overlap_loop_stream_scope(): + raise RuntimeError("loop failed") + assert executor._loop_on_execution_stream is False + + +@pytest.mark.parametrize( + "name, value", + [ + ("drafter", object()), + ("draft_model_engine", object()), + ("guided_decoder", object()), + ("kv_connector_manager", object()), + ("is_encoder_decoder", True), + ("_mm_encoder_item_scheduling_enabled", True), + ], +) +def test_overlap_loop_scope_stays_off_with_side_streams(monkeypatch, name, value): + executor = _scope_executor(armed=True, **{name: value}) + recorder = _StreamRecorder() + monkeypatch.setattr(torch.cuda, "stream", recorder) + + with executor._overlap_loop_stream_scope(): + assert executor._loop_on_execution_stream is False + assert recorder.entered == [] + + +# --------------------------------------------------------------------------------------------- +# With every hook off, every call site runs its original statements. +# +# The check below reads the engine files that consume a hook and evaluates each read of a hook +# at its all-off value. An ``if`` that reads a hook must be decided by that value; the branch it +# takes is checked like any other code and the other branch is not run, so it may use the hook +# freely. Outside such tests, live code may only bind a hook value to a local (or to an +# attribute that declares the hook) at its all-off value. Any other read is reported with its +# file and line. +# --------------------------------------------------------------------------------------------- + +_SEAM_CONSUMERS = ( + "pyexecutor/engine/runners/decoder/inputs.py", + "pyexecutor/py_executor.py", +) + +# Attributes that hold a hook, and their value when every hook is off. +_SEAM_ATTRS = { + "_executor_events": None, + "_loop_on_execution_stream": False, + "_execution_stream_loop_armed": False, + "runtime_hooks": DEFAULT_RUNTIME_HOOKS, + "decode_prologue": None, +} +# Locals the call sites hold a hook (or a decision on one) in. +_SEAM_LOCALS = ("prologue",) +_HOOK_FIELDS = {field.name for field in dataclasses.fields(RuntimeHooks)} + + +class _Unknown: + """A value the all-off evaluation cannot decide.""" + + +_UNKNOWN = _Unknown() + + +class _AllOffWalk: + """Walks one engine file with every hook at its all-off value.""" + + def __init__(self, rel: str) -> None: + self.rel = rel + self.problems = [] + self.decided = 0 + + def report(self, node: ast.AST, what: str) -> None: + self.problems.append(f"{self.rel}:{node.lineno}: {what}") + + # -- reads of a hook ------------------------------------------------------------------ + @staticmethod + def _is_seam_call(node: ast.AST) -> bool: + return ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "runtime_hooks_of" + ) + + def reads_seam(self, node: ast.AST, env: dict) -> bool: + for sub in ast.walk(node): + if isinstance(sub, ast.Attribute) and sub.attr in _SEAM_ATTRS: + return True + if isinstance(sub, ast.Name) and sub.id in env: + return True + if self._is_seam_call(sub): + return True + return False + + def value(self, node: ast.AST, env: dict): + """The all-off value of ``node``, or ``_UNKNOWN``.""" + if isinstance(node, ast.Constant): + return node.value + if isinstance(node, ast.Name): + return env.get(node.id, _UNKNOWN) + if isinstance(node, ast.Attribute): + receiver = self.value(node.value, env) + if receiver is DEFAULT_RUNTIME_HOOKS and node.attr in _HOOK_FIELDS: + return getattr(DEFAULT_RUNTIME_HOOKS, node.attr) + if node.attr in _SEAM_ATTRS: + return _SEAM_ATTRS[node.attr] + return _UNKNOWN + if self._is_seam_call(node): + return DEFAULT_RUNTIME_HOOKS + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not): + operand = self.value(node.operand, env) + return _UNKNOWN if operand is _UNKNOWN else not operand + if isinstance(node, ast.BoolOp): + values = [self.value(v, env) for v in node.values] + known = [v for v in values if v is not _UNKNOWN] + if isinstance(node.op, ast.And): + if any(not v for v in known): + return False + return True if len(known) == len(values) else _UNKNOWN + if any(known): + return True + return False if len(known) == len(values) else _UNKNOWN + if isinstance(node, ast.Compare) and len(node.ops) == 1: + left, right = self.value(node.left, env), self.value(node.comparators[0], env) + op = node.ops[0] + if isinstance(op, (ast.Is, ast.IsNot)) and left is not _UNKNOWN: + if isinstance(node.comparators[0], ast.Constant) and right is None: + return (left is None) == isinstance(op, ast.Is) + return _UNKNOWN + + # -- statements ------------------------------------------------------------------------- + def block(self, body, env: dict) -> bool: + """Walk live statements; True when the block always ends in return or raise.""" + for stmt in body: + if self.stmt(stmt, env): + return True + return False + + def bind(self, stmt, targets, value, env: dict) -> None: + reads = self.reads_seam(value, env) + names = [t.id for t in targets if isinstance(t, ast.Name)] + attrs = [t for t in targets if isinstance(t, ast.Attribute) and t.attr in _SEAM_ATTRS] + if not (reads or attrs or any(name in env or name in _SEAM_LOCALS for name in names)): + return + if len(names) + len(attrs) != len(targets): + self.report(stmt, "a hook value is bound to something other than a local") + return + off = self.value(value, env) + if off is _UNKNOWN: + if reads or attrs: + self.report(stmt, "a hook value with every hook off is not decided") + for name in names: + env.pop(name, None) + return + for name in names: + env[name] = off + for attr in attrs: + if off != _SEAM_ATTRS[attr.attr]: + self.report(stmt, f"{attr.attr} is set to a value other than its all-off value") + + def stmt(self, stmt: ast.stmt, env: dict) -> bool: + if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)): + self.block(stmt.body, dict(env)) + return False + if isinstance(stmt, ast.ClassDef): + self.block(stmt.body, {}) + return False + if isinstance(stmt, (ast.If, ast.While)): + if not self.reads_seam(stmt.test, env): + body_ends = self.block(stmt.body, env) + else_ends = self.block(stmt.orelse, env) + return isinstance(stmt, ast.If) and body_ends and else_ends + taken = self.value(stmt.test, env) + if taken is _UNKNOWN: + self.report(stmt, "a test of a hook is not decided with every hook off") + return False + self.decided += 1 + if isinstance(stmt, ast.While) and taken: + self.report(stmt, "a loop on a hook runs with every hook off") + return False + return self.block(stmt.body if taken else stmt.orelse, env) + if isinstance(stmt, ast.Assign): + self.bind(stmt, stmt.targets, stmt.value, env) + return False + if isinstance(stmt, ast.AnnAssign): + if stmt.value is not None: + self.bind(stmt, [stmt.target], stmt.value, env) + return False + if isinstance(stmt, (ast.With, ast.AsyncWith)): + for item in stmt.items: + if self.reads_seam(item.context_expr, env): + self.report(stmt, "a hook is entered as a context") + return self.block(stmt.body, env) + if isinstance(stmt, ast.Try): + ends = self.block(stmt.body, env) + for handler in stmt.handlers: + self.block(handler.body, env) + self.block(stmt.orelse, env) + return self.block(stmt.finalbody, env) or (ends and not stmt.handlers) + if isinstance(stmt, (ast.For, ast.AsyncFor)): + if self.reads_seam(stmt.iter, env): + self.report(stmt, "a loop iterates over a hook") + self.block(stmt.body, env) + self.block(stmt.orelse, env) + return False + if self.reads_seam(stmt, env): + self.report(stmt, "a hook is used outside a test decided with every hook off") + return isinstance(stmt, (ast.Return, ast.Raise)) + + +@pytest.mark.parametrize("rel", _SEAM_CONSUMERS) +def test_every_seam_read_has_a_native_branch(rel): + """With every hook off, each call site takes the branch with its original statements.""" + root = Path(tensorrt_llm.__file__).resolve().parent / "_torch" + walk = _AllOffWalk(rel) + walk.block(ast.parse((root / rel).read_text()).body, {}) + assert not walk.problems, "\n".join(walk.problems) + assert walk.decided, f"{rel} has no call site of a hook" + + +def test_the_seam_check_reports_undecided_reads(): + source = """ +def step(self, x): + events = self._executor_events + if events is not None: + events.after_warmup() + else: + x.copy_() + events.at_shutdown() + if x or events: + pass + self._executor_events = x +""" + walk = _AllOffWalk("sample.py") + walk.block(ast.parse(source).body, {}) + assert walk.decided == 1 + assert walk.problems == [ + "sample.py:8: a hook is used outside a test decided with every hook off", + "sample.py:9: a test of a hook is not decided with every hook off", + "sample.py:11: a hook value with every hook off is not decided", + ] + + +def test_the_decode_prologue_call_site_uses_its_protocol(): + """The input preparer calls the prologue with the protocol's positional and keyword parameters.""" + from tensorrt_llm._torch.pyexecutor.runtime_hooks import DecodePrologue + + signature = inspect.signature(DecodePrologue.__call__) + positional = [ + p.name + for p in signature.parameters.values() + if p.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD and p.name != "self" + ] + keywords = { + p.name for p in signature.parameters.values() if p.kind is inspect.Parameter.KEYWORD_ONLY + } + root = Path(tensorrt_llm.__file__).resolve().parent / "_torch" + tree = ast.parse((root / "pyexecutor/engine/runners/decoder/inputs.py").read_text()) + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "prologue" + ] + assert len(calls) == 1 + assert len(calls[0].args) == len(positional) + assert {keyword.arg for keyword in calls[0].keywords} == keywords + + +# -- the decode prologue ------------------------------------------------------------------------ + + +class _AttnMetadata: + def __init__(self) -> None: + self.kv_lens_updates = 0 + + def on_update_kv_lens(self) -> None: + self.kv_lens_updates += 1 + + +def _prologue_preparer(prologue): + """A ``DecoderInputPreparer`` with only what ``preprocess_inputs`` reads without spec decoding.""" + from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer + + preparer = DecoderInputPreparer.__new__(DecoderInputPreparer) + if prologue is not None: + preparer.decode_prologue = prologue + preparer.guided_decoder = None + preparer._config = types.SimpleNamespace(disable_overlap_scheduler=True) + preparer.mapping = types.SimpleNamespace(has_cp_helix=lambda: False) + return preparer + + +def test_preprocess_inputs_runs_the_native_prologue_without_a_hook(): + preparer = _prologue_preparer(None) + inputs = {"attn_metadata": _AttnMetadata()} + + assert preparer.decode_prologue is None + assert preparer.preprocess_inputs(inputs, enable_spec_decode=False, runtime_draft_len=0) is ( + inputs + ) + assert inputs["attn_metadata"].kv_lens_updates == 1 + + +@pytest.mark.parametrize("engaged", [False, True], ids=["declined", "engaged"]) +def test_preprocess_inputs_keeps_the_native_prologue_when_the_hook_declines(engaged): + calls = [] + + def prologue(preparer, inputs, *, enable_spec_decode, runtime_draft_len): + calls.append((preparer, inputs, enable_spec_decode, runtime_draft_len)) + return engaged + + preparer = _prologue_preparer(prologue) + inputs = {"attn_metadata": _AttnMetadata()} + + assert preparer.preprocess_inputs(inputs, enable_spec_decode=False, runtime_draft_len=3) is ( + inputs + ) + assert calls == [(preparer, inputs, False, 3)] + assert inputs["attn_metadata"].kv_lens_updates == (0 if engaged else 1) + + +# -- the execution-stream loop -------------------------------------------------------------------- + + +class _Stream: + def __init__(self, name: str, log: list) -> None: + self.name, self.log = name, log + + def wait_stream(self, other: "_Stream") -> None: + self.log.append(f"{self.name} waits for {other.name}") + + +def _forward_step_executor(monkeypatch, log: list): + """A ``PyExecutor`` with only what ``_forward_step`` reads, logging its stream work.""" + from tensorrt_llm._torch.pyexecutor import py_executor + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + current = _Stream("current", log) + + @contextlib.contextmanager + def stream(s): + log.append(f"enter {s.name}") + yield + log.append(f"leave {s.name}") + + monkeypatch.setattr(torch.cuda, "current_stream", lambda: current) + monkeypatch.setattr(torch.cuda, "stream", stream) + monkeypatch.setattr(py_executor, "nvtx_range", lambda message: (lambda fn: fn)) + + executor = PyExecutor.__new__(PyExecutor) + executor.iter_counter = 0 + executor.enable_spec_decode = False + executor.runtime_draft_len = 0 + executor.execution_stream = _Stream("execution", log) + executor.resource_manager = object() + executor.sampler = types.SimpleNamespace(get_cache_indirection=lambda: None) + executor.model_engine = types.SimpleNamespace( + forward=lambda *args, **kwargs: log.append("forward") or "outputs", + route_capture=None, + ) + executor._iter_adp_dummy_ctx_tokens = 0 + executor._iter_adp_dummy_gen_tokens = 0 + executor._compute_adp_dummy_tokens = lambda requests: (0, 0) + executor._maybe_record_hang_diagnostic_phase = lambda *args, **kwargs: None + executor._attach_encoder_output_to_execution_stream = lambda requests: None + executor._mark_cross_kv_projection_consumed = lambda requests: None + executor._kv_connector_wait_for_save = lambda: None + return executor + + +_NO_REQUESTS = types.SimpleNamespace( + context_requests=[], num_context_requests=0, num_generation_requests=0 +) + + +def test_forward_step_orders_the_streams_off_the_execution_stream_loop(monkeypatch): + log = [] + executor = _forward_step_executor(monkeypatch, log) + + assert executor._forward_step(_NO_REQUESTS) == "outputs" + assert log == [ + "execution waits for current", + "enter execution", + "forward", + "leave execution", + "current waits for execution", + ] + + +def test_forward_step_skips_the_waits_on_the_execution_stream_loop(monkeypatch): + log = [] + executor = _forward_step_executor(monkeypatch, log) + executor._loop_on_execution_stream = True + + assert executor._forward_step(_NO_REQUESTS) == "outputs" + assert log == ["enter execution", "forward", "leave execution"] + + +# -- executor events -------------------------------------------------------------------------- + + +def _cleanup_executor(events, completed: bool): + """A ``PyExecutor`` with only what ``_executor_loop_cleanup`` reads, without pipeline stages.""" + import threading + + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + executor = PyExecutor.__new__(PyExecutor) + executor.response_cv = threading.Condition() + executor.shutdown_event = threading.Event() + executor.num_micro_batches = 0 + executor._event_loop_completed = completed + if events is not None: + executor._executor_events = events + return executor + + +class _RecordingEvents(_Events): + def __init__(self) -> None: + self.calls = [] + + def at_shutdown(self) -> None: + self.calls.append("at_shutdown") + + +@pytest.mark.parametrize("completed", [False, True]) +def test_executor_loop_cleanup_without_events(completed): + executor = _cleanup_executor(None, completed) + + executor._executor_loop_cleanup() + assert executor.is_shutdown and executor.shutdown_event.is_set() + + +@pytest.mark.parametrize("completed", [False, True]) +def test_executor_loop_cleanup_calls_at_shutdown_after_a_completed_loop(completed): + events = _RecordingEvents() + executor = _cleanup_executor(events, completed) + + executor._executor_loop_cleanup() + assert executor.is_shutdown + assert events.calls == (["at_shutdown"] if completed else [])